From ccd805f7da8ecbdb17ce2c8561d97ff83492a604 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Mon, 24 Aug 2026 16:40:06 -0700 Subject: [PATCH 01/31] feat: add the MoeEp Python API Expose validated forward and backward contracts with lazy optional dependency loading so applications can configure expert-parallel execution without affecting existing imports. --- pyproject.toml | 15 +- python/cudnn/__init__.py | 34 +- python/cudnn/moe_ep/__init__.py | 22 + python/cudnn/moe_ep/_backend.py | 104 +++++ python/cudnn/moe_ep/_contracts.py | 78 ++++ python/cudnn/moe_ep/_tuning.py | 109 +++++ python/cudnn/moe_ep/_types.py | 302 ++++++++++++++ python/cudnn/moe_ep/_validation.py | 623 +++++++++++++++++++++++++++++ python/cudnn/moe_ep/api.py | 529 ++++++++++++++++++++++++ 9 files changed, 1813 insertions(+), 3 deletions(-) create mode 100644 python/cudnn/moe_ep/__init__.py create mode 100644 python/cudnn/moe_ep/_backend.py create mode 100644 python/cudnn/moe_ep/_contracts.py create mode 100644 python/cudnn/moe_ep/_tuning.py create mode 100644 python/cudnn/moe_ep/_types.py create mode 100644 python/cudnn/moe_ep/_validation.py create mode 100644 python/cudnn/moe_ep/api.py diff --git a/pyproject.toml b/pyproject.toml index 66d51bdd0..d40aa7987 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "nvidia-cudnn-frontend" dynamic = ["version"] -description = "NVIDIA cuDNN Frontend — Python and C++ Graph API with SOTA attention (SDPA / Flash Attention), MoE grouped GEMM fusions, and FP8/MXFP8 kernels for Hopper and Blackwell GPUs." +description = "NVIDIA cuDNN Frontend — Python and C++ Graph API with SOTA attention (SDPA / Flash Attention), MoE grouped GEMM fusions, and FP8/MXFP8 kernels for Hopper, Blackwell, and Rubin GPUs." readme = "README.md" requires-python = ">=3.9" license = {text = "Apache-2.0 AND MIT"} @@ -21,6 +21,7 @@ keywords = [ "fp8", "mxfp8", "blackwell", + "rubin", "hopper", "pytorch", "kernel", @@ -68,6 +69,14 @@ cutedsl = [ "cuda-python", "apache-tvm-ffi>=0.1.11", ] +moe_ep = [ + "nvidia-cutlass-dsl[cu13]>=4.8.0", + "nvshmem4py-cu13>=0.3.1", + "cuda-python", + "torch", + "apache-tvm-ffi>=0.1.11", + "torch-c-dlpack-ext", +] cutile = [ # The cuTile linear-attention engines. Base cuda-tile only -- its [tileiras] # extra pins cuda-toolkit>=13.2,<13.4, and that upper bound would cap the @@ -119,3 +128,7 @@ version = {attr = "cudnn.__version__"} [tool.setuptools.package-data] include = ["**/*"] +"cudnn.moe_ep._megamoe_backend.cutedsl_src" = [ + "LICENSE.Apache-2.0", + "VENDOR_INFO.md", +] diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 22ffbf7a8..065fd7767 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -305,10 +305,33 @@ def _dlopen_cudnn(): ) __all__ = [*_EAGER_PUBLIC_NAMES, "Graph", "wrapper"] -_OPTIONAL_DEPENDENCY_INSTALL_HINT = "Install with 'pip install nvidia-cudnn-frontend[cutedsl]'" +_CUTEDSL_INSTALL_HINT = ( + "Install with 'pip install nvidia-cudnn-frontend[cutedsl]'" +) +_MOE_EP_INSTALL_HINT = ( + "Install with 'pip install nvidia-cudnn-frontend[moe_ep]'" +) +_MOE_EP_OPTIONAL_IMPORTS = { + "moe_ep", + "BlockScaledTensor", + "MoeEp", + "MoeEpTuningConfig", + "MoeEpWgradForwardStash", + "MoeEpWgradOperands", + "MoeFormat", + "MoeTensor", +} _LAZY_OPTIONAL_IMPORTS = { "gnn": (".gnn", None), + "moe_ep": (".moe_ep", None), + "BlockScaledTensor": (".moe_ep", "BlockScaledTensor"), + "MoeEp": (".moe_ep", "MoeEp"), + "MoeEpTuningConfig": (".moe_ep", "MoeEpTuningConfig"), + "MoeEpWgradForwardStash": (".moe_ep", "MoeEpWgradForwardStash"), + "MoeEpWgradOperands": (".moe_ep", "MoeEpWgradOperands"), + "MoeFormat": (".moe_ep", "MoeFormat"), + "MoeTensor": (".moe_ep", "MoeTensor"), "BSA": (".block_sparse_attention", "BSA"), "block_sparse_attention_forward": (".block_sparse_attention", "block_sparse_attention_forward"), "block_sparse_attention_fp8_forward": (".block_sparse_attention", "block_sparse_attention_fp8_forward"), @@ -386,7 +409,14 @@ def _load_optional_symbol(name: str) -> Any: module = importlib.import_module(module_name, package=__name__) value = module if attr_name is None else getattr(module, attr_name) except Exception as e: - raise ImportError(f"{name} requires optional dependencies. {_OPTIONAL_DEPENDENCY_INSTALL_HINT}: {e}") from e + install_hint = ( + _MOE_EP_INSTALL_HINT + if name in _MOE_EP_OPTIONAL_IMPORTS + else _CUTEDSL_INSTALL_HINT + ) + raise ImportError( + f"{name} requires optional dependencies. {install_hint}: {e}" + ) from e globals()[name] = value return value diff --git a/python/cudnn/moe_ep/__init__.py b/python/cudnn/moe_ep/__init__.py new file mode 100644 index 000000000..ed160766f --- /dev/null +++ b/python/cudnn/moe_ep/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +from ._tuning import MoeEpTuningConfig +from ._types import ( + BlockScaledTensor, + MoeEpWgradForwardStash, + MoeEpWgradOperands, + MoeFormat, + MoeTensor, +) +from .api import MoeEp + +__all__ = [ + "BlockScaledTensor", + "MoeEp", + "MoeEpTuningConfig", + "MoeEpWgradForwardStash", + "MoeEpWgradOperands", + "MoeFormat", + "MoeTensor", +] diff --git a/python/cudnn/moe_ep/_backend.py b/python/cudnn/moe_ep/_backend.py new file mode 100644 index 000000000..4efa856ab --- /dev/null +++ b/python/cudnn/moe_ep/_backend.py @@ -0,0 +1,104 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lightweight private backend seam for :class:`cudnn.moe_ep.MoeEp`. + +Capability policy and the executable factory are imported lazily through this +backend-neutral seam. Importing :mod:`cudnn` still does not load CuTeDSL or +initialize CUDA. +""" + +from __future__ import annotations + +from typing import Protocol, Tuple, Union + +import torch + +from ._contracts import ( + ForwardConfig, + ValidatedBackwardRequest, + ValidatedForwardRequest, +) +from ._types import MoeEpWgradForwardStash, MoeEpWgradOperands, MoeTensor + + +ForwardResult = Union[ + MoeTensor, + Tuple[MoeTensor, torch.Tensor, torch.Tensor], + Tuple[ + MoeTensor, + torch.Tensor, + torch.Tensor, + MoeEpWgradForwardStash, + ], +] +BackwardResult = Union[ + Tuple[torch.Tensor, torch.Tensor], + Tuple[torch.Tensor, torch.Tensor, MoeEpWgradOperands], +] + + +class MoeEpBackend(Protocol): + """Instance-local backend created lazily for one static ``MoeEp`` config.""" + + def forward(self, request: ValidatedForwardRequest) -> ForwardResult: + """Execute one already-validated forward request.""" + + def backward(self, request: ValidatedBackwardRequest) -> BackwardResult: + """Execute one already-validated backward request.""" + + def close(self) -> None: + """Release backend-owned resources.""" + + +class BackendUnavailableError(RuntimeError): + """The requested supported path has no executable runtime backend yet.""" + + +def validate_config(config: ForwardConfig) -> None: + """Run the selected backend's static capability gate lazily.""" + + from ._megamoe_backend._capability import validate_config as validate + + validate(config) + + +def validate_request(request: ValidatedForwardRequest) -> None: + """Run the selected backend's per-request capability gate lazily.""" + + from ._megamoe_backend._capability import validate_request as validate + + validate(request) + + +def validate_backward_request(request: ValidatedBackwardRequest) -> None: + """Run the selected backend's backward capability gate lazily.""" + + from ._megamoe_backend._capability import ( + validate_backward_request as validate, + ) + + validate(request) + + +def create_backend( + config: ForwardConfig, + device: torch.device, +) -> MoeEpBackend: + """Create the default backend without an allocation-only fallback.""" + + from ._megamoe_backend.mxfp8._backend import Mxfp8Backend + + return Mxfp8Backend(config, device) + + +__all__ = [ + "BackwardResult", + "BackendUnavailableError", + "MoeEpBackend", + "ForwardResult", + "create_backend", + "validate_config", + "validate_backward_request", + "validate_request", +] diff --git a/python/cudnn/moe_ep/_contracts.py b/python/cudnn/moe_ep/_contracts.py new file mode 100644 index 000000000..bad8a0c60 --- /dev/null +++ b/python/cudnn/moe_ep/_contracts.py @@ -0,0 +1,78 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Private data contracts shared by the MoE EP API, validation, and backend.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Literal, Optional + +import torch + +from ._tuning import MoeEpTuningConfig +from ._types import MoeEpWgradForwardStash, MoeTensor + + +@dataclass(frozen=True) +class ForwardConfig: + """Static configuration snapshot for one ``MoeEp`` instance.""" + + num_experts: int + hidden_size: int + intermediate_size: int + top_k: int + experts_per_rank: int + ep_size: int + ep_rank: int + ep_group: Any + ep_global_ranks: tuple[int, ...] + max_tokens_per_rank: Optional[int] + output_format: str + combine_format: str + apply_topk_in_fc1: bool + gate_up_clamp: Optional[float] + generate_c: bool + token_padding_size: int + sf_padding_size: int + tuning: MoeEpTuningConfig + backward_wgrad_mode: Literal["none", "operands"] = "none" + + +@dataclass(frozen=True) +class ValidatedForwardRequest: + """Runtime inputs that have passed the public forward contract.""" + + config: ForwardConfig + activation: MoeTensor + fc1_weight: MoeTensor + fc2_weight: MoeTensor + topk_idx: torch.Tensor + topk_weights: torch.Tensor + token_count: int + device: torch.device + + +@dataclass(frozen=True) +class ValidatedBackwardRequest: + """Runtime inputs that have passed the public backward contract.""" + + config: ForwardConfig + grad_output: torch.Tensor + fc1_weight: MoeTensor + fc2_weight: MoeTensor + topk_idx: torch.Tensor + topk_weights: torch.Tensor + fc1_c: torch.Tensor + route_metadata: torch.Tensor + token_count: int + local_routes: int + device: torch.device + wgrad_forward_stash: Optional[MoeEpWgradForwardStash] = None + + +__all__ = [ + "ForwardConfig", + "ValidatedBackwardRequest", + "ValidatedForwardRequest", +] diff --git a/python/cudnn/moe_ep/_tuning.py b/python/cudnn/moe_ep/_tuning.py new file mode 100644 index 000000000..d9393c481 --- /dev/null +++ b/python/cudnn/moe_ep/_tuning.py @@ -0,0 +1,109 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Public, semantic-preserving performance tuning for :class:`MoeEp`.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + + +TokenBackMode = Literal[ + "epi_warps", + "standalone_warps", + "reuse_dispatch_warps", +] + +_TOKEN_BACK_MODES = frozenset( + { + "epi_warps", + "standalone_warps", + "reuse_dispatch_warps", + } +) +_EPI_FLAG_BATCHES = frozenset( + { + (4, 2), + (1, 1), + (1, 2), + (1, 4), + (2, 1), + (2, 2), + (2, 4), + (4, 4), + } +) +_TOKEN_IN_FLAG_BATCHES = frozenset({1, 2, 4, 8, 16}) +_GROUP_HINTS = frozenset({64, 128, 256, 512, 768, 1024}) + + +@dataclass(frozen=True, kw_only=True) +class MoeEpTuningConfig: + """Validated Rubin MegaMoE performance knobs. + + These fields select scheduling and transport implementations without + changing the public MoE mathematical contract. Every rank in an expert + parallel group must use the same configuration. + + ``group_hint=None`` preserves the default behavior: the backend uses the + number of hardware-resident CTA clusters. + """ + + token_back_mode: TokenBackMode = "epi_warps" + epi_flag_batch: tuple[int, int] = (1, 1) + token_in_flag_batch: int = 1 + group_hint: int | None = None + reduce_topk_in_kernel: bool = False + + def __post_init__(self) -> None: + if ( + not isinstance(self.token_back_mode, str) + or self.token_back_mode not in _TOKEN_BACK_MODES + ): + raise ValueError( + "token_back_mode must be one of " + f"{tuple(sorted(_TOKEN_BACK_MODES))}, got " + f"{self.token_back_mode!r}" + ) + if ( + not isinstance(self.epi_flag_batch, tuple) + or self.epi_flag_batch not in _EPI_FLAG_BATCHES + ): + raise ValueError( + "epi_flag_batch must be one of " + f"{tuple(sorted(_EPI_FLAG_BATCHES))}, got " + f"{self.epi_flag_batch!r}" + ) + if ( + isinstance(self.token_in_flag_batch, bool) + or self.token_in_flag_batch not in _TOKEN_IN_FLAG_BATCHES + ): + raise ValueError( + "token_in_flag_batch must be one of " + f"{tuple(sorted(_TOKEN_IN_FLAG_BATCHES))}, got " + f"{self.token_in_flag_batch!r}" + ) + if self.group_hint is not None and ( + isinstance(self.group_hint, bool) + or self.group_hint not in _GROUP_HINTS + ): + raise ValueError( + "group_hint must be None or one of " + f"{tuple(sorted(_GROUP_HINTS))}, got {self.group_hint!r}" + ) + if not isinstance(self.reduce_topk_in_kernel, bool): + raise ValueError( + "reduce_topk_in_kernel must be a bool, got " + f"{self.reduce_topk_in_kernel!r}" + ) + if ( + self.reduce_topk_in_kernel + and self.token_back_mode != "epi_warps" + ): + raise ValueError( + "reduce_topk_in_kernel requires " + "token_back_mode='epi_warps'" + ) + +__all__ = ["MoeEpTuningConfig"] diff --git a/python/cudnn/moe_ep/_types.py b/python/cudnn/moe_ep/_types.py new file mode 100644 index 000000000..1525060ab --- /dev/null +++ b/python/cudnn/moe_ep/_types.py @@ -0,0 +1,302 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lightweight public tensor and format types for :mod:`cudnn.moe_ep`.""" + +from __future__ import annotations + +import operator +from dataclasses import dataclass +from enum import Enum +from typing import Tuple, Union + +import torch + + +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: + """Normalize a public format value.""" + + 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 _normalize_axis(axis: int, ndim: int) -> int: + if isinstance(axis, bool): + raise ValueError(f"axis must be an integer, got {axis!r}") + try: + axis = operator.index(axis) + except TypeError as exc: + raise ValueError(f"axis must be an integer, got {axis!r}") from exc + normalized = axis + ndim if axis < 0 else axis + if normalized < 0 or normalized >= ndim: + raise ValueError(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: + if not isinstance(self.data, torch.Tensor): + raise ValueError( + f"data must be a torch.Tensor, got {type(self.data).__name__}" + ) + if not isinstance(self.scale, torch.Tensor): + raise ValueError( + f"scale must be a torch.Tensor, got {type(self.scale).__name__}" + ) + fmt = parse_format(self.format) + if fmt is MoeFormat.BF16: + raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") + if self.data.device != self.scale.device: + raise ValueError( + f"data device {self.data.device} does not match " + f"scale device {self.scale.device}" + ) + try: + raw_logical_shape = tuple(self.logical_shape) + except TypeError as exc: + raise ValueError( + "logical_shape must be an iterable of integers" + ) from exc + logical_shape = [] + for dim in raw_logical_shape: + if isinstance(dim, bool): + raise ValueError( + f"logical_shape dimensions must be integers, got {dim!r}" + ) + try: + dim = operator.index(dim) + except TypeError as exc: + raise ValueError( + f"logical_shape dimensions must be integers, got {dim!r}" + ) from exc + if dim < 0: + raise ValueError( + f"logical_shape dimensions must be non-negative, got {dim}" + ) + logical_shape.append(dim) + normalized_shape = tuple(logical_shape) + axis = _normalize_axis(self.axis, len(normalized_shape)) + logical_extent = normalized_shape[axis] + block_size = 32 if fmt is MoeFormat.MXFP8 else 16 + payload_extent = ( + logical_extent + if fmt is MoeFormat.MXFP8 + else (logical_extent + 1) // 2 + ) + scale_extent = (logical_extent + block_size - 1) // block_size + expected_data_shape = list(normalized_shape) + expected_data_shape[axis] = payload_extent + expected_scale_shape = list(normalized_shape) + expected_scale_shape[axis] = scale_extent + expected_data_shape = tuple(expected_data_shape) + expected_scale_shape = tuple(expected_scale_shape) + if tuple(self.data.shape) != expected_data_shape: + raise ValueError( + f"{fmt.value} data shape must be {expected_data_shape}, " + f"got {tuple(self.data.shape)}" + ) + if tuple(self.scale.shape) != expected_scale_shape: + raise ValueError( + f"{fmt.value} scale shape must be {expected_scale_shape}, " + f"got {tuple(self.scale.shape)}" + ) + e4m3_dtype = getattr(torch, "float8_e4m3fn", None) + if e4m3_dtype is None: + raise RuntimeError("this PyTorch build does not provide torch.float8_e4m3fn") + if fmt is MoeFormat.MXFP8: + expected_data_dtype = e4m3_dtype + expected_scale_dtype = getattr(torch, "float8_e8m0fnu", None) + if expected_scale_dtype is None: + raise RuntimeError( + "this PyTorch build does not provide torch.float8_e8m0fnu" + ) + else: + expected_data_dtype = torch.uint8 + expected_scale_dtype = e4m3_dtype + if self.data.dtype is not expected_data_dtype: + raise ValueError( + f"{fmt.value} data must have dtype {expected_data_dtype}, " + f"got {self.data.dtype}" + ) + if self.scale.dtype is not expected_scale_dtype: + raise ValueError( + f"{fmt.value} scale must have dtype {expected_scale_dtype}, " + f"got {self.scale.dtype}" + ) + object.__setattr__(self, "format", fmt) + object.__setattr__(self, "logical_shape", normalized_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) + + +@dataclass(frozen=True) +class MoeEpWgradForwardStash: + """Caller-owned forward state required to form expert-local wgrads. + + ``fc1_a`` and ``fc1_sfa`` represent the MXFP8 ``x.T`` operand. Valid + routes for each local expert occupy the beginning of its padded range; + ``expert_offsets`` contains cumulative padded end offsets and + ``valid_route_counts`` contains the corresponding unpadded row counts. + Scale factors use the blocked layout consumed by grouped wgrad, with + logical 1x32 scaling and physical 128x4 scale tiles. + ``route_metadata`` is the compact identity table returned by forward, + using ``(local_expert, src_rank, src_token, src_slot)`` rows. It validates + that the stash belongs to the matching routed call; it is not padded or + row-aligned with the operands' K dimension. + """ + + fc1_a: torch.Tensor + fc1_sfa: torch.Tensor + expert_offsets: torch.Tensor + valid_route_counts: torch.Tensor + route_metadata: torch.Tensor + + +@dataclass(frozen=True) +class MoeEpWgradOperands: + """Caller-owned MXFP8 operands for expert-local grouped wgrad GEMMs. + + The represented operations are ``dW1 = fc1_a @ fc1_b`` and + ``dW2 = fc2_a @ fc2_b``. For total padded route extent ``K``, their + logical shapes are ``fc1_a=(H,K)``, ``fc1_b=(K,2I)``, + ``fc2_a=(I,K)``, and ``fc2_b=(K,H)``. Each scale tensor uses grouped + wgrad's blocked 1x32 layout: ``(round_up(non-K,128), round_up(K/32,4))``. + The shared expert metadata has the same meaning as in + :class:`MoeEpWgradForwardStash`. + + Attributes: + fc1_a: E4M3 data for the FC1 A operand, logically ``x.T`` with shape + ``(H, K)``. ``x`` is the activation dispatched to each local + expert. The K dimension concatenates the experts' independently + padded route ranges. + fc1_sfa: E8M0 scales for ``fc1_a``. Each logical scale covers 32 + consecutive K elements of one hidden-feature row. + fc1_b: E4M3 data for the FC1 B operand, logically + ``dC=[d_gate | d_up]`` with shape ``(K, 2I)``. ``dC`` is the + gradient of the pre-SwiGLU FC1 accumulator; columns use the public + gate-then-up order rather than the kernel's internal strip + interleave. + fc1_sfb: E8M0 scales for ``fc1_b``. Each logical scale covers 32 + consecutive K rows for one gate/up feature column. + fc2_a: E4M3 data for the FC2 A operand, logically ``(p*h).T`` with + shape ``(I, K)``. ``h=SwiGLU(C)`` and ``p`` is the route's FP32 + router score, applied exactly once before column quantization. + fc2_sfa: E8M0 scales for ``fc2_a``. Each logical scale covers 32 + consecutive K elements of one intermediate-feature row. + fc2_b: E4M3 data for the FC2 B operand, logically unweighted ``dY`` + with shape ``(K, H)``. ``dY`` is the routed FC2 output gradient. + fc2_sfb: E8M0 scales for ``fc2_b``. Each logical scale covers 32 + consecutive K rows for one hidden-feature column. + expert_offsets: Int32 cumulative padded K-end offset for every local + expert. Adjacent equal offsets represent an empty expert. + valid_route_counts: Int32 unpadded route count for every local expert; + valid rows occupy the beginning of each padded expert range. + route_metadata: Compact Int32 route identity table with columns + ``(local_expert, src_rank, src_token, src_slot)``. It identifies + the routed call but is not padded or row-aligned with K. + """ + + fc1_a: torch.Tensor + fc1_sfa: torch.Tensor + fc1_b: torch.Tensor + fc1_sfb: torch.Tensor + fc2_a: torch.Tensor + fc2_sfa: torch.Tensor + fc2_b: torch.Tensor + fc2_sfb: torch.Tensor + expert_offsets: torch.Tensor + valid_route_counts: torch.Tensor + route_metadata: torch.Tensor + + +MoeTensor = Union[torch.Tensor, BlockScaledTensor] + + +__all__ = [ + "BlockScaledTensor", + "MoeEpWgradForwardStash", + "MoeEpWgradOperands", + "MoeFormat", + "MoeTensor", + "parse_format", +] diff --git a/python/cudnn/moe_ep/_validation.py b/python/cudnn/moe_ep/_validation.py new file mode 100644 index 000000000..9ee259fa3 --- /dev/null +++ b/python/cudnn/moe_ep/_validation.py @@ -0,0 +1,623 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Pure public-contract validation for :mod:`cudnn.moe_ep`. + +This module intentionally depends only on PyTorch and the lightweight public +API types. It must not import CuTeDSL, CUDA Python, NVSHMEM, or the private +MegaMoE runtime. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from ._contracts import ( + ForwardConfig, + ValidatedBackwardRequest, + ValidatedForwardRequest, +) +from ._types import ( + BlockScaledTensor, + MoeEpWgradForwardStash, + MoeFormat, + MoeTensor, +) + + +def _replace_axis(shape: Tuple[int, ...], axis: int, extent: int) -> Tuple[int, ...]: + result = list(shape) + result[axis] = extent + return tuple(result) + + +def _ceil_div(value: int, divisor: int) -> int: + return (value + divisor - 1) // divisor + + +def _round_up(value: int, multiple: int) -> int: + return _ceil_div(value, multiple) * multiple + + +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 _logical_shape(tensor: MoeTensor) -> Tuple[int, ...]: + if isinstance(tensor, BlockScaledTensor): + return tensor.logical_shape + if isinstance(tensor, torch.Tensor): + return tuple(tensor.shape) + raise ValueError( + f"expected torch.Tensor or BlockScaledTensor, got {type(tensor).__name__}" + ) + + +def _tensor_device(tensor: MoeTensor) -> torch.device: + if isinstance(tensor, BlockScaledTensor): + return tensor.device + if isinstance(tensor, torch.Tensor): + return tensor.device + raise ValueError( + f"expected torch.Tensor or BlockScaledTensor, got {type(tensor).__name__}" + ) + + +def _validate_strided(name: str, tensor: torch.Tensor) -> None: + if tensor.layout is not torch.strided: + raise ValueError(f"{name} must use torch.strided layout, got {tensor.layout}") + + +def _validate_tensor_representation( + name: str, + tensor: MoeTensor, + expected_logical_shape: Tuple[int, ...], +) -> None: + logical_shape = _logical_shape(tensor) + if logical_shape != expected_logical_shape: + raise ValueError( + f"{name} logical shape must be {expected_logical_shape}, got {logical_shape}" + ) + + if isinstance(tensor, torch.Tensor): + _validate_strided(name, tensor) + if not tensor.is_floating_point(): + raise ValueError(f"{name} must be floating point, got {tensor.dtype}") + return + + if not isinstance(tensor, BlockScaledTensor): + raise ValueError( + f"{name} must be a torch.Tensor or BlockScaledTensor, " + f"got {type(tensor).__name__}" + ) + + if tensor.axis != 1: + raise ValueError(f"{name} block-scaled axis must be 1, got {tensor.axis}") + _validate_strided(f"{name}.data", tensor.data) + _validate_strided(f"{name}.scale", tensor.scale) + + logical_extent = expected_logical_shape[tensor.axis] + if tensor.format is MoeFormat.MXFP8: + payload_extent = logical_extent + block_size = 32 + expected_data_dtype = _require_torch_dtype("float8_e4m3fn") + expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") + else: + payload_extent = _ceil_div(logical_extent, 2) + block_size = 16 + expected_data_dtype = torch.uint8 + expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") + + expected_data_shape = _replace_axis( + expected_logical_shape, + tensor.axis, + payload_extent, + ) + expected_scale_shape = _replace_axis( + expected_logical_shape, + tensor.axis, + _ceil_div(logical_extent, block_size), + ) + if tuple(tensor.data.shape) != expected_data_shape: + raise ValueError( + f"{name}.data shape must be {expected_data_shape}, " + f"got {tuple(tensor.data.shape)}" + ) + if tuple(tensor.scale.shape) != expected_scale_shape: + raise ValueError( + f"{name}.scale shape must be {expected_scale_shape}, " + f"got {tuple(tensor.scale.shape)}" + ) + if tensor.data.dtype != expected_data_dtype: + raise ValueError( + f"{name}.data must have dtype {expected_data_dtype}, " + f"got {tensor.data.dtype}" + ) + if tensor.scale.dtype != expected_scale_dtype: + raise ValueError( + f"{name}.scale must have dtype {expected_scale_dtype}, " + f"got {tensor.scale.dtype}" + ) + + +def _validate_expert_ids( + config: ForwardConfig, + topk_idx: torch.Tensor, +) -> None: + valid_experts = topk_idx.reshape(-1) + valid_experts = valid_experts[valid_experts != -1] + if valid_experts.numel() > 0 and bool( + ( + (valid_experts < 0) + | (valid_experts >= config.num_experts) + ).any().item() + ): + raise ValueError("topk_idx contains out-of-range expert ids") + + +def _validate_routes( + config: ForwardConfig, + token_count: int, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + validate_expert_ids: bool, +) -> None: + """Validate the public routing plane shared by forward and backward.""" + + if not isinstance(topk_idx, torch.Tensor): + raise ValueError( + f"topk_idx must be a torch.Tensor, got {type(topk_idx).__name__}" + ) + if not isinstance(topk_weights, torch.Tensor): + raise ValueError( + "topk_weights must be a torch.Tensor, " + f"got {type(topk_weights).__name__}" + ) + _validate_strided("topk_idx", topk_idx) + _validate_strided("topk_weights", topk_weights) + + route_shape = (token_count, config.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}, " + f"got {tuple(topk_weights.shape)}" + ) + if topk_idx.dtype not in (torch.int32, torch.int64): + raise ValueError( + "topk_idx must have dtype torch.int32 or torch.int64, " + f"got {topk_idx.dtype}" + ) + if not topk_weights.is_floating_point(): + raise ValueError( + f"topk_weights must be floating point, got {topk_weights.dtype}" + ) + if ( + config.max_tokens_per_rank is not None + and token_count > config.max_tokens_per_rank + ): + raise ValueError( + f"token count {token_count} exceeds " + f"max_tokens_per_rank={config.max_tokens_per_rank}" + ) + if validate_expert_ids: + _validate_expert_ids(config, topk_idx) + + +def validate_forward( + config: ForwardConfig, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + validate_expert_ids: bool = True, +) -> ValidatedForwardRequest: + """Validate public forward semantics without importing a device backend.""" + + activation_shape = _logical_shape(activation) + if len(activation_shape) != 2 or activation_shape[1] != config.hidden_size: + raise ValueError( + f"activation logical shape must be (T, {config.hidden_size}), " + f"got {activation_shape}" + ) + token_count = activation_shape[0] + expected_fc1 = ( + config.experts_per_rank, + config.hidden_size, + 2 * config.intermediate_size, + ) + expected_fc2 = ( + config.experts_per_rank, + config.intermediate_size, + config.hidden_size, + ) + _validate_tensor_representation("activation", activation, activation_shape) + _validate_tensor_representation("fc1_weight", fc1_weight, expected_fc1) + _validate_tensor_representation("fc2_weight", fc2_weight, expected_fc2) + + _validate_routes( + config, + token_count, + topk_idx, + topk_weights, + validate_expert_ids=False, + ) + + device = _tensor_device(activation) + for name, tensor in ( + ("fc1_weight", fc1_weight), + ("fc2_weight", fc2_weight), + ("topk_idx", topk_idx), + ("topk_weights", topk_weights), + ): + tensor_device = _tensor_device(tensor) + if tensor_device != device: + raise ValueError(f"{name} must be on {device}, got {tensor_device}") + + # Boolean compaction plus the host-visible ``item()`` below is not CUDA + # graph capturable. Eager calls (including the mandatory pre-capture + # warmup) retain strict validation. During capture/replay, callers must + # preserve that validated invariant: every route is -1 or a valid global + # expert ID. + if device.type == "cuda": + with torch.cuda.device(device): + capturing = torch.cuda.is_current_stream_capturing() + else: + capturing = False + if validate_expert_ids and not capturing: + _validate_expert_ids(config, topk_idx) + + return ValidatedForwardRequest( + config=config, + activation=activation, + fc1_weight=fc1_weight, + fc2_weight=fc2_weight, + topk_idx=topk_idx, + topk_weights=topk_weights, + token_count=token_count, + device=device, + ) + + +def _validate_wgrad_forward_stash( + config: ForwardConfig, + stash: MoeEpWgradForwardStash, + route_metadata: torch.Tensor, + device: torch.device, +) -> None: + """Validate caller-owned forward operands and their route identity.""" + + if not isinstance(stash, MoeEpWgradForwardStash): + raise TypeError( + "wgrad_forward_stash must be a MoeEpWgradForwardStash, " + f"got {type(stash).__name__}" + ) + + tensors = ( + ("fc1_a", stash.fc1_a), + ("fc1_sfa", stash.fc1_sfa), + ("expert_offsets", stash.expert_offsets), + ("valid_route_counts", stash.valid_route_counts), + ("route_metadata", stash.route_metadata), + ) + for name, tensor in tensors: + if not isinstance(tensor, torch.Tensor): + raise TypeError( + f"wgrad_forward_stash.{name} must be a torch.Tensor, " + f"got {type(tensor).__name__}" + ) + _validate_strided(f"wgrad_forward_stash.{name}", tensor) + if tensor.device != device: + raise ValueError( + f"wgrad_forward_stash.{name} must be on {device}, " + f"got {tensor.device}" + ) + + e4m3_dtype = _require_torch_dtype("float8_e4m3fn") + e8m0_dtype = _require_torch_dtype("float8_e8m0fnu") + if stash.fc1_a.dtype is not e4m3_dtype: + raise TypeError( + "wgrad_forward_stash.fc1_a must have dtype " + f"{e4m3_dtype}, got {stash.fc1_a.dtype}" + ) + if stash.fc1_sfa.dtype is not e8m0_dtype: + raise TypeError( + "wgrad_forward_stash.fc1_sfa must have dtype " + f"{e8m0_dtype}, got {stash.fc1_sfa.dtype}" + ) + + expert_shape = (config.experts_per_rank,) + for name, tensor in ( + ("expert_offsets", stash.expert_offsets), + ("valid_route_counts", stash.valid_route_counts), + ): + if tuple(tensor.shape) != expert_shape: + raise ValueError( + f"wgrad_forward_stash.{name} shape must be {expert_shape}, " + f"got {tuple(tensor.shape)}" + ) + if tensor.dtype is not torch.int32: + raise TypeError( + f"wgrad_forward_stash.{name} must have dtype torch.int32, " + f"got {tensor.dtype}" + ) + + if tuple(stash.route_metadata.shape) != tuple(route_metadata.shape): + raise ValueError( + "wgrad_forward_stash.route_metadata shape must match " + "route_metadata" + ) + if stash.route_metadata.dtype is not torch.int32: + raise TypeError( + "wgrad_forward_stash.route_metadata must have dtype torch.int32, " + f"got {stash.route_metadata.dtype}" + ) + if not torch.equal(stash.route_metadata, route_metadata): + raise ValueError( + "wgrad_forward_stash route identity does not match route_metadata" + ) + + offsets = [int(value) for value in stash.expert_offsets.cpu().tolist()] + counts = [int(value) for value in stash.valid_route_counts.cpu().tolist()] + previous = 0 + for expert, (offset, count) in enumerate(zip(offsets, counts)): + padded_routes = offset - previous + if offset < previous: + raise ValueError( + "wgrad_forward_stash.expert_offsets must be non-decreasing" + ) + if count < 0 or count > padded_routes: + raise ValueError( + "wgrad_forward_stash.valid_route_counts must fit each " + f"expert's padded range; expert {expert} has count={count} " + f"and capacity={padded_routes}" + ) + expected_padded_routes = _round_up( + count, + config.token_padding_size, + ) + if padded_routes != expected_padded_routes: + raise ValueError( + "wgrad_forward_stash expert ranges must use the canonical " + f"{config.token_padding_size}-row padding; expert {expert} " + f"has capacity={padded_routes}, expected=" + f"{expected_padded_routes}" + ) + previous = offset + + padded_route_count = offsets[-1] if offsets else 0 + expected_fc1_a = (config.hidden_size, padded_route_count) + if tuple(stash.fc1_a.shape) != expected_fc1_a: + raise ValueError( + "wgrad_forward_stash.fc1_a shape must be " + f"{expected_fc1_a}, got {tuple(stash.fc1_a.shape)}" + ) + if padded_route_count % 32: + raise ValueError( + "wgrad_forward_stash padded route count must be divisible by 32" + ) + expected_fc1_sfa = ( + _round_up(config.hidden_size, 128), + _round_up(padded_route_count // 32, 4), + ) + if tuple(stash.fc1_sfa.shape) != expected_fc1_sfa: + raise ValueError( + "wgrad_forward_stash.fc1_sfa shape must be " + f"{expected_fc1_sfa}, got {tuple(stash.fc1_sfa.shape)}" + ) + if padded_route_count and not stash.fc1_a.is_contiguous(): + raise ValueError( + "wgrad_forward_stash.fc1_a must use compact (K, 1) strides" + ) + if not stash.fc1_sfa.is_contiguous(): + raise ValueError( + "wgrad_forward_stash.fc1_sfa must be contiguous" + ) + for name, tensor, alignment in ( + ("fc1_a", stash.fc1_a, 16), + ("fc1_sfa", stash.fc1_sfa, 16), + ("expert_offsets", stash.expert_offsets, 4), + ("valid_route_counts", stash.valid_route_counts, 4), + ): + if tensor.data_ptr() % alignment: + raise ValueError( + f"wgrad_forward_stash.{name} must be " + f"{alignment}-byte aligned" + ) + + local_routes = int(route_metadata.shape[0]) + if sum(counts) != local_routes: + raise ValueError( + "wgrad_forward_stash.valid_route_counts must sum to the " + "route_metadata row count" + ) + if local_routes: + local_experts = route_metadata[:, 0].to(torch.int64) + if bool( + ( + (local_experts < 0) + | (local_experts >= config.experts_per_rank) + ).any().item() + ): + raise ValueError( + "route_metadata contains out-of-range local expert ids" + ) + metadata_counts = torch.bincount( + local_experts, + minlength=config.experts_per_rank, + ) + expected_counts = stash.valid_route_counts.to(torch.int64) + if not torch.equal(metadata_counts, expected_counts): + raise ValueError( + "wgrad_forward_stash.valid_route_counts do not match " + "route_metadata" + ) + expected_experts = torch.repeat_interleave( + torch.arange( + config.experts_per_rank, + dtype=torch.int64, + device=device, + ), + expected_counts, + output_size=local_routes, + ) + if not torch.equal(local_experts, expected_experts): + raise ValueError( + "route_metadata rows must be grouped by local expert" + ) + + src_ranks = route_metadata[:, 1] + src_tokens = route_metadata[:, 2] + src_slots = route_metadata[:, 3] + if bool(((src_ranks < 0) | (src_ranks >= config.ep_size)).any().item()): + raise ValueError("route_metadata contains out-of-range source ranks") + if bool((src_tokens < 0).any().item()): + raise ValueError("route_metadata contains negative source tokens") + if config.max_tokens_per_rank is not None and bool( + (src_tokens >= config.max_tokens_per_rank).any().item() + ): + raise ValueError("route_metadata contains out-of-range source tokens") + if bool(((src_slots < 0) | (src_slots >= config.top_k)).any().item()): + raise ValueError("route_metadata contains out-of-range source slots") + + +def validate_backward( + config: ForwardConfig, + grad_output: torch.Tensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + fc1_c: torch.Tensor, + route_metadata: torch.Tensor, + *, + wgrad_forward_stash: MoeEpWgradForwardStash | None = None, +) -> ValidatedBackwardRequest: + """Validate public backward semantics without importing a device backend.""" + + if not isinstance(grad_output, torch.Tensor): + raise TypeError( + "grad_output must be a torch.Tensor, " + f"got {type(grad_output).__name__}" + ) + _validate_strided("grad_output", grad_output) + if grad_output.ndim != 2 or grad_output.shape[1] != config.hidden_size: + raise ValueError( + f"grad_output shape must be (T, {config.hidden_size}), " + f"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}" + ) + token_count = int(grad_output.shape[0]) + expected_fc1 = ( + config.experts_per_rank, + config.hidden_size, + 2 * config.intermediate_size, + ) + expected_fc2 = ( + config.experts_per_rank, + config.intermediate_size, + config.hidden_size, + ) + _validate_tensor_representation("fc1_weight", fc1_weight, expected_fc1) + _validate_tensor_representation("fc2_weight", fc2_weight, expected_fc2) + + _validate_routes( + config, + token_count, + topk_idx, + topk_weights, + validate_expert_ids=True, + ) + + if not isinstance(route_metadata, torch.Tensor): + raise TypeError( + "route_metadata must be a torch.Tensor, " + f"got {type(route_metadata).__name__}" + ) + _validate_strided("route_metadata", route_metadata) + if route_metadata.ndim != 2 or route_metadata.shape[1] != 4: + raise ValueError( + "route_metadata shape must be (local_routes, 4), " + f"got {tuple(route_metadata.shape)}" + ) + if route_metadata.dtype is not torch.int32: + raise TypeError( + "route_metadata must have dtype torch.int32, " + f"got {route_metadata.dtype}" + ) + local_routes = int(route_metadata.shape[0]) + expected_fc1_c = (local_routes, 2 * config.intermediate_size) + if not isinstance(fc1_c, torch.Tensor): + raise TypeError( + f"fc1_c must be a torch.Tensor, got {type(fc1_c).__name__}" + ) + _validate_strided("fc1_c", fc1_c) + if tuple(fc1_c.shape) != expected_fc1_c: + raise ValueError( + f"fc1_c shape must be {expected_fc1_c}, got {tuple(fc1_c.shape)}" + ) + if fc1_c.dtype is not torch.bfloat16: + raise TypeError( + f"fc1_c must have dtype torch.bfloat16, got {fc1_c.dtype}" + ) + + device = grad_output.device + for name, tensor in ( + ("fc1_weight", fc1_weight), + ("fc2_weight", fc2_weight), + ("topk_idx", topk_idx), + ("topk_weights", topk_weights), + ("grad_output", grad_output), + ("fc1_c", fc1_c), + ("route_metadata", route_metadata), + ): + tensor_device = _tensor_device(tensor) + if tensor_device != device: + raise ValueError( + f"{name} must be on {device}, got {tensor_device}" + ) + + if config.backward_wgrad_mode == "operands": + _validate_wgrad_forward_stash( + config, + wgrad_forward_stash, + route_metadata, + device, + ) + elif wgrad_forward_stash is not None: + raise ValueError( + "wgrad_forward_stash is only accepted when " + "backward_wgrad_mode='operands'" + ) + + return ValidatedBackwardRequest( + config=config, + grad_output=grad_output, + fc1_weight=fc1_weight, + fc2_weight=fc2_weight, + topk_idx=topk_idx, + topk_weights=topk_weights, + fc1_c=fc1_c, + route_metadata=route_metadata, + token_count=token_count, + local_routes=local_routes, + device=device, + wgrad_forward_stash=wgrad_forward_stash, + ) + + +__all__ = ["validate_backward", "validate_forward"] diff --git a/python/cudnn/moe_ep/api.py b/python/cudnn/moe_ep/api.py new file mode 100644 index 000000000..900ab4b65 --- /dev/null +++ b/python/cudnn/moe_ep/api.py @@ -0,0 +1,529 @@ +# 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 API performs contract validation and dispatches through a private, +lazy backend seam. Device-runtime implementation details remain outside this +module. +""" + +from __future__ import annotations + +import contextlib +import math +import threading +import warnings +from numbers import Real +from typing import Literal, Optional, Tuple, Union + +import torch +import torch.distributed as dist + +from ._contracts import ForwardConfig, ValidatedBackwardRequest +from ._tuning import MoeEpTuningConfig +from ._types import ( + BlockScaledTensor, + MoeEpWgradForwardStash, + MoeEpWgradOperands, + MoeFormat, + MoeTensor, + parse_format as _parse_format, +) +from ._validation import validate_backward, validate_forward + +def _resolve_ep_topology( + ep_group: Optional[dist.ProcessGroup], +) -> tuple[int, int, tuple[int, ...]]: + """Return dense EP rank/size plus its ordered global-rank membership.""" + + if ep_group is None: + return 1, 0, () + 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 ep_size <= 0 or ep_rank < 0 or ep_rank >= ep_size: + raise ValueError("the current process must be a member of ep_group") + + ep_global_ranks = tuple( + dist.get_global_rank(ep_group, group_rank) + for group_rank in range(ep_size) + ) + if len(set(ep_global_ranks)) != ep_size: + raise RuntimeError("ep_group returned duplicate global ranks") + if ep_global_ranks[ep_rank] != dist.get_rank(): + raise RuntimeError( + "ep_group rank mapping is inconsistent with the current global rank" + ) + return ep_size, ep_rank, ep_global_ranks + + +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. + + The Rubin training-Mega backend accepts plain BF16/FP16/FP32 operands + (staged to MXFP8 E4M3) or MXFP8 ``BlockScaledTensor`` operands. Final + output is BF16. ``combine_format`` may be BF16 or MXFP8; forward MXFP8 + combine quantizes each FP32 route accumulator directly before top-k + reduction. The Rubin training backend requires + ``apply_topk_in_fc1=True``. + Native NVFP4 operands and NVFP4 combine/output are not executable. + + 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. + + With ``backward_wgrad_mode="operands"``, ``generate_c=True``, + ``token_padding_size=256``, and ``sf_padding_size=128`` are required. + Forward additionally returns a caller-owned + :class:`MoeEpWgradForwardStash`; backward accepts that exact routed-call + stash by keyword and additionally returns :class:`MoeEpWgradOperands`. + This opt-in path is available under the Rubin MXFP8 backward capability + gates documented below. + + The backend is created lazily on the first supported forward call. Valid + combinations outside the current backend capability matrix fail explicitly + instead of returning uninitialized storage. Once created, a backend and its + workspaces are bound to that call's device; use a separate ``MoeEp`` + instance for another device. + """ + + 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, + backward_wgrad_mode: Literal["none", "operands"] = "none", + token_padding_size: int = 128, + sf_padding_size: int = 128, + tuning: Optional[MoeEpTuningConfig] = None, + ) -> None: + self._lifecycle_lock = threading.RLock() + for name, value in ( + ("num_experts", num_experts), + ("hidden_size", hidden_size), + ("intermediate_size", intermediate_size), + ("top_k", top_k), + ): + if isinstance(value, bool) or 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 ( + isinstance(max_tokens_per_rank, bool) + or not isinstance(max_tokens_per_rank, int) + or max_tokens_per_rank < 0 + ): + raise ValueError( + "max_tokens_per_rank must be a non-negative integer or None" + ) + if not isinstance(apply_topk_in_fc1, bool): + raise ValueError("apply_topk_in_fc1 must be a bool") + if not isinstance(generate_c, bool): + raise ValueError("generate_c must be a bool") + if backward_wgrad_mode not in ("none", "operands"): + raise ValueError( + "backward_wgrad_mode must be 'none' or 'operands', " + f"got {backward_wgrad_mode!r}" + ) + if backward_wgrad_mode == "operands" and not generate_c: + raise ValueError( + "backward_wgrad_mode='operands' requires generate_c=True" + ) + for name, value in ( + ("token_padding_size", token_padding_size), + ("sf_padding_size", sf_padding_size), + ): + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError( + f"{name} must be a positive integer, got {value!r}" + ) + if backward_wgrad_mode == "operands" and token_padding_size != 256: + raise ValueError( + "backward_wgrad_mode='operands' requires " + "token_padding_size=256" + ) + if backward_wgrad_mode == "operands" and sf_padding_size != 128: + raise ValueError( + "backward_wgrad_mode='operands' requires " + "sf_padding_size=128" + ) + if sf_padding_size % 128: + raise ValueError( + "sf_padding_size must be a positive multiple of 128, " + f"got {sf_padding_size}" + ) + if tuning is not None and not isinstance(tuning, MoeEpTuningConfig): + raise TypeError( + "tuning must be a MoeEpTuningConfig or None, " + f"got {type(tuning).__name__}" + ) + if gate_up_clamp is not None: + if isinstance(gate_up_clamp, bool) or not isinstance(gate_up_clamp, Real): + raise ValueError("gate_up_clamp must be a finite real number or None") + gate_up_clamp = float(gate_up_clamp) + if not math.isfinite(gate_up_clamp): + raise ValueError("gate_up_clamp must be a finite real number or None") + + if ep_group is not None and not isinstance(ep_group, dist.ProcessGroup): + raise ValueError( + f"ep_group must be a torch.distributed.ProcessGroup or None, " + f"got {type(ep_group).__name__}" + ) + ep_size, ep_rank, ep_global_ranks = _resolve_ep_topology(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.ep_global_ranks = ep_global_ranks + 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 = apply_topk_in_fc1 + self.gate_up_clamp = None if gate_up_clamp is None else abs(gate_up_clamp) + self.generate_c = generate_c + self.backward_wgrad_mode = backward_wgrad_mode + self.token_padding_size = token_padding_size + self.sf_padding_size = sf_padding_size + self.tuning = MoeEpTuningConfig() if tuning is None else tuning + if self.tuning.reduce_topk_in_kernel and ( + self.combine_format is not MoeFormat.BF16 + or self.output_format is not MoeFormat.BF16 + or not self.apply_topk_in_fc1 + ): + raise ValueError( + "reduce_topk_in_kernel requires BF16 combine/output and " + "apply_topk_in_fc1=True" + ) + + 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}") + + self._forward_config = ForwardConfig( + num_experts=self.num_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + top_k=self.top_k, + experts_per_rank=self.experts_per_rank, + ep_size=self.ep_size, + ep_rank=self.ep_rank, + ep_group=self.ep_group, + ep_global_ranks=self.ep_global_ranks, + max_tokens_per_rank=self.max_tokens_per_rank, + output_format=self.output_format.value, + combine_format=self.combine_format.value, + apply_topk_in_fc1=self.apply_topk_in_fc1, + gate_up_clamp=self.gate_up_clamp, + generate_c=self.generate_c, + token_padding_size=self.token_padding_size, + sf_padding_size=self.sf_padding_size, + tuning=self.tuning, + backward_wgrad_mode=self.backward_wgrad_mode, + ) + self._forward_backend = None + self._forward_backend_device = None + self._validated_topk_idx = None + self._validated_topk_version = None + self._closed = False + + @staticmethod + def _tensor_version(tensor: torch.Tensor) -> int | None: + if not isinstance(tensor, torch.Tensor): + return None + try: + return tensor._version + except RuntimeError: + return None + + def _get_backend(self, request, *, backward: bool): + """Create and cache the private backend on first supported use.""" + + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + from . import _backend + + if ( + self._forward_backend is not None + and request.device != self._forward_backend_device + ): + raise ValueError( + f"MoeEp backend is bound to {self._forward_backend_device}; " + f"create a separate MoeEp instance for {request.device}" + ) + + _backend.validate_config(self._forward_config) + if backward: + _backend.validate_backward_request(request) + else: + _backend.validate_request(request) + + if self._forward_backend is None: + self._forward_backend = _backend.create_backend( + self._forward_config, + request.device, + ) + self._forward_backend_device = request.device + return self._forward_backend + + def _count_local_routes(self, request: ValidatedBackwardRequest) -> int: + """Number of valid routes this rank's experts receive. + + The request already passed expert-id validation. Data-dependent: + single-rank counts locally, while EP exchanges per-rank route counts + (the same exchange the device dispatch performs). + """ + + flat = request.topk_idx.reshape(-1).to(torch.int64) + expert = flat[flat != -1] + 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], + Tuple[ + MoeTensor, + torch.Tensor, + torch.Tensor, + MoeEpWgradForwardStash, + ], + ]: + """Validate and dispatch one fused MoE+EP forward call. + + 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 constructed with ``generate_c=True``. In + ``backward_wgrad_mode="operands"``, the latter tuple has a fourth + ``MoeEpWgradForwardStash`` item. + """ + + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + topk_version = self._tensor_version(topk_idx) + validate_expert_ids = not ( + self._validated_topk_idx is topk_idx + and topk_version is not None + and topk_version == self._validated_topk_version + ) + request = validate_forward( + self._forward_config, + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + validate_expert_ids=validate_expert_ids, + ) + version_after_validation = self._tensor_version(topk_idx) + if ( + topk_version is not None + and topk_version == version_after_validation + ): + self._validated_topk_idx = topk_idx + self._validated_topk_version = topk_version + else: + self._validated_topk_idx = None + self._validated_topk_version = None + return self._get_backend(request, backward=False).forward(request) + + def warmup( + self, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ) -> None: + """Prepare a forward plan for CUDA Graph capture. + + This runs one complete eager forward and synchronizes its CUDA device, + forcing runtime bootstrap, symmetric allocation, weight staging, JIT + compilation, and the first real kernel launch to finish before capture. + + For expert-parallel execution this method is collective by contract: + every rank in ``ep_group`` must call it concurrently with valid inputs. + It intentionally does not issue a process-group barrier; callers should + align all ranks after warmup and replay captured graphs in lockstep. + """ + + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + output = self( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ) + del output + device = activation.device + if device.type == "cuda": + torch.cuda.synchronize(device) + + def backward( + self, + grad_output: torch.Tensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + fc1_c: torch.Tensor, + route_metadata: torch.Tensor, + *, + wgrad_forward_stash: Optional[MoeEpWgradForwardStash] = None, + ) -> Union[ + Tuple[torch.Tensor, torch.Tensor], + Tuple[torch.Tensor, torch.Tensor, MoeEpWgradOperands], + ]: + """Validate and dispatch one device MoE backward call. + + Requires ``generate_c=True``; consumes the forward stash + (``fc1_c``, ``route_metadata``) plus the re-supplied weights and + routing inputs. Returns float32 + ``(grad_activation, grad_topk_weights)``. In + ``backward_wgrad_mode="operands"``, ``wgrad_forward_stash`` is + required and the return tuple has a third ``MoeEpWgradOperands`` item. + The Rubin MXFP8 device path supports BF16/MXFP8 combine and BF16 output + on EP1/EP2/EP4 under its documented capability gates. Forward and + backward both quantize each FP32 route accumulator directly to MXFP8 + before top-k reduction. + """ + + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + if not self.generate_c: + raise RuntimeError( + "backward requires the operator to be constructed with " + "generate_c=True" + ) + + backward_request = validate_backward( + self._forward_config, + grad_output, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + wgrad_forward_stash=wgrad_forward_stash, + ) + local_routes = self._count_local_routes(backward_request) + if backward_request.local_routes != local_routes: + raise ValueError( + "route_metadata row count must match the routes received " + "from the re-supplied topk_idx" + ) + return self._get_backend(backward_request, backward=True).backward( + backward_request + ) + + def close(self) -> None: + """Release compiled-backend instance resources; idempotent.""" + + with self._lifecycle_lock: + if self._closed: + return + if self._forward_backend is not None: + close_backend = getattr(self._forward_backend, "close", None) + if close_backend is not None: + close_backend() + self._forward_backend = None + self._forward_backend_device = None + self._validated_topk_idx = None + self._validated_topk_version = None + self._closed = True + + def __enter__(self) -> "MoeEp": + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + return self + + def __exit__(self, exc_type, exc_value, traceback) -> bool: + del exc_type, exc_value, traceback + self.close() + return False + + def __del__(self) -> None: + if not hasattr(self, "_closed"): + return + try: + self.close() + except Exception as exc: + # Explicit close propagates cleanup failures. During GC there is no + # safe global point to retry CUDA/NVSHMEM teardown, so report the + # failure without retaining the backend indefinitely. + with contextlib.suppress(Exception): + warnings.warn( + f"MoeEp finalizer could not release backend resources: {exc}", + ResourceWarning, + stacklevel=2, + ) + + +__all__ = [ + "BlockScaledTensor", + "MoeEp", + "MoeEpWgradForwardStash", + "MoeEpWgradOperands", + "MoeFormat", + "MoeTensor", +] From 17763a093266f503d775fb31bcc525164add86d2 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Mon, 24 Aug 2026 16:40:45 -0700 Subject: [PATCH 02/31] feat: vendor shared MegaMoE CuTeDSL sources Bring in the licensed communication, workspace, scheduling, and common kernel primitives required to host MegaMoE execution inside the frontend package. --- .../cutedsl_src/LICENSE.Apache-2.0 | 201 ++ .../cutedsl_src/VENDOR_INFO.md | 140 + .../_megamoe_backend/cutedsl_src/__init__.py | 1 + .../_megamoe_backend/cutedsl_src/api.py | 181 + .../cutedsl_src/communication/__init__.py | 22 + .../communication/nvlink_domain/__init__.py | 25 + .../nvlink_domain/symmetric_buffer.py | 257 ++ .../communication/nvlink_domain/token_comm.py | 2135 ++++++++++++ .../nvlink_domain/token_comm_deterministic.py | 2273 ++++++++++++ .../communication/token_protocol.py | 59 + .../cutedsl_src/helpers/__init__.py | 47 + .../cutedsl_src/helpers/constants.py | 25 + .../cutedsl_src/helpers/cute_py_helpers.py | 517 +++ .../cutedsl_src/helpers/device_workspace.py | 264 ++ .../cutedsl_src/helpers/dsl_helpers.py | 235 ++ .../cutedsl_src/helpers/flag_batch.py | 133 + .../cutedsl_src/helpers/iket_compat.py | 35 + .../cutedsl_src/helpers/ptx_helpers.py | 576 ++++ .../cutedsl_src/helpers/smem_workspace.py | 427 +++ .../cutedsl_src/helpers/software_sync.py | 276 ++ .../cutedsl_src/helpers/utils.py | 100 + .../cutedsl_src/kernel_src/__init__.py | 26 + .../block_scaled_swap_ab_fc12_epilogue.py | 3072 +++++++++++++++++ .../block_scaled_swap_ab_fc12_extension.py | 164 + .../blackwell/inference/mega/topk_reduce.py | 484 +++ .../kernel_src/function_mapping.py | 171 + .../kernel_src/schedulers/__init__.py | 29 + .../cutedsl_src/kernel_src/schedulers/base.py | 206 ++ .../kernel_src/schedulers/fc12_mapping.py | 1175 +++++++ .../kernel_src/schedulers/fc12_scheduler.py | 692 ++++ .../schedulers/non_clc_mixed_cga.py | 339 ++ .../kernel_src/schedulers/work_id_claim.py | 566 +++ .../_megamoe_backend/cutedsl_src/quant_def.py | 249 ++ 33 files changed, 15102 insertions(+) create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/LICENSE.Apache-2.0 create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR_INFO.md create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/api.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm_deterministic.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/token_protocol.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/constants.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/cute_py_helpers.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/dsl_helpers.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/flag_batch.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/iket_compat.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/software_sync.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/topk_reduce.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/function_mapping.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/base.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_scheduler.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/non_clc_mixed_cga.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/work_id_claim.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/quant_def.py diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/LICENSE.Apache-2.0 b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/LICENSE.Apache-2.0 new file mode 100644 index 000000000..261eeb9e9 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/LICENSE.Apache-2.0 @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR_INFO.md b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR_INFO.md new file mode 100644 index 000000000..7aab69c5a --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR_INFO.md @@ -0,0 +1,140 @@ +# Vendored CuTeDSL MegaMoE sources + +## Provenance + +- Source project: `cutedsl_megamoe` +- Source tree: `cutedsl_megamoe/next/sources` +- Base forward upstream revision: + `882c83e2ce4086c3cd4211fc5a2296143c5e2aea` +- Selected forward updates and backward dGLU upstream revision: + `92dd334af2eeedb36087834354b58ace08e880c6` +- Latest synchronized upstream revision: + `5a43c8523ea5215923c2fc8d0abae75bd6762011` (merge of source revision + `dc05bbdf38350a0eb67e9d9440e3c7c0e21e99fc`) +- Vendoring dates: 2026-08-11 (base), 2026-08-17 (selected updates), and + 2026-08-20 and 2026-08-24 (latest synchronizations). +- On 2026-08-24 every vendored Python source except the intentionally minimal + `kernel_src/rubin/training/__init__.py` was synchronized byte-for-byte with + the revision above. Other integration-specific behavior lives outside this + directory. +- License: synchronized Python sources retain their upstream BSD-3-Clause + SPDX identifiers. `LICENSE.Apache-2.0` remains as historical snapshot + metadata. + +The source repository URL is intentionally omitted because it is an internal +development location. The revisions above identify the upstream baselines; +the manifest and local-modification notes below describe the packaged snapshot. + +## Scope + +This directory contains the recursive Python import closures for these Rubin +SM107 products: + +- training MegaMoE forward GLU; +- optional forward MXFP8 column requantization (disabled by default); +- training MegaMoE backward dGLU. + +It preserves the `next/sources` package hierarchy and includes the shared API, +quantization, workspace, synchronization, NVLink token communication, +schedulers, TopK reduction, and Rubin helper modules required by those roots. +The Rubin training initializer remains a minimal package marker so importing +the MegaMoE products does not pull in the unused traditional-wgrad product. +The public backend exposes the backward dGLU product through a restricted Rubin +MXFP8 dgrad/dprob path. Unsupported formats and semantics retain explicit +capability gates; see the backend README. + +Complete Blackwell kernel products, runners, tests, and repository-only tooling +are excluded. Three architecture-neutral Blackwell donor modules remain because +the upstream Rubin source-copy shims import them. Imports of external CUTLASS +utility modules remain because they are CUTLASS helpers, not vendored kernel +support. + +## Manifest + +```text +LICENSE.Apache-2.0 +VENDOR_INFO.md +__init__.py +api.py +communication/__init__.py +communication/nvlink_domain/__init__.py +communication/nvlink_domain/symmetric_buffer.py +communication/nvlink_domain/token_comm.py +communication/nvlink_domain/token_comm_deterministic.py +communication/token_protocol.py +helpers/__init__.py +helpers/constants.py +helpers/cute_py_helpers.py +helpers/device_workspace.py +helpers/dsl_helpers.py +helpers/flag_batch.py +helpers/iket_compat.py +helpers/ptx_helpers.py +helpers/software_sync.py +helpers/smem_workspace.py +helpers/utils.py +kernel_src/__init__.py +kernel_src/function_mapping.py +kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py +kernel_src/blackwell/inference/mega/topk_reduce.py +kernel_src/rubin/__init__.py +kernel_src/rubin/training/__init__.py +kernel_src/rubin/training/mega/__init__.py +kernel_src/rubin/training/mega/bwd_dglu/__init__.py +kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py +kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py +kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py +kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py +kernel_src/rubin/training/mega/fwd_glu/__init__.py +kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py +kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py +kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_extension.py +kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_kernel.py +kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py +kernel_src/rubin/training/mega/helpers/__init__.py +kernel_src/rubin/training/mega/helpers/constants.py +kernel_src/rubin/training/mega/helpers/utils.py +kernel_src/rubin/training/mega/tmem_transpose.py +kernel_src/rubin/training/mega/topk_reduce.py +kernel_src/schedulers/__init__.py +kernel_src/schedulers/base.py +kernel_src/schedulers/fc12_mapping.py +kernel_src/schedulers/fc12_scheduler.py +kernel_src/schedulers/non_clc_mixed_cga.py +kernel_src/schedulers/work_id_claim.py +quant_def.py +``` + +## Integration boundary + +- Every `.py` file listed in the manifest except + `kernel_src/rubin/training/__init__.py` is a byte-for-byte copy of the same + relative path at revision `5a43c8523ea5215923c2fc8d0abae75bd6762011`. +- `kernel_src/rubin/training/__init__.py` is intentionally reduced to a + package marker. This avoids vendoring and eagerly importing the unused Rubin + traditional-wgrad product. +- `helpers/software_sync.py` replaces the former integration-only + `communication/nvlink_domain/software_sync.py` path. +- The upstream Rubin `topk_reduce.py` and `tmem_transpose.py` source-copy shims + require the three architecture-neutral Blackwell donor modules in the + manifest. Unrelated Blackwell and Rubin inference products remain excluded. +- Public API validation, symmetric-workspace ownership, overflow reporting, + staging, CUDA Graph handling, dprob materialization, and grouped-wgrad layout + conversion live in the parent `_megamoe_backend` package. +- The synchronized Rubin sources require a CUTLASS DSL distribution that + provides `cutlass.utils.rubin_helpers`. + +## Updating the snapshot + +1. Review source changes from the revision above. +2. Recompute recursive relative imports from all Rubin SM107 product roots, + exact package initializers, and shared product entry points. +3. Include only the exact donor and eager-import dependencies required by that + closure; do not add unrelated products. +4. Copy every selected Python source without modification and remove paths no + longer present in the selected upstream closure. +5. Update the revision, date, manifest, and integration-boundary notes. +6. Assert byte equality for every vendored Python source, then run compileall, + relative-import closure validation, package import smoke tests, and the + MoeEP regression suite. diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/__init__.py new file mode 100644 index 000000000..8feb6089f --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/__init__.py @@ -0,0 +1 @@ +"""Integration-ready CuTeDSL MegaMoE sources.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/api.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/api.py new file mode 100644 index 000000000..c3fd0af46 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/api.py @@ -0,0 +1,181 @@ +"""Stable construction API for composable kernel implementations.""" + +import types +from abc import ABC, ABCMeta, abstractmethod +from collections.abc import Iterator, Mapping +from types import MappingProxyType +from typing import Dict, Union, get_args, get_origin + +from cutlass.cute.typing import SymInt + + +RuntimeIntegerType = SymInt +StaticIntegerType = int +StaticOrRuntimeIntegerType = Union[int, SymInt] + + +class OptionalRequirement: + """Descriptor field that a component may consume conditionally.""" + + __slots__ = ("expected_type",) + + def __init__(self, expected_type) -> None: + self.expected_type = expected_type + + +Requirement = Union[type, OptionalRequirement] + + +def _required_type(requirement: Requirement): + return requirement.expected_type if isinstance(requirement, OptionalRequirement) else requirement + + +def _matches_type(value, expected_type) -> bool: + origin = get_origin(expected_type) + if origin in (Union, types.UnionType): + return any( + _matches_type(value, candidate) + for candidate in get_args(expected_type) + ) + return isinstance(value, expected_type) + + +class Desc(Mapping[str, object]): + """Immutable descriptor mapping validated against component schemas.""" + + def __init__(self, values: Mapping[str, object]) -> None: + self._values = MappingProxyType(dict(values)) + + def __getitem__(self, key: str): + return self._values[key] + + def __iter__(self) -> Iterator[str]: + return iter(self._values) + + def __len__(self) -> int: + return len(self._values) + + def validate( + self, + requirements: Dict[str, Requirement], + *, + component_name: str, + ) -> None: + for name, requirement in requirements.items(): + if name not in self._values: + if isinstance(requirement, OptionalRequirement): + continue + raise KeyError( + f"{component_name} requires descriptor field {name!r}." + ) + expected_type = _required_type(requirement) + value = self._values[name] + if not _matches_type(value, expected_type): + raise TypeError( + f"{component_name} requires {name!r} to have type " + f"{expected_type}, got {type(value)}." + ) + + +class ProblemDesc(Desc): + """Problem semantics shared by every component in one kernel.""" + + +class ImplDesc(Desc): + """Fully static implementation choices shared by kernel components.""" + + +class _KernelComponentMeta(ABCMeta): + def __call__(cls, *args, **kwargs): + instance = super().__call__(*args, **kwargs) + instance._validate_component_init() + return instance + + +class KernelComponent(ABC, metaclass=_KernelComponentMeta): + """Component that consumes descriptors only during construction.""" + + @classmethod + @abstractmethod + def problem_desc_require(cls) -> Dict[str, Requirement]: + ... + + @classmethod + @abstractmethod + def impl_desc_require(cls) -> Dict[str, Requirement]: + ... + + def _validate_desc_inputs( + self, + problem_desc: ProblemDesc, + impl_desc: ImplDesc, + ) -> None: + component_name = type(self).__name__ + overlap = ( + self.problem_desc_require().keys() + & self.impl_desc_require().keys() + ) + if overlap: + raise ValueError( + f"{component_name} requires fields from both descriptors: " + f"{sorted(overlap)}." + ) + problem_desc.validate( + self.problem_desc_require(), + component_name=component_name, + ) + impl_desc.validate( + self.impl_desc_require(), + component_name=component_name, + ) + + def _validate_component_init(self) -> None: + component_name = type(self).__name__ + requirements = { + **self.problem_desc_require(), + **self.impl_desc_require(), + } + for name, requirement in requirements.items(): + if not hasattr(self, name): + if isinstance(requirement, OptionalRequirement): + continue + raise RuntimeError( + f"{component_name} did not bind required field {name!r}." + ) + expected_type = _required_type(requirement) + value = getattr(self, name) + if not _matches_type(value, expected_type): + raise TypeError( + f"{component_name}.{name} must have type " + f"{expected_type}, got {type(value)}." + ) + for name, value in vars(self).items(): + if isinstance(value, Desc): + raise RuntimeError( + f"{component_name}.{name} retains a descriptor." + ) + + +class KernelClass(KernelComponent): + """Top-level host wrapper for one composable kernel implementation.""" + + @abstractmethod + def name(self) -> str: + ... + + @abstractmethod + def aot_compile(self): + ... + + +__all__ = [ + "Desc", + "ImplDesc", + "KernelClass", + "KernelComponent", + "OptionalRequirement", + "ProblemDesc", + "RuntimeIntegerType", + "StaticIntegerType", + "StaticOrRuntimeIntegerType", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/__init__.py new file mode 100644 index 000000000..cc1ef5158 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/__init__.py @@ -0,0 +1,22 @@ +"""Cross-rank communication protocols and implementations.""" + +from ..quant_def import CombineFormat, QuantKind +from .nvlink_domain.token_comm import ( + TokenBackScheduleMode, + TokenBackMode, + TokenCommArgs, + TokenCommNonDeterministic, +) +from .nvlink_domain.token_comm_deterministic import TokenCommDeterministic +from .token_protocol import TokenSrcMetadata + +__all__ = [ + "CombineFormat", + "QuantKind", + "TokenBackScheduleMode", + "TokenBackMode", + "TokenCommArgs", + "TokenCommDeterministic", + "TokenCommNonDeterministic", + "TokenSrcMetadata", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/__init__.py new file mode 100644 index 000000000..c72d7918a --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/__init__.py @@ -0,0 +1,25 @@ +"""NVLink-domain pointer mapping and token communication.""" + +from ...helpers.software_sync import NvlinkBarrier, SoftwareGridSync +from ...quant_def import QuantKind +from .symmetric_buffer import SymmetricBufferDevice, SymmetricBufferHost +from .token_comm import ( + TokenBackMode, + TokenBackScheduleMode, + TokenCommArgs, + TokenCommNonDeterministic, +) +from .token_comm_deterministic import TokenCommDeterministic + +__all__ = [ + "NvlinkBarrier", + "QuantKind", + "SoftwareGridSync", + "SymmetricBufferDevice", + "SymmetricBufferHost", + "TokenBackMode", + "TokenBackScheduleMode", + "TokenCommArgs", + "TokenCommDeterministic", + "TokenCommNonDeterministic", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py new file mode 100644 index 000000000..0fd98b9fb --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py @@ -0,0 +1,257 @@ +"""Symmetric-heap peer pointer mapping carried in kernel arguments.""" + +from dataclasses import dataclass +from typing import Any, Optional + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith, llvm +from cutlass.base_dsl.dsl import ( + extract_mlir_values, + get_mlir_types, + new_from_mlir_values, +) +from cutlass.base_dsl.runtime.jit_arg_adapters import JitArgAdapterRegistry +from cutlass.base_dsl.typing import get_c_pointers +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64, dsl_user_op + + +try: + from cutlass.base_dsl.typing import MLIR_DYNAMIC_INDEX +except ImportError: + MLIR_DYNAMIC_INDEX = -(2**31) + + +_byval_rank_limit = 16 + + +def _byval_struct_type() -> Any: + return ir.Type.parse( + f"!llvm.struct<(array<{_byval_rank_limit} x i64>)>" + ) + + +@dataclass(frozen=True) +class SymmetricBufferDevice: + """Device-side peer offset table in constant/by-value kernel arguments.""" + + value: Any + max_ranks: cutlass.Constexpr[int] + + def __extract_mlir_values__(self) -> list: + return [self.value] + + def __new_from_mlir_values__( + self, + values: list, + ) -> "SymmetricBufferDevice": + return SymmetricBufferDevice(values[0], self.max_ranks) + + def __get_mlir_types__(self) -> list: + if self.max_ranks <= _byval_rank_limit: + return [ir.Type.parse("!llvm.ptr")] + return [ir.Type.parse(f"vector<{self.max_ranks}xi64>")] + + def __extract_mlir_attributes__(self) -> list: + if self.max_ranks <= _byval_rank_limit: + return [ + ir.DictAttr.get( + { + "cute_nvgpu.grid_constant": ir.UnitAttr.get(), + "llvm.byval": ir.TypeAttr.get(_byval_struct_type()), + } + ) + ] + return [ir.DictAttr.get({})] + + @cute.jit + def map( + self, + local_address: Int64, + destination_rank: Int32, + byte_offset: Int64 = Int64(0), + ) -> Int64: + if cutlass.const_expr(self.max_ranks <= _byval_rank_limit): + i64_type = ir.Type.parse("i64") + offset_pointer = llvm.getelementptr( + ir.Type.parse("!llvm.ptr"), + self.value, + [destination_rank.ir_value()], + [MLIR_DYNAMIC_INDEX], + i64_type, + no_wrap_flags="None", + ) + peer_offset = Int64(llvm.load(i64_type, offset_pointer)) + else: + peer_offset = Int64( + llvm.extractelement( + self.value, + destination_rank.ir_value(), + ) + ) + return local_address + peer_offset + byte_offset + + @cute.jit + def map_pointer( + self, + pointer, + destination_rank: Int32, + byte_alignment: Optional[int] = None, + ): + if cutlass.const_expr(pointer.memspace != AddressSpace.gmem): + raise ValueError( + "Only GMEM pointers can be mapped to a symmetric peer." + ) + if cutlass.const_expr(byte_alignment is None): + byte_alignment = pointer.max_alignment + return cute.make_ptr( + pointer.dtype, + self.map(pointer.toint(), destination_rank), + pointer.memspace, + assumed_align=byte_alignment, + ) + + +@dataclass(frozen=True) +class SymmetricBufferHost: + """Host launch payload used to construct a SymmetricBufferDevice.""" + + base_address: Int64 + offsets: tuple + rank: Int32 + max_ranks: cutlass.Constexpr[int] + + @staticmethod + def _as_int64(value) -> Int64: + return value if isinstance(value, Int64) else Int64(int(value)) + + @dsl_user_op + def make_device_object(self, *, loc=None, ip=None) -> SymmetricBufferDevice: + offsets = tuple(self.offsets) + if len(offsets) != self.max_ranks: + raise ValueError( + f"Expected {self.max_ranks} peer offsets, got {len(offsets)}." + ) + + if self.max_ranks <= _byval_rank_limit: + pointer_type = ir.Type.parse("!llvm.ptr") + struct_type = _byval_struct_type() + i64_type = ir.Type.parse("i64") + one = arith.constant( + value=ir.IntegerAttr.get(i64_type, 1), + result=i64_type, + loc=loc, + ip=ip, + ) + buffer = llvm.alloca( + res=pointer_type, + elem_type=struct_type, + array_size=one, + alignment=64, + loc=loc, + ip=ip, + ) + for index, offset in enumerate(offsets): + slot = llvm.getelementptr( + pointer_type, + buffer, + [], + [index], + i64_type, + no_wrap_flags="None", + loc=loc, + ip=ip, + ) + llvm.store( + self._as_int64(offset).ir_value(), + slot, + loc=loc, + ip=ip, + ) + return SymmetricBufferDevice(buffer, self.max_ranks) + + i32_type = ir.Type.parse("i32") + vector_type = ir.Type.parse(f"vector<{self.max_ranks}xi64>") + vector = llvm.mlir_zero(vector_type, loc=loc, ip=ip) + for index, offset in enumerate(offsets): + element_index = arith.constant( + value=ir.IntegerAttr.get(i32_type, index), + result=i32_type, + loc=loc, + ip=ip, + ) + vector = llvm.insertelement( + vector, + self._as_int64(offset).ir_value(), + element_index, + loc=loc, + ip=ip, + ) + return SymmetricBufferDevice(vector, self.max_ranks) + + +@JitArgAdapterRegistry.register_jit_arg_adapter(SymmetricBufferHost) +class _SymmetricBufferHostAdapter: + def __init__(self, argument: SymmetricBufferHost) -> None: + self._argument = argument + offsets = tuple(argument.offsets) + if len(offsets) != int(argument.max_ranks): + raise ValueError( + f"Expected {int(argument.max_ranks)} peer offsets, " + f"got {len(offsets)}." + ) + self._fields = ( + Int64(argument.base_address), + *(Int64(offset) for offset in offsets), + Int32(argument.rank), + ) + + def __c_pointers__(self) -> list[Any]: + pointers: list[Any] = [] + for field in self._fields: + pointers.extend(get_c_pointers(field)) + return pointers + + def __get_mlir_types__(self) -> list[Any]: + types: list[Any] = [] + for field in self._fields: + types.extend(get_mlir_types(field)) + return types + + def __extract_mlir_values__(self) -> list[ir.Value]: + values: list[ir.Value] = [] + for field in self._fields: + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__( + self, + values: list[ir.Value], + ) -> SymmetricBufferHost: + value_index = 0 + rebuilt = [] + for field in self._fields: + field_value_count = len(get_mlir_types(field)) + rebuilt.append( + new_from_mlir_values( + field, + values[value_index : value_index + field_value_count], + ) + ) + value_index += field_value_count + if value_index != len(values): + raise ValueError( + f"Consumed {value_index} MLIR values, got {len(values)}." + ) + + result = object.__new__(SymmetricBufferHost) + object.__setattr__(result, "base_address", rebuilt[0]) + object.__setattr__(result, "offsets", tuple(rebuilt[1:-1])) + object.__setattr__(result, "rank", rebuilt[-1]) + object.__setattr__(result, "max_ranks", self._argument.max_ranks) + return result + + +__all__ = ["SymmetricBufferDevice", "SymmetricBufferHost"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm.py new file mode 100644 index 000000000..793ea2167 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm.py @@ -0,0 +1,2135 @@ +"""Metadata-push routing and fused non-deterministic token communication.""" + +import dataclasses +import os +from typing import Callable, ClassVar, Literal, Optional, Tuple, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 +from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF + +from ...api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.dsl_helpers import mark_alignment, smem_exclusive_prefix +from ...helpers.flag_batch import make_flag_batch_tracker +from ...helpers.iket_compat import iket +from ...helpers.software_sync import NvlinkBarrier +from ...helpers.ptx_helpers import ( + cp_async_bulk_s2g, + cp_reduce_async_bulk_add_bf16_s2g, + cp_reduce_async_bulk_add_u32_s2g, + nanosleep, + read_clock64, + red_add_relaxed_sys_s32, + stg_b64, + stg_f32, + tma_load_1d, +) +from ...helpers.smem_workspace import SmemWorkspace +from ...helpers.utils import ceil_div, round_up +from ...quant_def import CombineFormat, QuantKind +from ..token_protocol import TokenSrcMetadata +from .symmetric_buffer import SymmetricBufferDevice + + +TokenBackMode = Literal["epi_warps", "standalone_warps", "reuse_dispatch_warps"] +TokenBackScheduleMode = Literal["static", "atomic_counter"] + + +@dataclasses.dataclass(frozen=True) +class TokenCommArgs: + """Device views materialized inside the fused kernel region. + + ``pre_reduced_activation`` is the per-topk combine staging plane: TokenComm's own symmetric region, except + under in-kernel top-k reduction where it degenerates to a view of the caller's 2D output. That REDG + accumulates, so in that mode the incoming content is the caller's accumulation base -- zero, or a + shared-expert result. + """ + + activation: cute.Tensor + activation_sf: cute.Tensor + pre_reduced_activation: cute.Tensor + pre_reduced_activation_sf: Optional[cute.Tensor] + peer_rank_ptr_mapper: SymmetricBufferDevice + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "TokenCommArgs": + if values: + raise ValueError(f"TokenCommArgs expected no MLIR values, got {len(values)}.") + return self + + +@dataclasses.dataclass(frozen=True) +class _SortedElement: + flat_topk_index: Int32 + topk_score: Optional[cutlass.Float32] + + def pack(self) -> Union[Int64, Int32]: + if cutlass.const_expr(self.topk_score is None): + return self.flat_topk_index + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + scratch[0] = self.flat_topk_index + cute.recast_tensor(scratch, cutlass.Float32)[1] = self.topk_score + return cute.recast_tensor(scratch, cutlass.Int64)[0] + + @classmethod + def from_packed(cls, packed: Union[Int64, Int32]) -> "_SortedElement": + if cutlass.const_expr(type(packed).width == 32): + return cls(flat_topk_index=Int32(packed), topk_score=None) + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + cute.recast_tensor(scratch, cutlass.Int64)[0] = packed + return cls(flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1]) + + +@cute.jit +def _copy_atom(dtype, num_bits_per_copy: int): + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy) + + +class _MetadataPushRouter(KernelComponent): + """Sort and push routing metadata into each destination rank's final pool.""" + + router_smem_limit_bytes: ClassVar[int] = 227 * 1024 + router_warps_per_cta: ClassVar[int] = 16 + + sizes_by_rank_region = "nvlink.token_comm.sizes_by_rank" + sizes_region = "nvlink.token_comm.sizes" + sizes_ready_region = "nvlink.token_comm.sizes_ready" + metadata_ready_region = "nvlink.token_comm.metadata_ready" + sorted_metadata_region = "nvlink.token_comm.sorted_metadata" + sorted_scores_region = "nvlink.token_comm.sorted_scores" + pool_expert_base_region = "nvlink.token_comm.pool_expert_base" + token_src_metadata_region = "nvlink.token_comm.token_src_metadata" + fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" + source_expert_base_region = "nvlink.token_comm.source_expert_base" + push_destination_base_region = "nvlink.token_comm.push_destination_base" + sorted_metadata_ready_region = "nvlink.token_comm.sorted_metadata_ready" + push_table_ready_region = "nvlink.token_comm.push_table_ready" + router_size_counter_region = "nvlink.token_comm.router_size_counter" + router_histogram_done_region = "nvlink.token_comm.router_histogram_done" + source_base_ready_region = "nvlink.token_comm.source_base_ready" + + router_data_histogram_region = "nvlink.token_comm.router_smem.data_histogram" + router_data_prefix_region = "nvlink.token_comm.router_smem.data_prefix" + router_data_warp_totals_region = "nvlink.token_comm.router_smem.data_warp_totals" + router_data_sorted_region = "nvlink.token_comm.router_smem.data_sorted" + router_data_base_region = "nvlink.token_comm.router_smem.data_base" + router_helper_size_matrix_region = "nvlink.token_comm.router_smem.helper_size_matrix" + router_helper_totals_region = "nvlink.token_comm.router_smem.helper_totals" + router_helper_prefix_region = "nvlink.token_comm.router_smem.helper_prefix" + router_helper_warp_totals_region = "nvlink.token_comm.router_smem.helper_warp_totals" + router_helper_load_mbarrier_region = "nvlink.token_comm.router_smem.helper_load_mbarrier" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "token_padding_block": int, + "promised_launchable_sm_count": int, + "router_smem_limit_bytes": OptionalRequirement(int), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + + self.token_padding_block = impl_desc["token_padding_block"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.router_smem_limit_bytes = impl_desc.get("router_smem_limit_bytes", 227 * 1024) + + self._validate_router_configuration() + self.expert_count_padded = round_up(self.expert_count, 4) + self.expert_count_with_trash = self.expert_count_padded + 1 + self.router_elements_per_lane, self.router_data_cta_count = self._router_launch_configuration() + self.router_tokens_per_cta = self.router_elements_per_lane * self.router_warps_per_cta * 32 + self.router_push_cta_count = ceil_div(self.expert_count, self.router_warps_per_cta) + self.router_grid_cta_count = max(self.router_data_cta_count + 1, self.router_push_cta_count) + if self.router_grid_cta_count > self.promised_launchable_sm_count: + raise ValueError( + "Router grid exceeds promised_launchable_sm_count; all metadata-push CTAs must be concurrently resident." + ) + self.worst_case_token_count = self.worst_case_padded_tokens(self.token_padding_block) + self._router_smem_workspace = self._build_router_smem_workspace() + + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_grid_thread_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + + def _validate_router_configuration(self) -> None: + positive_fields = ( + "world_size", + "expert_count", + "topk", + "max_tokens_per_rank", + "token_padding_block", + "promised_launchable_sm_count", + "router_smem_limit_bytes", + ) + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.expert_count % self.world_size != 0: + raise ValueError( + f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." + ) + if self.expert_count > 16384: + raise NotImplementedError("TokenComm supports at most 16384 global experts.") + if self.topk > self.expert_count: + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + def worst_case_padded_tokens(self, block: int) -> int: + source_token_capacity = self.world_size * self.max_tokens_per_rank + routes_per_source_token = min(self.topk, self.experts_per_rank) + route_capacity = source_token_capacity * routes_per_source_token + active_expert_capacity = min(self.experts_per_rank, route_capacity) + route_budget_blocks = active_expert_capacity + (route_capacity - active_expert_capacity) // block + expert_bound_blocks = active_expert_capacity * int(ceil_div(source_token_capacity, block)) + return min(route_budget_blocks, expert_bound_blocks) * block + + def _router_launch_configuration(self) -> Tuple[int, int]: + def next_power_of_two(value: int) -> int: + return 1 << (max(value, 1) - 1).bit_length() + + routed_token_capacity = self.max_tokens_per_rank * self.topk + minimum_cta_capacity = 2048 + maximum_cta_capacity = 16384 + maximum_data_cta_count = 128 + maximum_supported_tokens = maximum_cta_capacity * maximum_data_cta_count + if routed_token_capacity > maximum_supported_tokens: + raise NotImplementedError(f"The router supports at most {maximum_supported_tokens} routed tokens per rank.") + cta_capacity = min(maximum_cta_capacity, next_power_of_two(max(routed_token_capacity, minimum_cta_capacity))) + elements_per_lane = cta_capacity // (self.router_warps_per_cta * 32) + data_cta_count = ceil_div(routed_token_capacity, cta_capacity) + return elements_per_lane, data_cta_count + + def _build_router_smem_workspace(self) -> SmemWorkspace: + workspace = SmemWorkspace() + workspace.register_mbarrier(self.router_helper_load_mbarrier_region, 1) + overlay = workspace.create_overlay("nvlink.token_comm.router_smem.role") + data_lifetime = overlay.add_lifetime("data_cta") + data_lifetime.register_tensor(self.router_data_histogram_region, cutlass.Int32, (self.expert_count_with_trash,)) + data_lifetime.register_tensor( + self.router_data_prefix_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 + ) + data_lifetime.register_tensor(self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,)) + data_lifetime.register_tensor( + self.router_data_sorted_region, + (cutlass.Int64 if self.apply_topk_at_fc1 else cutlass.Int32), + (self.router_tokens_per_cta,), + byte_alignment=16, + ) + if self.router_data_cta_count > 1: + data_lifetime.register_tensor( + self.router_data_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + + helper_lifetime = overlay.add_lifetime("helper_cta") + helper_lifetime.register_tensor( + self.router_helper_size_matrix_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + stride=(self.expert_count_padded, 1), + byte_alignment=16, + ) + helper_lifetime.register_tensor( + self.router_helper_totals_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_prefix_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 + ) + workspace.finalize(max_bytes=self.router_smem_limit_bytes) + return workspace + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router_smem_workspace + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + """Register router-private state and Router-to-Main outputs.""" + self._register_router_workspace(workspace) + + def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: + maximum_routed_tokens = self.max_tokens_per_rank * self.topk + workspace.register( + self.sizes_by_rank_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + buffer_space="shared", + stride=(self.expert_count_padded, 1), + ) + workspace.register( + self.sizes_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="shared", reset="tail_reset" + ) + workspace.register(self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.metadata_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.sorted_metadata_region, cutlass.Int64, (maximum_routed_tokens,), buffer_space="local") + if self.apply_topk_at_fc1: + workspace.register( + self.sorted_scores_region, cutlass.Float32, (maximum_routed_tokens,), buffer_space="local" + ) + workspace.register( + self.token_src_metadata_region, + cutlass.Int64, + (self.worst_case_token_count,), + buffer_space="shared", + byte_alignment=16, + ) + if self.apply_topk_at_fc1: + workspace.register( + self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" + ) + workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local") + workspace.register( + self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.push_destination_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.sorted_metadata_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register(self.push_table_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") + if self.router_data_cta_count > 1: + workspace.register( + self.router_size_counter_region, + cutlass.Int32, + (self.expert_count_with_trash,), + buffer_space="local", + reset="tail_reset", + ) + workspace.register( + self.router_histogram_done_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register( + self.source_base_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "_MetadataPushRouter": + if values: + raise ValueError("_MetadataPushRouter carries no MLIR values.") + return self + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + """Launch counting-sort DATA, size-exchange HELPER, and metadata PUSH roles.""" + if cutlass.const_expr(self.apply_topk_at_fc1 and topk_scores is None): + raise ValueError("apply_topk_at_fc1 requires router topk_scores.") + peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() + self._router_kernel( + topk_indices, + topk_scores, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper, + device_workspace, + ).launch( + grid=[self.router_grid_cta_count, 1, 1], + block=[self.router_warps_per_cta * 32, 1, 1], + min_blocks_per_mp=1, + stream=stream, + ) + + @cute.kernel + def _router_kernel( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper: SymmetricBufferDevice, + device_workspace: DeviceWorkspace, + ) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + linear_cta_idx, _, _ = cute.arch.block_idx() + cute.arch.griddepcontrol_launch_dependents() + block_thread_count = self.router_warps_per_cta * 32 + grid_thread_idx = thread_idx + linear_cta_idx * block_thread_count + warp_idx = cute.arch.make_warp_uniform(thread_idx // Int32(32)) + lane_idx = thread_idx % Int32(32) + + storage_type = self._router_smem_workspace.storage_class() + smem_allocator = cutlass.utils.SmemAllocator() + storage = smem_allocator.allocate(storage_type) + smem_base = storage.buffer.data_ptr() + + device_workspace.assign_device_members(local_workspace, shared_workspace) + self._device_workspace = device_workspace + self._peer_rank_ptr_mapper = peer_rank_ptr_mapper + self._router_local_rank = local_rank + self._router_thread_idx = thread_idx + self._router_linear_cta_idx = linear_cta_idx + self._router_grid_thread_idx = grid_thread_idx + self._router_warp_idx = warp_idx + self._router_lane_idx = lane_idx + + if cutlass.const_expr(self.router_data_cta_count == 1): + self._router_single_cta(topk_indices, topk_scores, smem_base) + else: + self._router_multiple_ctas(topk_indices, topk_scores, smem_base) + if linear_cta_idx < Int32(self.router_push_cta_count): + self._router_push_metadata() + + device_workspace.remove_device_members() + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_grid_thread_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + + @cute.jit + def _router_single_cta( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + + zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + expert = Int32(zero_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_with_trash): + histogram[expert] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + within_expert_indices = self._build_histogram(expert_registers, histogram) + iket.range_pop() + + iket.range_push("router.prefix_and_publish") + publish_sizes = self._broadcast_sizes_to_peers( + cute.make_tensor(histogram.iterator, cute.make_layout((self.expert_count_padded,))) + ) + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(histogram.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + publish_sizes() + iket.range_pop() + + iket.range_push("router.sort") + self._sort_router_elements( + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.write_out") + self._dump_contiguous_router_output(sorted_elements, total_valid_routes) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_single_cta(smem_base) + + @cute.jit + def _router_multiple_ctas( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + dump_base = self._router_smem_workspace.tensor(self.router_data_base_region, smem_base) + + zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + expert = Int32(zero_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_with_trash): + histogram[expert] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + within_expert_indices = self._build_histogram(expert_registers, histogram) + iket.range_pop() + + iket.range_push("router.reserve_and_prefix") + size_counter = self._device_workspace.ptr(self.router_size_counter_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + dump_base[expert] = Int32( + cute.arch.atomic_add(size_counter + expert, histogram[expert], sem="relaxed", scope="gpu") + ) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.router_histogram_done_region), Int32(1), sem="release", scope="gpu" + ) + + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(histogram.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.sort") + self._sort_router_elements( + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.wait_source_base") + source_base_ready = self._device_workspace.ptr(self.source_base_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.write_out") + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + dump_base[expert] = dump_base[expert] + source_expert_base[expert] + cute.arch.sync_threads() + self._dump_router_output_by_expert(histogram, prefix, dump_base, sorted_elements) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_multiple_ctas(smem_base) + + @cute.jit + def _router_helper_single_cta(self, smem_base: cute.Pointer) -> None: + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + size_counter = self._device_workspace.tensor(self.router_size_counter_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + + histogram_done = self._device_workspace.ptr(self.router_histogram_done_region) + iket.range_push("router.wait_histogram") + if self._router_thread_idx == Int32(0): + while cute.arch.load(histogram_done, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.broadcast_sizes") + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + totals[expert] = size_counter[expert] + cute.arch.sync_threads() + + publish_sizes = self._broadcast_sizes_to_peers(totals) + smem_exclusive_prefix( + totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.source_base_ready_region), Int32(1), sem="release", scope="gpu" + ) + publish_sizes() + iket.range_pop() + + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_push_metadata(self) -> None: + block_thread_count = self.router_warps_per_cta * 32 + sorted_metadata_ready = self._device_workspace.ptr(self.sorted_metadata_ready_region) + push_table_ready = self._device_workspace.ptr(self.push_table_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(sorted_metadata_ready, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + while cute.arch.load(push_table_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + + global_expert = self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx + if global_expert < Int32(self.expert_count): + sizes_by_rank = self._device_workspace.tensor(self.sizes_by_rank_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + route_count = sizes_by_rank[self._router_local_rank, global_expert] + source_begin = source_expert_base[global_expert] + destination_begin = push_destination_base[global_expert] + destination_rank = global_expert // Int32(self.experts_per_rank) + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) + destination_metadata_address = ( + self._device_workspace.ptr(self.token_src_metadata_region).toint() + peer_offset + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + source_scores = self._device_workspace.ptr(self.sorted_scores_region) + destination_scores_address = ( + self._device_workspace.ptr(self.fc1_topk_scores_region).toint() + peer_offset + ) + route_round_count = (route_count + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + if route < route_count: + source_position = source_begin + route + destination_position = destination_begin + route + metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) + stg_b64( + destination_metadata_address + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), + metadata, + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + score = cute.arch.load(source_scores + source_position, cutlass.Float32) + stg_f32(destination_scores_address + Int64(destination_position) * Int64(4), score) + + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + # Keep notifier threads behind the leader's system fence. + cute.arch.sync_threads() + metadata_ready_address = self._device_workspace.ptr(self.metadata_ready_region).toint() + rank_round_count = ceil_div(self.world_size, block_thread_count) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = Int32(rank_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(metadata_ready_address, destination_rank, Int64(0)), Int32(1) + ) + + @cute.jit + def _broadcast_sizes_to_peers(self, smem_expert_counts: cute.Tensor) -> Callable[[], None]: + block_thread_count = self.router_warps_per_cta * 32 + row_bytes = Int32(self.expert_count_padded * 4) + matrix_address = self._device_workspace.ptr(self.sizes_by_rank_region).toint() + total_address = self._device_workspace.ptr(self.sizes_region).toint() + rank_round_count = ceil_div(self.world_size, self.router_warps_per_cta) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = self._router_warp_idx + Int32(rank_round * self.router_warps_per_cta) + if destination_rank < Int32(self.world_size): + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + destination_row_address = ( + matrix_address + + peer_offset + + Int64(Int32(self._router_local_rank) * Int32(self.expert_count_padded)) * Int64(4) + ) + destination_row = cute.make_ptr( + cutlass.Int32, destination_row_address, AddressSpace.gmem, assumed_align=16 + ) + destination_total = cute.make_ptr( + cutlass.Int32, total_address + peer_offset, AddressSpace.gmem, assumed_align=16 + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_row, smem_expert_counts.iterator, row_bytes) + cp_reduce_async_bulk_add_u32_s2g(destination_total, smem_expert_counts.iterator, row_bytes) + cute.arch.cp_async_bulk_commit_group() + + def finalize() -> None: + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + cute.arch.sync_threads() + ready_address = self._device_workspace.ptr(self.sizes_ready_region).toint() + ready_round_count = ceil_div(self.world_size, block_thread_count) + for ready_round in cutlass.range_constexpr(ready_round_count): + destination_rank = Int32(ready_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(ready_address, destination_rank, Int64(0)), Int32(1) + ) + + return finalize + + @cute.jit + def _compute_push_tables(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + owner_expert_begin = Int32(self._router_local_rank) * Int32(self.experts_per_rank) + matrix_bytes = self.world_size * self.expert_count_padded * 4 + + size_matrix = self._router_smem_workspace.tensor(self.router_helper_size_matrix_region, smem_base) + padded_totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + load_mbarrier = self._router_smem_workspace.ptr(self.router_helper_load_mbarrier_region, smem_base) + sizes = self._device_workspace.tensor(self.sizes_region) + + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_init(load_mbarrier, 1) + + sizes_ready = self._device_workspace.ptr(self.sizes_ready_region) + iket.range_push("router.wait_sizes_ready") + if self._router_thread_idx == Int32(0): + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.world_size): + nanosleep(150) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.load_sizes_and_prefix") + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(load_mbarrier, Int32(matrix_bytes)) + tma_load_1d( + size_matrix.iterator, + self._device_workspace.ptr(self.sizes_by_rank_region), + load_mbarrier, + Int32(matrix_bytes), + ) + + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + expert_size = sizes[expert] + padded_totals[expert] = ( + (expert_size + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) + ) * Int32(self.token_padding_block) + cute.arch.sync_threads() + smem_exclusive_prefix( + padded_totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + pool_expert_base = self._device_workspace.tensor(self.pool_expert_base_region) + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + pool_expert_base[local_expert] = prefix[owner_expert_begin + local_expert] - prefix[owner_expert_begin] + + cute.arch.mbarrier_wait(load_mbarrier, 0) + iket.range_pop() + + iket.range_push("router.build_push_destinations") + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + global_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if global_expert < Int32(self.expert_count): + destination_rank = global_expert // Int32(self.experts_per_rank) + destination_expert_begin = destination_rank * Int32(self.experts_per_rank) + destination_pool_base = prefix[global_expert] - prefix[destination_expert_begin] + local_ring_position = ( + Int32(self._router_local_rank) - destination_rank + Int32(self.world_size) + ) % Int32(self.world_size) + source_ring_offset = Int32(0) + for ring_position in cutlass.range_constexpr(self.world_size): + source_rank = (destination_rank + Int32(ring_position)) % Int32(self.world_size) + if Int32(ring_position) < local_ring_position: + source_ring_offset = source_ring_offset + size_matrix[source_rank, global_expert] + push_destination_base[global_expert] = destination_pool_base + source_ring_offset + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.push_table_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + + @cute.jit + def _load_router_inputs( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor] + ) -> Tuple[cute.Tensor, Optional[cute.Tensor]]: + elements_per_vector = 128 // topk_indices.dtype.width + grid_thread_count = self.router_data_cta_count * self.router_warps_per_cta * 32 + tile_span = elements_per_vector * grid_thread_count + maximum_elements = self.max_tokens_per_rank * self.topk + actual_token_count = Int32(self.max_tokens_per_rank) + actual_elements = Int32(maximum_elements) + load_round_count = ceil_div(maximum_elements, tile_span) + elements_per_thread = load_round_count * elements_per_vector + + topk_flat = cute.make_tensor(topk_indices.iterator, cute.make_layout((maximum_elements,))) + topk_vectors = cute.logical_divide(cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None)) + load_atom = _copy_atom(topk_indices.dtype, 128) + expert_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Int32) + if cutlass.const_expr(topk_indices.dtype.width == 64): + raw_indices = cute.make_rmem_tensor((elements_per_thread,), topk_indices.dtype) + raw_vectors = cute.zipped_divide(raw_indices, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + load_atom, + mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), + raw_vectors[None, load_round], + ) + else: + expert_vectors = cute.zipped_divide(expert_registers, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + load_atom, + mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), + expert_vectors[None, load_round], + ) + + score_registers = None + if cutlass.const_expr(self.apply_topk_at_fc1): + score_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Float32) + scores_flat = cute.make_tensor(topk_scores.iterator, cute.make_layout((maximum_elements,))) + score_vectors = cute.logical_divide( + cute.zipped_divide(scores_flat, (tile_span,)), (elements_per_vector, None) + ) + score_atom = _copy_atom(cutlass.Float32, elements_per_vector * 32) + score_register_vectors = cute.zipped_divide(score_registers, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + score_atom, + mark_alignment(score_vectors[(None, self._router_grid_thread_idx), load_round], 16), + score_register_vectors[None, load_round], + ) + + expert_registers_u32 = cute.recast_tensor(expert_registers, cutlass.Uint32) + if cutlass.const_expr(topk_indices.dtype.width == 64): + raw_indices_i32 = cute.recast_tensor(raw_indices, cutlass.Int32) + for register_idx in cutlass.range_constexpr(elements_per_thread): + if cutlass.const_expr(topk_indices.dtype.width == 64): + expert_registers[register_idx] = raw_indices_i32[2 * register_idx] + token_idx, _ = self._router_value_coordinate(register_idx, topk_indices.dtype) + is_invalid = (expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count)) | ( + token_idx >= actual_token_count + ) + if is_invalid: + expert_registers[register_idx] = Int32(self.expert_count_padded) + return expert_registers, score_registers + + @cute.jit + def _build_histogram(self, expert_registers: cute.Tensor, histogram: cute.Tensor) -> cute.Tensor: + register_count = cute.size(expert_registers) + within_expert_indices = cute.make_rmem_tensor((register_count,), cutlass.Int32) + for register_idx in cutlass.range_constexpr(register_count): + within_expert_indices[register_idx] = Int32( + cute.arch.atomic_add( + histogram.iterator + expert_registers[register_idx], Int32(1), sem="relaxed", scope="cta" + ) + ) + cute.arch.sync_threads() + return within_expert_indices + + @cute.jit + def _sort_router_elements( + self, + expert_registers: cute.Tensor, + within_expert_indices: cute.Tensor, + score_registers: Optional[cute.Tensor], + sorted_elements: cute.Tensor, + expert_run_starts: cute.Tensor, + topk_index_type: type, + ) -> None: + register_count = cute.size(expert_registers) + for register_idx in cutlass.range_constexpr(register_count): + token_idx, topk_slot = self._router_value_coordinate(register_idx, topk_index_type) + flat_topk_index = token_idx * Int32(self.topk) + topk_slot + destination = expert_run_starts[expert_registers[register_idx]] + within_expert_indices[register_idx] + if cutlass.const_expr(self.apply_topk_at_fc1): + sorted_elements[destination] = _SortedElement(flat_topk_index, score_registers[register_idx]).pack() + else: + sorted_elements[destination] = _SortedElement(flat_topk_index, None).pack() + + @cute.jit + def _router_value_coordinate(self, register_idx: int, topk_index_type: type) -> Tuple[Int32, Int32]: + elements_per_vector = 128 // topk_index_type.width + tile_span = elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 + flat_index = Int32( + register_idx // elements_per_vector * tile_span + register_idx % elements_per_vector + ) + self._router_grid_thread_idx * Int32(elements_per_vector) + return (flat_index // Int32(self.topk), flat_index % Int32(self.topk)) + + @cute.jit + def _dump_contiguous_router_output(self, sorted_elements: cute.Tensor, total_valid_routes: Int32) -> None: + block_thread_count = self.router_warps_per_cta * 32 + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32(block_thread_count) + for dump_round in cutlass.range(dump_round_count, unroll=4): + position = Int32(dump_round * block_thread_count) + self._router_thread_idx + predicate = Int32(position < total_valid_routes) + element = _SortedElement.from_packed(sorted_elements[position]) + metadata = TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=(element.flat_topk_index // Int32(self.topk)), + src_topk=(element.flat_topk_index % Int32(self.topk)), + ) + stg_b64(metadata_address + Int64(position) * Int64(TokenSrcMetadata.nbytes), metadata.pack(), predicate) + if cutlass.const_expr(self.apply_topk_at_fc1): + stg_f32(score_address + Int64(position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def _dump_router_output_by_expert( + self, + histogram: cute.Tensor, + expert_run_starts: cute.Tensor, + expert_dump_bases: cute.Tensor, + sorted_elements: cute.Tensor, + ) -> None: + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + expert_round_count = ceil_div(self.expert_count_padded, self.router_warps_per_cta) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = self._router_warp_idx + Int32(expert_round * self.router_warps_per_cta) + if expert < Int32(self.expert_count_padded): + run_begin = expert_run_starts[expert] + run_length = histogram[expert] + dump_begin = expert_dump_bases[expert] + route_round_count = (run_length + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + predicate = Int32(route < run_length) + element = _SortedElement.from_packed(sorted_elements[predicate * (run_begin + route)]) + output_position = dump_begin + route + metadata = TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=(element.flat_topk_index // Int32(self.topk)), + src_topk=(element.flat_topk_index % Int32(self.topk)), + ) + stg_b64( + metadata_address + Int64(output_position) * Int64(TokenSrcMetadata.nbytes), + metadata.pack(), + predicate, + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + stg_f32(score_address + Int64(output_position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + """Return this rank's contiguous expert-size view.""" + sizes = device_workspace.tensor(self.sizes_region) + expert_begin = local_rank * Int32(self.experts_per_rank) + return cute.make_tensor(sizes.iterator + expert_begin, cute.make_layout((self.experts_per_rank,))) + + @property + def metadata_ready_target(self) -> int: + return self.router_push_cta_count * self.world_size + + @cute.jit + def sizes_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.sizes_region) + + @cute.jit + def pool_expert_base_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.pool_expert_base_region) + + @cute.jit + def token_src_metadata_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.token_src_metadata_region) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + sizes_ready = device_workspace.ptr(self.sizes_ready_region) + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.world_size): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def wait_for_metadata_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + metadata_ready = device_workspace.ptr(self.metadata_ready_region) + while cute.arch.load(metadata_ready, Int32, sem="acquire", scope="sys") != Int32( + self.metadata_ready_target + ): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.token_src_metadata_region) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.apply_topk_at_fc1): + return None + return device_workspace.tensor(self.fc1_topk_scores_region) + + +class TokenCommNonDeterministic(KernelComponent): + """Public fused-kernel token communication component.""" + + transfer_warp_count: ClassVar[int] = 4 + transfer_thread_count: ClassVar[int] = transfer_warp_count * 32 + standalone_chunk_bytes: ClassVar[int] = 2048 + minimum_pacing_window_cycles: ClassVar[int] = 512 + standalone_max_backoff_cycles: ClassVar[int] = 500 + adaptive_minimum_sleep_cycles: ClassVar[int] = 50 + transfer_lifetime_barrier_id: ClassVar[int] = 9 + grid_sync_barrier_id: ClassVar[int] = 10 + standalone_size_barrier_id: ClassVar[int] = 11 + token_in_size_barrier_id: ClassVar[int] = 12 + + fc1_ready_region = "nvlink.token_comm.fc1_ready" + fc1_activation_region = "nvlink.token_comm.fc1_activation" + fc1_activation_sf_region = "nvlink.token_comm.fc1_activation_sf" + fc2_done_region = "nvlink.token_comm.fc2_done" + fc2_activation_region = "nvlink.token_comm.fc2_activation" + fc2_activation_sf_region = "nvlink.token_comm.fc2_activation_sf" + pre_reduced_activation_region = "nvlink.token_comm.pre_reduced_activation" + pre_reduced_activation_sf_region = "nvlink.token_comm.pre_reduced_activation_sf" + token_back_schedule_region = "nvlink.token_comm.token_back_schedule" + + token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" + token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" + expert_sizes_smem_region = "nvlink.token_comm.main_smem.expert_sizes" + token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" + token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" + token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" + token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "hidden_size": int, + "quant_kind": str, + "combine_format": CombineFormat, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "token_padding_block": int, + "sf_padding_block": int, + "tokens_per_fc1_ready_slot": int, + "fc2_done_signals_per_token_tile": int, + "promised_launchable_sm_count": int, + "token_in_flag_batch": int, + "token_back_mode": str, + "token_back_schedule_mode": str, + "reduce_topk_in_kernel": bool, + "router_smem_limit_bytes": OptionalRequirement(int), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.hidden_size = problem_desc["hidden_size"] + self.quant_kind = QuantKind(problem_desc["quant_kind"]) + self.combine_format = problem_desc["combine_format"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.tokens_per_fc1_ready_slot = impl_desc["tokens_per_fc1_ready_slot"] + self.fc2_done_signals_per_token_tile = impl_desc["fc2_done_signals_per_token_tile"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.token_in_flag_batch = impl_desc["token_in_flag_batch"] + self.token_back_mode: TokenBackMode = impl_desc["token_back_mode"] + self.token_back_schedule_mode: TokenBackScheduleMode = impl_desc["token_back_schedule_mode"] + self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] + self.router_smem_limit_bytes = impl_desc.get("router_smem_limit_bytes", 227 * 1024) + + self._validate_configuration() + self._router = _MetadataPushRouter(problem_desc, impl_desc) + self._nvlink_barrier = NvlinkBarrier(world_size=self.world_size, barrier_id=self.grid_sync_barrier_id) + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + + def _validate_configuration(self) -> None: + positive_fields = ( + "world_size", + "expert_count", + "topk", + "max_tokens_per_rank", + "hidden_size", + "token_padding_block", + "sf_padding_block", + "tokens_per_fc1_ready_slot", + "promised_launchable_sm_count", + "router_smem_limit_bytes", + ) + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.expert_count % self.world_size != 0: + raise ValueError( + f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." + ) + if self.expert_count > 16384: + raise NotImplementedError("TokenComm supports at most 16384 global experts.") + if self.topk > self.expert_count: + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") + if self.token_back_mode not in ("epi_warps", "standalone_warps", "reuse_dispatch_warps"): + raise ValueError(f"Unsupported token_back_mode {self.token_back_mode!r}.") + if self.token_back_schedule_mode not in ("static", "atomic_counter"): + raise ValueError( + f"token_back_schedule_mode must be static or atomic_counter, got {self.token_back_schedule_mode!r}." + ) + if not 1 <= self.token_in_flag_batch <= 32: + raise ValueError(f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}.") + if self.tokens_per_fc1_ready_slot % self.token_padding_block != 0: + raise ValueError("tokens_per_fc1_ready_slot must be divisible by token_padding_block.") + if self.token_back_enabled and self.fc2_done_signals_per_token_tile <= 0: + raise ValueError("fc2_done_signals_per_token_tile must be positive when token-back is enabled.") + if self.reduce_topk_in_kernel and self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("In-kernel top-k reduction requires BF16 combine data.") + element_block = self.activation_sf_vector_size * 4 + if self.hidden_size % element_block != 0: + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + if self.sf_padding_block % 128 != 0: + raise ValueError("sf_padding_block must be a multiple of 128.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + @property + def activation_dtype(self) -> type: + return self.quant_kind.activation_dtype + + @property + def activation_sf_dtype(self) -> type: + return self.quant_kind.sf_dtype + + @property + def activation_sf_vector_size(self) -> int: + return self.quant_kind.sf_vec_size + + @property + def bytes_per_token(self) -> int: + return self.hidden_size * int(self.activation_dtype.width) // 8 + + @property + def activation_sf_hidden_padded(self) -> int: + valid_hidden = self.hidden_size // self.activation_sf_vector_size + elements_per_16_bytes = 128 // int(self.activation_sf_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def combine_sf_hidden_padded(self) -> int: + if not self.combine_format.is_quantized: + return 0 + valid_hidden = self.hidden_size // self.combine_format.scale_block + elements_per_16_bytes = 128 // int(self.combine_format.scale_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def token_back_push_data(self) -> bool: + return self.token_back_mode != "epi_warps" + + @property + def token_back_push_sf(self) -> bool: + return self.combine_format.is_quantized + + @property + def token_back_enabled(self) -> bool: + return self.token_back_push_data or self.token_back_push_sf + + @property + def worst_case_token_count(self) -> int: + return self._router.worst_case_token_count + + @property + def worst_case_sf_token_count(self) -> int: + return self._router.worst_case_padded_tokens(self.sf_padding_block) + + @property + def max_fc1_ready_slot_count(self) -> int: + return self._router.worst_case_padded_tokens(self.tokens_per_fc1_ready_slot) // self.tokens_per_fc1_ready_slot + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router.router_smem_workspace + + @property + def expert_count_padded(self) -> int: + return self._router.expert_count_padded + + @property + def expert_count_with_trash(self) -> int: + return self._router.expert_count_with_trash + + @property + def router_elements_per_lane(self) -> int: + return self._router.router_elements_per_lane + + @property + def router_warps_per_cta(self) -> int: + return self._router.router_warps_per_cta + + @property + def router_data_cta_count(self) -> int: + return self._router.router_data_cta_count + + @property + def router_tokens_per_cta(self) -> int: + return self._router.router_tokens_per_cta + + @property + def router_push_cta_count(self) -> int: + return self._router.router_push_cta_count + + @property + def router_grid_cta_count(self) -> int: + return self._router.router_grid_cta_count + + @property + def metadata_ready_target(self) -> int: + return self._router.metadata_ready_target + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + self._router.register_device_workspace(workspace) + self._register_main_workspace(workspace) + self._nvlink_barrier.register_device_workspace(workspace) + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + self._router.launch_router( + topk_indices, + topk_scores, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper_host, + device_workspace, + stream, + ) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + return self._router.local_expert_sizes(device_workspace, local_rank) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + self._router.wait_for_sizes_ready(device_workspace, sleep_cycles) + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return self._router.token_src_metadata_tensor(device_workspace) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return self._router.fc1_topk_scores_tensor(device_workspace) + + @cute.jit + def assign_device_members( + self, + *, + device_workspace: DeviceWorkspace, + token_comm_args: TokenCommArgs, + local_rank: Int32, + linear_cta_idx: Int32, + ) -> None: + self._device_workspace = device_workspace + self._token_comm_args = token_comm_args + self._local_rank = local_rank + self._linear_cta_idx = linear_cta_idx + thread_idx, _, _ = cute.arch.thread_idx() + transfer_thread_idx = thread_idx % Int32(self.transfer_thread_count) + self._transfer_warp_idx = cute.arch.make_warp_uniform(transfer_thread_idx // Int32(32)) + self._lane_idx = transfer_thread_idx % Int32(32) + self._nvlink_barrier.assign_device_members(device_workspace, token_comm_args.peer_rank_ptr_mapper) + + def remove_device_members(self) -> None: + self._nvlink_barrier.remove_device_members() + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "TokenCommNonDeterministic": + if values: + raise ValueError("TokenCommNonDeterministic carries no MLIR values.") + return self + + def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: + """Register the fused kernel's GMEM regions. Three groups, each stating its own existence condition. + + FC1 pools, always present: the dispatched payload ``token_in`` pulls in, addressed in POOL index space + (per-expert padded runs of tokens routed into this rank's experts, from every rank). + + Token-back machinery, iff ``token_back_enabled``: the transfer warps' counters plus the rank-local FC2 + staging they read, also in pool index space. ``epi_warps`` needs none of it, because there the FC2 + epilogue reaches the peers itself. The scale plane is registered separately from the data plane and is + staged locally in EVERY mode: pushing scales per token would scatter one warp's 32 lanes across up to 32 + ranks and explode the NVLink request count. + + Combine plane, iff ``not reduce_topk_in_kernel``: the symmetric per-topk landing zone peers deliver into, + addressed in (source token, source topk) space. It is the DESTINATION of the round trip whose source is + the staging above, so its condition is deliberately independent of who performs the transfer -- + ``epi_warps`` with a separate reduce registers this plane while registering no token-back machinery at + all. Peers reach it through the symmetric heap, so it cannot be a caller tensor: only this component + knows the wire dtype and the padded scale stride. Its ``data`` reset keeps it out of both the host zero + prefix and the per-launch tail reset, since every cell it exposes is rewritten each launch and the plane + is far too large to be worth clearing. + """ + workspace.register( + self.fc1_ready_region, + cutlass.Int32, + (self.max_fc1_ready_slot_count,), + buffer_space="local", + reset="tail_reset", + ) + activation_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc1_activation_region, + self.activation_dtype, + (activation_element_count,), + buffer_space="local", + byte_alignment=128, + ) + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded + workspace.register( + self.fc1_activation_sf_region, + self.activation_sf_dtype, + (activation_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + + if self.token_back_enabled: + workspace.register( + self.fc2_done_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" + ) + if self.token_back_push_data: + fc2_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc2_activation_region, + self.combine_format.act_dtype, + (fc2_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_push_sf: + fc2_sf_element_count = self.worst_case_token_count * self.combine_sf_hidden_padded + workspace.register( + self.fc2_activation_sf_region, + self.combine_format.scale_dtype, + (fc2_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": + workspace.register( + self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + + if not self.reduce_topk_in_kernel: + workspace.register( + self.pre_reduced_activation_region, + self.combine_format.act_dtype, + (self.max_tokens_per_rank, self.topk, self.hidden_size), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + if self.combine_format.is_quantized: + workspace.register( + self.pre_reduced_activation_sf_region, + self.combine_format.scale_dtype, + (self.max_tokens_per_rank, self.topk, self.combine_sf_hidden_padded), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + + def register_smem_regions(self, workspace: SmemWorkspace) -> None: + workspace.register_mbarrier(self.token_in_mbarrier_region, self.transfer_warp_count) + if self.token_back_enabled: + workspace.register_mbarrier(self.token_back_mbarrier_region, self.transfer_warp_count) + workspace.register_tensor( + self.expert_sizes_smem_region, cutlass.Int32, (self.experts_per_rank,), byte_alignment=16 + ) + transfer_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.transfer") + token_in_lifetime = transfer_overlay.add_lifetime("token_in") + token_in_lifetime.register_tensor( + self.token_in_activation_smem_region, + self.activation_dtype, + (self.transfer_warp_count, self.hidden_size), + byte_alignment=16, + ) + token_in_lifetime.register_tensor( + self.token_in_sf_smem_region, + self.activation_sf_dtype, + (self.transfer_warp_count, (self.activation_sf_vector_size, self.activation_sf_hidden_padded)), + stride=(self.activation_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + if not self.token_back_enabled: + return + + if self.token_back_mode == "standalone_warps": + token_back_lifetime = workspace.create_overlay( + "nvlink.token_comm.main_smem.standalone_token_back" + ).add_lifetime("token_back") + else: + token_back_lifetime = transfer_overlay.add_lifetime("token_back") + + if self.token_back_mode == "standalone_warps": + available_bytes_per_warp = self.standalone_chunk_bytes + else: + activation_bytes = self.bytes_per_token + sf_bytes = self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 + available_bytes_per_warp = activation_bytes + sf_bytes + + if self.token_back_push_data: + bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + if self.token_back_mode == "standalone_warps": + chunk_bytes = self.standalone_chunk_bytes + elif available_bytes_per_warp < bytes_per_output_token: + chunk_bytes = self.bytes_per_token + else: + chunk_bytes = bytes_per_output_token + if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: + raise ValueError("Token-back data chunk bytes must divide one row.") + chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) + token_back_lifetime.register_tensor( + self.token_back_activation_smem_region, + self.combine_format.act_dtype, + (self.transfer_warp_count, chunk_elements), + byte_alignment=16, + ) + if self.token_back_push_sf: + sf_row_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + if sf_row_bytes > available_bytes_per_warp: + raise ValueError("Token-back scale row exceeds its per-warp stage.") + token_back_lifetime.register_tensor( + self.token_back_sf_smem_region, + self.combine_format.scale_dtype, + (self.transfer_warp_count, (self.combine_format.scale_block, self.combine_sf_hidden_padded)), + stride=(self.combine_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + + @cute.jit + def fc1_ready_counter_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.fc1_ready_region) + + @cute.jit + def fc1_activation_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return cute.make_tensor( + device_workspace.ptr(self.fc1_activation_region), + cute.make_layout((self.worst_case_token_count, self.hidden_size), stride=(self.hidden_size, 1)), + ) + + @cute.jit + def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + layout = tile_atom_to_shape_SF( + (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size + ) + return cute.make_tensor(device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1])) + + @cute.jit + def fc2_done_counter_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_enabled): + return None + return device_workspace.tensor(self.fc2_done_region) + + @cute.jit + def fc2_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_push_data): + return None + return cute.make_tensor( + device_workspace.ptr(self.fc2_activation_region), + cute.make_layout( + (self.worst_case_token_count, 1, self.hidden_size), stride=(self.hidden_size, self.hidden_size, 1) + ), + ) + + @cute.jit + def fc2_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_push_sf): + return None + return cute.make_tensor( + device_workspace.ptr(self.fc2_activation_sf_region), + cute.make_layout( + ( + self.worst_case_token_count, + 1, + (self.combine_format.scale_block, self.hidden_size // self.combine_format.scale_block), + ), + stride=(self.combine_sf_hidden_padded, self.combine_sf_hidden_padded, (0, 1)), + ), + ) + + @cute.jit + def pre_reduced_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """The (tokens, topk, hidden) combine staging plane, or None under in-kernel top-k reduction.""" + if cutlass.const_expr(self.reduce_topk_in_kernel): + return None + return device_workspace.tensor(self.pre_reduced_activation_region) + + @cute.jit + def pre_reduced_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """The scale plane parallel to ``pre_reduced_activation_tensor``; only quantized wire formats carry one.""" + if cutlass.const_expr(self.reduce_topk_in_kernel or not self.combine_format.is_quantized): + return None + return device_workspace.tensor(self.pre_reduced_activation_sf_region) + + @cute.jit + def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Wait for pushed metadata, then pull activation payloads into local pools.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + global_warp_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_warp_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + + sizes = self._router.sizes_tensor(self._device_workspace) + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + + iket.range_push("token_in.wait_sizes_ready") + self._router.wait_for_sizes_ready(self._device_workspace) + iket.range_pop() + iket.range_push("token_in.stage_sizes") + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + owner_expert_begin = self._local_rank * Int32(self.experts_per_rank) + source_sizes = cute.make_tensor(sizes.iterator + owner_expert_begin, cute.make_layout((self.experts_per_rank,))) + copy_elements = 4 if self.experts_per_rank % 4 == 0 else 1 + source_size_vectors = cute.zipped_divide( + (mark_alignment(source_sizes, 16) if cutlass.const_expr(copy_elements == 4) else source_sizes), + (copy_elements,), + ) + destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) + size_vector_count = cute.size(destination_size_vectors, mode=[1]) + size_copy_atom = _copy_atom(cutlass.Int32, copy_elements * 32) + size_copy_rounds = ceil_div(size_vector_count, self.transfer_thread_count) + size_copy_registers = cute.make_rmem_tensor((copy_elements, size_copy_rounds), cutlass.Int32) + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + if vector_idx < Int32(size_vector_count): + cute.copy( + size_copy_atom, + source_size_vectors[None, vector_idx], + size_copy_registers[None, size_copy_round], + ) + iket.range_pop() + + iket.range_push("token_in.wait_metadata_ready") + self._router.wait_for_metadata_ready(self._device_workspace) + iket.range_pop() + + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + if vector_idx < Int32(size_vector_count): + cute.copy( + size_copy_atom, + size_copy_registers[None, size_copy_round], + destination_size_vectors[None, vector_idx], + ) + iket.range_push("token_in.size_barrier") + token_in_size_barrier = pipeline.NamedBarrier( + barrier_id=self.token_in_size_barrier_id, num_threads=self.transfer_thread_count + ) + token_in_size_barrier.arrive_and_wait() + iket.range_pop() + if cutlass.const_expr(self.token_back_mode == "standalone_warps"): + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive() + + iket.range_push("token_in.pull_payload") + token_in_mbarriers = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_mbarrier = token_in_mbarriers + transfer_warp_idx + warp_activation_stage = token_in_activation[transfer_warp_idx, None] + warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] + if lane_idx == Int32(0): + cute.arch.mbarrier_init(warp_mbarrier, 1) + cute.arch.sync_warp() + + fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) + fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) + fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) + activation_bytes = cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 + activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 + sf_copy_elements = 4 + source_sf_values = cute.slice_(warp_sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) + + next_dense_token = global_warp_idx + expert_valid_begin = Int32(0) + expert_sf_begin = Int32(0) + expert_ready_slot_begin = Int32(0) + pull_phase = Int32(0) + flag_tracker = make_flag_batch_tracker( + use_async=self.token_in_flag_batch == 1, + flag_address=Int64(0), + accumulated_flags=Int32(0), + phase=Int32(0), + thread_idx=lane_idx, + ) + + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + pull_count = Int32(0) + if next_dense_token < expert_valid_end: + pull_count = (expert_valid_end - next_dense_token + global_warp_count - Int32(1)) // global_warp_count + + for pull_round in cutlass.range(pull_count, unroll=1): + dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count + token_in_expert = dense_token_idx - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + sf_token_idx = expert_sf_begin + token_in_expert + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + remote_activation_address = ( + self._token_comm_args.activation.iterator.toint() + + peer_offset + + Int64(metadata.src_token) * Int64(self.bytes_per_token) + ) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, + ) + remote_sf_row = remote_sf[Int64(metadata.src_token), None] + + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) + ) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address), + warp_mbarrier, + Int32(activation_bytes), + ) + tma_load_1d( + warp_sf_stage.iterator, remote_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_activation, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + + destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] + destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) + destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) + sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) + for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): + sf_vector_idx = Int32(sf_round * 32) + lane_idx + if sf_vector_idx < Int32(sf_vector_count): + cute.copy( + sf_copy_atom, + source_sf_vectors[None, sf_vector_idx], + destination_sf_vectors[None, sf_vector_idx], + ) + + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + ready_address = (fc1_ready_counter + ready_slot_idx).toint() + flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) + cute.arch.sync_warp() + pull_phase = pull_phase ^ Int32(1) + + next_dense_token = next_dense_token + pull_count * global_warp_count + expert_valid_begin = expert_valid_end + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) + expert_ready_slot_begin = expert_ready_slot_begin + ( + (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot) + ) + local_expert = local_expert + Int32(1) + + flag_tracker.fire() + cute.arch.sync_warp() + iket.range_pop() + if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): + iket.range_push("token_in.transfer_barrier") + transfer_lifetime_barrier = pipeline.NamedBarrier( + barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count + ) + transfer_lifetime_barrier.arrive_and_wait() + iket.range_pop() + + @cute.jit + def _stateless_pace(self, reference_window: Int32, current_window: Int32) -> None: + sleep_cycles = Int32(0) + if current_window < reference_window: + sleep_cycles = reference_window - current_window + elif current_window > reference_window: + sleep_cycles = cutlass.min(current_window - reference_window, Int32(self.standalone_max_backoff_cycles)) + if sleep_cycles > Int32(0): + nanosleep(sleep_cycles) + + @cute.jit + def _adaptive_pace(self, average_window: Int32, current_window: Int32, low_window: int, high_window: int) -> Int32: + sleep_cycles = Int32(0) + if current_window > average_window: + sleep_cycles = current_window - average_window + average_window = average_window + ((current_window - average_window + Int32(3)) // Int32(4)) + if sleep_cycles > Int32(high_window): + sleep_cycles = Int32(high_window) + else: + sleep_cycles = average_window - current_window + average_window = average_window - ((average_window - current_window + Int32(3)) // Int32(4)) + # with cute.arch.elect_one(): + # cute.printf("avg_window: {}, current_window: {}, sleep_cycles: {}", average_window, current_window, sleep_cycles) + if sleep_cycles > Int32(self.adaptive_minimum_sleep_cycles): + nanosleep(sleep_cycles) + if average_window > Int32(high_window): + average_window = Int32(high_window) + if average_window < Int32(low_window): + average_window = Int32(low_window) + return average_window + + @cute.jit + def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Push completed FC2 data and scale rows to source ranks.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + if cutlass.const_expr(not self.token_back_enabled): + return + if cutlass.const_expr( + self.combine_format.is_quantized and self._token_comm_args.pre_reduced_activation_sf is None + ): + raise ValueError("Quantized token-back requires a scale destination.") + + global_worker_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_worker_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + token_back_mbarriers = smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) + worker_mbarrier = token_back_mbarriers + transfer_warp_idx + if cutlass.const_expr(self.token_back_push_data): + token_back_activation = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) + worker_activation_stage = token_back_activation[transfer_warp_idx, None] + activation_chunk_bytes = ( + cute.cosize(worker_activation_stage) * int(self.combine_format.act_dtype.width) // 8 + ) + if cutlass.const_expr(self.token_back_push_sf): + token_back_sf = smem_workspace.tensor(self.token_back_sf_smem_region, smem_base) + worker_sf_stage = token_back_sf[transfer_warp_idx, (None, None)] + sf_chunk_bytes = cute.cosize(worker_sf_stage) * int(self.combine_format.scale_dtype.width) // 8 + if lane_idx == Int32(0): + cute.arch.mbarrier_init(worker_mbarrier, 1) + cute.arch.sync_warp() + + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + if cutlass.const_expr(self.token_back_mode == "standalone_warps"): + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive_and_wait() + + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + fc2_done = self._device_workspace.ptr(self.fc2_done_region) + if cutlass.const_expr(self.token_back_push_data): + fc2_activation_pointer = self._device_workspace.ptr(self.fc2_activation_region) + output_token_bytes = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) + data_window_unit = activation_chunk_bytes * 2 + reuse_data_pacing_enabled = ( + self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles + ) + # Preserve the empirical low:initial:high ratio of 1:2.5:5. + data_average_window = Int32(data_window_unit) + data_low_window = data_window_unit * 2 // 5 + data_high_window = data_window_unit * 2 + if cutlass.const_expr(self.token_back_push_sf): + fc2_sf_pointer = self._device_workspace.ptr(self.fc2_activation_sf_region) + output_sf_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) + sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) + reuse_sf_pacing_enabled = ( + self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles + ) + sf_average_window = Int32(sf_window_unit) + sf_low_window = sf_window_unit * 2 // 5 + sf_high_window = sf_window_unit * 2 + + next_dense_token = global_worker_idx - global_worker_count + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): + next_dense_token = Int32(0) + next_dense_token = self.next_token(next_dense_token) + expert_valid_begin = Int32(0) + transfer_phase = Int32(0) + + iket.range_push("token_back.work") + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + if next_dense_token < expert_valid_end: + token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( + self.tokens_per_fc1_ready_slot + ) + completion_target = token_tile_count * Int32(self.fc2_done_signals_per_token_tile) + iket.range_push("token_back.wait_fc2") + while cute.arch.load(fc2_done + local_expert, Int32, sem="acquire", scope="gpu") < completion_target: + nanosleep(500) + iket.range_pop() + + while next_dense_token < expert_valid_end: + token_in_expert = next_dense_token - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + destination_topk = metadata.src_topk + if cutlass.const_expr(self.reduce_topk_in_kernel): + destination_topk = Int32(0) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + is_remote_token = metadata.src_rank != self._local_rank + + if cutlass.const_expr(self.token_back_push_data): + iket.range_push("token_back.push_data") + local_activation_address = fc2_activation_pointer.toint() + Int64(pool_token_idx) * Int64( + output_token_bytes + ) + remote_activation = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation.dtype, + self._token_comm_args.pre_reduced_activation.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation.layout, + ) + destination_row = remote_activation[Int64(metadata.src_token), destination_topk, None] + for chunk_idx in cutlass.range_constexpr(activation_chunk_count): + chunk_byte_offset = Int64(chunk_idx * activation_chunk_bytes) + chunk_bytes_this_round = min( + activation_chunk_bytes, output_token_bytes - chunk_idx * activation_chunk_bytes + ) + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_data_pacing_enabled = ( + self.token_back_mode != "reuse_dispatch_warps" + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) + cute.arch.sync_warp() + if cutlass.const_expr(reuse_data_pacing_enabled or stateless_data_pacing_enabled): + if is_remote_token: + round_start_clock = read_clock64() + else: + round_start_clock = round_start_clock + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) + tma_load_1d( + worker_activation_stage.iterator, + local_activation_address + chunk_byte_offset, + worker_mbarrier, + current_chunk_bytes, + ) + cute.arch.mbarrier_wait(worker_mbarrier, transfer_phase) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_row.iterator.toint() + chunk_byte_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + if cutlass.const_expr(self.reduce_topk_in_kernel): + cp_reduce_async_bulk_add_bf16_s2g( + destination_chunk, worker_activation_stage.iterator, current_chunk_bytes + ) + else: + cp_async_bulk_s2g( + destination_chunk, worker_activation_stage.iterator, current_chunk_bytes + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0) + transfer_phase = transfer_phase ^ Int32(1) + cute.arch.sync_warp() + if cutlass.const_expr(reuse_data_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + data_average_window = self._adaptive_pace( + data_average_window, current_window, data_low_window, data_high_window + ) + elif cutlass.const_expr(stateless_data_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + self._stateless_pace(Int32(current_window_unit), current_window) + iket.range_pop() + + if cutlass.const_expr(self.token_back_push_sf): + iket.range_push("token_back.push_sf") + local_sf_address = fc2_sf_pointer.toint() + Int64(pool_token_idx) * Int64(output_sf_bytes) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation_sf.dtype, + self._token_comm_args.pre_reduced_activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation_sf.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation_sf.layout, + ) + destination_sf_row = remote_sf[Int64(metadata.src_token), destination_topk, None] + for chunk_idx in cutlass.range_constexpr(sf_chunk_count): + chunk_byte_offset = Int64(chunk_idx * sf_chunk_bytes) + chunk_bytes_this_round = min(sf_chunk_bytes, output_sf_bytes - chunk_idx * sf_chunk_bytes) + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_sf_pacing_enabled = ( + self.token_back_mode != "reuse_dispatch_warps" + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) + cute.arch.sync_warp() + if cutlass.const_expr(reuse_sf_pacing_enabled or stateless_sf_pacing_enabled): + if is_remote_token: + round_start_clock = read_clock64() + else: + round_start_clock = round_start_clock + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) + tma_load_1d( + worker_sf_stage.iterator, + local_sf_address + chunk_byte_offset, + worker_mbarrier, + current_chunk_bytes, + ) + cute.arch.mbarrier_wait(worker_mbarrier, transfer_phase) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_sf_row.iterator.toint() + chunk_byte_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_chunk, worker_sf_stage.iterator, current_chunk_bytes) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0) + transfer_phase = transfer_phase ^ Int32(1) + cute.arch.sync_warp() + if cutlass.const_expr(reuse_sf_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + sf_average_window = self._adaptive_pace( + sf_average_window, current_window, sf_low_window, sf_high_window + ) + elif cutlass.const_expr(stateless_sf_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + self._stateless_pace(Int32(current_window_unit), current_window) + iket.range_pop() + + cute.arch.sync_warp() + next_dense_token = self.next_token(next_dense_token) + + expert_valid_begin = expert_valid_end + local_expert = local_expert + Int32(1) + iket.range_pop() + # with cute.arch.elect_one(): + # cute.printf(" final data_average_window: {} ", data_average_window) + + @cute.jit + def next_token(self, current_token: Int32) -> Int32: + global_worker_count = self.promised_launchable_sm_count * self.transfer_warp_count + schedule_counter = None + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): + schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): + claimed_token = Int32(0) + if self._lane_idx == Int32(0): + claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") + return Int32(cute.arch.shuffle_sync(claimed_token, Int32(0))) + return current_token + global_worker_count + + @cute.jit + def reset_tail(self) -> None: + """Reset communication state with the four token-in transfer warps.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + iket.range_push("tail.nvlink_drain") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + total_reset_threads = self.promised_launchable_sm_count * self.transfer_thread_count + global_reset_thread = self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx + iket.range_push("tail.reset_workspace") + self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) + self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) + iket.range_pop() + iket.range_push("tail.nvlink_publish") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + if cutlass.const_expr(os.environ.get("MEGA_USE_NCU", "0") == "1"): + iket.range_push("tail.ncu_finalize") + self._nvlink_barrier.finalize( + 2, + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + ) + iket.range_pop() + + +__all__ = ["TokenBackMode", "TokenBackScheduleMode", "TokenCommArgs", "TokenCommNonDeterministic"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm_deterministic.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm_deterministic.py new file mode 100644 index 000000000..05381406f --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm_deterministic.py @@ -0,0 +1,2273 @@ +"""Metadata-push routing with a fixed token sequence and fused communication.""" + +import dataclasses +import os +from typing import Callable, ClassVar, Optional, Tuple, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils +from cutlass._mlir.dialects import llvm +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 +from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF + +from ...api import ImplDesc, KernelComponent, ProblemDesc +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.dsl_helpers import smem_exclusive_prefix +from ...helpers.flag_batch import GpuReleaseFlagBatchTracker +from ...helpers.iket_compat import iket +from ...helpers.software_sync import NvlinkBarrier +from ...helpers.ptx_helpers import ( + cp_async_bulk_s2g, + cp_reduce_async_bulk_add_bf16_s2g, + cp_reduce_async_bulk_add_u32_s2g, + nanosleep, + read_clock64, + red_add_relaxed_sys_s32, + stg_b64, + stg_f32, + tma_load_1d, +) +from ...helpers.smem_workspace import SmemWorkspace +from ...helpers.utils import ceil_div, round_up +from ...quant_def import CombineFormat, QuantKind +from ..token_protocol import TokenSrcMetadata +from .symmetric_buffer import SymmetricBufferDevice +from .token_comm import TokenBackMode, TokenBackScheduleMode, TokenCommArgs + + +_quant_spec = { + "nvfp4": (cutlass.Float4E2M1FN, cutlass.Float8E4M3FN, 16), + "mxfp4": (cutlass.Float4E2M1FN, cutlass.Float8E8M0FNU, 32), + "mxfp8_e4m3": (cutlass.Float8E4M3FN, cutlass.Float8E8M0FNU, 32), + "mxfp8_e5m2": (cutlass.Float8E5M2, cutlass.Float8E8M0FNU, 32), +} + + +@dataclasses.dataclass(frozen=True) +class _ReceiveCapacity: + raw_route_count: int + logical_route_count: int + padded_route_count: int + + +def _compute_receive_capacity( + *, + world_size: int, + max_tokens_per_rank: int, + topk: int, + experts_per_rank: int, + max_recv_size_per_rank: int, + padding_block: int, +) -> _ReceiveCapacity: + raw_route_count = world_size * max_tokens_per_rank * topk + logical_route_count = min(max_recv_size_per_rank, raw_route_count) + active_expert_count = min(experts_per_rank, logical_route_count) + padded_block_count = active_expert_count + (logical_route_count - active_expert_count) // padding_block + return _ReceiveCapacity( + raw_route_count=raw_route_count, + logical_route_count=logical_route_count, + padded_route_count=padded_block_count * padding_block, + ) + + +@dataclasses.dataclass(frozen=True) +class _SortedElement: + flat_topk_index: Int32 + topk_score: Optional[cutlass.Float32] + + def pack(self) -> Union[Int64, Int32]: + if cutlass.const_expr(self.topk_score is None): + return self.flat_topk_index + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + scratch[0] = self.flat_topk_index + cute.recast_tensor(scratch, cutlass.Float32)[1] = self.topk_score + return cute.recast_tensor(scratch, cutlass.Int64)[0] + + @classmethod + def from_packed(cls, packed: Union[Int64, Int32]) -> "_SortedElement": + if cutlass.const_expr(type(packed).width == 32): + return cls(flat_topk_index=Int32(packed), topk_score=None) + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + cute.recast_tensor(scratch, cutlass.Int64)[0] = packed + return cls(flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1]) + + +@cute.jit +def _copy_atom(dtype, num_bits_per_copy: int): + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy) + + +@cute.jit +def _mark_alignment(tensor: cute.Tensor, byte_alignment: int) -> cute.Tensor: + pointer = tensor.iterator + return cute.make_tensor( + cute.make_ptr(pointer.dtype, pointer.toint(), pointer.memspace, assumed_align=byte_alignment), tensor.layout + ) + + +@cute.jit +def _device_trap() -> None: + """Terminate the current kernel after its failure state is published.""" + llvm.inline_asm( + res=None, + operands_=[], + asm_string="trap;", + constraints="", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +class _MetadataPushRouter(KernelComponent): + """Sort and push routing metadata into each destination rank's final pool.""" + + router_smem_limit_bytes: ClassVar[int] = 227 * 1024 + router_warps_per_cta: ClassVar[int] = 16 + + sizes_by_rank_region = "nvlink.token_comm.sizes_by_rank" + sizes_region = "nvlink.token_comm.sizes" + sizes_ready_region = "nvlink.token_comm.sizes_ready" + metadata_ready_region = "nvlink.token_comm.metadata_ready" + sorted_metadata_region = "nvlink.token_comm.sorted_metadata" + sorted_scores_region = "nvlink.token_comm.sorted_scores" + pool_expert_base_region = "nvlink.token_comm.pool_expert_base" + token_src_metadata_region = "nvlink.token_comm.token_src_metadata" + fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" + source_expert_base_region = "nvlink.token_comm.source_expert_base" + push_destination_base_region = "nvlink.token_comm.push_destination_base" + sorted_metadata_ready_region = "nvlink.token_comm.sorted_metadata_ready" + push_table_ready_region = "nvlink.token_comm.push_table_ready" + router_histogram_done_region = "nvlink.token_comm.router_histogram_done" + router_cta_histograms_region = "nvlink.token_comm.router_cta_histograms" + source_base_ready_region = "nvlink.token_comm.source_base_ready" + + router_data_histogram_region = "nvlink.token_comm.router_smem.data_histogram" + router_data_totals_region = "nvlink.token_comm.router_smem.data_totals" + router_data_prefix_region = "nvlink.token_comm.router_smem.data_prefix" + router_data_warp_totals_region = "nvlink.token_comm.router_smem.data_warp_totals" + router_data_sorted_region = "nvlink.token_comm.router_smem.data_sorted" + router_data_base_region = "nvlink.token_comm.router_smem.data_base" + router_helper_size_matrix_region = "nvlink.token_comm.router_smem.helper_size_matrix" + router_helper_totals_region = "nvlink.token_comm.router_smem.helper_totals" + router_helper_prefix_region = "nvlink.token_comm.router_smem.helper_prefix" + router_helper_warp_totals_region = "nvlink.token_comm.router_smem.helper_warp_totals" + router_helper_load_mbarrier_region = "nvlink.token_comm.router_smem.helper_load_mbarrier" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "max_recv_size_per_rank": int, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return {"token_padding_block": int, "promised_launchable_sm_count": int, "drop_on_overflow": bool} + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.max_recv_size_per_rank = min( + problem_desc["max_recv_size_per_rank"], self.world_size * self.max_tokens_per_rank * self.topk + ) + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + + self.token_padding_block = impl_desc["token_padding_block"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.drop_on_overflow = impl_desc["drop_on_overflow"] + + self._validate_router_configuration() + token_capacity = self.receive_capacity(self.token_padding_block) + self.raw_route_count = token_capacity.raw_route_count + self.logical_route_capacity = token_capacity.logical_route_count + self.worst_case_token_count = token_capacity.padded_route_count + self.expert_count_padded = round_up(self.expert_count, 4) + self.expert_count_with_trash = self.expert_count_padded + 1 + self.router_elements_per_lane, self.router_data_cta_count = self._router_launch_configuration() + self.router_tokens_per_cta = self.router_elements_per_lane * self.router_warps_per_cta * 32 + self.router_push_cta_count = ceil_div(self.expert_count, self.router_warps_per_cta) + self.router_grid_cta_count = max(self.router_data_cta_count + 1, self.router_push_cta_count) + if self.router_grid_cta_count > self.promised_launchable_sm_count: + raise ValueError( + "Router grid exceeds promised_launchable_sm_count; all metadata-push CTAs must be concurrently resident." + ) + self._router_smem_workspace = self._build_router_smem_workspace() + + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_grid_thread_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + self._overflow_flag = None + + def _validate_router_configuration(self) -> None: + if type(self.max_recv_size_per_rank) is not int: + raise TypeError(f"max_recv_size_per_rank must be an int, got {type(self.max_recv_size_per_rank).__name__}.") + if type(self.drop_on_overflow) is not bool: + raise TypeError(f"drop_on_overflow must be a bool, got {type(self.drop_on_overflow).__name__}.") + positive_fields = ( + "world_size", + "expert_count", + "topk", + "max_tokens_per_rank", + "max_recv_size_per_rank", + "token_padding_block", + "promised_launchable_sm_count", + ) + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.expert_count % self.world_size != 0: + raise ValueError( + f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." + ) + if self.expert_count > 16384: + raise NotImplementedError("TokenComm supports at most 16384 global experts.") + if self.topk > self.expert_count: + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + def receive_capacity(self, padding_block: int) -> _ReceiveCapacity: + return _compute_receive_capacity( + world_size=self.world_size, + max_tokens_per_rank=self.max_tokens_per_rank, + topk=self.topk, + experts_per_rank=self.experts_per_rank, + max_recv_size_per_rank=self.max_recv_size_per_rank, + padding_block=padding_block, + ) + + def _router_launch_configuration(self) -> Tuple[int, int]: + def next_power_of_two(value: int) -> int: + return 1 << (max(value, 1) - 1).bit_length() + + routed_token_capacity = self.max_tokens_per_rank * self.topk + minimum_cta_capacity = 2048 + maximum_cta_capacity = 16384 + maximum_data_cta_count = 128 + maximum_supported_tokens = maximum_cta_capacity * maximum_data_cta_count + if routed_token_capacity > maximum_supported_tokens: + raise NotImplementedError(f"The router supports at most {maximum_supported_tokens} routed tokens per rank.") + cta_capacity = min(maximum_cta_capacity, next_power_of_two(max(routed_token_capacity, minimum_cta_capacity))) + elements_per_lane = cta_capacity // (self.router_warps_per_cta * 32) + data_cta_count = ceil_div(routed_token_capacity, cta_capacity) + return elements_per_lane, data_cta_count + + def _build_router_smem_workspace(self) -> SmemWorkspace: + workspace = SmemWorkspace() + workspace.register_mbarrier(self.router_helper_load_mbarrier_region, 1) + overlay = workspace.create_overlay("nvlink.token_comm.router_smem.role") + data_lifetime = overlay.add_lifetime("data_cta") + data_lifetime.register_tensor( + self.router_data_histogram_region, + cutlass.Int32, + (self.expert_count_with_trash, self.router_warps_per_cta + 1), + stride=(self.router_warps_per_cta + 1, 1), + ) + data_lifetime.register_tensor( + self.router_data_totals_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 + ) + data_lifetime.register_tensor( + self.router_data_prefix_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 + ) + data_lifetime.register_tensor(self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,)) + data_lifetime.register_tensor( + self.router_data_sorted_region, + (cutlass.Int64 if self.apply_topk_at_fc1 else cutlass.Int32), + (self.router_tokens_per_cta,), + byte_alignment=16, + ) + if self.router_data_cta_count > 1: + data_lifetime.register_tensor( + self.router_data_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + + helper_lifetime = overlay.add_lifetime("helper_cta") + helper_lifetime.register_tensor( + self.router_helper_size_matrix_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + stride=(self.expert_count_padded, 1), + byte_alignment=16, + ) + helper_lifetime.register_tensor( + self.router_helper_totals_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_prefix_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 + ) + workspace.finalize(max_bytes=self.router_smem_limit_bytes) + return workspace + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router_smem_workspace + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + """Register router-private state and Router-to-Main outputs.""" + self._register_router_workspace(workspace) + + def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: + maximum_routed_tokens = self.max_tokens_per_rank * self.topk + workspace.register( + self.sizes_by_rank_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + buffer_space="shared", + stride=(self.expert_count_padded, 1), + ) + workspace.register( + self.sizes_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="shared", reset="tail_reset" + ) + workspace.register(self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.metadata_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.sorted_metadata_region, cutlass.Int64, (maximum_routed_tokens,), buffer_space="local") + if self.apply_topk_at_fc1: + workspace.register( + self.sorted_scores_region, cutlass.Float32, (maximum_routed_tokens,), buffer_space="local" + ) + workspace.register( + self.token_src_metadata_region, + cutlass.Int64, + (self.worst_case_token_count,), + buffer_space="shared", + byte_alignment=16, + ) + if self.apply_topk_at_fc1: + workspace.register( + self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" + ) + workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local") + workspace.register( + self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.push_destination_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.sorted_metadata_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register(self.push_table_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") + if self.router_data_cta_count > 1: + workspace.register( + self.router_cta_histograms_region, + cutlass.Int32, + (self.router_data_cta_count, self.expert_count_padded), + buffer_space="local", + stride=(self.expert_count_padded, 1), + ) + workspace.register( + self.router_histogram_done_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register( + self.source_base_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "_MetadataPushRouter": + if values: + raise ValueError("_MetadataPushRouter carries no MLIR values.") + return self + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + overflow_flag: cute.Tensor, + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + """Launch counting-sort DATA, size-exchange HELPER, and metadata PUSH roles.""" + if cutlass.const_expr(self.apply_topk_at_fc1 and topk_scores is None): + raise ValueError("apply_topk_at_fc1 requires router topk_scores.") + if cutlass.const_expr(overflow_flag.iterator.dtype is not cutlass.Int32): + raise TypeError("overflow_flag must use Int32 elements.") + if cutlass.const_expr(cute.size(overflow_flag) != 1): + raise ValueError("overflow_flag must contain exactly one element.") + peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() + self._router_kernel( + topk_indices, + topk_scores, + overflow_flag, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper, + device_workspace, + ).launch( + grid=[self.router_grid_cta_count, 1, 1], + block=[self.router_warps_per_cta * 32, 1, 1], + min_blocks_per_mp=1, + stream=stream, + ) + + @cute.kernel + def _router_kernel( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + overflow_flag: cute.Tensor, + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper: SymmetricBufferDevice, + device_workspace: DeviceWorkspace, + ) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + linear_cta_idx, _, _ = cute.arch.block_idx() + cute.arch.griddepcontrol_launch_dependents() + block_thread_count = self.router_warps_per_cta * 32 + grid_thread_idx = thread_idx + linear_cta_idx * block_thread_count + warp_idx = cute.arch.make_warp_uniform(thread_idx // Int32(32)) + lane_idx = thread_idx % Int32(32) + + storage_type = self._router_smem_workspace.storage_class() + smem_allocator = cutlass.utils.SmemAllocator() + storage = smem_allocator.allocate(storage_type) + smem_base = storage.buffer.data_ptr() + + device_workspace.assign_device_members(local_workspace, shared_workspace) + self._device_workspace = device_workspace + self._peer_rank_ptr_mapper = peer_rank_ptr_mapper + self._router_local_rank = local_rank + self._router_thread_idx = thread_idx + self._router_linear_cta_idx = linear_cta_idx + self._router_grid_thread_idx = grid_thread_idx + self._router_warp_idx = warp_idx + self._router_lane_idx = lane_idx + self._overflow_flag = overflow_flag + + if cutlass.const_expr(self.router_data_cta_count == 1): + self._router_single_cta(topk_indices, topk_scores, smem_base) + else: + self._router_multiple_ctas(topk_indices, topk_scores, smem_base) + if linear_cta_idx < Int32(self.router_push_cta_count): + self._router_push_metadata() + + device_workspace.remove_device_members() + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_grid_thread_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + self._overflow_flag = None + + @cute.jit + def _router_single_cta( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + totals = self._router_smem_workspace.tensor(self.router_data_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + + histogram_element_count = self.expert_count_with_trash * (self.router_warps_per_cta + 1) + histogram_flat = cute.make_tensor(histogram.iterator, cute.make_layout((histogram_element_count,))) + zero_round_count = ceil_div(histogram_element_count, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + slot = Int32(zero_round * block_thread_count) + self._router_thread_idx + if slot < Int32(histogram_element_count): + histogram_flat[slot] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + match_masks = self._build_histogram(expert_registers, histogram) + self._prefix_warp_histogram(histogram, totals) + iket.range_pop() + + iket.range_push("router.prefix_and_publish") + publish_sizes = self._broadcast_sizes_to_peers( + cute.make_tensor(totals.iterator, cute.make_layout((self.expert_count_padded,))) + ) + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(totals.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + publish_sizes() + iket.range_pop() + + iket.range_push("router.sort") + self._sort_router_elements( + expert_registers, match_masks, score_registers, sorted_elements, prefix, histogram, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.write_out") + self._dump_contiguous_router_output(sorted_elements, total_valid_routes) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_single_cta(smem_base) + + @cute.jit + def _router_multiple_ctas( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + totals = self._router_smem_workspace.tensor(self.router_data_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + dump_base = self._router_smem_workspace.tensor(self.router_data_base_region, smem_base) + + histogram_element_count = self.expert_count_with_trash * (self.router_warps_per_cta + 1) + histogram_flat = cute.make_tensor(histogram.iterator, cute.make_layout((histogram_element_count,))) + zero_round_count = ceil_div(histogram_element_count, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + slot = Int32(zero_round * block_thread_count) + self._router_thread_idx + if slot < Int32(histogram_element_count): + histogram_flat[slot] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + match_masks = self._build_histogram(expert_registers, histogram) + self._prefix_warp_histogram(histogram, totals) + iket.range_pop() + + iket.range_push("router.reserve_and_prefix") + cta_histograms = self._device_workspace.tensor(self.router_cta_histograms_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + cta_histograms[self._router_linear_cta_idx, expert] = totals[expert] + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.router_histogram_done_region), Int32(1), sem="release", scope="gpu" + ) + + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(totals.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.sort") + self._sort_router_elements( + expert_registers, match_masks, score_registers, sorted_elements, prefix, histogram, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.wait_source_base") + source_base_ready = self._device_workspace.ptr(self.source_base_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.write_out") + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + cta_histograms = self._device_workspace.tensor(self.router_cta_histograms_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + dump_base[expert] = source_expert_base[expert] + cta_histograms[self._router_linear_cta_idx, expert] + cute.arch.sync_threads() + self._dump_router_output_by_expert(totals, prefix, dump_base, sorted_elements) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_multiple_ctas(smem_base) + + @cute.jit + def _router_helper_single_cta(self, smem_base: cute.Pointer) -> None: + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + cta_histograms = self._device_workspace.tensor(self.router_cta_histograms_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + + histogram_done = self._device_workspace.ptr(self.router_histogram_done_region) + iket.range_push("router.wait_histogram") + if self._router_thread_idx == Int32(0): + while cute.arch.load(histogram_done, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.broadcast_sizes") + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + running = Int32(0) + for cta_idx in cutlass.range_constexpr(self.router_data_cta_count): + count = cta_histograms[cta_idx, expert] + cta_histograms[cta_idx, expert] = running + running = running + count + totals[expert] = running + cute.arch.sync_threads() + + publish_sizes = self._broadcast_sizes_to_peers(totals) + smem_exclusive_prefix( + totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.source_base_ready_region), Int32(1), sem="release", scope="gpu" + ) + publish_sizes() + iket.range_pop() + + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_push_metadata(self) -> None: + block_thread_count = self.router_warps_per_cta * 32 + sorted_metadata_ready = self._device_workspace.ptr(self.sorted_metadata_ready_region) + push_table_ready = self._device_workspace.ptr(self.push_table_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(sorted_metadata_ready, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + while cute.arch.load(push_table_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + + global_expert = self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx + if global_expert < Int32(self.expert_count): + sizes_by_rank = self._device_workspace.tensor(self.sizes_by_rank_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + route_count = sizes_by_rank[self._router_local_rank, global_expert] + source_begin = source_expert_base[global_expert] + destination_begin = push_destination_base[global_expert] + destination_rank = global_expert // Int32(self.experts_per_rank) + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) + destination_metadata_address = ( + self._device_workspace.ptr(self.token_src_metadata_region).toint() + peer_offset + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + source_scores = self._device_workspace.ptr(self.sorted_scores_region) + destination_scores_address = ( + self._device_workspace.ptr(self.fc1_topk_scores_region).toint() + peer_offset + ) + route_round_count = (route_count + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + if route < route_count: + source_position = source_begin + route + destination_position = destination_begin + route + physical_store = Int32(destination_position < Int32(self.worst_case_token_count)) + metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) + stg_b64( + destination_metadata_address + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), + metadata, + physical_store, + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + score = cute.arch.load(source_scores + source_position, cutlass.Float32) + stg_f32( + destination_scores_address + Int64(destination_position) * Int64(4), score, physical_store + ) + + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + # Keep notifier threads behind the leader's system fence. + cute.arch.sync_threads() + metadata_ready_address = self._device_workspace.ptr(self.metadata_ready_region).toint() + rank_round_count = ceil_div(self.world_size, block_thread_count) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = Int32(rank_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(metadata_ready_address, destination_rank, Int64(0)), Int32(1) + ) + + @cute.jit + def _broadcast_sizes_to_peers(self, smem_expert_counts: cute.Tensor) -> Callable[[], None]: + block_thread_count = self.router_warps_per_cta * 32 + row_bytes = Int32(self.expert_count_padded * 4) + matrix_address = self._device_workspace.ptr(self.sizes_by_rank_region).toint() + total_address = self._device_workspace.ptr(self.sizes_region).toint() + rank_round_count = ceil_div(self.world_size, self.router_warps_per_cta) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = self._router_warp_idx + Int32(rank_round * self.router_warps_per_cta) + if destination_rank < Int32(self.world_size): + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + destination_row_address = ( + matrix_address + + peer_offset + + Int64(Int32(self._router_local_rank) * Int32(self.expert_count_padded)) * Int64(4) + ) + destination_row = cute.make_ptr( + cutlass.Int32, destination_row_address, AddressSpace.gmem, assumed_align=16 + ) + destination_total = cute.make_ptr( + cutlass.Int32, total_address + peer_offset, AddressSpace.gmem, assumed_align=16 + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_row, smem_expert_counts.iterator, row_bytes) + cp_reduce_async_bulk_add_u32_s2g(destination_total, smem_expert_counts.iterator, row_bytes) + cute.arch.cp_async_bulk_commit_group() + + def finalize() -> None: + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + cute.arch.sync_threads() + ready_address = self._device_workspace.ptr(self.sizes_ready_region).toint() + ready_round_count = ceil_div(self.world_size, block_thread_count) + for ready_round in cutlass.range_constexpr(ready_round_count): + destination_rank = Int32(ready_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(ready_address, destination_rank, Int64(0)), Int32(1) + ) + + return finalize + + @cute.jit + def _compute_push_tables(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + owner_expert_begin = Int32(self._router_local_rank) * Int32(self.experts_per_rank) + matrix_bytes = self.world_size * self.expert_count_padded * 4 + + size_matrix = self._router_smem_workspace.tensor(self.router_helper_size_matrix_region, smem_base) + padded_totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + load_mbarrier = self._router_smem_workspace.ptr(self.router_helper_load_mbarrier_region, smem_base) + sizes = self._device_workspace.tensor(self.sizes_region) + + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_init(load_mbarrier, 1) + + sizes_ready = self._device_workspace.ptr(self.sizes_ready_region) + iket.range_push("router.wait_sizes_ready") + if self._router_thread_idx == Int32(0): + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.raw_sizes_ready_target): + nanosleep(150) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.load_sizes_and_prefix") + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(load_mbarrier, Int32(matrix_bytes)) + tma_load_1d( + size_matrix.iterator, + self._device_workspace.ptr(self.sizes_by_rank_region), + load_mbarrier, + Int32(matrix_bytes), + ) + + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + expert_size = sizes[expert] + padded_totals[expert] = ( + (expert_size + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) + ) * Int32(self.token_padding_block) + cute.arch.sync_threads() + smem_exclusive_prefix( + padded_totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + pool_expert_base = self._device_workspace.tensor(self.pool_expert_base_region) + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + pool_expert_base[local_expert] = prefix[owner_expert_begin + local_expert] - prefix[owner_expert_begin] + + cute.arch.mbarrier_wait(load_mbarrier, 0) + iket.range_pop() + + iket.range_push("router.build_push_destinations") + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + global_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if global_expert < Int32(self.expert_count): + destination_rank = global_expert // Int32(self.experts_per_rank) + destination_expert_begin = destination_rank * Int32(self.experts_per_rank) + destination_pool_base = prefix[global_expert] - prefix[destination_expert_begin] + local_ring_position = ( + Int32(self._router_local_rank) - destination_rank + Int32(self.world_size) + ) % Int32(self.world_size) + source_ring_offset = Int32(0) + for ring_position in cutlass.range_constexpr(self.world_size): + source_rank = (destination_rank + Int32(ring_position)) % Int32(self.world_size) + if Int32(ring_position) < local_ring_position: + source_ring_offset = source_ring_offset + size_matrix[source_rank, global_expert] + push_destination_base[global_expert] = destination_pool_base + source_ring_offset + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.apply_receive_limit") + if self._router_thread_idx == Int32(0): + warp_totals[0] = Int32(0) + cute.arch.sync_threads() + thread_local_total = Int32(0) + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + thread_local_total = thread_local_total + sizes[owner_expert_begin + local_expert] + if thread_local_total > Int32(0): + cute.arch.atomic_add(warp_totals.iterator, thread_local_total, scope="cta") + cute.arch.sync_threads() + raw_local_total = warp_totals[0] + did_overflow = Int32(raw_local_total > Int32(self.max_recv_size_per_rank)) + + if cutlass.const_expr(self.drop_on_overflow): + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + # Keep every raw row that still lands inside the statically allocated padded pool. + raw_size = sizes[owner_expert_begin + local_expert] + retained_size = Int32(0) + remaining_capacity = Int32(self.worst_case_token_count) - pool_expert_base[local_expert] + if remaining_capacity > Int32(0): + retained_size = cutlass.min(raw_size, remaining_capacity) + if did_overflow: + sizes[owner_expert_begin + local_expert] = retained_size + + if self._router_thread_idx == Int32(0): + cute.arch.store(self._overflow_flag.iterator, did_overflow, sem="relaxed", scope="sys") + cute.arch.sync_threads() + + if cutlass.const_expr(self.drop_on_overflow): + self._publish_push_table_and_sizes_ready(sizes_ready) + else: + if did_overflow: + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + _device_trap() + else: + self._publish_push_table_and_sizes_ready(sizes_ready) + iket.range_pop() + + @cute.jit + def _publish_push_table_and_sizes_ready(self, sizes_ready: cute.Pointer) -> None: + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_gpu() + cute.arch.atomic_add( + self._device_workspace.ptr(self.push_table_ready_region), Int32(1), sem="relaxed", scope="gpu" + ) + cute.arch.atomic_add(sizes_ready, Int32(1), sem="relaxed", scope="gpu") + + @cute.jit + def _load_router_inputs( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor] + ) -> Tuple[cute.Tensor, Optional[cute.Tensor]]: + elements_per_vector = 128 // topk_indices.dtype.width + grid_thread_count = self.router_data_cta_count * self.router_warps_per_cta * 32 + tile_span = elements_per_vector * grid_thread_count + maximum_elements = self.max_tokens_per_rank * self.topk + actual_token_count = Int32(self.max_tokens_per_rank) + actual_elements = Int32(maximum_elements) + load_round_count = ceil_div(maximum_elements, tile_span) + elements_per_thread = load_round_count * elements_per_vector + + topk_flat = cute.make_tensor(topk_indices.iterator, cute.make_layout((maximum_elements,))) + topk_vectors = cute.logical_divide(cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None)) + load_atom = _copy_atom(topk_indices.dtype, 128) + expert_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Int32) + if cutlass.const_expr(topk_indices.dtype.width == 64): + raw_indices = cute.make_rmem_tensor((elements_per_thread,), topk_indices.dtype) + raw_vectors = cute.zipped_divide(raw_indices, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + load_atom, + _mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), + raw_vectors[None, load_round], + ) + else: + expert_vectors = cute.zipped_divide(expert_registers, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + load_atom, + _mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), + expert_vectors[None, load_round], + ) + + score_registers = None + if cutlass.const_expr(self.apply_topk_at_fc1): + score_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Float32) + scores_flat = cute.make_tensor(topk_scores.iterator, cute.make_layout((maximum_elements,))) + score_vectors = cute.logical_divide( + cute.zipped_divide(scores_flat, (tile_span,)), (elements_per_vector, None) + ) + score_atom = _copy_atom(cutlass.Float32, elements_per_vector * 32) + score_register_vectors = cute.zipped_divide(score_registers, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + score_atom, + _mark_alignment(score_vectors[(None, self._router_grid_thread_idx), load_round], 16), + score_register_vectors[None, load_round], + ) + + expert_registers_u32 = cute.recast_tensor(expert_registers, cutlass.Uint32) + if cutlass.const_expr(topk_indices.dtype.width == 64): + raw_indices_i32 = cute.recast_tensor(raw_indices, cutlass.Int32) + for register_idx in cutlass.range_constexpr(elements_per_thread): + if cutlass.const_expr(topk_indices.dtype.width == 64): + expert_registers[register_idx] = raw_indices_i32[2 * register_idx] + token_idx, _ = self._router_value_coordinate(register_idx, topk_indices.dtype) + is_invalid = (expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count)) | ( + token_idx >= actual_token_count + ) + if is_invalid: + expert_registers[register_idx] = Int32(self.expert_count_padded) + return expert_registers, score_registers + + @cute.jit + def _build_histogram(self, expert_registers: cute.Tensor, histogram: cute.Tensor) -> cute.Tensor: + register_count = cute.size(expert_registers) + match_masks = cute.make_rmem_tensor((register_count,), cutlass.Int32) + lane_mask_less_than = (Int32(1) << self._router_lane_idx) - Int32(1) + for register_idx in cutlass.range_constexpr(register_count): + expert = expert_registers[register_idx] + match_mask = Int32(cute.arch.match_sync(0xFFFFFFFF, expert, kind="any")) + match_masks[register_idx] = match_mask + histogram_slot = expert * Int32(self.router_warps_per_cta + 1) + self._router_warp_idx + rank_in_group = Int32(cute.arch.popc(match_mask & lane_mask_less_than)) + if rank_in_group == Int32(0): + cute.arch.atomic_add( + histogram.iterator + histogram_slot, Int32(cute.arch.popc(match_mask)), sem="relaxed", scope="cta" + ) + cute.arch.sync_threads() + return match_masks + + @cute.jit + def _prefix_warp_histogram(self, histogram: cute.Tensor, totals: cute.Tensor) -> None: + """Turn per-warp counts into fixed warp bases and expert totals.""" + block_thread_count = self.router_warps_per_cta * 32 + expert_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_with_trash): + running = Int32(0) + for warp in cutlass.range_constexpr(self.router_warps_per_cta): + count = histogram[expert, warp] + histogram[expert, warp] = running + running = running + count + totals[expert] = running + cute.arch.sync_threads() + + @cute.jit + def _sort_router_elements( + self, + expert_registers: cute.Tensor, + match_masks: cute.Tensor, + score_registers: Optional[cute.Tensor], + sorted_elements: cute.Tensor, + expert_run_starts: cute.Tensor, + warp_histogram: cute.Tensor, + topk_index_type: type, + ) -> None: + """Stable CTA-local counting-sort scatter. + + The cursor for an expert is private to one warp. ``match.any`` gives + every equal-expert lane a fixed rank, and only that group's first lane + advances the cursor. Therefore the output order is a pure function of + warp id, register round, and lane id rather than atomic arrival order. + """ + register_count = cute.size(expert_registers) + lane_mask_less_than = (Int32(1) << self._router_lane_idx) - Int32(1) + for register_idx in cutlass.range_constexpr(register_count): + expert = expert_registers[register_idx] + match_mask = match_masks[register_idx] + group_size = Int32(cute.arch.popc(match_mask)) + rank_in_group = Int32(cute.arch.popc(match_mask & lane_mask_less_than)) + warp_base = warp_histogram[expert, self._router_warp_idx] + token_idx, topk_slot = self._router_value_coordinate(register_idx, topk_index_type) + flat_topk_index = token_idx * Int32(self.topk) + topk_slot + destination = expert_run_starts[expert] + warp_base + rank_in_group + + # Every lane in an equal-expert group reads the old cursor before + # its leader advances it for the next register round. + cute.arch.sync_warp() + if rank_in_group == Int32(0): + warp_histogram[expert, self._router_warp_idx] = warp_base + group_size + if cutlass.const_expr(self.apply_topk_at_fc1): + sorted_elements[destination] = _SortedElement(flat_topk_index, score_registers[register_idx]).pack() + else: + sorted_elements[destination] = _SortedElement(flat_topk_index, None).pack() + cute.arch.sync_warp() + + @cute.jit + def _router_value_coordinate(self, register_idx: int, topk_index_type: type) -> Tuple[Int32, Int32]: + elements_per_vector = 128 // topk_index_type.width + tile_span = elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 + flat_index = Int32( + register_idx // elements_per_vector * tile_span + register_idx % elements_per_vector + ) + self._router_grid_thread_idx * Int32(elements_per_vector) + return (flat_index // Int32(self.topk), flat_index % Int32(self.topk)) + + @cute.jit + def _dump_contiguous_router_output(self, sorted_elements: cute.Tensor, total_valid_routes: Int32) -> None: + block_thread_count = self.router_warps_per_cta * 32 + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32(block_thread_count) + for dump_round in cutlass.range(dump_round_count, unroll=4): + position = Int32(dump_round * block_thread_count) + self._router_thread_idx + predicate = Int32(position < total_valid_routes) + element = _SortedElement.from_packed(sorted_elements[position]) + metadata = TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=(element.flat_topk_index // Int32(self.topk)), + src_topk=(element.flat_topk_index % Int32(self.topk)), + ) + stg_b64(metadata_address + Int64(position) * Int64(TokenSrcMetadata.nbytes), metadata.pack(), predicate) + if cutlass.const_expr(self.apply_topk_at_fc1): + stg_f32(score_address + Int64(position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def _dump_router_output_by_expert( + self, + expert_totals: cute.Tensor, + expert_run_starts: cute.Tensor, + expert_dump_bases: cute.Tensor, + sorted_elements: cute.Tensor, + ) -> None: + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + expert_round_count = ceil_div(self.expert_count_padded, self.router_warps_per_cta) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = self._router_warp_idx + Int32(expert_round * self.router_warps_per_cta) + if expert < Int32(self.expert_count_padded): + run_begin = expert_run_starts[expert] + run_length = expert_totals[expert] + dump_begin = expert_dump_bases[expert] + route_round_count = (run_length + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + predicate = Int32(route < run_length) + element = _SortedElement.from_packed(sorted_elements[predicate * (run_begin + route)]) + output_position = dump_begin + route + metadata = TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=(element.flat_topk_index // Int32(self.topk)), + src_topk=(element.flat_topk_index % Int32(self.topk)), + ) + stg_b64( + metadata_address + Int64(output_position) * Int64(TokenSrcMetadata.nbytes), + metadata.pack(), + predicate, + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + stg_f32(score_address + Int64(output_position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + """Return this rank's contiguous expert-size view.""" + sizes = device_workspace.tensor(self.sizes_region) + expert_begin = local_rank * Int32(self.experts_per_rank) + return cute.make_tensor(sizes.iterator + expert_begin, cute.make_layout((self.experts_per_rank,))) + + @property + def metadata_ready_target(self) -> int: + return self.router_push_cta_count * self.world_size + + @property + def raw_sizes_ready_target(self) -> int: + return self.world_size + + @property + def published_sizes_ready_target(self) -> int: + return self.world_size + 1 + + @cute.jit + def sizes_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.sizes_region) + + @cute.jit + def pool_expert_base_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.pool_expert_base_region) + + @cute.jit + def token_src_metadata_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.token_src_metadata_region) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + sizes_ready = device_workspace.ptr(self.sizes_ready_region) + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="gpu") != Int32( + self.published_sizes_ready_target + ): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def wait_for_metadata_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + metadata_ready = device_workspace.ptr(self.metadata_ready_region) + while cute.arch.load(metadata_ready, Int32, sem="acquire", scope="sys") != Int32( + self.metadata_ready_target + ): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.token_src_metadata_region) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.apply_topk_at_fc1): + return None + return device_workspace.tensor(self.fc1_topk_scores_region) + + +class TokenCommDeterministic(KernelComponent): + """Fused communication with a fixed logical token-position sequence.""" + + transfer_warp_count: ClassVar[int] = 4 + transfer_thread_count: ClassVar[int] = transfer_warp_count * 32 + standalone_chunk_bytes: ClassVar[int] = 2048 + minimum_pacing_window_cycles: ClassVar[int] = 512 + standalone_max_backoff_cycles: ClassVar[int] = 500 + adaptive_minimum_sleep_cycles: ClassVar[int] = 50 + transfer_lifetime_barrier_id: ClassVar[int] = 9 + grid_sync_barrier_id: ClassVar[int] = 10 + standalone_size_barrier_id: ClassVar[int] = 11 + token_in_size_barrier_id: ClassVar[int] = 12 + + fc1_ready_region = "nvlink.token_comm.fc1_ready" + fc1_activation_region = "nvlink.token_comm.fc1_activation" + fc1_activation_sf_region = "nvlink.token_comm.fc1_activation_sf" + fc2_done_region = "nvlink.token_comm.fc2_done" + fc2_activation_region = "nvlink.token_comm.fc2_activation" + fc2_activation_sf_region = "nvlink.token_comm.fc2_activation_sf" + pre_reduced_activation_region = "nvlink.token_comm.pre_reduced_activation" + pre_reduced_activation_sf_region = "nvlink.token_comm.pre_reduced_activation_sf" + token_back_schedule_region = "nvlink.token_comm.token_back_schedule" + + token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" + token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" + expert_sizes_smem_region = "nvlink.token_comm.main_smem.expert_sizes" + token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" + token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" + token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" + token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "max_recv_size_per_rank": int, + "hidden_size": int, + "quant_kind": str, + "combine_format": CombineFormat, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "token_padding_block": int, + "sf_padding_block": int, + "tokens_per_fc1_ready_slot": int, + "fc2_done_signals_per_token_tile": int, + "promised_launchable_sm_count": int, + "token_in_flag_batch": int, + "token_back_mode": str, + "token_back_schedule_mode": str, + "reduce_topk_in_kernel": bool, + "drop_on_overflow": bool, + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.max_recv_size_per_rank = min( + problem_desc["max_recv_size_per_rank"], self.world_size * self.max_tokens_per_rank * self.topk + ) + self.hidden_size = problem_desc["hidden_size"] + self.quant_kind = problem_desc["quant_kind"] + self.combine_format = problem_desc["combine_format"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.tokens_per_fc1_ready_slot = impl_desc["tokens_per_fc1_ready_slot"] + self.fc2_done_signals_per_token_tile = impl_desc["fc2_done_signals_per_token_tile"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.token_in_flag_batch = impl_desc["token_in_flag_batch"] + self.token_back_mode: TokenBackMode = impl_desc["token_back_mode"] + self.token_back_schedule_mode: TokenBackScheduleMode = impl_desc["token_back_schedule_mode"] + self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] + self.drop_on_overflow = impl_desc["drop_on_overflow"] + + self._validate_configuration() + self._router = _MetadataPushRouter(problem_desc, impl_desc) + self._nvlink_barrier = NvlinkBarrier(world_size=self.world_size, barrier_id=self.grid_sync_barrier_id) + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + + def _validate_configuration(self) -> None: + positive_fields = ("hidden_size", "token_padding_block", "sf_padding_block", "tokens_per_fc1_ready_slot") + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.quant_kind not in _quant_spec: + raise ValueError(f"Unsupported quant_kind {self.quant_kind!r}.") + if self.token_back_mode not in ("epi_warps", "standalone_warps", "reuse_dispatch_warps"): + raise ValueError(f"Unsupported token_back_mode {self.token_back_mode!r}.") + if self.token_back_schedule_mode not in ("static", "atomic_counter"): + raise ValueError( + f"token_back_schedule_mode must be static or atomic_counter, got {self.token_back_schedule_mode!r}." + ) + if not 1 <= self.token_in_flag_batch <= 32: + raise ValueError(f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}.") + if self.tokens_per_fc1_ready_slot % self.token_padding_block != 0: + raise ValueError("tokens_per_fc1_ready_slot must be divisible by token_padding_block.") + _pad_lo = min(self.sf_padding_block, self.token_padding_block) + _pad_hi = max(self.sf_padding_block, self.token_padding_block) + if _pad_hi % _pad_lo != 0: + raise ValueError("sf_padding_block and token_padding_block must be power-of-two multiples of each other.") + if self.token_back_enabled and self.fc2_done_signals_per_token_tile <= 0: + raise ValueError("fc2_done_signals_per_token_tile must be positive when token-back is enabled.") + if self.reduce_topk_in_kernel and self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("In-kernel top-k reduction requires BF16 combine data.") + element_block = self.activation_sf_vector_size * 4 + if self.hidden_size % element_block != 0: + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + if self.sf_padding_block % 128 != 0: + raise ValueError("sf_padding_block must be a multiple of 128.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + @property + def activation_dtype(self) -> type: + return _quant_spec[self.quant_kind][0] + + @property + def activation_sf_dtype(self) -> type: + return _quant_spec[self.quant_kind][1] + + @property + def activation_sf_vector_size(self) -> int: + return _quant_spec[self.quant_kind][2] + + @property + def bytes_per_token(self) -> int: + return self.hidden_size * int(self.activation_dtype.width) // 8 + + @property + def activation_sf_hidden_padded(self) -> int: + valid_hidden = self.hidden_size // self.activation_sf_vector_size + elements_per_16_bytes = 128 // int(self.activation_sf_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def combine_sf_hidden_padded(self) -> int: + if not self.combine_format.is_quantized: + return 0 + valid_hidden = self.hidden_size // self.combine_format.scale_block + elements_per_16_bytes = 128 // int(self.combine_format.scale_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def token_back_push_data(self) -> bool: + return self.token_back_mode != "epi_warps" + + @property + def token_back_push_sf(self) -> bool: + return self.combine_format.is_quantized + + @property + def token_back_enabled(self) -> bool: + return self.token_back_push_data or self.token_back_push_sf + + @property + def worst_case_token_count(self) -> int: + return self._router.worst_case_token_count + + @property + def worst_case_sf_token_count(self) -> int: + # The SF pool has to be at least as tall as the data pool: once + # token_padding_block may exceed sf_padding_block, an expert's data + # cursor can walk past the SF-capacity-derived row count. + sf_capacity = self._router.receive_capacity(self.sf_padding_block).padded_route_count + return max(sf_capacity, self.worst_case_token_count) + + @property + def max_fc1_ready_slot_count(self) -> int: + return ( + self._router.receive_capacity(self.tokens_per_fc1_ready_slot).padded_route_count + // self.tokens_per_fc1_ready_slot + ) + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router.router_smem_workspace + + @property + def expert_count_padded(self) -> int: + return self._router.expert_count_padded + + @property + def expert_count_with_trash(self) -> int: + return self._router.expert_count_with_trash + + @property + def router_elements_per_lane(self) -> int: + return self._router.router_elements_per_lane + + @property + def router_warps_per_cta(self) -> int: + return self._router.router_warps_per_cta + + @property + def router_data_cta_count(self) -> int: + return self._router.router_data_cta_count + + @property + def router_tokens_per_cta(self) -> int: + return self._router.router_tokens_per_cta + + @property + def router_push_cta_count(self) -> int: + return self._router.router_push_cta_count + + @property + def router_grid_cta_count(self) -> int: + return self._router.router_grid_cta_count + + @property + def metadata_ready_target(self) -> int: + return self._router.metadata_ready_target + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + self._router.register_device_workspace(workspace) + self._register_main_workspace(workspace) + self._nvlink_barrier.register_device_workspace(workspace) + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + overflow_flag: cute.Tensor, + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + self._router.launch_router( + topk_indices, + topk_scores, + overflow_flag, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper_host, + device_workspace, + stream, + ) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + return self._router.local_expert_sizes(device_workspace, local_rank) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + self._router.wait_for_sizes_ready(device_workspace, sleep_cycles) + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return self._router.token_src_metadata_tensor(device_workspace) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return self._router.fc1_topk_scores_tensor(device_workspace) + + @cute.jit + def assign_device_members( + self, + *, + device_workspace: DeviceWorkspace, + token_comm_args: TokenCommArgs, + local_rank: Int32, + linear_cta_idx: Int32, + ) -> None: + self._device_workspace = device_workspace + self._token_comm_args = token_comm_args + self._local_rank = local_rank + self._linear_cta_idx = linear_cta_idx + thread_idx, _, _ = cute.arch.thread_idx() + transfer_thread_idx = thread_idx % Int32(self.transfer_thread_count) + self._transfer_warp_idx = cute.arch.make_warp_uniform(transfer_thread_idx // Int32(32)) + self._lane_idx = transfer_thread_idx % Int32(32) + self._nvlink_barrier.assign_device_members(device_workspace, token_comm_args.peer_rank_ptr_mapper) + + def remove_device_members(self) -> None: + self._nvlink_barrier.remove_device_members() + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "TokenCommDeterministic": + if values: + raise ValueError("TokenCommDeterministic carries no MLIR values.") + return self + + def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: + workspace.register( + self.fc1_ready_region, + cutlass.Int32, + (self.max_fc1_ready_slot_count,), + buffer_space="local", + reset="tail_reset", + ) + activation_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc1_activation_region, + self.activation_dtype, + (activation_element_count,), + buffer_space="local", + byte_alignment=128, + ) + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded + workspace.register( + self.fc1_activation_sf_region, + self.activation_sf_dtype, + (activation_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_enabled: + workspace.register( + self.fc2_done_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" + ) + if self.token_back_push_data: + fc2_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc2_activation_region, + self.combine_format.act_dtype, + (fc2_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_push_sf: + fc2_sf_element_count = self.worst_case_token_count * self.combine_sf_hidden_padded + workspace.register( + self.fc2_activation_sf_region, + self.combine_format.scale_dtype, + (fc2_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": + workspace.register( + self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + if not self.reduce_topk_in_kernel: + workspace.register( + self.pre_reduced_activation_region, + self.combine_format.act_dtype, + (self.max_tokens_per_rank, self.topk, self.hidden_size), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + if self.combine_format.is_quantized: + workspace.register( + self.pre_reduced_activation_sf_region, + self.combine_format.scale_dtype, + (self.max_tokens_per_rank, self.topk, self.combine_sf_hidden_padded), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + + def register_smem_regions(self, workspace: SmemWorkspace) -> None: + workspace.register_mbarrier(self.token_in_mbarrier_region, self.transfer_warp_count) + if self.token_back_enabled: + workspace.register_mbarrier(self.token_back_mbarrier_region, self.transfer_warp_count) + workspace.register_tensor( + self.expert_sizes_smem_region, cutlass.Int32, (self.experts_per_rank,), byte_alignment=16 + ) + transfer_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.transfer") + token_in_lifetime = transfer_overlay.add_lifetime("token_in") + token_in_lifetime.register_tensor( + self.token_in_activation_smem_region, + self.activation_dtype, + (self.transfer_warp_count, self.hidden_size), + byte_alignment=16, + ) + token_in_lifetime.register_tensor( + self.token_in_sf_smem_region, + self.activation_sf_dtype, + (self.transfer_warp_count, (self.activation_sf_vector_size, self.activation_sf_hidden_padded)), + stride=(self.activation_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + if not self.token_back_enabled: + return + + if self.token_back_mode == "standalone_warps": + token_back_lifetime = workspace.create_overlay( + "nvlink.token_comm.main_smem.standalone_token_back" + ).add_lifetime("token_back") + else: + token_back_lifetime = transfer_overlay.add_lifetime("token_back") + + if self.token_back_mode == "standalone_warps": + available_bytes_per_warp = self.standalone_chunk_bytes + else: + activation_bytes = self.bytes_per_token + sf_bytes = self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 + available_bytes_per_warp = activation_bytes + sf_bytes + + if self.token_back_push_data: + bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + if self.token_back_mode == "standalone_warps": + chunk_bytes = self.standalone_chunk_bytes + elif available_bytes_per_warp < bytes_per_output_token: + chunk_bytes = self.bytes_per_token + else: + chunk_bytes = bytes_per_output_token + if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: + raise ValueError("Token-back data chunk bytes must divide one row.") + chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) + token_back_lifetime.register_tensor( + self.token_back_activation_smem_region, + self.combine_format.act_dtype, + (self.transfer_warp_count, chunk_elements), + byte_alignment=16, + ) + if self.token_back_push_sf: + sf_row_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + if sf_row_bytes > available_bytes_per_warp: + raise ValueError("Token-back scale row exceeds its per-warp stage.") + token_back_lifetime.register_tensor( + self.token_back_sf_smem_region, + self.combine_format.scale_dtype, + (self.transfer_warp_count, (self.combine_format.scale_block, self.combine_sf_hidden_padded)), + stride=(self.combine_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + + @cute.jit + def fc1_ready_counter_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.fc1_ready_region) + + @cute.jit + def fc1_activation_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return cute.make_tensor( + device_workspace.ptr(self.fc1_activation_region), + cute.make_layout((self.worst_case_token_count, self.hidden_size), stride=(self.hidden_size, 1)), + ) + + @cute.jit + def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + layout = tile_atom_to_shape_SF( + (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size + ) + return cute.make_tensor(device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1])) + + @cute.jit + def fc2_done_counter_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_enabled): + return None + return device_workspace.tensor(self.fc2_done_region) + + @cute.jit + def fc2_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_push_data): + return None + return cute.make_tensor( + device_workspace.ptr(self.fc2_activation_region), + cute.make_layout( + (self.worst_case_token_count, 1, self.hidden_size), stride=(self.hidden_size, self.hidden_size, 1) + ), + ) + + @cute.jit + def fc2_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_push_sf): + return None + return cute.make_tensor( + device_workspace.ptr(self.fc2_activation_sf_region), + cute.make_layout( + ( + self.worst_case_token_count, + 1, + (self.combine_format.scale_block, self.hidden_size // self.combine_format.scale_block), + ), + stride=(self.combine_sf_hidden_padded, self.combine_sf_hidden_padded, (0, 1)), + ), + ) + + @cute.jit + def pre_reduced_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """Return the source-domain combine plane when top-k reduction is separate.""" + if cutlass.const_expr(self.reduce_topk_in_kernel): + return None + return device_workspace.tensor(self.pre_reduced_activation_region) + + @cute.jit + def pre_reduced_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """Return the source-domain combine scale plane when one is required.""" + if cutlass.const_expr(self.reduce_topk_in_kernel or not self.combine_format.is_quantized): + return None + return device_workspace.tensor(self.pre_reduced_activation_sf_region) + + @cute.jit + def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Wait for pushed metadata, then pull activation payloads into local pools.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + global_warp_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_warp_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + + sizes = self._router.sizes_tensor(self._device_workspace) + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + + iket.range_push("token_in.wait_sizes_ready") + self._router.wait_for_sizes_ready(self._device_workspace) + iket.range_pop() + iket.range_push("token_in.stage_sizes") + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + owner_expert_begin = self._local_rank * Int32(self.experts_per_rank) + source_sizes = cute.make_tensor(sizes.iterator + owner_expert_begin, cute.make_layout((self.experts_per_rank,))) + copy_elements = 4 if self.experts_per_rank % 4 == 0 else 1 + source_size_vectors = cute.zipped_divide( + (_mark_alignment(source_sizes, 16) if cutlass.const_expr(copy_elements == 4) else source_sizes), + (copy_elements,), + ) + destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) + size_vector_count = cute.size(destination_size_vectors, mode=[1]) + size_copy_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp( + cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL if copy_elements == 4 else cute.nvgpu.LoadCacheMode.ALWAYS + ), + cutlass.Int32, + num_bits_per_copy=copy_elements * 32, + ) + size_copy_rounds = ceil_div(size_vector_count, self.transfer_thread_count) + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + if vector_idx < Int32(size_vector_count): + cute.copy( + size_copy_atom, source_size_vectors[None, vector_idx], destination_size_vectors[None, vector_idx] + ) + cute.arch.cp_async_commit_group() + iket.range_pop() + + iket.range_push("token_in.wait_metadata_ready") + self._router.wait_for_metadata_ready(self._device_workspace) + iket.range_pop() + + cute.arch.cp_async_wait_group(0) + iket.range_push("token_in.size_barrier") + token_in_size_barrier = pipeline.NamedBarrier( + barrier_id=self.token_in_size_barrier_id, num_threads=self.transfer_thread_count + ) + token_in_size_barrier.arrive_and_wait() + iket.range_pop() + if cutlass.const_expr(self.token_back_mode == "standalone_warps"): + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive() + + iket.range_push("token_in.pull_payload") + token_in_mbarriers = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_mbarrier = token_in_mbarriers + transfer_warp_idx + warp_activation_stage = token_in_activation[transfer_warp_idx, None] + warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] + if lane_idx == Int32(0): + cute.arch.mbarrier_init(warp_mbarrier, 1) + cute.arch.sync_warp() + + fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) + fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) + fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) + activation_bytes = cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 + activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 + sf_copy_elements = 4 + source_sf_values = cute.slice_(warp_sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) + + next_dense_token = global_warp_idx + expert_valid_begin = Int32(0) + expert_sf_begin = Int32(0) + expert_ready_slot_begin = Int32(0) + pull_phase = Int32(0) + flag_tracker = GpuReleaseFlagBatchTracker( + flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(0), thread_idx=lane_idx + ) + + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + pull_count = Int32(0) + if next_dense_token < expert_valid_end: + pull_count = (expert_valid_end - next_dense_token + global_warp_count - Int32(1)) // global_warp_count + + for pull_round in cutlass.range(pull_count, unroll=1): + dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count + token_in_expert = dense_token_idx - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + sf_token_idx = expert_sf_begin + token_in_expert + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + remote_activation_address = ( + self._token_comm_args.activation.iterator.toint() + + peer_offset + + Int64(metadata.src_token) * Int64(self.bytes_per_token) + ) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, + ) + remote_sf_row = remote_sf[Int64(metadata.src_token), None] + + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) + ) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address), + warp_mbarrier, + Int32(activation_bytes), + ) + tma_load_1d( + warp_sf_stage.iterator, remote_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_activation, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + + destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] + destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) + destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) + sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) + for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): + sf_vector_idx = Int32(sf_round * 32) + lane_idx + if sf_vector_idx < Int32(sf_vector_count): + cute.copy( + sf_copy_atom, + source_sf_vectors[None, sf_vector_idx], + destination_sf_vectors[None, sf_vector_idx], + ) + + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + ready_address = (fc1_ready_counter + ready_slot_idx).toint() + flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) + cute.arch.sync_warp() + pull_phase = pull_phase ^ Int32(1) + + next_dense_token = next_dense_token + pull_count * global_warp_count + expert_valid_begin = expert_valid_end + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) + expert_ready_slot_begin = expert_ready_slot_begin + ( + (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot) + ) + local_expert = local_expert + Int32(1) + + flag_tracker.fire() + cute.arch.sync_warp() + iket.range_pop() + if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): + iket.range_push("token_in.transfer_barrier") + transfer_lifetime_barrier = pipeline.NamedBarrier( + barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count + ) + transfer_lifetime_barrier.arrive_and_wait() + iket.range_pop() + + @cute.jit + def _stateless_pace(self, reference_window: Int32, current_window: Int32) -> None: + sleep_cycles = Int32(0) + if current_window < reference_window: + sleep_cycles = reference_window - current_window + elif current_window > reference_window: + sleep_cycles = cutlass.min(current_window - reference_window, Int32(self.standalone_max_backoff_cycles)) + if sleep_cycles > Int32(0): + nanosleep(sleep_cycles) + + @cute.jit + def _adaptive_pace(self, average_window: Int32, current_window: Int32, low_window: int, high_window: int) -> Int32: + sleep_cycles = Int32(0) + if current_window > average_window: + average_window = average_window + ((current_window - average_window + Int32(3)) // Int32(4)) + sleep_cycles = current_window - average_window + if sleep_cycles > Int32(high_window): + sleep_cycles = Int32(high_window) + else: + average_window = average_window - ((average_window - current_window + Int32(3)) // Int32(4)) + sleep_cycles = average_window - current_window + if sleep_cycles > Int32(self.adaptive_minimum_sleep_cycles): + nanosleep(sleep_cycles) + if average_window > Int32(high_window): + average_window = Int32(high_window) + if average_window < Int32(low_window): + average_window = Int32(low_window) + return average_window + + @cute.jit + def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Push completed FC2 data and scale rows to source ranks.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + if cutlass.const_expr(not self.token_back_enabled): + return + if cutlass.const_expr( + self.combine_format.is_quantized and self._token_comm_args.pre_reduced_activation_sf is None + ): + raise ValueError("Quantized token-back requires a scale destination.") + + global_worker_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_worker_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + token_back_mbarriers = smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) + worker_mbarrier = token_back_mbarriers + transfer_warp_idx + if cutlass.const_expr(self.token_back_push_data): + token_back_activation = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) + worker_activation_stage = token_back_activation[transfer_warp_idx, None] + activation_chunk_bytes = ( + cute.cosize(worker_activation_stage) * int(self.combine_format.act_dtype.width) // 8 + ) + if cutlass.const_expr(self.token_back_push_sf): + token_back_sf = smem_workspace.tensor(self.token_back_sf_smem_region, smem_base) + worker_sf_stage = token_back_sf[transfer_warp_idx, (None, None)] + sf_chunk_bytes = cute.cosize(worker_sf_stage) * int(self.combine_format.scale_dtype.width) // 8 + if lane_idx == Int32(0): + cute.arch.mbarrier_init(worker_mbarrier, 1) + cute.arch.sync_warp() + + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + if cutlass.const_expr(self.token_back_mode == "standalone_warps"): + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive_and_wait() + + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + fc2_done = self._device_workspace.ptr(self.fc2_done_region) + if cutlass.const_expr(self.token_back_push_data): + fc2_activation_pointer = self._device_workspace.ptr(self.fc2_activation_region) + output_token_bytes = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) + data_window_unit = ceil_div(activation_chunk_bytes * 2, 3) + reuse_data_pacing_enabled = ( + self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles + ) + # Preserve the empirical low:initial:high ratio of 1:2.5:5. + data_average_window = Int32(data_window_unit) + data_low_window = data_window_unit * 2 // 5 + data_high_window = data_window_unit * 2 + print(f"[{data_window_unit}, {data_low_window}, {data_high_window}]") + if cutlass.const_expr(self.token_back_push_sf): + fc2_sf_pointer = self._device_workspace.ptr(self.fc2_activation_sf_region) + output_sf_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) + sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) + reuse_sf_pacing_enabled = ( + self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles + ) + sf_average_window = Int32(sf_window_unit) + sf_low_window = sf_window_unit * 2 // 5 + sf_high_window = sf_window_unit * 2 + + next_dense_token = global_worker_idx - global_worker_count + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): + next_dense_token = Int32(0) + next_dense_token = self.next_token(next_dense_token) + expert_valid_begin = Int32(0) + transfer_phase = Int32(0) + + iket.range_push("token_back.work") + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + if next_dense_token < expert_valid_end: + token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( + self.tokens_per_fc1_ready_slot + ) + completion_target = token_tile_count * Int32(self.fc2_done_signals_per_token_tile) + iket.range_push("token_back.wait_fc2") + while cute.arch.load(fc2_done + local_expert, Int32, sem="acquire", scope="gpu") < completion_target: + nanosleep(500) + iket.range_pop() + + while next_dense_token < expert_valid_end: + token_in_expert = next_dense_token - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + destination_topk = metadata.src_topk + if cutlass.const_expr(self.reduce_topk_in_kernel): + destination_topk = Int32(0) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + is_remote_token = metadata.src_rank != self._local_rank + + if cutlass.const_expr(self.token_back_push_data): + iket.range_push("token_back.push_data") + local_activation_address = fc2_activation_pointer.toint() + Int64(pool_token_idx) * Int64( + output_token_bytes + ) + remote_activation = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation.dtype, + self._token_comm_args.pre_reduced_activation.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation.layout, + ) + destination_row = remote_activation[Int64(metadata.src_token), destination_topk, None] + for chunk_idx in cutlass.range_constexpr(activation_chunk_count): + chunk_byte_offset = Int64(chunk_idx * activation_chunk_bytes) + chunk_bytes_this_round = min( + activation_chunk_bytes, output_token_bytes - chunk_idx * activation_chunk_bytes + ) + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_data_pacing_enabled = ( + self.token_back_mode != "reuse_dispatch_warps" + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) + if cutlass.const_expr(reuse_data_pacing_enabled or stateless_data_pacing_enabled): + if is_remote_token: + round_start_clock = read_clock64() + else: + round_start_clock = round_start_clock + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) + tma_load_1d( + worker_activation_stage.iterator, + local_activation_address + chunk_byte_offset, + worker_mbarrier, + current_chunk_bytes, + ) + cute.arch.mbarrier_wait(worker_mbarrier, transfer_phase) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_row.iterator.toint() + chunk_byte_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + if cutlass.const_expr(self.reduce_topk_in_kernel): + cp_reduce_async_bulk_add_bf16_s2g( + destination_chunk, worker_activation_stage.iterator, current_chunk_bytes + ) + else: + cp_async_bulk_s2g( + destination_chunk, worker_activation_stage.iterator, current_chunk_bytes + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + transfer_phase = transfer_phase ^ Int32(1) + if cutlass.const_expr(reuse_data_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + data_average_window = self._adaptive_pace( + data_average_window, current_window, data_low_window, data_high_window + ) + elif cutlass.const_expr(stateless_data_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + self._stateless_pace(Int32(current_window_unit), current_window) + iket.range_pop() + + if cutlass.const_expr(self.token_back_push_sf): + iket.range_push("token_back.push_sf") + local_sf_address = fc2_sf_pointer.toint() + Int64(pool_token_idx) * Int64(output_sf_bytes) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation_sf.dtype, + self._token_comm_args.pre_reduced_activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation_sf.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation_sf.layout, + ) + destination_sf_row = remote_sf[Int64(metadata.src_token), destination_topk, None] + for chunk_idx in cutlass.range_constexpr(sf_chunk_count): + chunk_byte_offset = Int64(chunk_idx * sf_chunk_bytes) + chunk_bytes_this_round = min(sf_chunk_bytes, output_sf_bytes - chunk_idx * sf_chunk_bytes) + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_sf_pacing_enabled = ( + self.token_back_mode != "reuse_dispatch_warps" + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) + if cutlass.const_expr(reuse_sf_pacing_enabled or stateless_sf_pacing_enabled): + if is_remote_token: + round_start_clock = read_clock64() + else: + round_start_clock = round_start_clock + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) + tma_load_1d( + worker_sf_stage.iterator, + local_sf_address + chunk_byte_offset, + worker_mbarrier, + current_chunk_bytes, + ) + cute.arch.mbarrier_wait(worker_mbarrier, transfer_phase) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_sf_row.iterator.toint() + chunk_byte_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_chunk, worker_sf_stage.iterator, current_chunk_bytes) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + transfer_phase = transfer_phase ^ Int32(1) + if cutlass.const_expr(reuse_sf_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + sf_average_window = self._adaptive_pace( + sf_average_window, current_window, sf_low_window, sf_high_window + ) + elif cutlass.const_expr(stateless_sf_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + self._stateless_pace(Int32(current_window_unit), current_window) + iket.range_pop() + + cute.arch.sync_warp() + next_dense_token = self.next_token(next_dense_token) + + expert_valid_begin = expert_valid_end + local_expert = local_expert + Int32(1) + iket.range_pop() + + @cute.jit + def next_token(self, current_token: Int32) -> Int32: + global_worker_count = self.promised_launchable_sm_count * self.transfer_warp_count + schedule_counter = None + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): + schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): + claimed_token = Int32(0) + if self._lane_idx == Int32(0): + claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") + return Int32(cute.arch.shuffle_sync(claimed_token, Int32(0))) + return current_token + global_worker_count + + @cute.jit + def reset_tail(self) -> None: + """Reset communication state with the four token-in transfer warps.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + iket.range_push("tail.nvlink_drain") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + total_reset_threads = self.promised_launchable_sm_count * self.transfer_thread_count + global_reset_thread = self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx + iket.range_push("tail.reset_workspace") + self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) + self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) + iket.range_pop() + iket.range_push("tail.nvlink_publish") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=False, + epilogue_grid_sync=True, + ) + iket.range_pop() + if cutlass.const_expr(os.environ.get("MEGA_USE_NCU", "0") == "1"): + iket.range_push("tail.ncu_finalize") + self._nvlink_barrier.finalize( + 2, + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + ) + iket.range_pop() + + +__all__ = ["TokenCommDeterministic"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/token_protocol.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/token_protocol.py new file mode 100644 index 000000000..289c5fe47 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/token_protocol.py @@ -0,0 +1,59 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Cross-rank token metadata protocol.""" + +import dataclasses +from typing import ClassVar, Union + +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 + +@dataclasses.dataclass(frozen=True) +class TokenSrcMetadata: + """One i64 routing record: rank:u16, topk:u16, token:u32.""" + + src_rank: Int32 + src_token: Int32 + src_topk: Int32 + + nbytes: ClassVar[int] = 8 + + def pack(self) -> Int64: + high = (Int64(self.src_rank) << Int64(16)) | Int64(self.src_topk) + return (high << Int64(32)) | ( + Int64(self.src_token) & Int64(0xFFFFFFFF) + ) + + @staticmethod + def _pointer(address: Union[cute.Pointer, Int64]) -> cute.Pointer: + raw_address = ( + address if isinstance(address, Int64) else address.toint() + ) + return cute.make_ptr( + Int64, + raw_address, + AddressSpace.gmem, + assumed_align=8, + ) + + def store(self, address: Union[cute.Pointer, Int64]) -> None: + cute.arch.store(self._pointer(address), self.pack(), scope="gpu") + + @classmethod + def load( + cls, + address: Union[cute.Pointer, Int64], + ) -> "TokenSrcMetadata": + packed = Int64( + cute.arch.load(cls._pointer(address), Int64, scope="gpu") + ) + high = packed >> Int64(32) + return cls( + src_rank=Int32((high >> Int64(16)) & Int64(0xFFFF)), + src_token=Int32(packed & Int64(0xFFFFFFFF)), + src_topk=Int32(high & Int64(0xFFFF)), + ) + + +__all__ = ["TokenSrcMetadata"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/__init__.py new file mode 100644 index 000000000..0e7d0e526 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/__init__.py @@ -0,0 +1,47 @@ +"""Low-level workspace, synchronization, and PTX helpers.""" + +from .device_workspace import DeviceWorkspace +from .dsl_helpers import spin_peek, spin_wait +from .flag_batch import GpuAsyncReleaseFlagBatchTracker, GpuReleaseFlagBatchTracker, make_flag_batch_tracker +from .iket_compat import iket +from .ptx_helpers import ( + cvt_f32_to_fp8_to_f32, + cvt_f32x4_to_f8x4_pack_i32, + stg_e8m0_from_f32, + stg_e8m0x8_from_f32, +) +from .smem_workspace import SmemWorkspace +from .software_sync import NvlinkBarrier, SoftwareGridSync +from .utils import ( + IntegerType, + ceil_div, + cosize_from_shape_stride_tuples, + product, + row_major_stride, + round_up, + validate_static_integer_tuple, +) + +__all__ = [ + "DeviceWorkspace", + "GpuAsyncReleaseFlagBatchTracker", + "GpuReleaseFlagBatchTracker", + "IntegerType", + "NvlinkBarrier", + "SmemWorkspace", + "SoftwareGridSync", + "ceil_div", + "cosize_from_shape_stride_tuples", + "cvt_f32_to_fp8_to_f32", + "cvt_f32x4_to_f8x4_pack_i32", + "make_flag_batch_tracker", + "product", + "row_major_stride", + "round_up", + "iket", + "spin_peek", + "spin_wait", + "stg_e8m0_from_f32", + "stg_e8m0x8_from_f32", + "validate_static_integer_tuple", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/constants.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/constants.py new file mode 100644 index 000000000..7870293a0 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/constants.py @@ -0,0 +1,25 @@ +"""Specification-defined numeric constants shared across kernel components.""" + + +Log2E = 1.4426950408889634 +Fp32Max = 3.40282346638528859812e38 + +Nvfp4E2M1Max = 6.0 +Fp8E4M3FNMax = 448.0 +Fp8E5M2Max = 57344.0 + +Nvfp4E2M1RcpLimit = 1.0 / Nvfp4E2M1Max +Fp8E4M3RcpLimit = 1.0 / Fp8E4M3FNMax +Fp8E5M2RcpLimit = 1.0 / Fp8E5M2Max + + +__all__ = [ + "Fp32Max", + "Fp8E4M3FNMax", + "Fp8E4M3RcpLimit", + "Fp8E5M2Max", + "Fp8E5M2RcpLimit", + "Log2E", + "Nvfp4E2M1Max", + "Nvfp4E2M1RcpLimit", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/cute_py_helpers.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/cute_py_helpers.py new file mode 100644 index 000000000..8cf623735 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/cute_py_helpers.py @@ -0,0 +1,517 @@ +from dataclasses import dataclass +from math import gcd +from typing import Optional, Tuple, Union + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import OperandMajorMode + +from .smem_workspace import SmemRegion + + +# Every block-scaled tcgen05 MMA kind accumulates in F32 and no other accumulator is legal, so +# this belongs to the instruction family rather than being a caller choice. +tcgen05_block_scaled_acc_dtype = cutlass.Float32 + + +@dataclass(frozen=True) +class Tcgen05MmaInstruction: + a_type: type[cutlass.Numeric] + b_type: type[cutlass.Numeric] + instruction_mnk: Tuple[int, int, int] + participates: int + acc_type: type[cutlass.Numeric] = cutlass.Float32 + sfa_type: Optional[type[cutlass.Numeric]] = None + sfb_type: Optional[type[cutlass.Numeric]] = None + sf_vec_size: Optional[int] = None + + +@dataclass(frozen=True) +class Tcgen05TmemPlan: + """Column-level TMEM placement and accumulator staging contract.""" + + allocation_columns: int + accumulator_columns: int + accumulator_stage_columns: int + accumulator_stage_count: int + accumulator_stage_stride_columns: int + accumulator_pipeline_stages: int + sfa_columns: int + sfb_columns: int + + +def make_tcgen05_tmem_plan( + mma_instruction: Tcgen05MmaInstruction, arch: str, mma_tiler_mnk: Tuple[int, int, int] +) -> Tcgen05TmemPlan: + """Plan TMEM for canonical dense or block-scaled TCGen05 MMA. + + This planner requires A and B to reside in SMEM, per-CTA M to equal 128, + and tile M to equal instruction M. A-from-TMEM, sparse MMA, B-reuse, and + custom atom layouts or permutations are outside its contract. + SM100/SM103 use the canonical two-stage overlap for 256-column block-scaled + accumulators; other supported cases maximize disjoint accumulator stages. + """ + if mma_instruction.participates not in (1, 2): + raise ValueError(f"TCGen05 MMA participates must be one or two, got {mma_instruction.participates}.") + if len(mma_instruction.instruction_mnk) != 3 or len(mma_tiler_mnk) != 3: + raise ValueError("instruction_mnk and mma_tiler_mnk must each contain three dimensions.") + if any(dimension <= 0 for dimension in (*mma_instruction.instruction_mnk, *mma_tiler_mnk)): + raise ValueError("MMA instruction and tiler dimensions must be positive.") + + instruction_m, instruction_n, instruction_k = mma_instruction.instruction_mnk + tile_m, tile_n, tile_k = mma_tiler_mnk + if instruction_m % mma_instruction.participates != 0 or instruction_n % mma_instruction.participates != 0: + raise ValueError("MMA instruction M and N must be divisible by participates.") + if tile_m % instruction_m != 0 or tile_n % instruction_n != 0 or tile_k % instruction_k != 0: + raise ValueError("MMA instruction dimensions must divide mma_tiler_mnk.") + if tile_m != instruction_m: + raise ValueError("TCGen05 TMEM planning does not support M repetition or B-reuse.") + if tile_m // mma_instruction.participates != 128: + raise ValueError("TCGen05 TMEM planning requires per-CTA M to equal 128.") + + sf_fields = (mma_instruction.sfa_type, mma_instruction.sfb_type, mma_instruction.sf_vec_size) + has_scale_factors = any(value is not None for value in sf_fields) + if has_scale_factors and any(value is None for value in sf_fields): + raise ValueError("sfa_type, sfb_type, and sf_vec_size must be provided together.") + + sfa_columns = 0 + sfb_columns = 0 + if has_scale_factors: + sf_vec_size = mma_instruction.sf_vec_size + if not isinstance(sf_vec_size, int) or isinstance(sf_vec_size, bool) or sf_vec_size <= 0: + raise ValueError("sf_vec_size must be a positive Python int.") + if instruction_k % sf_vec_size != 0: + raise ValueError("Instruction K must be divisible by sf_vec_size.") + if tile_k % (sf_vec_size * 4) != 0: + raise ValueError("Tile K must contain complete block-scaled basic chunks.") + sfa_columns = tile_k // sf_vec_size + sfb_columns = max(tile_n // 128, 1) * tile_k // sf_vec_size + + tmem_column_capacity = cute.arch.get_max_tmem_alloc_cols(arch) + accumulator_stage_columns = tile_n + scale_factor_columns = sfa_columns + sfb_columns + accumulator_stage_count = (tmem_column_capacity - scale_factor_columns) // accumulator_stage_columns + if accumulator_stage_count < 1: + raise ValueError("Scale-factor TMEM leaves no accumulator stage.") + + accumulator_stage_stride_columns = accumulator_stage_columns + accumulator_pipeline_stages = accumulator_stage_count + accumulator_columns = accumulator_stage_columns * accumulator_stage_count + + arch_number = _parse_arch_number(arch) + use_sm100_overlap = ( + arch_number in (100, 103) + and has_scale_factors + and accumulator_stage_columns == 256 + and accumulator_stage_count == 1 + and scale_factor_columns <= 64 + ) + if use_sm100_overlap: + accumulator_stage_count = 2 + accumulator_stage_stride_columns = accumulator_stage_columns - scale_factor_columns + accumulator_pipeline_stages = 1 + accumulator_columns = accumulator_stage_columns + accumulator_stage_stride_columns + + used_columns = accumulator_columns + scale_factor_columns + allocation_columns = _round_tmem_allocation_columns(used_columns, tmem_column_capacity) + return Tcgen05TmemPlan( + allocation_columns=allocation_columns, + accumulator_columns=accumulator_columns, + accumulator_stage_columns=accumulator_stage_columns, + accumulator_stage_count=accumulator_stage_count, + accumulator_stage_stride_columns=accumulator_stage_stride_columns, + accumulator_pipeline_stages=accumulator_pipeline_stages, + sfa_columns=sfa_columns, + sfb_columns=sfb_columns, + ) + + +def _parse_arch_number(arch: str) -> int: + if not isinstance(arch, str): + raise TypeError(f"arch must be a string, got {type(arch)}.") + normalized = arch.lower() + if normalized.startswith("sm_"): + normalized = normalized[3:] + elif normalized.startswith("sm"): + normalized = normalized[2:] + digits = [] + for character in normalized: + if not character.isdigit(): + break + digits.append(character) + if not digits: + raise ValueError(f"Cannot parse architecture {arch!r}.") + return int("".join(digits)) + + +def _round_tmem_allocation_columns(used_columns: int, capacity_columns: int) -> int: + if used_columns <= 0: + raise ValueError("TMEM usage must be positive.") + if used_columns <= 512: + allocation_columns = max(32, 1 << (used_columns - 1).bit_length()) + else: + allocation_columns = _round_up(used_columns, 32) + if allocation_columns > capacity_columns: + raise ValueError(f"TMEM plan needs {allocation_columns} columns, exceeding {capacity_columns}.") + return allocation_columns + + +def tcgen05_smem_alloc_type( + dtype: type[cutlass.Numeric], peer_dtype: type[cutlass.Numeric], arch: str +) -> type[cutlass.Numeric]: + """SMEM container type for one block-scaled TCGen05 operand. + + Blackwell mixed-width MMA consumes a uniform byte-per-element SMEM image, so its narrow operand + arrives through U4_UNPACK_U8. Rubin consumes mixed FP4 directly from packed SMEM. A 6-bit + operand always uses U6_UNPACK_U8 because no packed U6 TMA format exists. + """ + arch_number = _parse_arch_number(arch) + if arch_number not in (100, 103, 107): + raise ValueError(f"Unsupported TCGen05 architecture {arch!r}.") + needs_mixed_width_unpack = arch_number in (100, 103) and dtype.width != peer_dtype.width + needs_unpack = needs_mixed_width_unpack or 6 in (dtype.width, peer_dtype.width) + return cutlass.Int8 if (needs_unpack and dtype.width < 8) else dtype + + +def make_smem_layouts( + mma_inst: Tcgen05MmaInstruction, + mma_tiler_mnk: Tuple[int, int, int], + stages: Union[int, Tuple[int, ...]], + ab_gmem_major_modes: Tuple[OperandMajorMode, OperandMajorMode], + arch: str, +) -> Union[Tuple[SmemRegion, SmemRegion], Tuple[SmemRegion, SmemRegion, SmemRegion, SmemRegion]]: + """Derive workspace-ready TCGen05 operand regions without an MLIR context.""" + has_scale_factors = any(value is not None for value in (mma_inst.sfa_type, mma_inst.sfb_type, mma_inst.sf_vec_size)) + if has_scale_factors and any( + value is None for value in (mma_inst.sfa_type, mma_inst.sfb_type, mma_inst.sf_vec_size) + ): + raise ValueError("sfa_type, sfb_type, and sf_vec_size must be provided together.") + + operand_count = 4 if has_scale_factors else 2 + if isinstance(stages, int): + operand_stages = (stages,) * operand_count + else: + operand_stages = stages + if len(operand_stages) != operand_count: + raise ValueError(f"Expected {operand_count} stage counts, got {len(operand_stages)}.") + if any(stage <= 0 for stage in operand_stages): + raise ValueError(f"All stage counts must be positive, got {operand_stages}.") + + if mma_inst.participates not in (1, 2): + raise ValueError(f"TCGen05 MMA participates must be one or two, got {mma_inst.participates}.") + if len(mma_inst.instruction_mnk) != 3 or len(mma_tiler_mnk) != 3: + raise ValueError("instruction_mnk and mma_tiler_mnk must each contain three dimensions.") + if any(dimension <= 0 for dimension in (*mma_inst.instruction_mnk, *mma_tiler_mnk)): + raise ValueError("MMA instruction and tiler dimensions must be positive.") + + inst_m, inst_n, inst_k = mma_inst.instruction_mnk + tile_m, tile_n, tile_k = mma_tiler_mnk + if inst_m % mma_inst.participates != 0 or inst_n % mma_inst.participates != 0: + raise ValueError("MMA instruction M and N must be divisible by participates.") + if tile_m % inst_m != 0 or tile_n % inst_n != 0 or tile_k % inst_k != 0: + raise ValueError("MMA instruction dimensions must divide mma_tiler_mnk.") + + a_region = _make_ab_region( + mma_inst.a_type, + tcgen05_smem_alloc_type(mma_inst.a_type, mma_inst.b_type, arch), + tile_m // mma_inst.participates, + tile_k, + (inst_m // mma_inst.participates, inst_k), + operand_stages[0], + ab_gmem_major_modes[0], + ) + b_region = _make_ab_region( + mma_inst.b_type, + tcgen05_smem_alloc_type(mma_inst.b_type, mma_inst.a_type, arch), + tile_n // mma_inst.participates, + tile_k, + (inst_n // mma_inst.participates, inst_k), + operand_stages[1], + ab_gmem_major_modes[1], + ) + if not has_scale_factors: + return a_region, b_region + + sfa_region = _make_sf_region( + mma_inst.sfa_type, + inst_m // mma_inst.participates // 128, + tile_m // inst_m, + tile_k // inst_k, + inst_k, + mma_inst.sf_vec_size, + operand_stages[2], + ) + sfb_region = _make_sf_region( + mma_inst.sfb_type, + _round_up(inst_n, 128) // 128, + tile_n // inst_n, + tile_k // inst_k, + inst_k, + mma_inst.sf_vec_size, + operand_stages[3], + ) + return a_region, b_region, sfa_region, sfb_region + + +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +def _canonical_mode(shape, stride): + if not isinstance(shape, tuple): + return (1, 0) if shape == 1 else (shape, stride) + + result_shape = [] + result_stride = [] + for current_shape, current_stride in zip(shape, stride): + if current_shape == 1: + continue + if result_shape and result_shape[-1] * result_stride[-1] == current_stride: + result_shape[-1] *= current_shape + else: + result_shape.append(current_shape) + result_stride.append(current_stride) + if not result_shape: + return 1, 0 + if len(result_shape) == 1: + return result_shape[0], result_stride[0] + return tuple(result_shape), tuple(result_stride) + + +def _shape_size(shape) -> int: + if isinstance(shape, tuple): + result = 1 + for child_shape in shape: + result *= _shape_size(child_shape) + return result + return shape + + +def _prefix_tile_profile(shape, tile_extent: int): + """Represent a scalar prefix tile using the source mode boundaries.""" + if tile_extent <= 0 or _shape_size(shape) % tile_extent != 0: + raise ValueError(f"Tile extent {tile_extent} must divide mode {shape}.") + if not isinstance(shape, tuple): + return tile_extent + + tile_modes = [] + remaining_extent = tile_extent + for child_shape in shape: + if remaining_extent == 1: + break + + child_size = _shape_size(child_shape) + if remaining_extent >= child_size and remaining_extent % child_size == 0: + tile_modes.append(child_shape) + remaining_extent //= child_size + elif child_size % remaining_extent == 0: + tile_modes.append(_prefix_tile_profile(child_shape, remaining_extent)) + remaining_extent = 1 + else: + raise ValueError(f"Tile extent {tile_extent} does not divide a prefix of mode {shape}.") + + if remaining_extent != 1: + raise ValueError(f"Tile extent {tile_extent} exceeds the prefix of mode {shape}.") + if not tile_modes: + return 1 + if len(tile_modes) == 1: + return tile_modes[0] + return tuple(tile_modes) + + +def _flatten_mode(shape, stride): + if isinstance(shape, tuple): + result = [] + for child_shape, child_stride in zip(shape, stride): + result.extend(_flatten_mode(child_shape, child_stride)) + return result + return [(shape, stride)] + + +def _flatten_shape(shape): + if isinstance(shape, tuple): + result = [] + for child_shape in shape: + result.extend(_flatten_shape(child_shape)) + return result + return [shape] + + +def _rebuild_stride_like(shape, flat_strides): + stride_iterator = iter(flat_strides) + + def rebuild(current_shape): + if isinstance(current_shape, tuple): + return tuple(rebuild(child_shape) for child_shape in current_shape) + return next(stride_iterator) + + return rebuild(shape) + + +def _divide_mode_by_tile(shape, stride, tile_shape): + """Apply a static tile profile and return its stride and canonical rest.""" + remaining_modes = _flatten_mode(shape, stride) + tile_strides = [] + for tile_extent in _flatten_shape(tile_shape): + if not remaining_modes: + raise ValueError(f"Tile {tile_shape} exceeds mode {shape}.") + + mode_extent, mode_stride = remaining_modes[0] + if tile_extent <= 0 or mode_extent % tile_extent != 0: + raise ValueError(f"Tile mode {tile_extent} must divide source mode {mode_extent}.") + + tile_strides.append(mode_stride if tile_extent > 1 else 0) + rest_extent = mode_extent // tile_extent + if rest_extent > 1: + remaining_modes[0] = (rest_extent, mode_stride * tile_extent) + else: + remaining_modes.pop(0) + + rest_shape, rest_stride = _canonical_mode( + tuple(mode_extent for mode_extent, _ in remaining_modes), + tuple(mode_stride for _, mode_stride in remaining_modes), + ) + return (_rebuild_stride_like(tile_shape, tile_strides), rest_shape, rest_stride) + + +def _tiled_divide_2d(shape, stride, tile_shape): + """Tile a static rank-2 layout and group value modes before rest modes.""" + if ( + not isinstance(shape, tuple) + or not isinstance(stride, tuple) + or not isinstance(tile_shape, tuple) + or len(shape) != 2 + or len(stride) != 2 + or len(tile_shape) != 2 + ): + raise ValueError("A 2-D tiled divide requires rank-2 shape, stride, and tile.") + + mn_tile_stride, mn_rest_shape, mn_rest_stride = _divide_mode_by_tile(shape[0], stride[0], tile_shape[0]) + k_tile_stride, k_rest_shape, k_rest_stride = _divide_mode_by_tile(shape[1], stride[1], tile_shape[1]) + return ( + ((tile_shape[0], tile_shape[1]), mn_rest_shape, k_rest_shape), + ((mn_tile_stride, k_tile_stride), mn_rest_stride, k_rest_stride), + ) + + +def _make_ab_region( + dtype: type[cutlass.Numeric], + alloc_dtype: type[cutlass.Numeric], + mn_extent: int, + k_extent: int, + value_shape: Tuple[int, int], + stages: int, + major_mode: OperandMajorMode, +) -> SmemRegion: + """Plan one operand's SMEM region. + + ``dtype`` is the logical element type the MMA sees; ``alloc_dtype`` is the container it + occupies in SMEM. The two differ only for a sub-byte operand loaded through the unpacking TMA + (see ``tcgen05_smem_alloc_type``), where the layout must be sized and swizzled for 1-byte + containers while the major-mode rule still follows the logical type. + """ + if dtype.width in (4, 6) and major_mode != OperandMajorMode.K: + raise ValueError(f"{dtype} TCGen05 operands require K-major SMEM.") + + leading_extent = k_extent if major_mode == OperandMajorMode.K else mn_extent + leading_bits = leading_extent * alloc_dtype.width + if leading_bits % 8 != 0: + raise ValueError("The leading dimension must occupy a whole number of bytes.") + swizzle_bytes = gcd(leading_bits // 8, 128) + swizzle_by_bytes = {16: (0, 4, 3), 32: (1, 4, 3), 64: (2, 4, 3), 128: (3, 4, 3)} + if swizzle_bytes not in swizzle_by_bytes: + raise ValueError(f"Unsupported leading dimension size {leading_bits // 8} bytes.") + swizzle = swizzle_by_bytes[swizzle_bytes] + if major_mode == OperandMajorMode.MN and alloc_dtype.width == 32 and swizzle_bytes == 128: + swizzle = (2, 5, 2) + + chunk_elements = swizzle_bytes * 8 // alloc_dtype.width + if leading_extent % chunk_elements != 0: + raise ValueError(f"Leading extent {leading_extent} must be divisible by {chunk_elements}.") + repeats = leading_extent // chunk_elements + stage_stride = mn_extent * k_extent + + if major_mode == OperandMajorMode.K: + mn_shape = mn_extent + mn_stride = chunk_elements if repeats > 1 else k_extent + if repeats > 1: + k_shape = (chunk_elements, repeats) + k_stride = (1, mn_extent * chunk_elements) + else: + k_shape = k_extent + k_stride = 1 + else: + k_shape = k_extent + k_stride = chunk_elements if repeats > 1 else mn_extent + if repeats > 1: + mn_shape = (chunk_elements, repeats) + mn_stride = (1, k_extent * chunk_elements) + else: + mn_shape = mn_extent + mn_stride = 1 + + mma_mn_shape = _prefix_tile_profile(mn_shape, value_shape[0]) + mma_k_shape = _prefix_tile_profile(k_shape, value_shape[1]) + mma_shape, mma_stride = _tiled_divide_2d((mn_shape, k_shape), (mn_stride, k_stride), (mma_mn_shape, mma_k_shape)) + return SmemRegion( + name="", + kind="tensor", + dtype=alloc_dtype, + shape=(*mma_shape, stages), + stride=(*mma_stride, stage_stride if stages > 1 else 0), + swizzle=swizzle, + byte_alignment=128, + ) + + +def _make_sf_region( + dtype: type[cutlass.Numeric], + instruction_mn_blocks: int, + mn_iterations: int, + k_iterations: int, + instruction_k: int, + sf_vec_size: int, + stages: int, +) -> SmemRegion: + if instruction_mn_blocks <= 0: + raise ValueError("A scale-factor instruction must cover at least one 128-element MN block.") + if instruction_k % sf_vec_size != 0: + raise ValueError("instruction K must be divisible by sf_vec_size.") + + cta_mn_blocks = instruction_mn_blocks * mn_iterations + cta_k = instruction_k * k_iterations + basic_chunk_k = sf_vec_size * 4 + if cta_k % basic_chunk_k != 0: + raise ValueError("The CTA K extent must contain complete block-scaled basic chunks.") + basic_chunk_repetitions = cta_k // basic_chunk_k + + full_mn_shape = ((32, 4), cta_mn_blocks) + full_mn_stride = ((16, 4), basic_chunk_repetitions * 512 if cta_mn_blocks > 1 else 0) + full_k_shape = ((sf_vec_size, 4), basic_chunk_repetitions) + full_k_stride = ((0, 1), 512 if basic_chunk_repetitions > 1 else 0) + + mma_mn_shape = ((32, 4), instruction_mn_blocks) + mma_k_shape = (sf_vec_size, _prefix_tile_profile((4, basic_chunk_repetitions), instruction_k // sf_vec_size)) + mma_shape, mma_stride = _tiled_divide_2d( + (full_mn_shape, full_k_shape), (full_mn_stride, full_k_stride), (mma_mn_shape, mma_k_shape) + ) + stage_stride = cta_mn_blocks * basic_chunk_repetitions * 512 + return SmemRegion( + name="", + kind="tensor", + dtype=dtype, + shape=(*mma_shape, stages), + stride=(*mma_stride, stage_stride if stages > 1 else 0), + swizzle=(0, 4, 3), + byte_alignment=128, + ) + + +__all__ = [ + "Tcgen05MmaInstruction", + "Tcgen05TmemPlan", + "make_smem_layouts", + "make_tcgen05_tmem_plan", + "tcgen05_block_scaled_acc_dtype", + "tcgen05_smem_alloc_type", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py new file mode 100644 index 000000000..4956e6bd2 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py @@ -0,0 +1,264 @@ +"""Declarative GMEM workspace shared by host layout and device access.""" + +import dataclasses +from typing import Any, Dict, List, Literal, Optional, Tuple, Type + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 + +from .utils import ( + ceil_div, + cosize_from_shape_stride_tuples, + is_nested_shape, + is_power_of_two, + ordered_stride, + row_major_stride, + round_up, + validate_static_integer_tuple, +) + + +BufferResetAttr = Literal["data", "zero_on_first_allocate", "tail_reset"] +BufferSpace = Literal["local", "shared"] + +_reset_order = {"tail_reset": 0, "zero_on_first_allocate": 1, "data": 2} + + +@dataclasses.dataclass(frozen=True) +class DeviceRegion: + """One typed region in a local or symmetric GMEM workspace.""" + + name: str + dtype: Type[cutlass.Numeric] + shape: Tuple + buffer_space: BufferSpace + stride: Optional[Tuple] = None + mem_order: Optional[Tuple[int, ...]] = None + byte_alignment: int = 16 + reset: BufferResetAttr = "data" + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("A workspace region needs a non-empty name.") + if self.buffer_space not in ("local", "shared"): + raise ValueError(f"Invalid buffer space {self.buffer_space!r}.") + if self.reset not in _reset_order: + raise ValueError(f"Invalid reset policy {self.reset!r}.") + if not is_power_of_two(self.byte_alignment): + raise ValueError(f"Region {self.name!r} alignment must be a positive power of two.") + validate_static_integer_tuple(self.shape, field_name=f"{self.name}.shape") + if self.stride is not None and self.mem_order is not None: + raise ValueError(f"Region {self.name!r} accepts either stride or mem_order, not both.") + if self.stride is None and self.mem_order is None: + if is_nested_shape(self.shape) or len(self.shape) != 1: + raise ValueError(f"Region {self.name!r} needs an explicit layout.") + object.__setattr__(self, "stride", row_major_stride(self.shape)) + if is_nested_shape(self.shape) and self.mem_order is not None: + raise ValueError(f"Nested region {self.name!r} needs an explicit stride.") + if self.stride is not None: + validate_static_integer_tuple(self.stride, field_name=f"{self.name}.stride") + if self.mem_order is not None: + validate_static_integer_tuple(self.mem_order, field_name=f"{self.name}.mem_order") + expected = tuple(range(len(self.shape))) + if tuple(sorted(self.mem_order)) != expected: + raise ValueError( + f"Region {self.name!r} mem_order must be a permutation of {expected}, got {self.mem_order}." + ) + resolved_stride, _ = ordered_stride(self.shape, self.mem_order) + object.__setattr__(self, "stride", resolved_stride) + + +class DeviceWorkspace: + """Single source of truth for GMEM region layout and device pointer derivation.""" + + def __init__(self) -> None: + self._registered: Dict[BufferSpace, List[DeviceRegion]] = {"local": [], "shared": []} + self._region_by_name: Dict[str, DeviceRegion] = {} + self._offset: Dict[str, int] = {} + self._stride: Dict[str, Tuple] = {} + self._cosize: Dict[str, int] = {} + self._nbytes: Dict[str, int] = {} + self._byte_alignment: Dict[str, int] = {} + self._total: Dict[BufferSpace, int] = {"local": 0, "shared": 0} + self._zero_leading: Dict[BufferSpace, int] = {"local": 0, "shared": 0} + self._tail_leading: Dict[BufferSpace, int] = {"local": 0, "shared": 0} + self._base: Dict[BufferSpace, Any] = {"local": None, "shared": None} + self._finalized = False + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "DeviceWorkspace": + return self + + @property + def finalized(self) -> bool: + return self._finalized + + def register( + self, + name: str, + dtype: Type[cutlass.Numeric], + shape: Tuple, + *, + buffer_space: BufferSpace, + stride: Optional[Tuple] = None, + mem_order: Optional[Tuple[int, ...]] = None, + byte_alignment: int = 16, + reset: BufferResetAttr = "data", + ) -> None: + if self._finalized: + raise RuntimeError("Cannot register a region after finalize().") + if name in self._region_by_name or any( + region.name == name for regions in self._registered.values() for region in regions + ): + raise ValueError(f"Duplicate workspace region {name!r}.") + self._registered[buffer_space].append( + DeviceRegion( + name=name, + dtype=dtype, + shape=shape, + buffer_space=buffer_space, + stride=stride, + mem_order=mem_order, + byte_alignment=byte_alignment, + reset=reset, + ) + ) + + def finalize(self) -> None: + if self._finalized: + raise RuntimeError("DeviceWorkspace.finalize() may only be called once.") + for buffer_space in ("local", "shared"): + ordered = sorted(self._registered[buffer_space], key=lambda region: _reset_order[region.reset]) + cursor = 0 + for region_index, region in enumerate(ordered): + stride = region.stride + if stride is None: + raise RuntimeError(f"Region {region.name!r} stride was not resolved.") + cosize = cosize_from_shape_stride_tuples(region.shape, stride) + nbytes = (cosize * int(region.dtype.width) + 7) // 8 + if region.reset == "tail_reset" and ( + region_index == 0 or ordered[region_index - 1].reset != "tail_reset" + ): + cursor = round_up(cursor, 16) + byte_alignment = region.byte_alignment + cursor = round_up(cursor, byte_alignment) + self._region_by_name[region.name] = region + self._offset[region.name] = cursor + self._stride[region.name] = stride + self._cosize[region.name] = cosize + self._nbytes[region.name] = nbytes + self._byte_alignment[region.name] = byte_alignment + cursor += nbytes + is_last_tail_reset_region = region.reset == "tail_reset" and ( + region_index + 1 == len(ordered) or ordered[region_index + 1].reset != "tail_reset" + ) + if is_last_tail_reset_region: + cursor = round_up(cursor, 16) + if region.reset != "data": + self._zero_leading[buffer_space] = cursor + if is_last_tail_reset_region: + self._tail_leading[buffer_space] = cursor + self._total[buffer_space] = round_up(cursor, 16) + self._finalized = True + + def regions(self, buffer_space: BufferSpace) -> Tuple[DeviceRegion, ...]: + return tuple(self._registered[buffer_space]) + + def region(self, name: str) -> DeviceRegion: + self._require_finalized() + return self._region_by_name[name] + + def offset(self, name: str) -> int: + self._require_finalized() + return self._offset[name] + + def stride(self, name: str) -> Tuple: + self._require_finalized() + return self._stride[name] + + def cosize(self, name: str) -> int: + self._require_finalized() + return self._cosize[name] + + def nbytes(self, name: str) -> int: + self._require_finalized() + return self._nbytes[name] + + def byte_alignment(self, name: str) -> int: + self._require_finalized() + return self._byte_alignment[name] + + def total_bytes(self, buffer_space: BufferSpace) -> int: + self._require_finalized() + return self._total[buffer_space] + + @property + def local_and_shared_bytes(self) -> Tuple[int, int]: + self._require_finalized() + return self._total["local"], self._total["shared"] + + def zero_on_allocate_bytes(self, buffer_space: BufferSpace) -> int: + self._require_finalized() + return self._zero_leading[buffer_space] + + def tail_reset_bytes(self, buffer_space: BufferSpace) -> int: + self._require_finalized() + return self._tail_leading[buffer_space] + + @property + def require_zero_workspace_leading_bytes(self) -> Tuple[int, int]: + self._require_finalized() + return self._zero_leading["local"], self._zero_leading["shared"] + + @cute.jit + def assign_device_members( + self, local_workspace: cute.Pointer, shared_workspace: Optional[cute.Pointer] = None + ) -> None: + self._base["local"] = local_workspace + self._base["shared"] = shared_workspace + + def remove_device_members(self) -> None: + self._base = {"local": None, "shared": None} + + @cute.jit + def ptr(self, name: str) -> cute.Pointer: + region = self._region_by_name[name] + base = self._base[region.buffer_space] + address = base.toint() + Int64(self._offset[name]) + return cute.make_ptr(region.dtype, address, AddressSpace.gmem, assumed_align=self._byte_alignment[name]) + + @cute.jit + def tensor(self, name: str) -> cute.Tensor: + region = self._region_by_name[name] + return cute.make_tensor(self.ptr(name), cute.make_layout(region.shape, stride=self._stride[name])) + + @cute.jit + def reset_tail(self, tid: Int32, total_threads: int) -> None: + self.reset_tail_space("local", tid, total_threads) + self.reset_tail_space("shared", tid, total_threads) + + @cute.jit + def reset_tail_space(self, buffer_space: BufferSpace, tid: Int32, total_threads: int) -> None: + num_vectors = self._tail_leading[buffer_space] // 16 + if cutlass.const_expr(num_vectors > 0): + vectors = cute.make_tensor( + cute.make_ptr(Int32, self._base[buffer_space].toint(), AddressSpace.gmem, assumed_align=16), + cute.make_layout((num_vectors, 4), stride=(4, 1)), + ) + zero = cute.make_rmem_tensor((4,), Int32) + for element in cutlass.range_constexpr(4): + zero[element] = Int32(0) + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Int32, num_bits_per_copy=128) + reset_round_count = ceil_div(num_vectors, total_threads) + for reset_round in cutlass.range_constexpr(reset_round_count): + vector_index = Int32(reset_round * total_threads) + tid + if vector_index < Int32(num_vectors): + cute.copy(store_atom, zero, vectors[vector_index, None]) + + def _require_finalized(self) -> None: + if not self._finalized: + raise RuntimeError("DeviceWorkspace must be finalized first.") diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/dsl_helpers.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/dsl_helpers.py new file mode 100644 index 000000000..807d85c27 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/dsl_helpers.py @@ -0,0 +1,235 @@ +"""General-purpose CuTe DSL helpers.""" + +from typing import Callable, Literal, Optional + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import Pointer +from cutlass.cutlass_dsl import Boolean, Int16, Int32 + +from .ptx_helpers import nanosleep +from .utils import ceil_div + + +@cute.jit +def smem_exclusive_prefix( + input_tensor: cute.Tensor, + output_tensor: cute.Tensor, + warp_totals: cute.Tensor, + block_thread_count: int, + thread_idx: Int32, + lane_idx: Int32, + warp_idx: Int32, +) -> Int32: + """Compute a CTA-wide exclusive prefix over an Int32 SMEM tensor.""" + num_elements = cute.size(input_tensor) + scan_rows = num_elements // 4 + num_warps = block_thread_count // 32 + input_vectors = cute.make_tensor(input_tensor.iterator, cute.make_layout((scan_rows, 4), stride=(4, 1))) + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128) + + values = cute.make_rmem_tensor((4,), cutlass.Int32) + carry = Int32(0) + for segment in cutlass.range_constexpr(ceil_div(scan_rows, block_thread_count)): + row = Int32(segment * block_thread_count) + thread_idx + if row < Int32(scan_rows): + row_slice = input_vectors[row, None] + pointer = row_slice.iterator + aligned_row_slice = cute.make_tensor( + cute.make_ptr(pointer.dtype, pointer.toint(), pointer.memspace, assumed_align=16), + row_slice.layout, + ) + cute.copy(load_atom, aligned_row_slice, values) + else: + for element in cutlass.range_constexpr(4): + values[element] = Int32(0) + + local_prefix = (Int32(0), values[0], values[0] + values[1], values[0] + values[1] + values[2]) + lane_total = local_prefix[3] + values[3] + inclusive = lane_total + for step_log in cutlass.range_constexpr(5): + step = Int32(1 << step_log) + previous = Int32(cute.arch.shuffle_sync(inclusive, lane_idx - step)) + if lane_idx >= step: + inclusive = inclusive + previous + lane_base = inclusive - lane_total + warp_total = Int32(cute.arch.shuffle_sync(inclusive, Int32(31))) + if lane_idx == Int32(0): + warp_totals[warp_idx] = warp_total + cute.arch.sync_threads() + + region_total = Int32(0) + if lane_idx < Int32(num_warps): + region_total = warp_totals[lane_idx] + inclusive_region = region_total + for step_log in cutlass.range_constexpr(5): + step = Int32(1 << step_log) + previous = Int32(cute.arch.shuffle_sync(inclusive_region, lane_idx - step)) + if lane_idx >= step: + inclusive_region = inclusive_region + previous + warp_base = Int32(cute.arch.shuffle_sync(inclusive_region - region_total, warp_idx)) + segment_total = Int32(cute.arch.shuffle_sync(inclusive_region, Int32(31))) + base = carry + warp_base + lane_base + if row < Int32(scan_rows): + first_element = row * Int32(4) + for element in cutlass.range_constexpr(4): + output_tensor[first_element + Int32(element)] = base + local_prefix[element] + carry = carry + segment_total + cute.arch.sync_threads() + return carry + + +@cute.jit +def mark_alignment(tensor: cute.Tensor, byte_alignment: int) -> cute.Tensor: + pointer = tensor.iterator + return cute.make_tensor( + cute.make_ptr(pointer.dtype, pointer.toint(), pointer.memspace, assumed_align=byte_alignment), tensor.layout + ) + + +@cute.jit +def spin_peek(pointer: Pointer, condition: Callable[[Int32], Boolean], scope: str = "gpu") -> Boolean: + """Perform one acquire load and test its value.""" + value = cute.arch.load(pointer, pointer.dtype, sem="acquire", scope=scope) + return Boolean(condition(value)) + + +@cute.jit +def spin_wait( + pointer: Pointer, + condition: Callable[[Int32], Boolean], + scope: str = "gpu", + sleep_cycles: int = 150, + peek_status: Optional[Boolean] = None, +) -> None: + """Wait until an acquire-loaded value satisfies the condition.""" + wait_required = Boolean(True) + if cutlass.const_expr(peek_status is not None): + wait_required = not peek_status + + if wait_required: + value = cute.arch.load(pointer, pointer.dtype, sem="acquire", scope=scope) + while not condition(value): + nanosleep(sleep_cycles) + value = cute.arch.load(pointer, pointer.dtype, sem="acquire", scope=scope) + + +def _tma_multicast_pattern( + cluster_mn: tuple[int, int], mma_cta_count: int, tensor_role: Literal["a", "b", "sfa", "sfb"] +) -> int: + cluster_m, cluster_n = cluster_mn + if tensor_role in ("a", "sfa"): + return sum(1 << (cluster_n_index * cluster_m) for cluster_n_index in range(cluster_n)) + if tensor_role == "b": + return sum(1 << (cluster_m_index * mma_cta_count) for cluster_m_index in range(cluster_m // mma_cta_count)) + if tensor_role == "sfb": + return (1 << cluster_m) - 1 + raise ValueError(f"Unsupported TMA tensor role {tensor_role!r}.") + + +@cute.jit +def tma_multicast_mask( + preferred_cluster_mn: tuple[int, int], + fallback_cluster_mn: Optional[tuple[int, int]], + cta_coord_in_cluster: cute.Coord, + is_preferred: Optional[Boolean], + is_2cta: bool, + tensor_role: Literal["a", "b", "sfa", "sfb"], +) -> Int16: + """Build a preferred/fallback TMA multicast mask.""" + preferred_m, preferred_n = preferred_cluster_mn + if cutlass.const_expr(preferred_m <= 0 or preferred_n <= 0 or preferred_m * preferred_n > 16): + raise ValueError(f"Invalid preferred cluster shape {preferred_cluster_mn}.") + + mma_cta_count = 2 if cutlass.const_expr(is_2cta) else 1 + if cutlass.const_expr(preferred_m % mma_cta_count != 0): + raise ValueError("Preferred cluster M must be divisible by the MMA CTA count.") + + preferred_pattern = _tma_multicast_pattern(preferred_cluster_mn, mma_cta_count, tensor_role) + cta_m = Int32(cta_coord_in_cluster[0]) + if cutlass.const_expr(fallback_cluster_mn is None): + if cutlass.const_expr(tensor_role in ("a", "sfa")): + result = cute.arch.inline_ptx( + f"shl.b16 {{$w0}}, 0x{preferred_pattern:04x}, {{$r0}};", + write_only_types=[Int16], + read_only_args=[cta_m], + ) + else: + cta_n = Int32(cta_coord_in_cluster[1]) + mma_cta_index = Int32(0) if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) else cta_m % 2 + result = cute.arch.inline_ptx( + "{\n\t" + ".reg .u32 offset;\n\t" + f"mad.lo.u32 offset, {{$r0}}, {preferred_m}, {{$r1}};\n\t" + f"shl.b16 {{$w0}}, 0x{preferred_pattern:04x}, offset;\n\t" + "}", + write_only_types=[Int16], + read_only_args=[cta_n, mma_cta_index], + ) + return Int16(result) + + fallback_m, fallback_n = fallback_cluster_mn + if cutlass.const_expr(fallback_m <= 0 or fallback_n <= 0 or fallback_m * fallback_n > 16): + raise ValueError(f"Invalid fallback cluster shape {fallback_cluster_mn}.") + if cutlass.const_expr(preferred_m % fallback_m != 0 or preferred_n % fallback_n != 0): + raise ValueError("Preferred cluster dimensions must be divisible by fallback dimensions.") + if cutlass.const_expr(fallback_m % mma_cta_count != 0): + raise ValueError("Fallback cluster M must be divisible by the MMA CTA count.") + + fallback_pattern = _tma_multicast_pattern(fallback_cluster_mn, mma_cta_count, tensor_role) + if cutlass.const_expr(preferred_pattern == fallback_pattern): + if cutlass.const_expr(tensor_role in ("a", "sfa")): + result = cute.arch.inline_ptx( + f"shl.b16 {{$w0}}, 0x{preferred_pattern:04x}, {{$r0}};", + write_only_types=[Int16], + read_only_args=[cta_m], + ) + else: + cta_n = Int32(cta_coord_in_cluster[1]) + mma_cta_index = Int32(0) if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) else cta_m % 2 + result = cute.arch.inline_ptx( + "{\n\t" + ".reg .u32 offset;\n\t" + f"mad.lo.u32 offset, {{$r0}}, {preferred_m}, {{$r1}};\n\t" + f"shl.b16 {{$w0}}, 0x{preferred_pattern:04x}, offset;\n\t" + "}", + write_only_types=[Int16], + read_only_args=[cta_n, mma_cta_index], + ) + return Int16(result) + + if cutlass.const_expr(is_preferred is None): + raise ValueError("is_preferred is required when the preferred and fallback multicast patterns differ.") + + if cutlass.const_expr(tensor_role in ("a", "sfa")): + result = cute.arch.inline_ptx( + "{\n\t" + f"mov.b16 {{$w0}}, 0x{preferred_pattern:04x};\n\t" + f"@!{{$r0}} mov.b16 {{$w0}}, 0x{fallback_pattern:04x};\n\t" + "shl.b16 {$w0}, {$w0}, {$r1};\n\t" + "}", + write_only_types=[Int16], + read_only_args=[is_preferred, cta_m], + ) + else: + if cutlass.const_expr(fallback_pattern & preferred_pattern != fallback_pattern): + raise ValueError("Fallback B/SFB multicast pattern must be a subset of the preferred pattern.") + cta_n = Int32(cta_coord_in_cluster[1]) + mma_cta_index = Int32(0) if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) else cta_m % 2 + result = cute.arch.inline_ptx( + "{\n\t" + ".reg .u32 cluster_m, offset;\n\t" + f"mov.u32 cluster_m, {preferred_m};\n\t" + f"@!{{$r0}} mov.u32 cluster_m, {fallback_m};\n\t" + "mad.lo.u32 offset, {$r1}, cluster_m, {$r2};\n\t" + f"mov.b16 {{$w0}}, 0x{preferred_pattern:04x};\n\t" + f"@!{{$r0}} and.b16 {{$w0}}, {{$w0}}, 0x{fallback_pattern:04x};\n\t" + "shl.b16 {$w0}, {$w0}, offset;\n\t" + "}", + write_only_types=[Int16], + read_only_args=[is_preferred, cta_n, mma_cta_index], + ) + return Int16(result) + + +__all__ = ["mark_alignment", "smem_exclusive_prefix", "spin_peek", "spin_wait", "tma_multicast_mask"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/flag_batch.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/flag_batch.py new file mode 100644 index 000000000..28088ec8f --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/flag_batch.py @@ -0,0 +1,133 @@ +"""Lane- and warp-distributed GPU release-counter batching.""" + +import dataclasses +from typing import Any, Union + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 + +from .ptx_helpers import red_add_release_gpu_s32, red_async_add_release_gpu_s32 + + +@dataclasses.dataclass(frozen=True) +class GpuReleaseFlagBatchTracker: + """Lane-distributed delayed publication state for synchronous GPU release counters.""" + + flag_address: Int64 + accumulated_flags: Int32 + phase: Int32 + thread_idx: Int32 + + @cute.jit + def _make(self, flag_address: Int64, accumulated_flags: Int32, phase: Int32) -> "GpuReleaseFlagBatchTracker": + return type(self)( + flag_address=flag_address, accumulated_flags=accumulated_flags, phase=phase, thread_idx=self.thread_idx + ) + + @cute.jit + def fire(self) -> None: + if self.flag_address != Int64(0): + pointer = cute.make_ptr(cutlass.Int32, self.flag_address, AddressSpace.gmem, assumed_align=4) + red_add_release_gpu_s32(pointer, Int32(1)) + + @cute.jit + def accumulate( + self, next_phase: Any, flush_threshold: int, flag_address: Int64, no_fire: bool = False + ) -> "GpuReleaseFlagBatchTracker": + if cutlass.const_expr(flush_threshold == 1): + if cutlass.const_expr(not no_fire): + lane_address = Int64(0) + if self.thread_idx == Int32(0): + lane_address = flag_address + self._make(flag_address=lane_address, accumulated_flags=Int32(1), phase=self.phase).fire() + return self._make(flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(next_phase)) + + current_address = self.flag_address + accumulated_flags = self.accumulated_flags + if self.thread_idx == accumulated_flags: + current_address = flag_address + accumulated_flags = accumulated_flags + Int32(1) + + if accumulated_flags == Int32(flush_threshold) or next_phase != self.phase: + if cutlass.const_expr(not no_fire): + self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=self.phase).fire() + accumulated_flags = Int32(0) + current_address = Int64(0) + + return self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=Int32(next_phase)) + + +@dataclasses.dataclass(frozen=True) +class GpuAsyncReleaseFlagBatchTracker: + """Loop-carried warp-uniform state for asynchronous GPU release counters.""" + + flag_address: Int64 + accumulated_flags: Int32 + phase: Int32 + warp_idx: Int32 + + @cute.jit + def _make(self, flag_address: Int64, accumulated_flags: Int32, phase: Int32) -> "GpuAsyncReleaseFlagBatchTracker": + return type(self)( + flag_address=flag_address, accumulated_flags=accumulated_flags, phase=phase, warp_idx=self.warp_idx + ) + + @cute.jit + def fire(self) -> None: + if self.flag_address != Int64(0): + pointer = cute.make_ptr(cutlass.Int32, self.flag_address, AddressSpace.gmem, assumed_align=4) + with cute.arch.elect_one(): + red_async_add_release_gpu_s32(pointer, Int32(1)) + + @cute.jit + def accumulate( + self, next_phase: Any, flush_threshold: int, flag_address: Int64, no_fire: bool = False + ) -> "GpuAsyncReleaseFlagBatchTracker": + if cutlass.const_expr(flush_threshold == 1): + if cutlass.const_expr(not no_fire): + warp_address = Int64(0) + if self.warp_idx == Int32(0): + warp_address = flag_address + self._make(flag_address=warp_address, accumulated_flags=Int32(1), phase=self.phase).fire() + return self._make(flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(next_phase)) + + current_address = self.flag_address + accumulated_flags = self.accumulated_flags + if self.warp_idx == accumulated_flags: + current_address = flag_address + accumulated_flags = accumulated_flags + Int32(1) + + if accumulated_flags == Int32(flush_threshold) or next_phase != self.phase: + if cutlass.const_expr(not no_fire): + self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=self.phase).fire() + accumulated_flags = Int32(0) + current_address = Int64(0) + + return self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=Int32(next_phase)) + + +@cute.jit +def make_flag_batch_tracker( + use_async: bool, *, flag_address: Int64, accumulated_flags: Int32, phase: Int32, thread_idx: Int32 +) -> Union[GpuReleaseFlagBatchTracker, GpuAsyncReleaseFlagBatchTracker]: + """Construct a lane-batched synchronous or warp-batched asynchronous tracker. + + ``thread_idx`` must be zero-based within the caller's cooperating publisher + group. The async tracker maps each contiguous group of 32 thread indices to + one warp-uniform publisher. + """ + if cutlass.const_expr(use_async): + return GpuAsyncReleaseFlagBatchTracker( + flag_address=flag_address, + accumulated_flags=accumulated_flags, + phase=phase, + warp_idx=cute.arch.make_warp_uniform(thread_idx // Int32(32)), + ) + return GpuReleaseFlagBatchTracker( + flag_address=flag_address, accumulated_flags=accumulated_flags, phase=phase, thread_idx=thread_idx + ) + + +__all__ = ["GpuAsyncReleaseFlagBatchTracker", "GpuReleaseFlagBatchTracker", "make_flag_batch_tracker"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/iket_compat.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/iket_compat.py new file mode 100644 index 000000000..d4cfd4c97 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/iket_compat.py @@ -0,0 +1,35 @@ +"""Compatibility wrapper for optional in-kernel event tracing support.""" + +try: + from cutlass.cute.experimental import iket +except (ImportError, NotImplementedError): + try: + from cutlass.cute import iket # type: ignore + except (ImportError, NotImplementedError): + class _IketShim: + """No-op IKET interface for toolchains without the dialect.""" + + @staticmethod + def range_push(_name, *_args, **_kwargs): + return None + + @staticmethod + def range_pop(*_args, **_kwargs): + return None + + @staticmethod + def range_start(_name, *_args, **_kwargs): + return None + + @staticmethod + def range_end(_token=None, *_args, **_kwargs): + return None + + @staticmethod + def mark(_name, *_args, **_kwargs): + return None + + iket = _IketShim() # type: ignore + + +__all__ = ["iket"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py new file mode 100644 index 000000000..56d81bc87 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py @@ -0,0 +1,576 @@ +"""Minimal inline-PTX primitives required by the greenfield NVFP4 kernels.""" + +from typing import Optional + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith, llvm, vector +from cutlass.cutlass_dsl import Float32, Int32, Int64, T, dsl_user_op + + +TmaCacheHintEvictFirst = 0x12F0000000000000 + + +def _address_value(pointer_or_address, *, loc=None, ip=None): + if isinstance(pointer_or_address, Int64): + return pointer_or_address.ir_value() + return pointer_or_address.toint(loc=loc, ip=ip).ir_value() + + +@dsl_user_op +def nanosleep(sleep_cycles: int, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> None: + """Suspend the calling thread for up to the requested clock cycles.""" + if cutlass.const_expr(hasattr(cute.arch, "nanosleep")): + cute.arch.nanosleep(sleep_time=sleep_cycles, loc=loc, ip=ip) + return + + llvm.inline_asm( + res=None, + operands_=[Int32(sleep_cycles).ir_value(loc=loc, ip=ip)], + asm_string="nanosleep.u32 $0;", + constraints="r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def read_clock64(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> Int64: + """Read the per-SM 64-bit cycle counter.""" + return Int64( + llvm.inline_asm( + T.i64(), [], "mov.u64 $0, %clock64;", "=l", has_side_effects=True, asm_dialect=0, loc=loc, ip=ip + ) + ) + + +@cute.jit +def movmatrix_b16(input_regs: cute.Tensor) -> cute.Tensor: + """Transpose every packed m8n8 b16 register fragment across the warp.""" + if cutlass.const_expr(input_regs.element_type.width != 32): + raise TypeError(f"movmatrix_b16 expects packed 32-bit registers, got {input_regs.element_type}.") + + input_words = cute.coalesce(cute.flatten(cute.recast_tensor(input_regs, Int32))) + output_words = cute.make_rmem_tensor((cute.size(input_words),), Int32) + for word_idx in cutlass.range_constexpr(cute.size(input_words)): + output_words[word_idx] = Int32( + llvm.inline_asm( + T.i32(), + [Int32(input_words[word_idx]).ir_value()], + "movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;", + "=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + output_regs = cute.make_rmem_tensor(input_regs.layout, input_regs.element_type) + output_regs_words = cute.coalesce(cute.flatten(cute.recast_tensor(output_regs, Int32))) + for word_idx in cutlass.range_constexpr(cute.size(output_words)): + output_regs_words[word_idx] = output_words[word_idx] + return output_regs + + +@dsl_user_op +def cvt_f32_to_fp8_to_f32( + value, fp8_type, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> Float32: + """Round one f32 through the selected FP8 format and widen it back.""" + if cutlass.const_expr(fp8_type is cutlass.Float8E8M0FNU): + downcast_instruction = "cvt.rp.satfinite.ue8m0x2.f32" + upcast_instruction = "cvt.rn.bf16x2.ue8m0x2" + elif cutlass.const_expr(fp8_type is cutlass.Float8E4M3FN): + downcast_instruction = "cvt.rn.satfinite.e4m3x2.f32" + upcast_instruction = "cvt.rn.bf16x2.e4m3x2" + elif cutlass.const_expr(fp8_type is cutlass.Float8E5M2): + downcast_instruction = "cvt.rn.satfinite.e5m2x2.f32" + upcast_instruction = "cvt.rn.bf16x2.e5m2x2" + else: + raise ValueError(f"Unsupported FP8 type {fp8_type}.") + + packed_bf16 = llvm.inline_asm( + T.i32(), + [Float32(value).ir_value(loc=loc, ip=ip)], + "{\n" + " .reg .b16 converted;\n" + f" {downcast_instruction} converted, 0f00000000, $1;\n" + f" {upcast_instruction} $0, converted;\n" + "}", + "=r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + bf16_pair_type = ir.Type.parse("vector<2xbf16>") + bf16_pair = llvm.bitcast(bf16_pair_type, packed_bf16, loc=loc, ip=ip) + rounded_bf16 = vector.extract(bf16_pair, [], [0], loc=loc, ip=ip) + return Float32(arith.extf(Float32.mlir_type, rounded_bf16, loc=loc, ip=ip)) + + +@dsl_user_op +def tma_load_1d(destination_smem, source_gmem, mbarrier_smem, num_bytes, *, loc=None, ip=None) -> None: + """Issue a cache-hinted 1D GMEM-to-SMEM bulk copy.""" + llvm.inline_asm( + None, + [ + destination_smem.toint(loc=loc, ip=ip).ir_value(), + _address_value(source_gmem, loc=loc, ip=ip), + num_bytes.ir_value(), + mbarrier_smem.toint(loc=loc, ip=ip).ir_value(), + Int64(TmaCacheHintEvictFirst).ir_value(), + ], + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [$0], [$1], $2, [$3], $4;", + "r,l,r,r,l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_i32_to_peer_cluster_smem_async( + smem_pointer, + value: Int32, + mbarrier_pointer, + destination_cta_rank, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + """Store one Int32 to peer SMEM and complete its transaction barrier.""" + smem_address = llvm.ptrtoint(T.i32(), smem_pointer.llvm_ptr, loc=loc, ip=ip) + mbarrier_address = llvm.ptrtoint(T.i32(), mbarrier_pointer.llvm_ptr, loc=loc, ip=ip) + llvm.inline_asm( + res=None, + operands_=[ + smem_address, + value.ir_value(loc=loc, ip=ip), + mbarrier_address, + Int32(destination_cta_rank).ir_value(loc=loc, ip=ip), + ], + asm_string="""{{ + .reg .u32 remote_addr; + .reg .u32 remote_mbar; + mapa.shared::cluster.u32 remote_addr, $0, $3; + mapa.shared::cluster.u32 remote_mbar, $2, $3; + st.async.shared::cluster.mbarrier::complete_tx::bytes.u32 [remote_addr], $1, [remote_mbar]; + }}""", + constraints="r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def mbarrier_arrive_expect_tx_on_peer( + mbarrier_pointer, + transaction_bytes: Int32, + destination_cta_rank, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + """Declare an expected peer-CTA SMEM transaction.""" + mbarrier_address = llvm.ptrtoint(T.i32(), mbarrier_pointer.llvm_ptr, loc=loc, ip=ip) + llvm.inline_asm( + res=None, + operands_=[ + mbarrier_address, + Int32(destination_cta_rank).ir_value(loc=loc, ip=ip), + transaction_bytes.ir_value(loc=loc, ip=ip), + ], + asm_string="""{{ + .reg .u32 remote_mbar; + mapa.shared::cluster.u32 remote_mbar, $0, $1; + mbarrier.arrive.expect_tx.shared::cluster.b64 _, [remote_mbar], $2; + }}""", + constraints="r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cp_async_bulk_s2g( + destination_gmem, + source_smem, + num_bytes, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + llvm.inline_asm( + None, + [ + destination_gmem.toint(loc=loc, ip=ip).ir_value(), + source_smem.toint(loc=loc, ip=ip).ir_value(), + num_bytes.ir_value(), + ], + "cp.async.bulk.global.shared::cta.bulk_group [$0], [$1], $2;", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cp_reduce_async_bulk_add_bf16_s2g( + destination_gmem, + source_smem, + num_bytes, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + llvm.inline_asm( + None, + [ + destination_gmem.toint(loc=loc, ip=ip).ir_value(), + source_smem.toint(loc=loc, ip=ip).ir_value(), + num_bytes.ir_value(), + ], + "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.noftz.bf16 [$0], [$1], $2;", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cp_reduce_async_bulk_add_u32_s2g( + destination_gmem, + source_smem, + num_bytes, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + llvm.inline_asm( + None, + [ + destination_gmem.toint(loc=loc, ip=ip).ir_value(), + source_smem.toint(loc=loc, ip=ip).ir_value(), + num_bytes.ir_value(), + ], + "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.u32 [$0], [$1], $2;", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def lds128_v4_b32(smem_pointer, *, loc=None, ip=None): + result = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 4), + [smem_pointer.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return tuple(Int32(llvm.extractvalue(T.i32(), result, [index])) for index in range(4)) + + +@dsl_user_op +def stg_f32(address: Int64, value: Float32, predicate: Optional[Int32] = None, *, loc=None, ip=None) -> None: + if predicate is None: + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value()], + "st.global.f32 [$0], $1;", + "l,f", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value(), predicate.ir_value()], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.f32 [$0], $1;\n\t}", + "l,f,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stg_b64(address: Int64, value: Int64, predicate: Optional[Int32] = None, *, loc=None, ip=None) -> None: + if predicate is None: + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value()], + "st.global.u64 [$0], $1;", + "l,l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + return + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value(), predicate.ir_value()], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.u64 [$0], $1;\n\t}", + "l,l,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add_relaxed_sys_s32(address: Int64, value: Int32, *, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value()], + "red.relaxed.sys.global.add.s32 [$0], $1;", + "l,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add_relaxed_sys_f32( + address, value: Float32, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> None: + llvm.inline_asm( + None, + [address.toint(loc=loc, ip=ip).ir_value(), value.ir_value(loc=loc, ip=ip)], + "red.relaxed.sys.global.add.f32 [$0], $1;", + "l,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add_release_sys_s32(address: Int64, value: Int32, *, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value()], + "red.release.sys.global.add.s32 [$0], $1;", + "l,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add_release_gpu_s32( + counter_pointer, value: Int32, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> None: + llvm.inline_asm( + None, + [counter_pointer.toint(loc=loc, ip=ip).ir_value(), value.ir_value()], + "red.release.gpu.global.add.s32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_async_add_release_gpu_s32( + counter_pointer, value: Int32, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> None: + llvm.inline_asm( + None, + [counter_pointer.toint(loc=loc, ip=ip).ir_value(), value.ir_value()], + "red.async.release.gpu.global.add.s32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add_relaxed_sys_v2_bf16x2( + address, value0, value1, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> None: + llvm.inline_asm( + None, + [address.toint(loc=loc, ip=ip).ir_value(), value0.ir_value(), value1.ir_value()], + "red.relaxed.sys.global.add.noftz.v2.bf16x2 [$0], {$1, $2};", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def cvt_f32x4_to_f8x4_pack_i32(fp32x4: cute.Tensor, fp8_type, *, loc=None, ip=None) -> Int32: + """Round four f32 lanes to the selected FP8 format and pack them into one i32.""" + fp32x4 = fp32x4.load() + src_vec4 = fp32x4.ir_value(loc=loc, ip=ip) if hasattr(fp32x4, "ir_value") else fp32x4 + + src0 = Float32(vector.extract(src_vec4, [], [0])).ir_value(loc=loc, ip=ip) + src1 = Float32(vector.extract(src_vec4, [], [1])).ir_value(loc=loc, ip=ip) + src2 = Float32(vector.extract(src_vec4, [], [2])).ir_value(loc=loc, ip=ip) + src3 = Float32(vector.extract(src_vec4, [], [3])).ir_value(loc=loc, ip=ip) + + if cutlass.const_expr(fp8_type is cutlass.Float8E8M0FNU): + cvt_instruction = "cvt.rp.satfinite.ue8m0x2.f32" + elif cutlass.const_expr(fp8_type is cutlass.Float8E4M3FN): + cvt_instruction = "cvt.rn.satfinite.e4m3x2.f32" + elif cutlass.const_expr(fp8_type is cutlass.Float8E5M2): + cvt_instruction = "cvt.rn.satfinite.e5m2x2.f32" + else: + raise ValueError(f"Unsupported FP8 type {fp8_type}.") + + packed_i32 = llvm.inline_asm( + T.i32(), + [src0, src1, src2, src3], + "{\n" + " .reg .b16 lo;\n" + " .reg .b16 hi;\n" + f" {cvt_instruction} lo, $2, $1;\n" + f" {cvt_instruction} hi, $4, $3;\n" + " mov.b32 $0, {lo, hi};\n" + "}", + "=r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return Int32(packed_i32) + + +@dsl_user_op +def stg_e8m0_from_f32(addr: Int64, fp32_val: Float32, *, loc=None, ip=None) -> None: + """Convert ``fp32_val`` to E8M0 via PTX and store the 1-byte result to global memory. + + Uses ``cvt.rp.satfinite.ue8m0x2.f32`` -- the correct fp32 -> E8M0 path -- rather + than the DSL's generic ``.to(Float8E8M0FNU)``, which does not lower correctly for + the non-IEEE-754 E8M0 type. + """ + llvm.inline_asm( + None, + [addr.ir_value(), fp32_val.ir_value()], + "{\n" + " .reg .b16 bf_lo;\n" + " .reg .u32 tmp;\n" + " cvt.rp.satfinite.ue8m0x2.f32 bf_lo, 0f00000000, $1;\n" + " cvt.u32.u16 tmp, bf_lo;\n" + " st.global.b8 [$0], tmp;\n" + "}", + "l,f", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stg_e8m0x8_from_f32( + addr: Int64, + v0: Float32, v1: Float32, v2: Float32, v3: Float32, + v4: Float32, v5: Float32, v6: Float32, v7: Float32, + *, loc=None, ip=None, +) -> None: + """Convert 8 fp32 values to E8M0 and store them as 8 contiguous bytes in one shot. + + Batched form of ``stg_e8m0_from_f32``: four ``cvt.rp.satfinite.ue8m0x2.f32`` each + pack two E8M0 bytes, the four ``.b16`` results are assembled into two ``.b32`` words, + and a single ``st.global.v2.u32`` writes all 8 bytes. ``addr`` must be 8-byte aligned; + output byte ``k`` holds E8M0(``v{k}``). + """ + llvm.inline_asm( + None, + [ + addr.ir_value(), + v0.ir_value(), v1.ir_value(), v2.ir_value(), v3.ir_value(), + v4.ir_value(), v5.ir_value(), v6.ir_value(), v7.ir_value(), + ], + "{\n" + " .reg .b16 p0, p1, p2, p3;\n" + " .reg .b32 w0, w1;\n" + " cvt.rp.satfinite.ue8m0x2.f32 p0, $2, $1;\n" + " cvt.rp.satfinite.ue8m0x2.f32 p1, $4, $3;\n" + " cvt.rp.satfinite.ue8m0x2.f32 p2, $6, $5;\n" + " cvt.rp.satfinite.ue8m0x2.f32 p3, $8, $7;\n" + " mov.b32 w0, {p0, p1};\n" + " mov.b32 w1, {p2, p3};\n" + " st.global.v2.u32 [$0], {w0, w1};\n" + "}", + "l,f,f,f,f,f,f,f,f", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +__all__ = [ + "TmaCacheHintEvictFirst", + "cp_async_bulk_s2g", + "cp_reduce_async_bulk_add_bf16_s2g", + "cp_reduce_async_bulk_add_u32_s2g", + "cvt_f32_to_fp8_to_f32", + "cvt_f32x4_to_f8x4_pack_i32", + "lds128_v4_b32", + "mbarrier_arrive_expect_tx_on_peer", + "movmatrix_b16", + "nanosleep", + "read_clock64", + "red_add_relaxed_sys_f32", + "red_add_relaxed_sys_s32", + "red_add_relaxed_sys_v2_bf16x2", + "red_async_add_release_gpu_s32", + "red_add_release_gpu_s32", + "red_add_release_sys_s32", + "store_i32_to_peer_cluster_smem_async", + "stg_b64", + "stg_e8m0_from_f32", + "stg_e8m0x8_from_f32", + "stg_f32", + "tma_load_1d", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py new file mode 100644 index 000000000..ce412673e --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py @@ -0,0 +1,427 @@ +"""Pure-Python SMEM declarations with lifetime-aware overlay placement.""" + +import dataclasses +from typing import Dict, List, Literal, Optional, Tuple, Type, Union + +import cutlass +import cutlass.cute as cute + +from .utils import ( + cosize_from_shape_stride_tuples, + is_power_of_two, + row_major_stride, + round_up, + validate_static_integer_tuple, +) + + +SmemRegionKind = Literal["mbarrier", "tensor"] +SwizzleSpec = Tuple[int, int, int] + + +def _swizzle_alignment(swizzle: Optional[SwizzleSpec]) -> Optional[int]: + if swizzle is None or swizzle[0] == 0: + return None + num_bits, num_base, _ = swizzle + return 1 << (num_base + num_bits) + + +@dataclasses.dataclass(frozen=True) +class SmemRegion: + """One logical SMEM tensor with no assigned byte offset.""" + + name: str + kind: SmemRegionKind + dtype: Type[cutlass.Numeric] + shape: Tuple + stride: Tuple + swizzle: Optional[SwizzleSpec] + byte_alignment: int + + @property + def cosize(self) -> int: + return int(cosize_from_shape_stride_tuples(self.shape, self.stride)) + + @property + def nbytes(self) -> int: + return (self.cosize * int(self.dtype.width) + 7) // 8 + + +class SmemLifetime: + """One mutually exclusive use of an overlay's physical storage.""" + + def __init__( + self, + workspace: "SmemWorkspace", + overlay: "SmemOverlay", + name: str, + ) -> None: + self._workspace = workspace + self._overlay = overlay + self.name = name + self._regions: List[SmemRegion] = [] + + @property + def regions(self) -> Tuple[SmemRegion, ...]: + return tuple(self._regions) + + def register_tensor( + self, + name: str, + dtype: Type[cutlass.Numeric], + shape: Tuple, + *, + stride: Optional[Tuple] = None, + swizzle: Optional[SwizzleSpec] = None, + byte_alignment: Optional[int] = None, + ) -> SmemRegion: + region = self._workspace._make_region( + name=name, + kind="tensor", + dtype=dtype, + shape=shape, + stride=stride, + swizzle=swizzle, + byte_alignment=byte_alignment, + ) + self._workspace._claim_region_name(region) + self._regions.append(region) + return region + + +class SmemOverlay: + """One physical allocation shared by mutually exclusive lifetimes.""" + + def __init__(self, workspace: "SmemWorkspace", name: str) -> None: + self._workspace = workspace + self.name = name + self._lifetimes: List[SmemLifetime] = [] + self._lifetime_names: set[str] = set() + + @property + def lifetimes(self) -> Tuple[SmemLifetime, ...]: + return tuple(self._lifetimes) + + def add_lifetime(self, name: str) -> SmemLifetime: + if self._workspace.finalized: + raise RuntimeError("Cannot add an SMEM lifetime after finalize().") + if not name: + raise ValueError("An SMEM lifetime needs a non-empty name.") + if name in self._lifetime_names: + raise ValueError( + f"Duplicate lifetime {name!r} in overlay {self.name!r}." + ) + lifetime = SmemLifetime(self._workspace, self, name) + self._lifetimes.append(lifetime) + self._lifetime_names.add(name) + return lifetime + + +_TopLevelDeclaration = Union[SmemRegion, SmemOverlay] + + +class SmemWorkspace: + """Collect static SMEM declarations and finalize one physical placement.""" + + def __init__( + self, + *, + base_alignment: int = 1024, + total_alignment: int = 16, + ) -> None: + if not is_power_of_two(base_alignment): + raise ValueError("base_alignment must be a positive power of two.") + if not is_power_of_two(total_alignment): + raise ValueError("total_alignment must be a positive power of two.") + self.base_alignment = base_alignment + self.total_alignment = total_alignment + self._mbarriers: List[SmemRegion] = [] + self._declarations: List[_TopLevelDeclaration] = [] + self._region_by_name: Dict[str, SmemRegion] = {} + self._overlay_names: set[str] = set() + self._offset: Dict[str, int] = {} + self._total_bytes = 0 + self._finalized = False + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "SmemWorkspace": + return self + + @property + def finalized(self) -> bool: + return self._finalized + + @property + def total_bytes(self) -> int: + self._require_finalized() + return self._total_bytes + + def register_mbarrier( + self, + name: str, + count: int, + *, + byte_alignment: Optional[int] = None, + ) -> SmemRegion: + if count <= 0: + raise ValueError(f"Mbarrier region {name!r} needs a positive count.") + region = self._make_region( + name=name, + kind="mbarrier", + dtype=cutlass.Int64, + shape=(count,), + stride=(1,), + swizzle=None, + byte_alignment=byte_alignment, + ) + self._claim_region_name(region) + self._mbarriers.append(region) + return region + + def register_tensor( + self, + name: str, + dtype: Type[cutlass.Numeric], + shape: Tuple, + *, + stride: Optional[Tuple] = None, + swizzle: Optional[SwizzleSpec] = None, + byte_alignment: Optional[int] = None, + ) -> SmemRegion: + region = self._make_region( + name=name, + kind="tensor", + dtype=dtype, + shape=shape, + stride=stride, + swizzle=swizzle, + byte_alignment=byte_alignment, + ) + self._claim_region_name(region) + self._declarations.append(region) + return region + + def create_overlay(self, name: str) -> SmemOverlay: + if self._finalized: + raise RuntimeError("Cannot create an SMEM overlay after finalize().") + if not name: + raise ValueError("An SMEM overlay needs a non-empty name.") + if name in self._overlay_names or name in self._region_by_name: + raise ValueError(f"Duplicate SMEM overlay {name!r}.") + overlay = SmemOverlay(self, name) + self._overlay_names.add(name) + self._declarations.append(overlay) + return overlay + + def _make_region( + self, + *, + name: str, + kind: SmemRegionKind, + dtype: Type[cutlass.Numeric], + shape: Tuple, + stride: Optional[Tuple], + swizzle: Optional[SwizzleSpec], + byte_alignment: Optional[int], + ) -> SmemRegion: + if self._finalized: + raise RuntimeError("Cannot register an SMEM region after finalize().") + if not name: + raise ValueError("An SMEM region needs a non-empty name.") + validate_static_integer_tuple(shape, field_name=f"{name}.shape") + if stride is None: + stride = row_major_stride(shape) + validate_static_integer_tuple(stride, field_name=f"{name}.stride") + if len(shape) != len(stride): + raise ValueError( + f"SMEM region {name!r} shape and stride ranks differ." + ) + if swizzle is not None: + if len(swizzle) != 3 or not all( + isinstance(parameter, int) for parameter in swizzle + ): + raise TypeError( + f"SMEM region {name!r} swizzle must be three Python ints." + ) + if swizzle[0] < 0 or swizzle[1] < 0: + raise ValueError( + f"SMEM region {name!r} swizzle bits/base must be non-negative." + ) + + natural_alignment = max(1, (int(dtype.width) + 7) // 8) + explicit_alignment = ( + natural_alignment if byte_alignment is None else byte_alignment + ) + if not is_power_of_two(explicit_alignment): + raise ValueError( + f"SMEM region {name!r} alignment must be a positive power of two." + ) + swizzle_alignment = _swizzle_alignment(swizzle) + effective_alignment = max( + explicit_alignment, + natural_alignment, + 1 if swizzle_alignment is None else swizzle_alignment, + ) + if effective_alignment > self.base_alignment: + raise ValueError( + f"SMEM region {name!r} needs {effective_alignment}B alignment, " + f"but the workspace base only promises {self.base_alignment}B." + ) + return SmemRegion( + name=name, + kind=kind, + dtype=dtype, + shape=shape, + stride=stride, + swizzle=swizzle, + byte_alignment=effective_alignment, + ) + + def _claim_region_name(self, region: SmemRegion) -> None: + if ( + region.name in self._region_by_name + or region.name in self._overlay_names + ): + raise ValueError(f"Duplicate SMEM region {region.name!r}.") + self._region_by_name[region.name] = region + + def estimate_total_bytes(self) -> int: + """Exact size of what is registered so far, priced by running the real placement. + + ``_build_placement`` is side-effect free -- ``finalize`` is what stores its result -- so the layout can be + measured without consuming the workspace. Summing region sizes instead would miss the alignment padding + that only exists once regions are placed, and a budget derived from that undercount overspends the + workspace: ``finalize`` then rejects a plan the host arithmetic had already accepted. + """ + return self._build_placement()[1] + + def finalize(self, *, max_bytes: Optional[int] = None) -> None: + if self._finalized: + raise RuntimeError("SmemWorkspace.finalize() may only be called once.") + offsets, total_bytes = self._build_placement() + if max_bytes is not None and total_bytes > max_bytes: + raise ValueError( + f"SMEM plan needs {total_bytes} bytes, exceeding {max_bytes} bytes." + ) + self._offset = offsets + self._total_bytes = total_bytes + self._finalized = True + + def _build_placement(self) -> Tuple[Dict[str, int], int]: + offsets: Dict[str, int] = {} + cursor = 0 + for mbarrier in self._mbarriers: + cursor = round_up(cursor, mbarrier.byte_alignment) + offsets[mbarrier.name] = cursor + cursor += mbarrier.nbytes + for declaration in self._declarations: + if isinstance(declaration, SmemRegion): + cursor = round_up(cursor, declaration.byte_alignment) + offsets[declaration.name] = cursor + cursor += declaration.nbytes + continue + relative_offsets, overlay_alignment, overlay_bytes = ( + self._layout_overlay(declaration) + ) + cursor = round_up(cursor, overlay_alignment) + for region_name, relative_offset in relative_offsets.items(): + offsets[region_name] = cursor + relative_offset + cursor += overlay_bytes + return offsets, int(round_up(cursor, self.total_alignment)) + + def _layout_overlay( + self, + overlay: SmemOverlay, + ) -> Tuple[Dict[str, int], int, int]: + if not overlay.lifetimes: + raise ValueError(f"SMEM overlay {overlay.name!r} has no lifetimes.") + relative_offsets: Dict[str, int] = {} + overlay_alignment = 1 + overlay_bytes = 0 + for lifetime in overlay.lifetimes: + if not lifetime.regions: + raise ValueError( + f"SMEM lifetime {overlay.name}.{lifetime.name} has no regions." + ) + lifetime_cursor = 0 + for region in lifetime.regions: + overlay_alignment = max( + overlay_alignment, + region.byte_alignment, + ) + lifetime_cursor = round_up( + lifetime_cursor, + region.byte_alignment, + ) + relative_offsets[region.name] = lifetime_cursor + lifetime_cursor += region.nbytes + overlay_bytes = max(overlay_bytes, lifetime_cursor) + return relative_offsets, overlay_alignment, overlay_bytes + + def regions(self) -> Tuple[SmemRegion, ...]: + return tuple(self._region_by_name.values()) + + def region(self, name: str) -> SmemRegion: + return self._region_by_name[name] + + def offset(self, name: str) -> int: + self._require_finalized() + return self._offset[name] + + def nbytes(self, name: str) -> int: + return self._region_by_name[name].nbytes + + def byte_alignment(self, name: str) -> int: + return self._region_by_name[name].byte_alignment + + def storage_class(self) -> type: + self._require_finalized() + storage_bytes = max(self._total_bytes, 1) + base_alignment = self.base_alignment + + @cute.struct + class SmemStorage: + buffer: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int8, storage_bytes], + base_alignment, + ] + + return SmemStorage + + @cute.jit + def ptr(self, name: str, smem_base: cute.Pointer) -> cute.Pointer: + region = self._region_by_name[name] + swizzle = ( + None + if region.swizzle is None + else cute.make_swizzle(*region.swizzle) + ) + return cute.make_ptr( + region.dtype, + smem_base.toint() + self._offset[name], + smem_base.memspace, + assumed_align=region.byte_alignment, + swizzle_=swizzle, + ) + + @cute.jit + def tensor(self, name: str, smem_base: cute.Pointer) -> cute.Tensor: + region = self._region_by_name[name] + layout = cute.make_layout(region.shape, stride=region.stride) + return cute.make_tensor(self.ptr(name, smem_base), layout) + + def _require_finalized(self) -> None: + if not self._finalized: + raise RuntimeError("SmemWorkspace must be finalized first.") + + +__all__ = [ + "SmemLifetime", + "SmemOverlay", + "SmemRegion", + "SmemRegionKind", + "SmemWorkspace", + "SwizzleSpec", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/software_sync.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/software_sync.py new file mode 100644 index 000000000..9b1ef4010 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/software_sync.py @@ -0,0 +1,276 @@ +"""Software grid and NVLink synchronization for persistent kernels.""" + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import Int32, Int64 + +from .device_workspace import DeviceWorkspace +from .ptx_helpers import red_add_relaxed_sys_s32 + + +class SoftwareGridSync: + """Reusable device-local barrier over a phase-flipping GMEM counter.""" + + finish_sum_tag = 0x80000000 + grid_counter_region = "software_grid_sync.counter" + + def __init__(self, *, barrier_id: int) -> None: + self.barrier_id = barrier_id + self._grid_counter = None + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "SoftwareGridSync": + if values: + raise ValueError(f"SoftwareGridSync expected no MLIR values, got {len(values)}.") + return self + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + workspace.register( + self.grid_counter_region, + cutlass.Int32, + (1,), + buffer_space="local", + byte_alignment=16, + reset="zero_on_first_allocate", + ) + + @cute.jit + def assign_device_members(self, workspace: DeviceWorkspace) -> None: + self._grid_counter = workspace.ptr(self.grid_counter_region) + + def remove_device_members(self) -> None: + self._grid_counter = None + + @cute.jit + def _cta_rendezvous(self, participating_threads: int) -> None: + cute.arch.barrier(barrier_id=self.barrier_id, number_of_threads=participating_threads) + + @cute.jit + def sync( + self, + participating_threads: int, + actual_cta_count: Int32, + linear_cta_idx: Int32, + thread_idx_in_group: Int32, + ) -> None: + self._cta_rendezvous(participating_threads) + leader_delta = Int32(-self.finish_sum_tag) - (actual_cta_count - Int32(1)) + _inline_grid_sync( + self._grid_counter, + linear_cta_idx, + leader_delta, + Int32(1), + thread_idx_in_group, + ) + self._cta_rendezvous(participating_threads) + + +class NvlinkBarrier(SoftwareGridSync): + """Sense-reversing all-rank barrier layered over software grid sync.""" + + period = 4 + grid_counter_region = "nvlink.token_comm.grid_sync_counter" + phase_counter_region = "nvlink.token_comm.nvlink_phase_counter" + signal_region = "nvlink.token_comm.nvlink_signal" + + def __init__( + self, + *, + world_size: int, + barrier_id: int, + ) -> None: + super().__init__(barrier_id=barrier_id) + self.world_size = world_size + self._phase_counter = None + self._signal = None + self._peer_rank_ptr_mapper = None + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + super().register_device_workspace(workspace) + workspace.register( + self.phase_counter_region, + cutlass.Int32, + (1,), + buffer_space="local", + byte_alignment=16, + reset="zero_on_first_allocate", + ) + workspace.register( + self.signal_region, + cutlass.Int32, + (2,), + buffer_space="shared", + byte_alignment=16, + reset="zero_on_first_allocate", + ) + + @cute.jit + def assign_device_members( + self, + workspace: DeviceWorkspace, + peer_rank_ptr_mapper, + ) -> None: + super().assign_device_members(workspace) + self._phase_counter = workspace.ptr(self.phase_counter_region) + self._signal = workspace.ptr(self.signal_region) + self._peer_rank_ptr_mapper = peer_rank_ptr_mapper + + def remove_device_members(self) -> None: + super().remove_device_members() + self._phase_counter = None + self._signal = None + self._peer_rank_ptr_mapper = None + + @cute.jit + def arrive_and_wait( + self, + participating_threads: int, + actual_cta_count: Int32, + linear_cta_idx: Int32, + thread_idx_in_group: Int32, + *, + prologue_grid_sync: bool, + epilogue_grid_sync: bool, + ) -> None: + if cutlass.const_expr(prologue_grid_sync): + self.sync( + participating_threads, + actual_cta_count, + linear_cta_idx, + thread_idx_in_group, + ) + + if linear_cta_idx == Int32(0): + status = cute.arch.load( + self._phase_counter, + Int32, + sem="relaxed", + scope="gpu", + ) & Int32(3) + signal_phase = status & Int32(1) + signal_direction = status >> Int32(1) + signal_delta = Int32(1) + signal_target = Int32(self.world_size) + if signal_direction != Int32(0): + signal_delta = Int32(-1) + signal_target = Int32(0) + + self._cta_rendezvous(participating_threads) + if thread_idx_in_group == Int32(0): + cute.arch.fence_acq_rel_sys() + self._cta_rendezvous(participating_threads) + + rank_round_count = ( + self.world_size + participating_threads - 1 + ) // participating_threads + signal_base_address = self._signal.toint() + signal_byte_offset = Int64(signal_phase * Int32(4)) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = ( + Int32(rank_round * participating_threads) + + thread_idx_in_group + ) + if destination_rank < Int32(self.world_size): + destination_address = self._peer_rank_ptr_mapper.map( + signal_base_address, + destination_rank, + signal_byte_offset, + ) + red_add_relaxed_sys_s32( + destination_address, + signal_delta, + ) + + self._cta_rendezvous(participating_threads) + if thread_idx_in_group == Int32(0): + cute.arch.atomic_add( + self._phase_counter, + Int32(1), + sem="relaxed", + scope="gpu", + ) + local_signal = self._signal + signal_phase + while cute.arch.load( + local_signal, + Int32, + sem="acquire", + scope="sys", + ) != signal_target: + pass + + if cutlass.const_expr(epilogue_grid_sync): + self.sync( + participating_threads, + actual_cta_count, + linear_cta_idx, + thread_idx_in_group, + ) + + @cute.jit + def finalize( + self, + completed_calls: int, + participating_threads: int, + actual_cta_count: Int32, + linear_cta_idx: Int32, + thread_idx_in_group: Int32, + ) -> None: + padding_calls = (-completed_calls) % self.period + for _ in cutlass.range_constexpr(padding_calls): + self.arrive_and_wait( + participating_threads, + actual_cta_count, + linear_cta_idx, + thread_idx_in_group, + prologue_grid_sync=True, + epilogue_grid_sync=True, + ) + + +@cute.jit +def _inline_grid_sync( + counter, + linear_cta_idx, + leader_delta, + other_delta, + thread_idx_in_group, +) -> None: + llvm.inline_asm( + None, + [ + counter.toint().ir_value(), + Int32(linear_cta_idx).ir_value(), + leader_delta.ir_value(), + other_delta.ir_value(), + Int32(thread_idx_in_group).ir_value(), + ], + ( + "{\n\t" + ".reg .b32 %delta; .reg .b32 %old; .reg .b32 %current;\n\t" + ".reg .pred %not_leader; .reg .pred %is_cta0; " + ".reg .pred %waiting;\n\t" + "setp.ne.u32 %not_leader, $4, 0;\n\t" + "@%not_leader bra DONE;\n\t" + "setp.eq.u32 %is_cta0, $1, 0;\n\t" + "selp.b32 %delta, $2, $3, %is_cta0;\n\t" + "atom.release.gpu.global.add.u32 %old, [$0], %delta;\n\t" + "SPIN:\n\t" + "ld.relaxed.gpu.global.b32 %current, [$0];\n\t" + "xor.b32 %current, %current, %old;\n\t" + "and.b32 %current, %current, 0x80000000;\n\t" + "setp.eq.u32 %waiting, %current, 0;\n\t" + "@%waiting bra SPIN;\n\t" + "fence.acq_rel.gpu;\n\t" + "DONE:\n\t" + "}" + ), + "l,r,r,r,r", + has_side_effects=True, + asm_dialect=0, + ) + + +__all__ = ["NvlinkBarrier", "SoftwareGridSync"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py new file mode 100644 index 000000000..b14dae2cd --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py @@ -0,0 +1,100 @@ +"""Small integer and layout helpers shared by workspace implementations.""" + +from typing import Iterable, List, Tuple, Union + +import cutlass + + +IntegerType = Union[int, cutlass.Int32, cutlass.Int64, cutlass.Uint32, cutlass.Uint64] + + +def round_up(value: IntegerType, alignment: IntegerType) -> IntegerType: + return ((value + alignment - 1) // alignment) * alignment + + +def ceil_div(value: IntegerType, divisor: IntegerType) -> IntegerType: + return (value + divisor - 1) // divisor + + +def is_power_of_two(value: int) -> bool: + return value > 0 and (value & (value - 1)) == 0 + + +def is_nested_shape(shape: Tuple) -> bool: + return any(isinstance(dimension, tuple) for dimension in shape) + + +def validate_static_integer_tuple(value: Tuple, *, field_name: str) -> None: + for element in value: + if isinstance(element, tuple): + validate_static_integer_tuple(element, field_name=field_name) + elif not isinstance(element, int): + raise TypeError(f"{field_name} must contain Python ints, got {type(element)}.") + + +def flatten_shape_stride(shape: Tuple, stride: Tuple) -> List[Tuple[IntegerType, IntegerType]]: + pairs: List[Tuple[IntegerType, IntegerType]] = [] + for size, step in zip(shape, stride): + if isinstance(size, tuple): + pairs.extend(flatten_shape_stride(size, step)) + else: + pairs.append((size, step)) + return pairs + + +def strides_equal_ignoring_singletons(shape, lhs_stride, rhs_stride) -> bool: + """Compare strides while ignoring leaves whose logical extent is one.""" + if isinstance(shape, tuple): + if not isinstance(lhs_stride, tuple) or not isinstance(rhs_stride, tuple): + return False + if len(shape) != len(lhs_stride) or len(shape) != len(rhs_stride): + return False + return all( + strides_equal_ignoring_singletons(child_shape, lhs_step, rhs_step) + for child_shape, lhs_step, rhs_step in zip(shape, lhs_stride, rhs_stride) + ) + return shape == 1 or lhs_stride == rhs_stride + + +def ordered_stride(shape: Tuple[int, ...], mem_order: Tuple[int, ...]) -> Tuple[Tuple[int, ...], int]: + stride = [0] * len(shape) + cosize = 1 + for mode in sorted(range(len(shape)), key=lambda index: mem_order[index]): + stride[mode] = cosize + cosize *= shape[mode] + return tuple(stride), cosize + + +def row_major_stride(shape: Tuple[int, ...]) -> Tuple[int, ...]: + if is_nested_shape(shape): + raise ValueError("A nested shape needs an explicit stride.") + stride, _ = ordered_stride(shape, tuple(reversed(range(len(shape))))) + return stride + + +def cosize_from_shape_stride_tuples(shape: Tuple, stride: Tuple) -> IntegerType: + leaf_pairs = flatten_shape_stride(shape, stride) if shape else [] + return 1 + sum((size - 1) * step for size, step in leaf_pairs) + + +def product(values: Iterable[IntegerType]) -> IntegerType: + result: IntegerType = 1 + for value in values: + result = result * value + return result + + +__all__ = [ + "IntegerType", + "ceil_div", + "cosize_from_shape_stride_tuples", + "flatten_shape_stride", + "is_nested_shape", + "is_power_of_two", + "ordered_stride", + "product", + "row_major_stride", + "round_up", + "strides_equal_ignoring_singletons", + "validate_static_integer_tuple", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/__init__.py new file mode 100644 index 000000000..337547020 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/__init__.py @@ -0,0 +1,26 @@ +"""Kernel source modules independent of repository-only runners.""" + +from .schedulers import ( + BlackwellFusedFc12Scheduler, + BlockPhase, + Fc12WorkTileState, + NonSwapAbFc12WorkTileInfo, + SchedulerBase, + SchedulerConsumer, + SchedulerWorkTileBase, + SwapAbFc12WorkTileInfo, + WorkIdAcquisitionMode, +) + + +__all__ = [ + "BlackwellFusedFc12Scheduler", + "BlockPhase", + "Fc12WorkTileState", + "NonSwapAbFc12WorkTileInfo", + "SchedulerBase", + "SchedulerConsumer", + "SchedulerWorkTileBase", + "SwapAbFc12WorkTileInfo", + "WorkIdAcquisitionMode", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py new file mode 100644 index 000000000..cedc19ea4 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py @@ -0,0 +1,3072 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Autonomous epilogue for the fused FC1+FC2 swap-AB MegaMoE kernel. + +Per-thread RMEM tensors flow between the transpose / SwiGLU / quantize / fc2 +store steps as bare ``cute.Tensor`` fragments; their thread distribution is a +fixed physical property of the surrounding atom sequence and is documented in +local comments. FC2 store mappings are finite ``FunctionMapping`` objects +evaluated at runtime to drive metadata lookup and destination pointer math. +""" + +import dataclasses +import math +from typing import Any, Callable, ClassVar, List, Literal, Optional, Tuple, Type, Union + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Int64, T +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import AddressSpace +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils + +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, vector + +from .....communication.nvlink_domain.symmetric_buffer import SymmetricBufferDevice +from .....communication.token_protocol import TokenSrcMetadata +from .....quant_def import CombineFormat, QuantKind +from .....api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from .....helpers.constants import Nvfp4E2M1RcpLimit, Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Fp32Max +from .....helpers.cute_py_helpers import tcgen05_block_scaled_acc_dtype +from .....helpers.dsl_helpers import mark_alignment +from .....helpers.flag_batch import GpuReleaseFlagBatchTracker +from .....helpers.ptx_helpers import ( + cp_async_bulk_s2g, + cp_reduce_async_bulk_add_bf16_s2g, + cvt_f32_to_fp8_to_f32, + movmatrix_b16, + red_add_relaxed_sys_v2_bf16x2, +) +from .....helpers.smem_workspace import SmemRegion, SmemWorkspace +from ....function_mapping import CoordinateSpace, FunctionMapping +from ....schedulers import BlockPhase, SchedulerConsumer, SwapAbFc12WorkTileInfo +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension + + +@dataclasses.dataclass(frozen=True) +class QuantImpl: + """Register-level block quantizer shared by the fc1 / fc2 epilogues. + + Returns ``(data_regs, sf_regs)`` ONLY: the caller pre-multiplies the topk + weight / global scale into ``prequant_reg`` beforehand and owns the data / + sf plane stores afterwards. ``sf_vec_direction`` selects the per-block amax + reduction: + + * ``regs_in_thread`` -- a block is ``sf_vec`` contiguous regs of + one thread; amax is thread-local (packed bf16x2 abs-max for combine, + fp32 ``fmax`` for orthodox). + * ``threads_with_the_same_reg`` -- a block is one reg across ``sf_vec`` warp + lanes; amax is a warp CREDUX (full warp for vec=32, lane-predicated + halves for vec=16) and is fp32-only, so bf16 is upconverted first. + * ``regs_in_pair_threads`` -- a block is ``sf_vec / 2`` regs in each of + two paired warps; amax is thread-local over the half, then exchanged with + the partner warp through SMEM. Orthodox mx only: fc1's TMEM transpose + hands one warp exactly 16 intermediate values per token, so a 32-wide + scale block necessarily straddles warps ``w`` and ``w ^ 1``. + ``prequant_reg`` must be 1D with size divisible by the per-thread block share + (``sf_vec``, or ``sf_vec / 2`` for the paired direction). Combine inputs are + bf16 (fc2's bf16 reorder regs); orthodox input is fp32 (swiglu). + """ + + quant_kind: Union[QuantKind, CombineFormat] + sf_vec_direction: Literal["regs_in_thread", "threads_with_the_same_reg", "regs_in_pair_threads"] + lane_idx: Optional[Any] = None + warp_idx: Optional[Any] = None + pair_exchange_barrier: Optional[Any] = None + + _directions: ClassVar[Tuple[str, ...]] = ("regs_in_thread", "threads_with_the_same_reg", "regs_in_pair_threads") + + # -- config / validation -------------------------------------------------- + + def __post_init__(self): + if isinstance(self.quant_kind, CombineFormat): + if not self.quant_kind.is_quantized: + raise ValueError(f"QuantImpl combine path needs a quantized CombineFormat, got {self.quant_kind}.") + elif not isinstance(self.quant_kind, QuantKind): + raise ValueError(f"quant_kind must be a QuantKind or a CombineFormat, got {self.quant_kind!r}.") + if self.sf_vec_direction not in self._directions: + raise ValueError(f"sf_vec_direction must be one of {self._directions}, got {self.sf_vec_direction!r}.") + if self.sf_vec_direction == "threads_with_the_same_reg" and self.lane_idx is None: + raise ValueError("across-lane quant needs lane_idx for the CREDUX half-warp predicate.") + if self.sf_vec_direction == "regs_in_pair_threads": + if isinstance(self.quant_kind, CombineFormat): + raise ValueError("The paired direction is orthodox-only; combine blocks never straddle warps.") + if self.sf_vec_size % 2 != 0: + raise ValueError(f"The paired direction needs an even sf_vec, got {self.sf_vec_size}.") + if self.lane_idx is None or self.warp_idx is None or self.pair_exchange_barrier is None: + raise ValueError("The paired direction needs lane_idx, warp_idx and pair_exchange_barrier.") + elif not isinstance(self.quant_kind, CombineFormat) and self.sf_vec_size != 16: + # A wider orthodox block cannot be reduced inside one thread: see the paired direction. + raise ValueError(f"Orthodox sf_vec {self.sf_vec_size} needs the paired direction.") + + @property + def data_dtype(self): + if isinstance(self.quant_kind, CombineFormat): + return self.quant_kind.act_dtype + # Orthodox output feeds fc2 as its activation, so it is the kind's activation type. + return self.quant_kind.activation_dtype + + @property + def scale_dtype(self): + if isinstance(self.quant_kind, CombineFormat): + return self.quant_kind.scale_dtype + return self.quant_kind.sf_dtype + + @property + def sf_vec_size(self) -> int: + if isinstance(self.quant_kind, CombineFormat): + return self.quant_kind.scale_block + return self.quant_kind.sf_vec_size + + @property + def _is_combine(self) -> bool: + return isinstance(self.quant_kind, CombineFormat) + + @property + def _data_rcp_limit(self) -> float: + # 1 / max representable magnitude of the data element type. + dt = self.data_dtype + if dt is cutlass.Float4E2M1FN: + return Nvfp4E2M1RcpLimit # 1/6 + if dt is cutlass.Float8E4M3FN: + return Fp8E4M3RcpLimit # 1/448 + return Fp8E5M2RcpLimit # 1/57344 + + @property + def _block_share_per_thread(self) -> int: + """How many of a block's elements this thread holds.""" + return self.sf_vec_size // 2 if self.sf_vec_direction == "regs_in_pair_threads" else self.sf_vec_size + + # -- dispatch ------------------------------------------------------------- + + @cute.jit + def __call__(self, prequant_reg: cute.Tensor, *, norm_const=None, smem_intermediate=None): + if cutlass.const_expr(cute.size(prequant_reg) % self._block_share_per_thread != 0): + raise ValueError("prequant_reg size must be divisible by this thread's share of a block.") + # Combine quantizes fc2's bf16 reorder regs; orthodox the fp32 swiglu. + expected_in = cutlass.BFloat16 if self._is_combine else cutlass.Float32 + if cutlass.const_expr(prequant_reg.element_type is not expected_in): + raise TypeError( + f"QuantImpl({self.quant_kind}) expects {expected_in} prequant input, got {prequant_reg.element_type}." + ) + needs_smem = cutlass.const_expr(self.sf_vec_direction == "regs_in_pair_threads") + if cutlass.const_expr(needs_smem != (smem_intermediate is not None)): + raise ValueError("smem_intermediate must be supplied for the paired direction and only for it.") + if cutlass.const_expr(not self._is_combine): + if cutlass.const_expr(self.sf_vec_direction == "regs_in_pair_threads"): + return self.mx_quant_regs_in_pair_threads_impl(prequant_reg, smem_intermediate) + return self.nvfp4_quant_impl(prequant_reg, norm_const=norm_const) + if cutlass.const_expr(self.data_dtype is cutlass.Float4E2M1FN): + if cutlass.const_expr(self.sf_vec_direction == "regs_in_thread"): + return self.nvfp4_combine_quant_regs_in_thread_impl(prequant_reg) + return self.nvfp4_combine_quant_threads_with_the_same_reg_impl(prequant_reg) + if cutlass.const_expr(self.sf_vec_direction == "regs_in_thread"): + return self.mxfp8_combine_quant_regs_in_thread_impl(prequant_reg) + return self.mxfp8_combine_quant_threads_with_the_same_reg_impl(prequant_reg) + + # -- impls ---------------------------------------------------------------- + + # regs_in_thread only; fc1 promises the vec direction via its TMEM transpose. + @cute.jit + def nvfp4_quant_impl( + self, prequant_reg: cute.Tensor, *, norm_const: Optional[cutlass.Float32] = None + ) -> Tuple[cute.Tensor, cute.Tensor]: + # fp32 in -> e2m1 data + e4m3 sfc. Mirrors the prior nvfp4_quant scale + # math (sfc -> capped/masked acc_scale); topk pre-mult + sf store are the + # caller's job now. + vec = self.sf_vec_size + n_blocks = cute.size(prequant_reg) // vec + data = cute.make_rmem_tensor((cute.size(prequant_reg),), cutlass.Float4E2M1FN) + sf = cute.make_rmem_tensor((n_blocks,), cutlass.Float8E4M3FN) + in_blocks = cute.zipped_divide(prequant_reg, (vec,)) # ((vec,), (n_blocks,)) + data_blocks = [] + sf_values = [] + rcp_limit = cutlass.Float32(self._data_rcp_limit) + for vec_block_idx in cutlass.range_constexpr(n_blocks): + block = in_blocks[None, vec_block_idx] + amax = self._amax_thread_fp32(block) + if cutlass.const_expr(norm_const is not None): + sfc_fp32 = amax * rcp_limit * norm_const + else: + sfc_fp32 = amax * rcp_limit + sfc_e4m3 = sfc_fp32.to(cutlass.Float8E4M3FN) + sfc_rt = cutlass.Float32(sfc_e4m3) + if cutlass.const_expr(norm_const is not None): + acc_scale = norm_const * cute.arch.rcp_approx(sfc_rt) + else: + acc_scale = cute.arch.rcp_approx(sfc_rt) + acc_scale = cute.arch.fmin(acc_scale, Fp32Max) + mask = cute.arch.fmin(sfc_rt * cutlass.Float32(1e30), cutlass.Float32(1.0)) + acc_scale = acc_scale * mask + sf_values.append(sfc_e4m3) + data_blocks.append(self._scale_to_data_ssa(block, acc_scale)) + self._store_packed_blocks(data, data_blocks) + sf.store(self._values_to_ssa(sf_values, cutlass.Float8E4M3FN)) + return data, sf + + @cute.jit + def mx_quant_regs_in_pair_threads_impl( + self, + prequant_reg: cute.Tensor, # (token_2_intermeidate_x) + smem_intermediate: cute.Tensor, # (token_blocks, epi_threads) + ) -> Tuple[cute.Tensor, cute.Tensor]: + half = self._block_share_per_thread + n_blocks = cute.size(prequant_reg) // half + data = cute.make_rmem_tensor((cute.size(prequant_reg),), self.data_dtype) + sf = cute.make_rmem_tensor((n_blocks,), self.scale_dtype) + in_blocks = cute.zipped_divide(prequant_reg, (half,)) # ((half,), (n_blocks,)) + + if cutlass.const_expr(cute.size(smem_intermediate, mode=[0]) != n_blocks): + raise ValueError( + f"The paired amax exchange needs {n_blocks} rows, got {cute.size(smem_intermediate, mode=[0])}." + ) + if cutlass.const_expr(smem_intermediate.stride[0] != 1): + # Thread-major would still be correct but would silently split the exchange into + # per-block scalar accesses. + raise ValueError("The paired amax exchange must be block-major to stay one vector access.") + + exchange_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=n_blocks * cutlass.Float32.width + ) + # Which warps pair up is epilogue knowledge: warp w owns one contiguous run of + # intermediate outputs, so a block that is twice that run wide joins adjacent warps. + partner_warp = self.warp_idx ^ cutlass.Int32(1) + own_thread = self.warp_idx * cutlass.Int32(32) + self.lane_idx + partner_thread = partner_warp * cutlass.Int32(32) + self.lane_idx + + half_amax = cute.make_rmem_tensor((n_blocks,), cutlass.Float32) + for vec_block_idx in cutlass.range_constexpr(n_blocks): + half_amax[vec_block_idx] = self._amax_thread_fp32(in_blocks[None, vec_block_idx]) + cute.copy(exchange_atom, cute.coalesce(half_amax), cute.coalesce(smem_intermediate[None, own_thread])) + + self.pair_exchange_barrier.arrive_and_wait() + + partner_amax = cute.make_rmem_tensor((n_blocks,), cutlass.Float32) + cute.copy(exchange_atom, cute.coalesce(smem_intermediate[None, partner_thread]), cute.coalesce(partner_amax)) + + data_blocks = [] + sf_values = [] + for vec_block_idx in cutlass.range_constexpr(n_blocks): + block_amax = cute.arch.fmax(half_amax[vec_block_idx], partner_amax[vec_block_idx]) + scale_e8m0, scale_f32 = self._e8m0(block_amax) + sf_values.append(scale_e8m0) + data_blocks.append(self._scale_to_data_ssa(in_blocks[None, vec_block_idx], self._enc_mxfp8(scale_f32))) + self._store_packed_blocks(data, data_blocks) + sf.store(self._values_to_ssa(sf_values, self.scale_dtype)) + return data, sf + + @cute.jit + def nvfp4_combine_quant_regs_in_thread_impl(self, prequant_reg: cute.Tensor) -> Tuple[cute.Tensor, cute.Tensor]: + # bf16 in -> e2m1 data + per-16 bf16 amax. amax found on bf16 (packed). + vec = self.sf_vec_size + n_blocks = cute.size(prequant_reg) // vec + data = cute.make_rmem_tensor((cute.size(prequant_reg),), cutlass.Float4E2M1FN) + sf = cute.make_rmem_tensor((n_blocks,), cutlass.BFloat16) + in_blocks = cute.zipped_divide(prequant_reg, (vec,)) # ((vec,), (n_blocks,)) + data_blocks = [] + sf_values = [] + for vec_block_idx in cutlass.range_constexpr(n_blocks): + block = in_blocks[None, vec_block_idx] + amax = self._amax_thread_bf16(block) + sf_values.append(amax) + decode_scale = cutlass.Float32(amax) * cutlass.Float32(self._data_rcp_limit) + data_blocks.append(self._scale_to_data_ssa(block, self._enc_nvfp4(decode_scale))) + self._store_packed_blocks(data, data_blocks) + sf.store(self._values_to_ssa(sf_values, cutlass.BFloat16)) + return data, sf + + # Mapping: (lane_idx, selected_sf_idx) -> (token_64, hidden_32) + # token_idx = lane_idx % 16 + selected_sf_idx * 16 + # hidden_idx = lane_idx // 16 * 16 + @cute.jit + def nvfp4_combine_quant_threads_with_the_same_reg_impl( + self, prequant_reg: cute.Tensor + ) -> Tuple[cute.Tensor, cute.Tensor]: + # UBLK has lane == hidden, so the warp's 32 lanes are 32 consecutive + # hidden. A scale block = sf_vec hidden, so the lanes split along hidden + # into 32 // sf_vec blocks of sf_vec lanes each (warp = blocks_per_warp * + # lanes_per_block, the EP x TP split). Only the sf_vec lanes inside a + # block share its CREDUX scale, so they pool the subtile tokens: sf_vec + # == 32 pools the whole warp, sf_vec < 32 pools fewer (more per lane). + lanes_per_block = self.sf_vec_size + n_tokens = cute.size(prequant_reg) + lane_in_block = self.lane_idx % cutlass.Int32(lanes_per_block) + data = cute.make_rmem_tensor((n_tokens,), cutlass.Float4E2M1FN) + selected_sf = cute.make_rmem_tensor((n_tokens // lanes_per_block,), cutlass.BFloat16) + scaled_vec = cute.full((n_tokens,), cutlass.Float32(0.0), cutlass.Float32) + for token_idx in cutlass.range_constexpr(n_tokens): + value = cutlass.Float32(prequant_reg[token_idx]) + amax_bf16 = self._amax_lane(value).to(cutlass.BFloat16) + slot = token_idx // lanes_per_block + if (token_idx % lanes_per_block) == lane_in_block: + selected_sf[slot] = amax_bf16 + else: + selected_sf[slot] = selected_sf[slot] + decode_scale = cutlass.Float32(amax_bf16) * cutlass.Float32(self._data_rcp_limit) + scaled_value = value * self._enc_nvfp4(decode_scale) + scaled_vec = cute.TensorSSA( + vector.insert(scaled_value.ir_value(), scaled_vec.ir_value(), [], [token_idx]), + (n_tokens,), + cutlass.Float32, + ) + self._store_packed_data(data, scaled_vec.to(cutlass.Float4E2M1FN)) + return data, selected_sf + + @cute.jit + def mxfp8_combine_quant_regs_in_thread_impl(self, prequant_reg: cute.Tensor) -> Tuple[cute.Tensor, cute.Tensor]: + # bf16 in -> e4m3/e5m2 data + per-32 e8m0. amax found on bf16 (packed). + vec = self.sf_vec_size + n_blocks = cute.size(prequant_reg) // vec + data = cute.make_rmem_tensor((cute.size(prequant_reg),), self.data_dtype) + sf = cute.make_rmem_tensor((n_blocks,), cutlass.Float8E8M0FNU) + in_blocks = cute.zipped_divide(prequant_reg, (vec,)) # ((vec,), (n_blocks,)) + data_blocks = [] + sf_values = [] + for vec_block_idx in cutlass.range_constexpr(n_blocks): + block = in_blocks[None, vec_block_idx] + # widen the native-bf16 amax to fp32 for the e8m0 round-up math. + scale_e8m0, scale_f32 = self._e8m0(cutlass.Float32(self._amax_thread_bf16(block))) + sf_values.append(scale_e8m0) + data_blocks.append(self._scale_to_data_ssa(block, self._enc_mxfp8(scale_f32))) + self._store_packed_blocks(data, data_blocks) + sf.store(self._values_to_ssa(sf_values, cutlass.Float8E8M0FNU)) + return data, sf + + # Mapping: (lane_idx, selected_sf_idx) -> (token_64, hidden_32) + # token_idx = lane_idx + selected_sf_idx * 32 + # hidden_idx = 0 + @cute.jit + def mxfp8_combine_quant_threads_with_the_same_reg_impl( + self, prequant_reg: cute.Tensor + ) -> Tuple[cute.Tensor, cute.Tensor]: + # UBLK has lane == hidden, so the warp's 32 lanes are 32 consecutive + # hidden. A scale block = sf_vec hidden, so the lanes split along hidden + # into 32 // sf_vec blocks of sf_vec lanes each (warp = blocks_per_warp * + # lanes_per_block, the EP x TP split). Only the sf_vec lanes inside a + # block share its CREDUX scale, so they pool the subtile tokens. mxfp8 + # sf_vec == 32 -> the whole warp is one block, all 32 lanes pool. + lanes_per_block = self.sf_vec_size + n_tokens = cute.size(prequant_reg) + lane_in_block = self.lane_idx % cutlass.Int32(lanes_per_block) + data = cute.make_rmem_tensor((n_tokens,), self.data_dtype) + selected_sf = cute.make_rmem_tensor((n_tokens // lanes_per_block,), cutlass.Float8E8M0FNU) + scaled_vec = cute.full((n_tokens,), cutlass.Float32(0.0), cutlass.Float32) + for token_idx in cutlass.range_constexpr(n_tokens): + value = cutlass.Float32(prequant_reg[token_idx]) + scale_e8m0, scale_f32 = self._e8m0(self._amax_lane(value)) + slot = token_idx // lanes_per_block + if (token_idx % lanes_per_block) == lane_in_block: + selected_sf[slot] = scale_e8m0 + else: + selected_sf[slot] = selected_sf[slot] + scaled_value = value * self._enc_mxfp8(scale_f32) + scaled_vec = cute.TensorSSA( + vector.insert(scaled_value.ir_value(), scaled_vec.ir_value(), [], [token_idx]), + (n_tokens,), + cutlass.Float32, + ) + self._store_packed_data(data, scaled_vec.to(self.data_dtype)) + return data, selected_sf + + # -- shared sub-steps ----------------------------------------------------- + + @cute.jit + def _scale_to_data_ssa(self, block: cute.Tensor, enc: cutlass.Float32) -> cute.TensorSSA: + block_f32 = block.load().to(cutlass.Float32) + enc_vec = cute.full_like(block_f32, enc, cutlass.Float32) + return (block_f32 * enc_vec).to(self.data_dtype) + + @cute.jit + def _concat_blocks_ssa(self, blocks, dtype: Type[cutlass.Numeric]) -> cute.TensorSSA: + values = [] + for block_idx in cutlass.range_constexpr(len(blocks)): + block = blocks[block_idx] + for elem_idx in cutlass.range_constexpr(cute.size(block.shape)): + values.append(block[elem_idx].ir_value()) + vec = vector.from_elements(T.vector(len(values), dtype.mlir_type), values) + return cute.TensorSSA(vec, (len(values),), dtype) + + @cute.jit + def _values_to_ssa(self, values, dtype: Type[cutlass.Numeric]) -> cute.TensorSSA: + vec = vector.from_elements( + T.vector(len(values), dtype.mlir_type), [values[i].ir_value() for i in range(len(values))] + ) + return cute.TensorSSA(vec, (len(values),), dtype) + + @cute.jit + def _concat_i32_blocks_ssa(self, blocks) -> cute.TensorSSA: + values = [] + for block_idx in cutlass.range_constexpr(len(blocks)): + packed_block = blocks[block_idx].bitcast(cutlass.Int32) + for elem_idx in cutlass.range_constexpr(cute.size(packed_block.shape)): + values.append(packed_block[elem_idx].ir_value()) + vec = vector.from_elements(T.vector(len(values), cutlass.Int32.mlir_type), values) + return cute.TensorSSA(vec, (len(values),), cutlass.Int32) + + @cute.jit + def _store_packed_blocks(self, data: cute.Tensor, blocks) -> None: + packed_data = cute.recast_tensor(data, cutlass.Int32) + packed_data.store(self._concat_i32_blocks_ssa(blocks)) + + @cute.jit + def _store_packed_data(self, data: cute.Tensor, data_ssa: cute.TensorSSA) -> None: + packed_data = cute.recast_tensor(data, cutlass.Int32) + packed_data.store(data_ssa.bitcast(cutlass.Int32)) + + @cute.jit + def _amax_thread_fp32(self, block: cute.Tensor) -> cutlass.Float32: + # max.xorsign.abs reduces |.| in one op per element; the result sign is + # the xor of the inputs (junk for an amax), so clear it at the end. + def max_abs(lhs: cutlass.Float32, rhs: cutlass.Float32) -> cutlass.Float32: + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [cutlass.Float32(lhs).ir_value(), cutlass.Float32(rhs).ir_value()], + "max.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + acc = block[0] + for elem_idx in cutlass.range_constexpr(1, cute.size(block)): + acc = max_abs(acc, block[elem_idx]) + mag_bits = cutlass.Int32(llvm.bitcast(T.i32(), cutlass.Float32(acc).ir_value())) & cutlass.Int32(0x7FFFFFFF) + return cutlass.Float32(llvm.bitcast(T.f32(), mag_bits.ir_value())) + + @cute.jit + def _amax_thread_bf16(self, block: cute.Tensor) -> cutlass.BFloat16: + # Packed bf16x2 abs-max: tree-reduce the pairs, then fold the survivor's + # two halves (high shifted into low). max.xorsign.abs leaves a junk sign, + # so the low bf16 is masked before being read back. The amax is natively + # bf16 -- exactly what the wire format stores. + def max_abs(lhs: cutlass.Int32, rhs: cutlass.Int32) -> cutlass.Int32: + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [cutlass.Int32(lhs).ir_value(), cutlass.Int32(rhs).ir_value()], + "max.xorsign.abs.bf16x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + pairs = cute.recast_tensor(block, cutlass.Int32) # (vec/2,) bf16x2 + acc = cutlass.Int32(pairs[0]) + for pair_idx in cutlass.range_constexpr(1, cute.size(pairs)): + acc = max_abs(acc, pairs[pair_idx]) + acc = max_abs(acc, acc >> cutlass.Int32(16)) + amax_bits = cute.make_rmem_tensor((1,), cutlass.Int32) + amax_bits[0] = acc & cutlass.Int32(0x7FFF) + return cute.recast_tensor(amax_bits, cutlass.BFloat16)[0] + + @cute.jit + def _amax_lane(self, v: cutlass.Float32) -> cutlass.Float32: + if cutlass.const_expr(self.sf_vec_size == 32): + return cute.arch.warp_redux_sync(v, "fmax", abs=True) + first_half = (self.lane_idx % cutlass.Int32(32)) < cutlass.Int32(16) + vsel = cutlass.Float32(0.0) + if first_half: + vsel = v + amax = cute.arch.warp_redux_sync(vsel, "fmax", abs=True) + if not first_half: + amax = cute.arch.warp_redux_sync(v, "fmax", abs=True) + return amax + + @cute.jit + def _e8m0(self, amax: cutlass.Float32) -> Tuple[cutlass.Float8E8M0FNU, cutlass.Float32]: + candidate = amax * cutlass.Float32(self._data_rcp_limit) + scale_f32 = cutlass.Float32(cvt_f32_to_fp8_to_f32(candidate, cutlass.Float8E8M0FNU)) + return scale_f32.to(cutlass.Float8E8M0FNU), scale_f32 + + @cute.jit + def _enc_nvfp4(self, decode_scale: cutlass.Float32) -> cutlass.Float32: + # rcp.approx.ftz with the fc1 cap+mask idiom (amax==0 -> 0, no inf*0 NaN). + enc = cute.arch.fmin(cute.arch.rcp_approx(decode_scale), Fp32Max) + mask = cute.arch.fmin(decode_scale * cutlass.Float32(1e30), cutlass.Float32(1.0)) + return enc * mask + + @cute.jit + def _enc_mxfp8(self, scale_f32: cutlass.Float32) -> cutlass.Float32: + # Skip nan + enc = cute.arch.fmin(cute.arch.rcp_approx(scale_f32), Fp32Max) + mask = cute.arch.fmin(scale_f32 * cutlass.Float32(1e30), cutlass.Float32(1.0)) + return enc * mask + + +# ============================================================================= +# Region tag +# ============================================================================= + + +class Region: + """Codegen-time region tag for a 16x32 sub-region within a 32x32 tile.""" + + Top = 0 + Bottom = 1 + + +# ============================================================================= +# TmemTranspose16x32 +# ============================================================================= + + +class _TmemTranspose16x32Core: + """Physical implementation of the 16x32 -> 32x16 TMEM in-place transpose. + + The transpose is a fixed sequence of tcgen05 32-bit element atoms; each + 32-bit slot is an fp32 SwiGLU-fold value for FC1. The (thread, reg) -> + (tmem_dp, tmem_col) input / output mapping is documented on the + ``TmemTranspose16x32`` subclass, which is the public entry point. + + Per-thread RMEM coordinate convention: + + - ``lane_idx`` -- warp lane id (= thread index within warp), in [0, 32). + - ``elem_idx`` -- per-thread reg index, in [0, 16). + """ + + _PermR1 = (0, 8, 2, 10, 4, 12, 6, 14, 1, 9, 3, 11, 5, 13, 7, 15) + _PermR3 = (0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15) + _PermR4 = (0, 8, 2, 10, 4, 12, 6, 14, 1, 9, 3, 11, 5, 13, 7, 15) + + _TmemRowStride = 1 << 16 + _io_dtype = cutlass.Float32 + + @staticmethod + def _tmem_layout(num_lanes: int, num_cols: int) -> cute.Layout: + return cute.make_layout( + (((num_lanes, num_cols), 1),), stride=(((_TmemTranspose16x32Core._TmemRowStride, 1), 0),) + ) + + @staticmethod + def _rmem_copy_view(rmem: cute.Tensor, num_regs: int, offset: int = 0) -> cute.Tensor: + return cute.make_tensor(rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),))) + + @staticmethod + def load_subtile_raw_acc( + tmem_subtile_tensor: cute.Tensor, + ) -> Tuple[cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor]: + """LDTM the entire 32-lane x 64-col raw acc region of one epi + subtile into 4 independent (16,) fp32 RMEM tensors. + + Used by the FC1 overlap-acc unroll path to extract all raw acc data + of the first 2 subtiles up front, so that the acc TMEM can be released + after the first subtile's 4 LDTMs. + + ``tmem_subtile_tensor`` is the (32 lanes, 64 cols) view onto a + single epi subtile's acc TMEM region (already offset by + ``warp_lane_offset + acc_stage_col_offset + subtile_col_offset``; + see ``SwapABGatedActEpilogue._subtile_local_tmem_tensor``). + + Returns a 4-tuple of (16,) fp32 RMEM tensors carrying the FC1 raw + LDTM distribution: + + [0] gate_lo / first-half top -- subtile cols 0..31, lanes 0..15 + [1] up_lo / first-half bot -- subtile cols 0..31, lanes 16..31 + [2] raw_top / second-half top -- subtile cols 32..63, lanes 0..15 + [3] raw_bot / second-half bot -- subtile cols 32..63, lanes 16..31 + + 4 atom calls of ``Ld16x64bOp(Repetition.x16) Float32`` -- the same + atom used by the per-subtile entry LDTM. Each output is in the + raw-LDTM input distribution consumed by ``TmemTranspose16x32``. + """ + atom_ld16x64 = cute.make_copy_atom( + tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), _TmemTranspose16x32Core._io_dtype + ) + + ptr = tmem_subtile_tensor.iterator + half_lane_off = 16 * _TmemTranspose16x32Core._TmemRowStride + + # 4 source 16-lane x 32-col views over the (32, 64) subtile region: + # first half (cols 0..31): top lanes 0..15 / bot lanes 16..31 + # second half (cols 32..63): top lanes 0..15 / bot lanes 16..31 + # All offsets are Python ints (compile-time const) so cute can + # const-fold them and infer the correct (>= 8 B / 2 col) ptr + # alignment that the LDTM atom requires. Using ``cutlass.Int32`` + # offsets here would wrap them as SSA values that cute treats as + # alignment-unknown, tripping the atom's verifier. + first_top_view = cute.make_tensor(ptr, _TmemTranspose16x32Core._tmem_layout(16, 32)) + first_bot_view = cute.make_tensor(ptr + half_lane_off, _TmemTranspose16x32Core._tmem_layout(16, 32)) + second_top_view = cute.make_tensor(ptr + 32, _TmemTranspose16x32Core._tmem_layout(16, 32)) + second_bot_view = cute.make_tensor(ptr + 32 + half_lane_off, _TmemTranspose16x32Core._tmem_layout(16, 32)) + + first_top = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) + first_bot = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) + second_top = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) + second_bot = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) + + cute.copy(atom_ld16x64, first_top_view, _TmemTranspose16x32Core._rmem_copy_view(first_top, 16)) + cute.copy(atom_ld16x64, first_bot_view, _TmemTranspose16x32Core._rmem_copy_view(first_bot, 16)) + cute.copy(atom_ld16x64, second_top_view, _TmemTranspose16x32Core._rmem_copy_view(second_top, 16)) + cute.copy(atom_ld16x64, second_bot_view, _TmemTranspose16x32Core._rmem_copy_view(second_bot, 16)) + + return (first_top, first_bot, second_top, second_bot) + + def __init__(self, tmem_ptr, region: int, reg_tensor: Optional[cute.Tensor] = None) -> None: + # The whole transpose is built from 32-bit element atoms; _io_dtype + # drives _src_regs / output / every LDTM/STTM atom below, so guard the + # invariant once here (tautological today, defensive against future + # dtype edits). + if cutlass.const_expr(self._io_dtype.width != 32): + raise TypeError( + f"{type(self).__name__} requires a 32-bit _io_dtype (the " + f"transpose uses 32-bit element atoms), got {self._io_dtype} " + f"(width {self._io_dtype.width})." + ) + + half_lane_off = 16 * self._TmemRowStride + if region == Region.Top: + src_ptr = tmem_ptr + dst_ptr = tmem_ptr + elif region == Region.Bottom: + src_ptr = tmem_ptr + half_lane_off + dst_ptr = tmem_ptr + 16 + else: + raise ValueError("region must be Region.Top or Region.Bottom") + + self.region = region + + self._tmem_src_full = cute.make_tensor(src_ptr, self._tmem_layout(16, 32)) + self._tmem_dst_full = cute.make_tensor(dst_ptr, self._tmem_layout(32, 16)) + self._tmem_dst_top = cute.make_tensor(dst_ptr, self._tmem_layout(16, 16)) + self._tmem_dst_bot = cute.make_tensor(dst_ptr + half_lane_off, self._tmem_layout(16, 16)) + + self._atom_ld16x64 = cute.make_copy_atom(tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), self._io_dtype) + self._atom_st16x128 = cute.make_copy_atom(tcgen05.St16x128bOp(tcgen05.Repetition.x8), self._io_dtype) + self._atom_st32x32 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), self._io_dtype) + self._atom_ld16x256 = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), self._io_dtype) + self._atom_ld16x128 = cute.make_copy_atom(tcgen05.Ld16x128bOp(tcgen05.Repetition.x4), self._io_dtype) + + self._src_regs = cute.make_rmem_tensor((16,), self._io_dtype) + # ``output`` is a bare (16,) RMEM fragment; its (lane_idx, elem_idx) + # distribution after all four rounds is the transpose output mapping + # documented on ``TmemTranspose16x32``. + self.output = cute.make_rmem_tensor((16,), self._io_dtype) + + # skip-R1.Load mode: ``reg_tensor`` must already be in the transpose + # input distribution (see ``TmemTranspose16x32`` / produced by + # ``load_subtile_raw_acc``); we copy it in lieu of the R1 LDTM. + # Weak entry guard (replaces the removed input contract): the transpose + # atoms are 32-bit element atoms over exactly 16 regs/lane, so the fed + # tensor must be a 32-bit element type of size 16. + self._reg_tensor = reg_tensor + if reg_tensor is not None: + if cutlass.const_expr(reg_tensor.element_type.width != 32): + raise TypeError( + f"{type(self).__name__} reg_tensor must be a 32-bit element " + f"type, got element type " + f"{reg_tensor.element_type} (width {reg_tensor.element_type.width})." + ) + if cutlass.const_expr(cute.size(reg_tensor) != 16): + raise ValueError( + f"{type(self).__name__} reg_tensor must hold exactly 16 elements, got {cute.size(reg_tensor)}." + ) + for r in range(16): + self._src_regs[r] = reg_tensor[r] + + # -- R1 ------------------------------------------------------------------ + + def r1_load(self) -> None: + """LDTM src region -> ``_src_regs``. No-op in skip-R1.Load mode.""" + if self._reg_tensor is not None: + return + cute.copy(self._atom_ld16x64, self._tmem_src_full, self._rmem_copy_view(self._src_regs, 16)) + + def r1_perm(self) -> None: + for r in range(16): + self.output[r] = self._src_regs[self._PermR1[r]] + + def r1_store(self) -> None: + cute.copy(self._atom_st16x128, self._rmem_copy_view(self.output, 16), self._tmem_src_full) + + # -- R2 ------------------------------------------------------------------ + + def r2_load(self) -> None: + cute.copy(self._atom_ld16x64, self._tmem_src_full, self._rmem_copy_view(self._src_regs, 16)) + + def r2_store(self) -> None: + cute.copy(self._atom_st32x32, self._rmem_copy_view(self._src_regs, 16), self._tmem_dst_full) + + # -- R3 ------------------------------------------------------------------ + + def r3_load_top(self) -> None: + cute.copy(self._atom_ld16x256, self._tmem_dst_top, self._rmem_copy_view(self._src_regs, 8, offset=0)) + + def r3_load_bot(self) -> None: + cute.copy(self._atom_ld16x256, self._tmem_dst_bot, self._rmem_copy_view(self._src_regs, 8, offset=8)) + + def r3_perm(self) -> None: + for r in range(16): + self.output[r] = self._src_regs[self._PermR3[r]] + + def r3_store(self) -> None: + cute.copy(self._atom_st32x32, self._rmem_copy_view(self.output, 16), self._tmem_dst_full) + + # -- R4 ------------------------------------------------------------------ + + def r4_load_top(self) -> None: + cute.copy(self._atom_ld16x128, self._tmem_dst_top, self._rmem_copy_view(self._src_regs, 8, offset=0)) + + def r4_load_bot(self) -> None: + cute.copy(self._atom_ld16x128, self._tmem_dst_bot, self._rmem_copy_view(self._src_regs, 8, offset=8)) + + def r4_perm(self) -> None: + for r in range(16): + self.output[r] = self._src_regs[self._PermR4[r]] + + def r4_store(self) -> None: + cute.copy(self._atom_st32x32, self._rmem_copy_view(self.output, 16), self._tmem_dst_full) + + def from_r1_perm_until_last_store(self) -> cute.Tensor: + self.r1_perm() + self.r1_store() + self.r2_load() + self.r2_store() + self.r3_load_top() + self.r3_load_bot() + self.r3_perm() + self.r3_store() + self.r4_load_top() + self.r4_load_bot() + self.r4_perm() + return self.output + + +class TmemTranspose16x32(_TmemTranspose16x32Core): + """FC1 16x32 -> 32x16 TMEM in-place transpose. + + The per-thread RMEM ``(lane_idx, elem_idx) -> (tmem_dp, tmem_col)`` mapping + is fixed by the underlying atom sequence. Each slot is an fp32 SwiGLU-fold + value and ``tmem_col`` is the intermediate-output index. + + Input distribution -- what each (lane_idx, elem_idx) reg holds on entry + (i.e. straight after the 16-dp x 32-col source LDTM, or as fed in via + ``reg_tensor`` / ``load_subtile_raw_acc`` for skip-R1.Load mode): + + tmem_dp = elem_idx * 2 + (lane_idx // 2) % 2 # in [0, 32) + tmem_col = (lane_idx % 2) * 8 + lane_idx // 4 # in [0, 16) + + Output distribution -- after all four rounds, the 32-dp x 16-col result has + each lane owning one full dp-row of 16 cols: + + tmem_dp = lane_idx # in [0, 32) + tmem_col = elem_idx # in [0, 16) + """ + + +# ============================================================================= +# TmemTranspose32x32Inplace +# ============================================================================= + + +class TmemTranspose32x32Inplace: + """fc1 epi 32x32 in-place TMEM transpose: two ``TmemTranspose16x32`` + sub-instances (``top`` = lanes 0..15, ``bot`` = lanes 16..31). + + Optional ``reg_tensor_top`` / ``reg_tensor_bot`` enable skip-R1.Load mode + for both halves; they must be provided or omitted together. + """ + + def __init__( + self, tmem_ptr, reg_tensor_top: Optional[cute.Tensor] = None, reg_tensor_bot: Optional[cute.Tensor] = None + ) -> None: + if (reg_tensor_top is None) != (reg_tensor_bot is None): + raise ValueError( + "TmemTranspose32x32Inplace: reg_tensor_top and reg_tensor_bot " + "must be provided or omitted together (both halves either " + "skip-R1.Load or do R1.Load)." + ) + self.top = TmemTranspose16x32(tmem_ptr, Region.Top, reg_tensor=reg_tensor_top) + self.bot = TmemTranspose16x32(tmem_ptr, Region.Bottom, reg_tensor=reg_tensor_bot) + + def from_r1_perm_until_last_store(self) -> Tuple[cute.Tensor, cute.Tensor]: + self.bot.r1_perm() + self.top.r1_perm() + self.bot.r1_store() + self.top.r1_store() + + self.bot.r2_load() + self.top.r2_load() + self.top.r2_store() + self.bot.r2_store() + + self.top.r3_load_top() + self.top.r3_load_bot() + self.bot.r3_load_top() + self.bot.r3_load_bot() + self.top.r3_perm() + self.bot.r3_perm() + self.top.r3_store() + self.bot.r3_store() + + self.top.r4_load_top() + self.top.r4_load_bot() + self.bot.r4_load_top() + self.bot.r4_load_bot() + self.top.r4_perm() + self.bot.r4_perm() + return self.top.output, self.bot.output + + +class TmemTranspose32x64B16Movm: + """FC2 warp-local 32-hidden x 64-token BF16 transpose using MOVM. + + Input is the flat ``[top, bottom]`` distribution produced by two + 16dp256bit accumulator loads followed by ``fc2_f2fp``. Output + ``(lane_idx, elem_idx)`` coordinates are: + + token = lane_idx + 32 * (elem_idx // 32) + hidden = elem_idx % 32 + + The fixed register permutation between MOVM and STTM is an SSA rename. It + keeps each thread's two complete hidden-32 rows without any lane exchange. + """ + + _tmem_row_stride = 1 << 16 + _store_reg_source_indices = ( + 0, + 2, + 1, + 3, + 16, + 18, + 17, + 19, + 8, + 10, + 9, + 11, + 24, + 26, + 25, + 27, + 4, + 6, + 5, + 7, + 20, + 22, + 21, + 23, + 12, + 14, + 13, + 15, + 28, + 30, + 29, + 31, + ) + + @staticmethod + def _tmem_layout(num_lanes: int, num_cols: int) -> cute.Layout: + return cute.make_layout( + (((num_lanes, num_cols), 1),), stride=(((TmemTranspose32x64B16Movm._tmem_row_stride, 1), 0),) + ) + + @staticmethod + def _rmem_copy_view(rmem: cute.Tensor, num_regs: int, offset: int = 0) -> cute.Tensor: + return cute.make_tensor(rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),))) + + @cute.jit + def __init__(self, tmem_ptr, reg_tensor: cute.Tensor) -> None: + if cutlass.const_expr(reg_tensor.element_type is not cutlass.BFloat16): + raise TypeError(f"{type(self).__name__} expects BF16 input after f2fp, got {reg_tensor.element_type}.") + if cutlass.const_expr(cute.size(reg_tensor) != 64): + raise ValueError(f"{type(self).__name__} expects 64 BF16 elements, got {cute.size(reg_tensor)}.") + + movm_words = movmatrix_b16(cute.recast_tensor(reg_tensor, cutlass.Int32)) + self._store_words = cute.make_rmem_tensor(movm_words.layout, movm_words.element_type) + for store_reg in cutlass.range_constexpr(32): + self._store_words[store_reg] = movm_words[self._store_reg_source_indices[store_reg]] + + half_lane_offset = 16 * self._tmem_row_stride + self._tmem_top = cute.make_tensor(tmem_ptr, self._tmem_layout(16, 32)) + self._tmem_bottom = cute.make_tensor(tmem_ptr + half_lane_offset, self._tmem_layout(16, 32)) + self._tmem_full = cute.make_tensor(tmem_ptr, self._tmem_layout(32, 32)) + self._store_atom = cute.make_copy_atom(tcgen05.St16x128bOp(tcgen05.Repetition.x8), cutlass.Float32) + self._load_atom = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) + + @cute.jit + def __call__(self) -> cute.Tensor: + movm_words_f32 = cute.recast_tensor(self._store_words, cutlass.Float32) + cute.copy(self._store_atom, self._rmem_copy_view(movm_words_f32, 16), self._tmem_top) + cute.copy(self._store_atom, self._rmem_copy_view(movm_words_f32, 16, offset=16), self._tmem_bottom) + + output_words = cute.make_rmem_tensor((32,), cutlass.Float32) + cute.copy(self._load_atom, self._tmem_full, self._rmem_copy_view(output_words, 32)) + return cute.recast_tensor(output_words, cutlass.BFloat16) + + +@dataclasses.dataclass(frozen=True) +class GatedActEpilogueArgs: + """Optional runtime tensors used by the gated-activation epilogue.""" + + fc1_alpha: Optional[cute.Tensor] + fc2_alpha: Optional[cute.Tensor] + fc1_norm_const: Optional[cute.Tensor] + # ----------------------------------- + # MoE domain (token, topk), deepgemm graph only? for transformer graph, we want reduce kernel to perform the score mul. + topk_scores: Optional[cute.Tensor] + + +class SwapABGatedActEpilogue(KernelComponent): + """Autonomous epilogue for the swap-AB SwiGLU NVFP4 kernel. + + ``run()`` is the single entry point the kernel calls inside the epi + warp body. The kernel's responsibility is reduced to: + + - allocate / free TMEM and build ``acc_tensor`` + - construct the AB / acc pipelines + - obtain the scheduler consumer + + Everything else (acc consumer state, task-tile loop, overlap rotation, + early release, TMA store commit / drain, per-subtile dispatch) lives + inside this class. + """ + + _EpilogueSyncWaitBarId = 1 # Arrive and wait only + _EpilogueAsyncBarIdBase = 4 # Some arrive, the others arrive and wait + _EpilogueFc1GateUpInterleave = 16 + _EpilogueTokenTileSize = 64 # Fundamentally the epi_tile_n + _EpilogueFc1IntermediateGateUpTileSize = 128 # Fundamentally epi_tile_m + _EpilogueFc1IntermediateDownTileSize = 64 # Fundamentally epi_tile_m // 2 + _EpilogueFc2HiddenTileSize = 128 # Fundamentally epi_tile_m + _EpilogueWarpCnt = 4 + # One warp owns this many intermediate_down outputs per token: the TMEM transpose gives each + # warp 32 accumulator rows, which the gate/up interleave halves. + _EpilogueFc1IntermediateDownPerWarp = 16 + smem_scratch_overlay: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.scratch" + fc1_staging_region: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.fc1_staging" + fc2_staging_region: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.fc2_staging" + # A 144 B FP8 stride rotates each row by four SMEM banks. + _Fc2UblkFp8RowStrideBytes = 144 + _ScratchByteAlignment = 128 + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + # The accumulator follows from the instruction family; see __init__. + "quant_kind": str, + "hidden_size": StaticOrRuntimeIntegerType, + "intermediate_gateup_size": StaticOrRuntimeIntegerType, + "combine_format": CombineFormat, + "gate_up_clamp": Optional[float], + # SiTU (Kimi K3) selects a different gated-activation core; optional so + # existing SwiGLU descriptors stay valid unchanged. See _resolve_situ_betas. + "situ_beta": OptionalRequirement(Optional[float]), + "situ_linear_beta": OptionalRequirement(Optional[float]), + } + + @classmethod + def impl_desc_require(cls) -> dict[str, object]: + return { + "mma_tiler_mnk": tuple, + "cluster_shape_mn": tuple, + "use_2cta_instrs": bool, + "fc2_use_bulk": bool, + "communication_enabled": bool, + "fc1_epi_flag_batch": int, + "fc2_epi_flag_batch": int, + "fc2_tma_stages": OptionalRequirement(int), + "reduce_topk_in_kernel": OptionalRequirement(bool), + "token_back_push_data": OptionalRequirement(bool), + } + + @staticmethod + def _resolve_situ_betas(problem_desc: ProblemDesc) -> Tuple[Optional[float], Optional[float]]: + """Resolve the SiTU (Kimi K3) betas from the ProblemDesc; ``(None, None)`` means SwiGLU. + + Both betas must be given together, and SiTU excludes ``gate_up_clamp`` -- matching DeepGEMM's + ``DG_HOST_ASSERT(not use_situ or not activation_clamp_opt.has_value())``. They are baked in at + codegen time exactly like ``gate_up_clamp``, so the enclosing KernelClass must also fold them + into its ``name()`` cache key. + """ + beta = problem_desc.get("situ_beta") + linear_beta = problem_desc.get("situ_linear_beta") + if (beta is None) != (linear_beta is None): + raise ValueError(f"situ_beta and situ_linear_beta must be set together, got {beta} and {linear_beta}.") + if beta is None: + return None, None + if beta <= 0 or linear_beta <= 0: + raise ValueError(f"SiTU beta parameters must be positive, got {beta} and {linear_beta}.") + if problem_desc["gate_up_clamp"] is not None: + raise ValueError("SiTU does not support gate_up_clamp; the two activation variants are exclusive.") + return float(beta), float(linear_beta) + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.quant_kind = QuantKind(problem_desc["quant_kind"]) + self.acc_dtype = tcgen05_block_scaled_acc_dtype + self.hidden_size = problem_desc["hidden_size"] + self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] + self.combine_format = problem_desc["combine_format"] + self.gate_up_clamp = problem_desc["gate_up_clamp"] + self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) + self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + self.cluster_shape_mn = impl_desc["cluster_shape_mn"] + self.use_2cta_instrs = impl_desc["use_2cta_instrs"] + self.fc2_use_bulk = impl_desc["fc2_use_bulk"] + self.communication_enabled = impl_desc["communication_enabled"] + self.fc1_epi_flag_batch = impl_desc["fc1_epi_flag_batch"] + self.fc2_epi_flag_batch = impl_desc["fc2_epi_flag_batch"] + + if self.communication_enabled: + for field_name in ("reduce_topk_in_kernel", "token_back_push_data"): + if field_name not in impl_desc: + raise KeyError(f"Communication-enabled Epilogue requires ImplDesc field {field_name!r}.") + self.reduce_topk_in_kernel = impl_desc.get("reduce_topk_in_kernel", False) + self.token_back_push_data = impl_desc.get("token_back_push_data", False) + if not self.communication_enabled and (self.reduce_topk_in_kernel or self.token_back_push_data): + raise ValueError("A communication-disabled Epilogue cannot enable communication policies.") + + # FC1 emits what FC2 consumes as its activation, in the kind's own scale format. + self.fc1_output_dtype = self.quant_kind.activation_dtype + self.fc1_output_sf_dtype = self.quant_kind.sf_dtype + self.sf_vec_size = self.quant_kind.sf_vec_size + # A 32-wide scale block spans two epilogue warps; see QuantImpl's paired direction. + self.needs_pair_amax_exchange = self.sf_vec_size > self._EpilogueFc1IntermediateDownPerWarp + self.token_back_push_sf = self.communication_enabled and self.combine_format.is_quantized + self.token_back_enabled = self.token_back_push_data or self.token_back_push_sf + self.fc2_output_is_local = not self.communication_enabled or self.token_back_push_data + self.fc2_use_tma = self.fc2_use_bulk and self.fc2_output_is_local + self.fc2_use_ublk = self.fc2_use_bulk and not self.fc2_output_is_local + if self.fc2_use_ublk and self.combine_format.act_dtype.width == 4: + raise ValueError("FC2 UBLK does not support an FP4 combine payload.") + if self.reduce_topk_in_kernel and self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("In-kernel top-k reduction requires a BF16 combine format.") + self.reduce_topk_in_epilogue = self.reduce_topk_in_kernel and not self.token_back_push_data + if not 1 <= self.fc1_epi_flag_batch <= 32 or not 1 <= self.fc2_epi_flag_batch <= 32: + raise ValueError("Epilogue flag batch sizes must be in [1, 32].") + self.cluster_tile_intermediate_downproj = self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] + + atom_thr_size = 2 if self.use_2cta_instrs else 1 + self.cta_tile_m = self._EpilogueFc2HiddenTileSize + self.cta_tile_n = self.mma_tiler_mnk[1] + self.cta_tile_k = self.mma_tiler_mnk[2] + assert self.mma_tiler_mnk[0] // atom_thr_size == self.cta_tile_m + assert self.cta_tile_n % self._EpilogueTokenTileSize == 0 + tmem_plan = impl_desc["tmem_plan"] + self.num_sfa_tmem_cols = tmem_plan.sfa_columns + self.num_sfb_tmem_cols = tmem_plan.sfb_columns + self.num_sf_tmem_cols = tmem_plan.sfa_columns + tmem_plan.sfb_columns + self.num_tmem_alloc_cols = tmem_plan.allocation_columns + self.num_accumulator_stages = tmem_plan.accumulator_stage_count + self.num_accumulator_pipeline_stages = tmem_plan.accumulator_pipeline_stages + self.accumulator_overlap_columns = ( + tmem_plan.accumulator_stage_columns - tmem_plan.accumulator_stage_stride_columns + ) + self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns + self.overlapping_accum = self.accumulator_overlap_columns > 0 + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) + self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) + + if isinstance(self.hidden_size, int) and self.hidden_size % (self.cta_tile_m * self.cluster_shape_mn[0]) == 0: + self.fc2_hidden_needs_predicate: bool = False + else: + self.fc2_hidden_needs_predicate: bool = True + + if isinstance(self.intermediate_gateup_size, int): + self.intermediate_downproj: Optional[int] = self.intermediate_gateup_size // 2 + else: + self.intermediate_downproj: Optional[int] = None + + self.subtile_cnt = self.cta_tile_n // self._EpilogueTokenTileSize + + # One staging stage per token subtile, each an (epi_tile_n, epi_tile_m // 2) quantized tile. + self.fc1_staging_stage_bytes = ( + self._EpilogueTokenTileSize * self._EpilogueFc1IntermediateDownTileSize * self.fc1_output_dtype.width // 8 + ) + self.fc1_staging_bytes = self.subtile_cnt * self.fc1_staging_stage_bytes + # The amax exchange plane is (token chunk, epilogue thread), block-major so each thread's + # column is one vector access. A thread holds one token per 32-lane chunk of the subtile. + # These slots borrow the current subtile's staging stage rather than costing their own + # bytes -- see the lifetime argument in fc1_quant. + self.fc1_amax_token_chunks = self._EpilogueTokenTileSize // 32 + self.fc1_amax_slot_count = ( + self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 if self.needs_pair_amax_exchange else 0 + ) + if self.fc1_amax_slot_count * 4 > self.fc1_staging_stage_bytes: + raise ValueError("The paired amax exchange does not fit in one fc1 staging stage.") + + requested_fc2_tma_stages = impl_desc.get("fc2_tma_stages") + if requested_fc2_tma_stages is not None and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt: + raise ValueError(f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}.") + if self.fc2_use_bulk: + single_stage_region = self._make_fc2_single_stage_region() + if self.fc2_use_ublk and single_stage_region.nbytes % 16 != 0: + raise ValueError("Each FC2 UBLK staging stage must occupy a multiple of 16 bytes.") + # Additional stages trade mainloop SMEM for store overlap. + self.fc2_tma_stages = requested_fc2_tma_stages if requested_fc2_tma_stages is not None else 1 + self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region(self.fc2_tma_stages) + else: + self.fc2_tma_stages = 0 + self.fc2_staging_spec = None + + @classmethod + def epilogue_sync_barrier(cls) -> pipeline.NamedBarrier: + """The one barrier every epilogue rendezvous uses: all four warps, arrive-and-wait. + + Reused rather than split per purpose because the participant set is always the same 128 + threads and no two uses are ever in flight together -- the tile-boundary rendezvous sits + outside the subtile loop, FC1's amax exchange and FC2's bulk-store handshake belong to + different work tiles. + """ + return pipeline.NamedBarrier(barrier_id=cls._EpilogueSyncWaitBarId, num_threads=32 * cls._EpilogueWarpCnt) + + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: + """Declare the epilogue scratch as one allocation shared by two exclusive lifetimes. + + FC1's staging tile and FC2's store tile belong to different work tiles, separated by the + tile-boundary TMA drain and rendezvous in ``run()``, so they never coexist. + """ + overlay = smem_workspace.create_overlay(self.smem_scratch_overlay) + overlay.add_lifetime("fc1_staging").register_tensor( + self.fc1_staging_region, cutlass.Int8, (self.fc1_staging_bytes,), byte_alignment=128 + ) + if self.fc2_staging_spec is not None: + overlay.add_lifetime("fc2_staging").register_tensor( + self.fc2_staging_region, + self.fc2_staging_spec.dtype, + self.fc2_staging_spec.shape, + stride=self.fc2_staging_spec.stride, + swizzle=self.fc2_staging_spec.swizzle, + byte_alignment=self.fc2_staging_spec.byte_alignment, + ) + + def fc1_staged_smem_layout( + self, n_stages: int, without_stage_mode: bool = False + ) -> Union[cute.Layout, cute.ComposedLayout]: + layout = sm100_utils.make_smem_layout_epi( + self.fc1_output_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize), + n_stages, + ) + if without_stage_mode: + return cute.select(layout, mode=[0, 1]) + return layout + + def fc2_tma_staged_smem_spec(self, n_stages: int) -> Tuple[Tuple, Tuple, Tuple[int, int, int]]: + """Return the bank-conflict-free token-major FC2 TMA staging layout.""" + stage_stride = self._EpilogueTokenTileSize * self._EpilogueFc2HiddenTileSize + wire_dtype = self.combine_format.act_dtype + if wire_dtype is cutlass.BFloat16: + shape = ((32, 2), (64, 2), n_stages) + stride = ((64, 2048), (1, 4096), stage_stride if n_stages > 1 else 0) + swizzle = (3, 4, 3) + elif wire_dtype.width == 8: + shape = ((32, 2), 128, n_stages) + stride = ((128, 4096), 1, stage_stride if n_stages > 1 else 0) + swizzle = (3, 4, 3) + elif wire_dtype.width == 4: + shape = ((32, 2), 128, n_stages) + stride = ((128, 4096), 1, stage_stride if n_stages > 1 else 0) + swizzle = (2, 4, 3) + else: + raise ValueError(f"Unsupported FC2 TMA staging dtype {wire_dtype}.") + return shape, stride, swizzle + + def fc2_tma_staged_smem_layout(self, n_stages: int, without_stage_mode: bool = False) -> cute.ComposedLayout: + shape, stride, swizzle = self.fc2_tma_staged_smem_spec(n_stages) + layout = cute.make_composed_layout(cute.make_swizzle(*swizzle), 0, cute.make_layout(shape, stride=stride)) + if without_stage_mode: + return cute.select(layout, mode=[0, 1]) + return layout + + def _make_fc2_single_stage_region(self) -> SmemRegion: + """Describe one FC2 bulk-staging subtile for TMA or UBLK.""" + if self.fc2_use_tma: + shape, stride, swizzle = self.fc2_tma_staged_smem_spec(1) + return SmemRegion( + name="", + kind="tensor", + dtype=self.combine_format.act_dtype, + shape=shape[:-1], + stride=stride[:-1], + swizzle=swizzle, + byte_alignment=128, + ) + row_stride_elements = ( + self._Fc2UblkFp8RowStrideBytes + if self.combine_format.act_dtype.width == 8 + else self._EpilogueFc2HiddenTileSize + ) + return SmemRegion( + name="", + kind="tensor", + dtype=self.combine_format.act_dtype, + shape=(self._EpilogueTokenTileSize, self._EpilogueFc2HiddenTileSize), + stride=(row_stride_elements, 1), + swizzle=None, + byte_alignment=16, + ) + + def _make_fc2_staging_region(self, stage_count: int) -> SmemRegion: + single_stage_region = self._make_fc2_single_stage_region() + return SmemRegion( + name="", + kind="tensor", + dtype=single_stage_region.dtype, + shape=(*single_stage_region.shape, stage_count), + stride=(*single_stage_region.stride, single_stage_region.cosize if stage_count > 1 else 0), + swizzle=single_stage_region.swizzle, + byte_alignment=single_stage_region.byte_alignment, + ) + + def prepare_tma_store_params( + self, fc1_output_template: cute.Tensor, fc2_output_template: cute.Tensor + ) -> Tuple[cute.CopyAtom, cute.Tensor, Optional[cute.CopyAtom], Optional[cute.Tensor]]: + """Build the FC1 TMA store and the optional local FC2 TMA store.""" + fc1_operation = cpasync.CopyBulkTensorTileS2GOp() + fc1_smem_layout = self.fc1_staged_smem_layout(1, without_stage_mode=True) + fc1_tile = (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize) + fc1_atom, fc1_tensor = cpasync.make_tiled_tma_atom( + fc1_operation, fc1_output_template, fc1_smem_layout, fc1_tile + ) + + if cutlass.const_expr(self.fc2_use_tma): + # Keep tiled rest modes dynamic so a single static tile cannot collapse to stride zero. + runtime_token_extent = cutlass.Int32(fc2_output_template.shape[0]) + runtime_hidden_extent = cutlass.Int32(fc2_output_template.shape[2]) + fc2_token_major_template = cute.make_tensor( + fc2_output_template.iterator, + cute.make_layout( + (runtime_token_extent, runtime_hidden_extent, cutlass.Int32(1)), + stride=(fc2_output_template.stride[0], fc2_output_template.stride[2], 0), + ), + ) + fc2_operation = cpasync.CopyBulkTensorTileS2GOp() + fc2_smem_layout = self.fc2_tma_staged_smem_layout(1, without_stage_mode=True) + fc2_tile = (self._EpilogueTokenTileSize, self._EpilogueFc2HiddenTileSize) + fc2_atom, fc2_tensor = cpasync.make_tiled_tma_atom( + fc2_operation, fc2_token_major_template, fc2_smem_layout, fc2_tile + ) + else: + fc2_atom = None + fc2_tensor = None + return fc1_atom, fc1_tensor, fc2_atom, fc2_tensor + + @cute.jit + def run( + self, + smem_workspace: SmemWorkspace, + smem_base: cute.Pointer, + tmem_ptr: cute.Pointer, + acc_pipeline, + # ── Sched ──────────────────────────────────────────────────────── + sched_consumer: SchedulerConsumer, + kernel_extension: BlockScaledSwapAbFc12Extension, + # ── tensors ────────────────────────────────── + tma_atom_fc1_output: cute.CopyAtom, + fc1_output: cute.Tensor, # Domain of fake (m, n, l) + fc1_output_sf: cute.Tensor, # Domain of fake (m, n, l) + tma_atom_fc2_output: Optional[cute.CopyAtom], + fc2_tma_output: Optional[cute.Tensor], # Domain (physical_token, hidden, l=1) + fc2_output: cute.Tensor, # MoE domain (token, topk, hidden) + fc1_done_counter: cute.Tensor, # 1D tensor + tidx: cutlass.Int32, + token_src_metadata: Optional[cute.Tensor], + fc2_done_counter: Optional[cute.Tensor], + fc2_output_sf: Optional[cute.Tensor], + peer_rank_ptr_mapper: Optional[SymmetricBufferDevice], + optional_epi_args: Optional[GatedActEpilogueArgs] = None, + ): + if cutlass.const_expr(not smem_workspace.finalized): + raise RuntimeError("SwapABGatedActEpilogue.run requires a finalized SmemWorkspace.") + if cutlass.const_expr(optional_epi_args is None): + optional_epi_args = GatedActEpilogueArgs( + fc1_alpha=None, fc2_alpha=None, fc1_norm_const=None, topk_scores=None + ) + fc1_staging_pointer = smem_workspace.ptr(self.fc1_staging_region, smem_base) + if cutlass.const_expr(self.fc2_tma_stages > 0): + fc2_smem_tensor = smem_workspace.tensor(self.fc2_staging_region, smem_base) + else: + fc2_smem_tensor = None + if cutlass.const_expr(self.fc2_use_tma and (tma_atom_fc2_output is None or fc2_tma_output is None)): + raise ValueError("FC2 TMA store requires a TMA atom and token-major output tensor.") + if cutlass.const_expr( + self.communication_enabled and (token_src_metadata is None or peer_rank_ptr_mapper is None) + ): + raise ValueError("Communication-enabled Epilogue requires token metadata and a peer pointer mapper.") + if cutlass.const_expr(self.token_back_enabled and fc2_done_counter is None): + raise ValueError("Token-back requires an FC2 done counter.") + if cutlass.const_expr(self.token_back_push_sf and fc2_output_sf is None): + raise ValueError("Quantized token-back requires an FC2 output scale tensor.") + tmem_acc = cute.make_tensor( + cute.recast_ptr(tmem_ptr, dtype=cutlass.Float32), + cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride), + ) + + fc1_epi = SwapABFc1Epilogue( + self, + tidx, + fc1_staging_pointer, + kernel_extension, + tma_atom_fc1_output, + fc1_output, + fc1_output_sf, + fc1_done_counter, + optional_epi_args, + ) + fc2_epi = SwapABFc2Epilogue( + self, + tidx, + fc2_smem_tensor, + tma_atom_fc2_output, + fc2_tma_output, + fc2_output, + token_src_metadata, + fc2_done_counter, + fc2_output_sf, + peer_rank_ptr_mapper, + optional_epi_args, + ) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_accumulator_pipeline_stages + ) + wait_only_named_barrier = self.epilogue_sync_barrier() + is_odd_turn = cutlass.Int32(1) + work_tile_info = sched_consumer.consume_work() + + flag_tracker = GpuReleaseFlagBatchTracker( + flag_address=Int64(0), + accumulated_flags=cutlass.Int32(0), + phase=cutlass.Int32(work_tile_info.phase), + thread_idx=tidx % (self._EpilogueWarpCnt * 32), + ) + + while work_tile_info.is_valid_tile: + if cutlass.const_expr(self.overlapping_accum): + tmem_stage_idx = acc_consumer_state.phase + else: + tmem_stage_idx = acc_consumer_state.index + tmem_acc_current = tmem_acc[None, None, tmem_stage_idx] + if work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1): + # The __call__ args should only take the while loop args, leave all loop irrevalent args to the init. + fc1_epi( + work_tile_info=work_tile_info, + tmem_acc_tensor=tmem_acc_current, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + is_odd_turn=is_odd_turn, + ) + else: + # The __call__ args should only take the while loop args, leave all loop irrevalent args to the init. + fc2_epi( + work_tile_info=work_tile_info, + tmem_acc_tensor=tmem_acc_current, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + is_odd_turn=is_odd_turn, + ) + prev_work_tile_info = work_tile_info + cur_was_linear1 = prev_work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + + acc_consumer_state.advance() + if cutlass.const_expr(self.overlapping_accum): + is_odd_turn = cutlass.Int32(1) - is_odd_turn + + work_tile_info = sched_consumer.consume_work() + + # Every asynchronous store commits at issue; drain before completion or scratch reuse. + cute.arch.cp_async_bulk_wait_group(0) + # _fence_rel_gpu() + wait_only_named_barrier.arrive_and_wait() + + # Publish completion for the work tile snapshotted above. + if cur_was_linear1: + flag_tracker = fc1_epi.signal_fc1_done(prev_work_tile_info, work_tile_info, flag_tracker) + else: + flag_tracker = fc2_epi.signal_fc2_done(prev_work_tile_info, work_tile_info, flag_tracker) + # Tail flush + flag_tracker.fire() + + +class _ImmutableAfterInit: + """Froze at the point calling `_freeze()`""" + + def __setattr__(self, name, value): + if self.__dict__.get("_frozen_", False): + raise AttributeError(f"{type(self).__name__} is immutable after __init__ (cannot set {name!r}).") + object.__setattr__(self, name, value) + + def _freeze(self) -> None: + object.__setattr__(self, "_frozen_", True) + + +# Device only object +class SwapABFc1Epilogue(_ImmutableAfterInit): + def __init__( + self, + base: SwapABGatedActEpilogue, + tidx: cutlass.Int32, + staging_pointer: cute.Pointer, + kernel_extension: BlockScaledSwapAbFc12Extension, + tma_atom_fc1_output: cute.CopyAtom, + fc1_output: cute.Tensor, # fake (m,n,l) domain + fc1_output_sf: cute.Tensor, # fake (m,n,l) domain + fc1_done_counter: cute.Tensor, # 1D tensor + optional_epi_args: GatedActEpilogueArgs, + ): + self.base = base + self.tidx = tidx % (base._EpilogueWarpCnt * 32) + self.warp_idx = self.tidx // 32 + self.lane_idx = self.tidx % 32 + # (token64, intermediate, stage). The swizzle travels with the layout instead of being + # spelled out here, so an 8-bit fc1 output picks up its own atom without a code change. + staged_layout = base.fc1_staged_smem_layout(base.subtile_cnt) + self.smem_tensor = cute.make_tensor( + cute.recast_ptr(staging_pointer, staged_layout.inner, dtype=base.fc1_output_dtype), staged_layout.outer + ) + # Kept unswizzled and untyped so fc1_quant can carve an fp32 amax-exchange view out of one + # staging stage without going through the staging tensor's swizzle. + self.staging_pointer = staging_pointer + self.kernel_extension = kernel_extension + self.fc1_tma_atom = tma_atom_fc1_output + self.fc1_output = fc1_output + self.fc1_output_sf = fc1_output_sf + self.fc1_done_counter = fc1_done_counter + self.optional_epi_args = optional_epi_args + self._freeze() + + def __getattr__(self, name): + return getattr(object.__getattribute__(self, "base"), name) + + def __extract_mlir_values__(self) -> List[ir.Value]: + # This object is a loop-invariant Python context wrapper, not a + # dynamic value. Keep it out of scf.while iter_args and reconstruct by + # identity across region boundaries. Any field that becomes a + # loop-carried SSA value must be passed explicitly to __call__ instead + # of being stored here. + return [] + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "SwapABFc1Epilogue": + assert len(values) == 0 + return self + + @cute.jit + def signal_fc1_done(self, work_tile_info, next_work_tile_info, flag_tracker): + # Only in-bound intermediate_downproj tiles signal; OOB -> null slot. + needs_intermediate_guard = ( + self.intermediate_downproj is None + or self.intermediate_downproj % self.cluster_tile_intermediate_downproj != 0 + ) + if cutlass.const_expr(needs_intermediate_guard): + in_bound = work_tile_info.tile_m_idx * self._EpilogueFc1IntermediateDownTileSize < self.fc1_output.shape[1] + else: + in_bound = True + slot = work_tile_info.cumulative_token_block_count + work_tile_info.tile_n_idx + flag_address = Int64(0) + if in_bound: + flag_address = (self.fc1_done_counter.iterator + slot).toint() + return flag_tracker.accumulate(next_work_tile_info.phase, self.fc1_epi_flag_batch, flag_address) + + @cute.jit + def __call__( + self, + work_tile_info: SwapAbFc12WorkTileInfo, + tmem_acc_tensor: cute.Tensor, # (cta_tile_m, cta_tile_n) + acc_pipeline, + acc_consumer_state, + is_odd_turn: cutlass.Int32, + ): + # (tokens_this_expert, intermediate_down, 1) + real_fc1_output, _ = self.kernel_extension.get_gmem_tensor("c", self.fc1_output, work_tile_info) + # (tokens_this_expert, intermediate_down, 1) + real_fc1_output_sf, _ = self.kernel_extension.get_gmem_tensor("sfc", self.fc1_output_sf, work_tile_info) + # subtile-irrevalent hoist out here. + if cutlass.const_expr(self.optional_epi_args.fc1_alpha is not None): + alpha_val = self.optional_epi_args.fc1_alpha[work_tile_info.expert_idx] + else: + alpha_val = None + if cutlass.const_expr(self.optional_epi_args.fc1_norm_const is not None): + norm_const = self.optional_epi_args.fc1_norm_const[work_tile_info.expert_idx] + else: + norm_const = None + # (cta_tile_m, cta_tile_n) -> (epi_tile_m, epi_tile_n, iters) + tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( + tmem_acc_tensor, (self._EpilogueFc1IntermediateGateUpTileSize, self._EpilogueTokenTileSize) + )[None, None, 0, None] + + acc_pipeline.consumer_wait(acc_consumer_state) + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + + # Overlap path preloads two subtiles before releasing acc TMEM. + unroll_tile_cnt = 2 if cutlass.const_expr(self.overlapping_accum) else 0 + remain_subtile_cnt = self.subtile_cnt - unroll_tile_cnt + + if cutlass.const_expr(unroll_tile_cnt > 0): + subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + subtile_idx_second = (cutlass.Int32(self.subtile_cnt + 1) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + + # preload_subtile_first: subtile_idx_first's raw PRE-transpose acc, LDTM'd by + # all 128 epi threads into 4 reg tensors == the 4 quadrants of the subtile's + # (128 tmem_dp x 64 tmem_col) footprint. Only these raw-TMEM offsets are + # guaranteed: + # reg[0]/reg[1], reg[2]/reg[3] : top vs bot -> 16 apart in tmem_dp + # reg[0]/reg[2], reg[1]/reg[3] : 1st vs 2nd half -> 32 apart in tmem_col + # (so reg[0..1] = the first 128x32, reg[2..3] = the second 128x32 of the 128x64.) + # The per-lane (lane_idx, elem_idx) -> (tmem_dp, tmem_col) layout INSIDE each + # reg tensor is opaque -- do not assume it; it only becomes well-defined once + # the tmem transpose consumes them. + preload_subtile_first: Tuple[cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor] = ( + _TmemTranspose16x32Core.load_subtile_raw_acc( + tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_first] + ) + ) + + # Release acc to next MMA unconditionally. + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + # preload_subtile_second: same 128 tmem_dp x 64 tmem_col footprint, but for + # subtile_idx_second (the other token subtile, not the 2nd col-half). Same + # quadrant/offset invariants and opaque per-lane layout as above. + preload_subtile_second: Tuple[cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor] = ( + _TmemTranspose16x32Core.load_subtile_raw_acc( + tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second] + ) + ) + + # Both unrolled subtiles borrow tmem_subtile_second as workspace. + preload_pair = (preload_subtile_first, preload_subtile_second) + subtile_idx_pair = (subtile_idx_first, subtile_idx_second) + for i in cutlass.range_constexpr(unroll_tile_cnt): + if subtile_idx_pair[i] * cutlass.Int32(self._EpilogueTokenTileSize) < valid_tokens: + self.run_subtile( + work_tile_info=work_tile_info, + subtile_idx=subtile_idx_pair[i], + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second], + preload_acc=preload_pair[i], + fc1_output=real_fc1_output, + fc1_output_sf=real_fc1_output_sf, + alpha_val=alpha_val, + norm_const=norm_const, + ) + + for i in cutlass.range(remain_subtile_cnt, unroll=1): + real_i = i + unroll_tile_cnt + if cutlass.const_expr(self.overlapping_accum): + subtile_idx = (cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + else: + subtile_idx = cutlass.Int32(real_i) + + if subtile_idx * cutlass.Int32(self._EpilogueTokenTileSize) < valid_tokens: + self.run_subtile( + work_tile_info=work_tile_info, + subtile_idx=subtile_idx, + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx], + preload_acc=None, + fc1_output=real_fc1_output, + fc1_output_sf=real_fc1_output_sf, + alpha_val=alpha_val, + norm_const=norm_const, + ) + + # Non-overlap-path release: at the natural task-tile boundary. + if cutlass.const_expr(not self.overlapping_accum): + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + @cute.jit + def run_subtile( + self, + work_tile_info: SwapAbFc12WorkTileInfo, + subtile_idx: cutlass.Int32, + # (intermedaite_gateup_tile, token_subtile), fundamentally (epi_tile_m, epi_tile_n) + tmem_subtile_tensor: cute.Tensor, + # Rmems preloaded from tmem, contract with downstream tmem trans. Do not assume mapping here. + preload_acc: Tuple[cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor], + # (tokens_this_expert, intermediate_down, 1) + fc1_output: cute.Tensor, + fc1_output_sf: cute.Tensor, + alpha_val: Optional[cutlass.Float32], + norm_const: Optional[cutlass.Float32], + ): + if cutlass.const_expr(self.optional_epi_args.topk_scores is not None): + # This means we need to perform DeepGEMM computation graph, topk_score at fc1 pre-quant + topk_score_tensor, _ = self.kernel_extension.get_gmem_tensor( + "topk", self.optional_epi_args.topk_scores, work_tile_info + ) # (tokens_this_expert) + else: + topk_score_tensor = None + + # Mapping of the transposed accumulator for orthodox NVFP4 output: + # (epi_tid, val_id) -> (token_idx, intermediate_down_idx) + # token_idx = epi_tid % 32 + val_id // 16 * 32 + # intermediate_down_idx = val_id % 16 + epi_tid // 32 * 16 + # Each thread holds (intermediate_down_16, token_2):(1, 16) + + # Step -1: preload topk scores. + current_two_token_idices = ( + work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + self.lane_idx, + work_tile_info.tile_n_idx * self.cta_tile_n + + subtile_idx * self._EpilogueTokenTileSize + + self.lane_idx + + 32, + ) + if cutlass.const_expr(topk_score_tensor is not None): + topk_scores = ( + topk_score_tensor[current_two_token_idices[0]], + topk_score_tensor[current_two_token_idices[1]], + ) + else: + topk_scores = None + + # Step 0: load tmem + if cutlass.const_expr(preload_acc is not None): + gate_token_0_32, up_token_0_32, gate_token_32_64, up_token_32_64 = preload_acc + else: + gate_token_0_32 = cute.make_rmem_tensor((16,), cutlass.Float32) + up_token_0_32 = cute.make_rmem_tensor((16,), cutlass.Float32) + gate_token_32_64 = cute.make_rmem_tensor((16,), cutlass.Float32) + up_token_32_64 = cute.make_rmem_tensor((16,), cutlass.Float32) + # Although hardcode is not right, but since the whole tmem transpose is too tricky, I have to hardcode... + # (epi_tile_m, epi_tile_n) -> (warp_local_epi_tile_m, epi_tile_n) + # tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, self.warp_idx), None] + tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, 0), None] + # (warp_local_epi_tile_m, epi_tile_n) -> (((16, 32), 1), (2, 2)) + tmem_subtile_tensor_in_first_load_view = cute.logical_divide( + cute.zipped_divide(tmem_subtile_tensor_per_warp, (16, 32)), ((16, 32), 1) + ) + atom = cute.make_copy_atom(tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), cutlass.Float32) + cute.copy( + atom, + wrap_into_copy_standard_layout(tmem_subtile_tensor_in_first_load_view[None, 0]), + wrap_into_copy_standard_layout(gate_token_0_32), + ) + cute.copy( + atom, + wrap_into_copy_standard_layout(tmem_subtile_tensor_in_first_load_view[None, 1]), + wrap_into_copy_standard_layout(up_token_0_32), + ) + cute.copy( + atom, + wrap_into_copy_standard_layout(tmem_subtile_tensor_in_first_load_view[None, 2]), + wrap_into_copy_standard_layout(gate_token_32_64), + ) + cute.copy( + atom, + wrap_into_copy_standard_layout(tmem_subtile_tensor_in_first_load_view[None, 3]), + wrap_into_copy_standard_layout(up_token_32_64), + ) + + # Step 1: perform swiglu on the first part, interleave with the second's 32x32 tmem transpose. + token_0_32_pre_quant_pre_trans = self.alpha_swiglu_clamp(gate_token_0_32, up_token_0_32, alpha_val) + + # gate_token_32_64 / up_token_32_64 are already in the transpose input + # distribution (see TmemTranspose16x32 / load_subtile_raw_acc). + token_32_64_tmem_trans = TmemTranspose32x32Inplace( + tmem_subtile_tensor.iterator, reg_tensor_top=gate_token_32_64, reg_tensor_bot=up_token_32_64 + ) + + # Transpose output: each lane holds (token_1, intermediate_16); tmem_dp + # = lane_idx (token), tmem_col = elem_idx (intermediate output idx). + gate_token_32_64_trans_pre_act, up_token_32_64_trans_pre_act = ( + token_32_64_tmem_trans.from_r1_perm_until_last_store() + ) + + token_32_64_pre_quant = self.alpha_swiglu_clamp( + gate_token_32_64_trans_pre_act, up_token_32_64_trans_pre_act, alpha_val + ) + + token_0_32_tmem_trans = TmemTranspose16x32( + tmem_subtile_tensor.iterator, Region.Top, reg_tensor=token_0_32_pre_quant_pre_trans + ) + token_0_32_pre_quant = token_0_32_tmem_trans.from_r1_perm_until_last_store() + + # Step 2: Quant + self.fc1_quant( + work_tile_info=work_tile_info, + two_token=(token_0_32_pre_quant, token_32_64_pre_quant), + topk_scores=topk_scores, + norm_const=norm_const, + intermediate_output_size=cute.size(fc1_output, 1), + fc1_output_sf=fc1_output_sf, + subtile_idx=subtile_idx, + ) + + # Step 3: TMASTG + # (token_64, intermeidate_64) + fc1_smem = self.smem_tensor[None, None, subtile_idx] + # (token, intermediate_down, l=1) -> (cta_token, cta_intermediate_down) + fc1_gmem_cta_view = cute.flat_divide(fc1_output, (self.cta_tile_n, self.cta_tile_m // 2))[ + None, None, work_tile_info.tile_n_idx, work_tile_info.tile_m_idx, 0 + ] + # (cta_token, cta_intermediate_down) -> (token_64, intermediate_64) + fc1_gmem_subtile_view = cute.flat_divide( + fc1_gmem_cta_view, (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize) + )[None, None, subtile_idx, 0] + tma_smem_src, tma_gmem_dst = cpasync.tma_partition( + self.fc1_tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(fc1_smem, 0, 2), + cute.group_modes(fc1_gmem_subtile_view, 0, 2), + ) + + subtile_bar_id = subtile_idx + cutlass.Int32(SwapABGatedActEpilogue._EpilogueAsyncBarIdBase) + tma_ready_to_read_smem_named_barrier = pipeline.NamedBarrier( + barrier_id=subtile_bar_id, num_threads=self._EpilogueWarpCnt * 32 + ) + cute.arch.fence_proxy("async.shared", space="cta") + if self.warp_idx == subtile_idx: + tma_ready_to_read_smem_named_barrier.arrive_and_wait() + with cute.arch.elect_one(): + # if work_tile_info.tile_m_idx * (self.cta_tile_m // 2) < cute.size(fc1_output, 1): + cute.copy(self.fc1_tma_atom, tma_smem_src, tma_gmem_dst) + cute.arch.cp_async_bulk_commit_group() + else: + tma_ready_to_read_smem_named_barrier.arrive() + + @cute.jit + def alpha_swiglu_clamp( + self, + gate_rmem: cute.Tensor, # Raw fc1 acc (pre-dequant); even-size 1D fp32 rmem + up_rmem: cute.Tensor, # Raw fc1 acc (pre-dequant); even-size 1D fp32 rmem + alpha_val: Optional[cutlass.Float32], + ) -> cute.Tensor: + # ── Input contract checks (compile-time): fp32, 1D, even-count, rmem ── + # Wrapped in const_expr so the DSL evaluates them at trace time and the + # raise fires during compilation rather than emitting a runtime branch. + for _name, _t in (("gate_rmem", gate_rmem), ("up_rmem", up_rmem)): + if cutlass.const_expr(_t.element_type is not cutlass.Float32): + raise TypeError(f"alpha_swiglu_clamp: {_name} must be Float32, got {_t.element_type}") + if cutlass.const_expr(_t.memspace != AddressSpace.rmem): + raise ValueError( + f"alpha_swiglu_clamp: {_name} must be a register (rmem) tensor, got address space {_t.memspace}" + ) + if cutlass.const_expr(cute.rank(_t) != 1): + raise ValueError(f"alpha_swiglu_clamp: {_name} must be 1D, got rank {cute.rank(_t)}") + if cutlass.const_expr(cute.size(_t) % 2 != 0): + raise ValueError(f"alpha_swiglu_clamp: {_name} element count must be even, got {cute.size(_t)}") + if cutlass.const_expr(cute.size(gate_rmem) != cute.size(up_rmem)): + raise ValueError( + "alpha_swiglu_clamp: gate_rmem and up_rmem must have equal size, got " + f"{cute.size(gate_rmem)} vs {cute.size(up_rmem)}" + ) + + # gate_rmem / up_rmem are the RAW fc1 fp32 accumulator (pre-dequant). + # Order follows the NVFP4 -> fp32 -> SwiGLU contract and MUST be: + # + # 1. dequant: gate = alpha * gate_raw ; up = alpha * up_raw + # (alpha = expert-wise global scale on the acc; None => alpha == 1.) + # 2. clamp the DEQUANTED (real) values, gpt-oss ``_apply_gate`` style: + # gate = min(gate, +limit) (upper bound only) + # up = clamp(up, -limit, +limit) (symmetric) + # 3. gated activation, either SwiGLU or SiTU (mutually exclusive; SiTU has no clamp): + # SwiGLU: out = up * gate * sigmoid(gate) + # SiTU: out = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) + # sigmoid(x) = rcp(1 + exp2(-x * log2e)) + # + # The symmetric up-clamp is a single ``min.xorsign.abs.f32`` (magnitude + # min(|up|, limit), sign = sign(up)^sign(limit) = sign(up) since limit>=0); + # the gate-clamp is a plain ``min.f32``. ``.xorsign.abs`` has no f32x2 form, + # so dequant+clamp run scalar while the swiglu core stays packed f32x2. + n = cute.size(gate_rmem) + out = cute.make_rmem_tensor((n,), cutlass.Float32) + log2_e = 1.4426950408889634 + + neg_log2e_pair = (cutlass.Float32(-log2_e), cutlass.Float32(-log2_e)) + one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) + if cutlass.const_expr(self.gate_up_clamp is not None): + limit = cutlass.Float32(self.gate_up_clamp) + + for i in cutlass.range_constexpr(0, n, 2): + g0 = gate_rmem[i] + g1 = gate_rmem[i + 1] + u0 = up_rmem[i] + u1 = up_rmem[i + 1] + + # 1) dequant raw acc to real values (skip entirely when alpha is None). + if cutlass.const_expr(alpha_val is not None): + alpha_pair = (alpha_val, alpha_val) + g0, g1 = cute.arch.mul_packed_f32x2((g0, g1), alpha_pair) + u0, u1 = cute.arch.mul_packed_f32x2((u0, u1), alpha_pair) + + # 2) clamp the real values (skip when no clamp configured). + if cutlass.const_expr(self.gate_up_clamp is not None): + # gate upper-clamp: min(gate, +limit) + g0 = cutlass.Float32( + llvm.inline_asm( + cutlass.Float32.mlir_type, + [g0.ir_value(), limit.ir_value()], + "min.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + g1 = cutlass.Float32( + llvm.inline_asm( + cutlass.Float32.mlir_type, + [g1.ir_value(), limit.ir_value()], + "min.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + # up symmetric-clamp: clamp(up, -limit, +limit) in one instruction + u0 = cutlass.Float32( + llvm.inline_asm( + cutlass.Float32.mlir_type, + [u0.ir_value(), limit.ir_value()], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + u1 = cutlass.Float32( + llvm.inline_asm( + cutlass.Float32.mlir_type, + [u1.ir_value(), limit.ir_value()], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + # 3) gated activation on the dequanted (and clamped) real values. + # sigmoid(x) = rcp(1 + exp2(-x * log2e)) -- shared by both cores. + def _sigmoid(p0, p1): + neg = cute.arch.mul_packed_f32x2((p0, p1), neg_log2e_pair) + e = (cute.math.exp2(neg[0], fastmath=True), cute.math.exp2(neg[1], fastmath=True)) + d = cute.arch.add_packed_f32x2(e, one_pair) + return (cute.arch.rcp_approx(d[0]), cute.arch.rcp_approx(d[1])) + + sigmoid_pair = _sigmoid(g0, g1) + + if cutlass.const_expr(self.situ_beta is None): + # SwiGLU: out = up * gate * sigmoid(gate) + ug = cute.arch.mul_packed_f32x2((u0, u1), (g0, g1)) + out_pair = cute.arch.mul_packed_f32x2(ug, sigmoid_pair) + else: + # SiTU (Kimi K3), matching HF ``modeling_kimi.py``: + # situ_gate = beta * tanh(gate / beta) * sigmoid(gate) + # situ_up = linear_beta * tanh(up / linear_beta) + # out = situ_gate * situ_up + # + # ``tanh(z) = 2 * sigmoid(2z) - 1`` keeps the whole core on the packed f32x2 path -- + # there is no packed tanh, so calling one would force this loop back to scalar. + # + # beta * tanh(x/beta) = beta * (2*sigmoid(2x/beta) - 1) = 2*beta*sigmoid(2x/beta) - beta + # so the reciprocals and the 2*beta factors fold at trace time. + inv_2beta = cutlass.Float32(2.0 / self.situ_beta) + two_beta = cutlass.Float32(2.0 * self.situ_beta) + neg_beta = cutlass.Float32(-self.situ_beta) + inv_2lbeta = cutlass.Float32(2.0 / self.situ_linear_beta) + two_lbeta = cutlass.Float32(2.0 * self.situ_linear_beta) + neg_lbeta = cutlass.Float32(-self.situ_linear_beta) + + gs = _sigmoid(*cute.arch.mul_packed_f32x2((g0, g1), (inv_2beta, inv_2beta))) + tanh_g = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(gs, (two_beta, two_beta)), (neg_beta, neg_beta) + ) + + us = _sigmoid(*cute.arch.mul_packed_f32x2((u0, u1), (inv_2lbeta, inv_2lbeta))) + tanh_u = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(us, (two_lbeta, two_lbeta)), (neg_lbeta, neg_lbeta) + ) + + situ_gate = cute.arch.mul_packed_f32x2(tanh_g, sigmoid_pair) + out_pair = cute.arch.mul_packed_f32x2(situ_gate, tanh_u) + + out[i] = out_pair[0] + out[i + 1] = out_pair[1] + + return out + + @cute.jit + def fc1_quant( + self, + work_tile_info: SwapAbFc12WorkTileInfo, + two_token: Tuple[cute.Tensor, cute.Tensor], # two rmem tensor, each fp32 @ (token_1, intermediate_16) + topk_scores: Optional[Tuple[cutlass.Float32, cutlass.Float32]], + norm_const: Optional[cutlass.Float32], + intermediate_output_size: cutlass.Int32, + fc1_output_sf: cute.Tensor, # MoE domain (token_this_rank, intermediate_down, 1) + subtile_idx: cutlass.Int32, + ): + # ``two_token`` are the two post-swiglu, transposed token rmem tensors; each lane holds one + # token's 16 intermediate-output values. half 0 -> token (lane), half 1 -> (lane+32). + # + # Those 16 values are a whole scale block at sf_vec 16 (nvfp4) but only half of one at + # sf_vec 32 (mx), where the other half sits in the paired warp -- hence the two directions + # below. Both are handed the two tokens in one call so the paired path needs a single + # exchange barrier per subtile rather than one per token. + # + # Per token (ported from PostSwigluHalf._gen_sfc_quantize + stg_sfc + r2s): + # 1. (Path A) pre-multiply topk weight into the values, if present. + # 2. block quant -> data regs + one scale factor (QuantImpl's job). + # 3. write the scale factor to fc1_output_sf[token, intermediate_idx, 0] + # (plain scalar store; predicated unless statically in-bound). + # 4. STS the quantized values into this subtile's shared output stage. + # norm_const is treated like alpha_val: None => behaves as 1.0 (factors const-elided, not + # multiplied by 1.0). It only exists for nvfp4; an e8m0 scale absorbs the rescale itself. + values_per_token = cute.size(two_token[0]) + if cutlass.const_expr(self.needs_pair_amax_exchange): + # The exchange slots live in THIS subtile's staging stage. That stage is dead right + # now -- the previous tile's copy of it was drained at the tile boundary, and this + # tile writes it only after the retire barrier below -- whereas every other stage may + # still have a TMA store reading it. Borrowing any other stage would corrupt it. + amax_exchange = mark_alignment( + cute.make_tensor( + cute.recast_ptr(self.staging_pointer, dtype=cutlass.Float32) + + subtile_idx * cutlass.Int32(self.fc1_staging_stage_bytes // 4), + cute.make_layout( + (self.fc1_amax_token_chunks, self._EpilogueWarpCnt * 32), stride=(1, self.fc1_amax_token_chunks) + ), + ), + 16, + ) + quant = QuantImpl( + self.quant_kind, + "regs_in_pair_threads", + lane_idx=self.lane_idx, + warp_idx=self.warp_idx, + pair_exchange_barrier=SwapABGatedActEpilogue.epilogue_sync_barrier(), + ) + else: + amax_exchange = None + quant = QuantImpl(self.quant_kind, "regs_in_thread") + + # Both warps of a pair derive the same scale, so only the even one stores it. The scale + # plane folds any coordinate inside a block onto that block's slot, so the per-warp base + # addresses the right entry for either vec size. + intermediate_idx = ( + work_tile_info.tile_m_idx * (self.cta_tile_m // 2) + + self.warp_idx * self._EpilogueFc1IntermediateDownPerWarp + ) + if cutlass.const_expr(self.needs_pair_amax_exchange): + stores_scale_factor = self.warp_idx % 2 == 0 + else: + stores_scale_factor = True + subtile_token_start = work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + token_idx_pair = (subtile_token_start + self.lane_idx, subtile_token_start + self.lane_idx + 32) + + # This subtile's (token, intermediate) shared output stage, tiled into (1, 16) blocks so + # each thread's cells slice out directly (zipped_divide + slice; avoids the ambiguous + # local_tile surface). + smem_stage = self.smem_tensor[None, None, subtile_idx] + # (token_64, intermediate_down_64) -> ((1, 16), (token_tile_size, warp_cnt)) + smem_tiled = cute.zipped_divide(smem_stage, (1, values_per_token)) + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.fc1_output_dtype, + num_bits_per_copy=values_per_token * self.fc1_output_dtype.width, + ) + + # 1) topk-weight pre-multiply (Path A) into a weighted scratch, both tokens back to back. + weighted = cute.make_rmem_tensor((2 * values_per_token,), cutlass.Float32) + for half in cutlass.range_constexpr(2): + tok = two_token[half] + base = half * values_per_token + if cutlass.const_expr(topk_scores is not None): + topk_pair = (topk_scores[half], topk_scores[half]) + for i in cutlass.range_constexpr(0, values_per_token, 2): + w0, w1 = cute.arch.mul_packed_f32x2((tok[i], tok[i + 1]), topk_pair) + weighted[base + i] = w0 + weighted[base + i + 1] = w1 + else: + for i in cutlass.range_constexpr(0, values_per_token): + weighted[base + i] = tok[i] + + # 2) Core block quant. One scale factor per token either way; the paired direction spends + # its exchange barrier inside this call. + data_regs, sf_regs = quant(weighted, norm_const=norm_const, smem_intermediate=amax_exchange) + data_by_token = cute.zipped_divide(data_regs, (values_per_token,)) + + # Retire the exchange: its slots are this stage's own bytes, so no thread may start + # writing the stage until every thread has read its partner's amax. Sitting after the + # quantization rather than before it puts the sync latency behind work that is already in + # flight (the e8m0 path alone goes through MUFU). + if cutlass.const_expr(self.needs_pair_amax_exchange): + SwapABGatedActEpilogue.epilogue_sync_barrier().arrive_and_wait() + + for half in cutlass.range_constexpr(2): + # 3) scale-factor store (predicate const-elided when statically in-bound, mirroring + # signal_fc1_done's intermediate predicate). + if stores_scale_factor: + if cutlass.const_expr( + self.intermediate_downproj is None + or self.intermediate_downproj % self.cluster_tile_intermediate_downproj != 0 + ): + if intermediate_idx < intermediate_output_size: + fc1_output_sf[token_idx_pair[half], intermediate_idx, 0] = sf_regs[half] + else: + fc1_output_sf[token_idx_pair[half], intermediate_idx, 0] = sf_regs[half] + + # 4) STS the quantized values into this subtile's shared output stage. + # ((1, 16), (token_tile_size, warp_cnt)) -> (16) + smem_thread_row = smem_tiled[(0, None), (self.lane_idx + 32 * half, self.warp_idx)] + cute.copy(store_atom, cute.coalesce(data_by_token[None, half]), cute.coalesce(smem_thread_row)) + + +@dataclasses.dataclass(frozen=True) +class Fc2ProcessPipeline: + tmem_acc_load: Callable + f2fp: Callable + post_f2fp_reorder: Callable + store_function: Callable + # Kept as a finer-grained, elem-level reading aid for the store-out layout + # (never evaluated); ``store_out_mapping`` is the per-issue form that the + # router actually evaluates at runtime to drive metadata / pointer math. + fc2_cta_tile_mapping: FunctionMapping + store_out_mapping: FunctionMapping + require_tmem_trans: bool + # SF plane per-issue mapping; None for the bf16 (unquantized) paths. + sf_store_out_mapping: Optional[FunctionMapping] = None + + +# Device only object +class SwapABFc2Epilogue(_ImmutableAfterInit): + def __init__( + self, + base: SwapABGatedActEpilogue, + tidx: cutlass.Int32, + smem_tensor: Optional[cute.Tensor], + tma_atom_fc2_output: Optional[cute.CopyAtom], + fc2_tma_output: Optional[cute.Tensor], + fc2_output: cute.Tensor, # MoE domain (token, topk, hidden) + token_src_metadata: Optional[cute.Tensor], + fc2_done_counter: Optional[cute.Tensor], + fc2_output_sf: Optional[cute.Tensor], + peer_rank_ptr_mapper: Optional[SymmetricBufferDevice], + optional_epi_args: GatedActEpilogueArgs, + ): + self.base = base + self.tidx = tidx % (base._EpilogueWarpCnt * 32) + self.warp_idx = self.tidx // 32 + self.lane_idx = self.tidx % 32 + self.tma_atom_fc2_output = tma_atom_fc2_output + self.fc2_tma_output = fc2_tma_output + self.fc2_output = fc2_output + self.token_src_metadata = token_src_metadata + self.fc2_done_counter = fc2_done_counter + self.fc2_output_sf = fc2_output_sf + self.peer_rank_ptr_mapper = peer_rank_ptr_mapper + self.optional_epi_args = optional_epi_args + if cutlass.const_expr(base.fc2_use_tma): + self.smem_tensor = smem_tensor + self.process_pipeline = make_fc2_tma_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ) + elif cutlass.const_expr(base.fc2_use_ublk): + self.smem_tensor = smem_tensor + self.process_pipeline = make_fc2_ublk_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ) + else: + self.smem_tensor = None + if cutlass.const_expr(base.reduce_topk_in_epilogue): + self.process_pipeline = make_fc2_redg_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ) + else: + self.process_pipeline = make_fc2_stg_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ) + self._freeze() + + def __getattr__(self, name): + return getattr(object.__getattribute__(self, "base"), name) + + def __extract_mlir_values__(self) -> List[ir.Value]: + # See SwapABFc1Epilogue.__extract_mlir_values__: this helper carries + # only loop-invariant Python context. It intentionally serializes no + # MLIR values, so changing it to store loop-carried state would be a + # correctness bug. + return [] + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "SwapABFc2Epilogue": + assert len(values) == 0 + return self + + @cute.jit + def signal_fc2_done(self, work_tile_info, next_work_tile_info, flag_tracker): + publish: cutlass.Constexpr = self.token_back_enabled + flag_address = Int64(0) + if cutlass.const_expr(publish): + flag_address = (self.fc2_done_counter.iterator + work_tile_info.expert_idx).toint() + no_fire: cutlass.Constexpr = not publish + return flag_tracker.accumulate(next_work_tile_info.phase, self.fc2_epi_flag_batch, flag_address, no_fire) + + @cute.jit + def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2OutputRouter": + task_tile_data_row_start = ( + work_tile_info.cumulative_data_physical_row + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) + ) + hidden_base_this_cta_tile = work_tile_info.tile_m_idx * cutlass.Int32(self.cta_tile_m) + valid_hidden_this_cta_tile = cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile + if valid_hidden_this_cta_tile < 0: + valid_hidden_this_cta_tile = 0 + if valid_hidden_this_cta_tile > self._EpilogueFc2HiddenTileSize: + valid_hidden_this_cta_tile = self._EpilogueFc2HiddenTileSize + + metadata = None + peer_rank_ptr_mapper = None + data_token_base = task_tile_data_row_start + if cutlass.const_expr(self.token_src_metadata is not None and not self.token_back_push_data): + metadata = cute.domain_offset((task_tile_data_row_start,), self.token_src_metadata) + peer_rank_ptr_mapper = self.peer_rank_ptr_mapper + data_token_base = None + + if cutlass.const_expr(self.combine_format.is_quantized): + base_outputs = (self.fc2_output, self.fc2_output_sf) + token_bases = (data_token_base, task_tile_data_row_start) + output_mappings = (self.process_pipeline.store_out_mapping, self.process_pipeline.sf_store_out_mapping) + else: + base_outputs = self.fc2_output + token_bases = data_token_base + output_mappings = self.process_pipeline.store_out_mapping + + return Fc2OutputRouter( + metadata=metadata, + token_bases=token_bases, + base_outputs=base_outputs, + hidden_base_this_cta_tile=hidden_base_this_cta_tile, + peer_rank_ptr_mapper=peer_rank_ptr_mapper, + valid_tokens_this_cta_tile=work_tile_info.valid_tokens_in_cta_tile, + valid_hidden_this_cta_tile=valid_hidden_this_cta_tile, + reduce_topk_in_epilogue=self.reduce_topk_in_epilogue, + output_mappings=output_mappings, + epi_tid=self.tidx, + combine_format=self.combine_format, + ).prefetch() + + @cute.jit + def __call__( + self, + work_tile_info: SwapAbFc12WorkTileInfo, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + acc_consumer_state, + is_odd_turn: cutlass.Int32, + ): + # subtile-irrelevant hoist: fc2 alpha scales raw fc2 accumulators before f2fp. + if cutlass.const_expr(self.optional_epi_args.fc2_alpha is not None): + alpha_val = self.optional_epi_args.fc2_alpha[work_tile_info.expert_idx] + else: + alpha_val = None + acc_ready = False + if not work_tile_info.peek_ready: + acc_ready = True + acc_pipeline.consumer_wait(acc_consumer_state) + fc2_output_router = self._make_output_router(work_tile_info) + # (cta_tile_m, cta_tile_n) -> (epi_tile_m, epi_tile_n, iters) + tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( + tmem_acc_tensor, (self._EpilogueFc2HiddenTileSize, self._EpilogueTokenTileSize) + )[None, None, 0, None] + + acc_pipeline.consumer_wait(acc_consumer_state, acc_ready) + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + + # Overlap path preloads two subtiles before releasing acc TMEM. + unroll_tile_cnt = ( + 2 if cutlass.const_expr(self.overlapping_accum and self.process_pipeline.require_tmem_trans) else 0 + ) + remain_subtile_cnt = self.subtile_cnt - unroll_tile_cnt + + if cutlass.const_expr(unroll_tile_cnt > 0): + subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + subtile_idx_second = (cutlass.Int32(self.subtile_cnt + 1) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + + # Each warp preloads its local 32-hidden x 64-token accumulator as two + # (32,) FP32 tensors: + # + # preload[0]: hidden 0..15 x token 0..63 + # preload[1]: hidden 16..31 x token 0..63 + # + # One LDTM.16dp256bit.x8 produces each tensor. Within it, every + # consecutive register pair holds adjacent token columns at one + # hidden coordinate. After F2FP, that pair is one BF16x2 MOVM input. + preload_subtile_first: Tuple[cute.Tensor, ...] = self.process_pipeline.tmem_acc_load( + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_first], epi=self + ) + + # Release acc to next MMA unconditionally. + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + # Same two-half mapping for the other token subtile. Both preloaded + # subtiles can now use this second subtile's TMEM as transpose workspace. + preload_subtile_second: Tuple[cute.Tensor, ...] = self.process_pipeline.tmem_acc_load( + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second], epi=self + ) + + # Both unrolled subtiles borrow tmem_subtile_second as workspace. + preload_pair = (preload_subtile_first, preload_subtile_second) + subtile_idx_pair = (subtile_idx_first, subtile_idx_second) + for i in cutlass.range_constexpr(unroll_tile_cnt): + if subtile_idx_pair[i] * cutlass.Int32(self._EpilogueTokenTileSize) < valid_tokens: + self.run_subtile( + work_tile_info=work_tile_info, + epilogue_iter_idx=cutlass.Int32(i), + subtile_idx=subtile_idx_pair[i], + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second], + preload_acc=preload_pair[i], + fc2_output_router=fc2_output_router, + alpha_val=alpha_val, + release_after_ldtm=False, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + ) + + if cutlass.const_expr(self.overlapping_accum and unroll_tile_cnt == 0): + release_after_ldtm = True + else: + release_after_ldtm = False + for i in cutlass.range(remain_subtile_cnt, unroll=1): + # for i in cutlass.range_constexpr(remain_subtile_cnt): + real_i = i + unroll_tile_cnt + if cutlass.const_expr(self.overlapping_accum): + subtile_idx = (cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + else: + subtile_idx = cutlass.Int32(real_i) + + if subtile_idx * cutlass.Int32(self._EpilogueTokenTileSize) < valid_tokens: + self.run_subtile( + work_tile_info=work_tile_info, + epilogue_iter_idx=real_i, + subtile_idx=subtile_idx, + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx], + preload_acc=None, + fc2_output_router=fc2_output_router, + alpha_val=alpha_val, + release_after_ldtm=release_after_ldtm, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + ) + release_after_ldtm = False + + # Non-overlap-path release: at the natural task-tile boundary. + if cutlass.const_expr(not self.overlapping_accum): + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + @cute.jit + def run_subtile( + self, + work_tile_info: SwapAbFc12WorkTileInfo, + epilogue_iter_idx: cutlass.Int32, + subtile_idx: cutlass.Int32, + # (hidden_tile, token_subtile), fundamentally (epi_tile_m, epi_tile_n) + tmem_subtile_tensor: cute.Tensor, + preload_acc: Optional[Tuple[cute.Tensor, ...]], + fc2_output_router: "Fc2OutputRouter", + alpha_val: Optional[cutlass.Float32], + release_after_ldtm: Union[cutlass.Boolean, bool], + acc_pipeline, + acc_consumer_state, + ): + process_pipeline = self.process_pipeline + if cutlass.const_expr(preload_acc is None): + loaded = process_pipeline.tmem_acc_load(tmem_subtile_tensor=tmem_subtile_tensor, epi=self) + if release_after_ldtm: + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + else: + loaded = preload_acc + + casted = process_pipeline.f2fp(*loaded, alpha_val=alpha_val) + # reorder returns a bare RMEM fragment in the store's expected pre-store + # distribution; reorder + store are paired 1:1 inside the pipeline. + pre_store = process_pipeline.post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_tensor) + process_pipeline.store_function( + epi=self, + subtile=pre_store, + work_tile_info=work_tile_info, + epilogue_iter_idx=epilogue_iter_idx, + subtile_idx=subtile_idx, + fc2_output_router=fc2_output_router, + ) + + +@dataclasses.dataclass(frozen=True) +class Fc2OutputRouter: + # One packed i64 TokenSrcMetadata record per pool token. None means a local write. + metadata: Optional[cute.Tensor] + # token + possible sf + token_bases: Union[Tuple[Optional[cutlass.Int32], cutlass.Int32], Optional[cutlass.Int32]] + base_outputs: Union[Tuple[cute.Tensor, cute.Tensor], cute.Tensor] # (token, topk, hidden) + hidden_base_this_cta_tile: Union[cutlass.Int32, int] + peer_rank_ptr_mapper: Optional[SymmetricBufferDevice] + valid_tokens_this_cta_tile: cutlass.Int32 + valid_hidden_this_cta_tile: Union[cutlass.Int32, int] + reduce_topk_in_epilogue: bool + # Per-issue (epi_tid, iter_idx) -> (token_cta_tile, hidden_cta_tile). Data + # mapping, or (data mapping, sf mapping) when quantized. + output_mappings: Union[Tuple[FunctionMapping, FunctionMapping], FunctionMapping] + epi_tid: cutlass.Int32 + combine_format: CombineFormat + # After metadata prefetch + dst_ptrs: Optional[cute.Tensor] = None # i64 x (copy_iters_this_thread_cta_tile), fundamentally the pointers. + valid: Optional[cute.Tensor] = None # (copy_iters_this_thread_cta_tile) + + @property + def data_output(self) -> cute.Tensor: + return self.base_outputs[0] if isinstance(self.base_outputs, tuple) else self.base_outputs + + @property + def sf_output(self) -> Optional[cute.Tensor]: + # Present iff quantized; (pool_token, 1, hidden // sf_vec) rank-local. + return self.base_outputs[1] if isinstance(self.base_outputs, tuple) else None + + @property + def data_token_base(self) -> Optional[cutlass.Int32]: + return self.token_bases[0] if isinstance(self.token_bases, tuple) else self.token_bases + + @property + def sf_token_base(self) -> Optional[cutlass.Int32]: + return self.token_bases[1] if isinstance(self.token_bases, tuple) else None + + @property + def data_mapping(self) -> FunctionMapping: + return self.output_mappings[0] if isinstance(self.output_mappings, tuple) else self.output_mappings + + @property + def sf_mapping(self) -> Optional[FunctionMapping]: + return self.output_mappings[1] if isinstance(self.output_mappings, tuple) else None + + def __post_init__(self) -> None: + if (self.metadata is None) == (self.data_token_base is None): + raise ValueError("Fc2OutputRouter requires exactly one of metadata or a (data) token base.") + if (self.metadata is None) != (self.peer_rank_ptr_mapper is None): + raise ValueError("Fc2OutputRouter requires peer_rank_ptr_mapper iff metadata is set.") + if self.reduce_topk_in_epilogue and self.metadata is None: + raise ValueError("Fc2OutputRouter reduction requires metadata routing.") + + @cute.jit + def prefetch(self) -> "Fc2OutputRouter": + # Only the metadata (comm) path prefetches a pointer array: its + # metadata-derived address has long-latency LDGs worth issuing early. + # The local (no-comm) path computes its affine address on demand in + # get_dst() -- no array, hence no runtime-indexed local-memory spill. + if cutlass.const_expr(self.metadata is None): + return self + copy_iters: cutlass.Constexpr[int] = self.data_mapping.domain.axis_size("iter_idx") + + valid = cute.make_rmem_tensor((copy_iters,), cutlass.Int32) + dst_ptrs = cute.make_rmem_tensor((copy_iters,), cutlass.Int64) + + # Compiler should be able to optimize the same token_copy_group's offset add. (Fundamental cse + strength_reduce) + # We should check the SASS to ensure this happens. + for iter_idx in cutlass.range_constexpr(copy_iters): + coord = self.data_mapping.evaluate(epi_tid=self.epi_tid, iter_idx=iter_idx) + token_in_tile = cutlass.Int32(coord["token_in_cta_tile"]) + hidden_in_tile = cutlass.Int32(coord["hidden_in_cta_tile"]) + + valid[iter_idx] = cutlass.Int32(0) + dst_ptrs[iter_idx] = cutlass.Int64(0) + + token_valid = token_in_tile < self.valid_tokens_this_cta_tile + hidden_valid = hidden_in_tile < cutlass.Int32(self.valid_hidden_this_cta_tile) + if token_valid and hidden_valid: + valid[iter_idx] = cutlass.Int32(1) + if cutlass.const_expr(self.metadata is None): + dst_tokens = self.data_token_base + token_in_tile + dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile + # Int64 token coord: dst_tokens*K*H overflows int32 once + # T*K*H exceeds 2^31 (data_output is (token, topk, hidden)). + dst_ptrs[iter_idx] = self.data_output[Int64(dst_tokens), None, dst_hidden].iterator.toint() + + else: + md = TokenSrcMetadata.load( + self.metadata.iterator.toint() + Int64(token_in_tile) * Int64(TokenSrcMetadata.nbytes) + ) + dst_rank = md.src_rank + dst_token = md.src_token + dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile + if cutlass.const_expr(not self.reduce_topk_in_epilogue): + dst_topk = md.src_topk + else: + dst_topk = 0 + # Int64 token coord: domain_offset on (token, topk, hidden) + # computes dst_token*K*H, which overflows int32 once T*K*H > 2^31. + dst_ptrs[iter_idx] = self.peer_rank_ptr_mapper.map_pointer( + cute.domain_offset((Int64(dst_token), dst_topk, dst_hidden), self.data_output).iterator, + dst_rank, + byte_alignment=32, + ).toint() + + return dataclasses.replace(self, dst_ptrs=dst_ptrs, valid=valid) + + @cute.jit + def get_data_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, cutlass.Int32]: + """Per-issue DATA destination: gmem pointer + validity predicate. + + The router owns ``data_output`` so the caller never re-assembles a + pointer from a raw int; it just builds its own copy tensor (STG) or + feeds the pointer to inline asm (REDG/UBLK). + + Alignment is unified at 32 B: only STG feeds this pointer to a real + ``cute.copy`` (256 b vector store, genuinely 32 B aligned); REDG/UBLK + only ``ptrtoint`` it for inline-asm issue, where the hint is inert. + """ + if cutlass.const_expr(self.metadata is None): + # no-comm: on-demand affine address (no prefetched array). The + # invariant base hoists out of the caller's loop via CSE; a + # constexpr iter folds the per-issue offset into the store. + coord = self.data_mapping.evaluate(epi_tid=self.epi_tid, iter_idx=iter_idx) + token_in_tile = cutlass.Int32(coord["token_in_cta_tile"]) + hidden_in_tile = cutlass.Int32(coord["hidden_in_cta_tile"]) + pred = cutlass.Int32(0) + addr = cutlass.Int64(0) + if token_in_tile < self.valid_tokens_this_cta_tile and hidden_in_tile < cutlass.Int32( + self.valid_hidden_this_cta_tile + ): + pred = cutlass.Int32(1) + dst_tokens = self.data_token_base + token_in_tile + dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile + # Int64 token coord: dst_tokens*K*H overflows int32 once T*K*H > 2^31. + addr = self.data_output[Int64(dst_tokens), None, dst_hidden].iterator.toint() + else: + # comm: read the pointer / validity prefetched by prefetch(). + addr = self.dst_ptrs[iter_idx] + pred = self.valid[iter_idx] + ptr = cute.make_ptr(self.data_output.element_type, addr, AddressSpace.gmem, assumed_align=32) + return ptr, pred + + @cute.jit + def get_sf_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, cutlass.Int32]: + """Per-issue SF destination: rank-local gmem pointer + validity predicate. + + SF never goes to a peer (it is staged locally and pushed token-contiguously + by the dispatch / standalone warps), so this is always the affine local + address -- no metadata routing, no prefetch. ``sf_output`` is the broadcast + plane ``(pool_token, 1, (sf_vec, hidden//sf_vec)):(., ., (0, 1))``, so the + logical hidden coordinate folds to its scale block on indexing. + """ + coord = self.sf_mapping.evaluate(epi_tid=self.epi_tid, iter_idx=iter_idx) + token_in_tile = cutlass.Int32(coord["token_in_cta_tile"]) + hidden_in_tile = cutlass.Int32(coord["hidden_in_cta_tile"]) + pred = cutlass.Int32(0) + addr = cutlass.Int64(0) + if token_in_tile < self.valid_tokens_this_cta_tile and hidden_in_tile < cutlass.Int32( + self.valid_hidden_this_cta_tile + ): + pred = cutlass.Int32(1) + sf_row = self.sf_token_base + token_in_tile + sf_hidden = hidden_in_tile + self.hidden_base_this_cta_tile + addr = self.sf_output[Int64(sf_row), None, sf_hidden].iterator.toint() + # Per-block scale offsets are element-granular; claim the scale dtype's + # natural element alignment (e8m0 1 B / bf16 2 B). + sf_ptr = cute.make_ptr(self.sf_output.element_type, addr, AddressSpace.gmem, assumed_align=4) + return sf_ptr, pred + + +def make_fc2_stg_cta_store_out_mapping( + combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int +): + assert cta_hidden_tile_size == 128 + assert cta_token_tile_size % 64 == 0 + wire_dtype = combine_format.act_dtype + assert wire_dtype.width in (4, 8, 16), "fc2 STG wire dtype must be fp4/fp8/bf16." + elems_per_stg = min(256 // wire_dtype.width, 32) + stgs_per_hidden32 = 32 // elems_per_stg + fundamental_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "elem_idx"), (128, cta_token_tile_size)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) + ), + function=lambda epi_tid, elem_idx: { + "token_in_cta_tile": epi_tid % 32 + elem_idx // 32 * 32, + "hidden_in_cta_tile": elem_idx % 32 + epi_tid // 32 * 32, + }, + ) + store_out_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, stgs_per_hidden32 * cta_token_tile_size // 32)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) + ), + function=lambda epi_tid, iter_idx: { + "token_in_cta_tile": epi_tid % 32 + iter_idx // stgs_per_hidden32 * 32, + "hidden_in_cta_tile": (iter_idx % stgs_per_hidden32) * elems_per_stg + epi_tid // 32 * 32, + }, + ) + sf_store_out_mapping = None + if combine_format.is_quantized: + + def stg_sf_mapping(epi_tid, iter_idx): + lane = epi_tid % 32 + warp = epi_tid // 32 + return {"token_in_cta_tile": lane + iter_idx * 32, "hidden_in_cta_tile": warp * 32} + + sf_store_out_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, cta_token_tile_size // 32)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) + ), + function=stg_sf_mapping, + ) + return store_out_mapping, sf_store_out_mapping, fundamental_mapping + + +def make_fc2_redg_cta_store_out_mapping( + combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int +): + assert cta_hidden_tile_size == 128 + assert cta_token_tile_size % 64 == 0 + # In-kernel reduce is bf16-only and never quantized, so there is no SF plane. + assert combine_format.act_dtype.width == 16 + assert not combine_format.is_quantized + + fundamental_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "elem_idx"), (128, cta_token_tile_size)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) + ), + function=lambda epi_tid, elem_idx: { + "token_in_cta_tile": ( + ((elem_idx // 4) // 16) * 64 + + (((elem_idx // 4) % 16) // 8) * 32 + + (((elem_idx // 4) % 8) // 4) * 16 + + (((elem_idx // 4) % 4) % 2) * 8 + + (epi_tid % 32) // 4 + ), + "hidden_in_cta_tile": ( + (epi_tid // 32) * 32 + (epi_tid % 4) * 4 + (((elem_idx // 4) % 4) // 2) * 16 + elem_idx % 4 + ), + }, + ) + # SIMT REDG emits one 8B red.v2.bf16x2 per 4 hidden elements. Each + # 64-token subtile contributes two token rows per lane and 8 hidden + # segments per token row. + store_out_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, cta_token_tile_size // 64 * 16)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) + ), + function=lambda epi_tid, iter_idx: { + "token_in_cta_tile": ( + (iter_idx // 16) * 64 + + ((iter_idx % 16) // 8) * 32 + + ((iter_idx % 8) // 4) * 16 + + ((iter_idx % 4) % 2) * 8 + + (epi_tid % 32) // 4 + ), + "hidden_in_cta_tile": ((epi_tid // 32) * 32 + (epi_tid % 4) * 4 + ((iter_idx % 4) // 2) * 16), + }, + ) + return store_out_mapping, None, fundamental_mapping + + +def make_fc2_ublk_store_out_mapping(combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int): + assert cta_hidden_tile_size == 128 + assert cta_token_tile_size % 64 == 0 + # UBLK pushes whole hidden rows by byte count, so the token/hidden mapping + # is element-indexed and dtype-independent (wire dtype only sets copy bytes). + assert combine_format.act_dtype.width in (4, 8, 16), "fc2 UBLK wire dtype must be fp4/fp8/bf16." + assert cta_token_tile_size <= 256 + max_token_cta_tile = 256 + fundamental_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "elem_idx"), (128, cta_token_tile_size)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (max_token_cta_tile, cta_hidden_tile_size) + ), + function=lambda epi_tid, elem_idx: { + "token_in_cta_tile": elem_idx // cta_hidden_tile_size * 32 + + epi_tid % 8 + + epi_tid // 32 * 8 + + ((epi_tid % 32) // 8) * 64, + "hidden_in_cta_tile": elem_idx % cta_hidden_tile_size, + }, + ) + # One row per thread covers two 64-row subtiles per iteration. + copy_iters = (cta_token_tile_size + 127) // 128 + store_out_mapping = FunctionMapping( + domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, copy_iters)), + codomain=CoordinateSpace( + ("token_in_cta_tile", "hidden_in_cta_tile"), (max_token_cta_tile, cta_hidden_tile_size) + ), + function=lambda epi_tid, iter_idx: { + "token_in_cta_tile": (iter_idx * 128 + ((epi_tid % 32) // 16) * 64 + (epi_tid // 32) * 16 + epi_tid % 16), + "hidden_in_cta_tile": 0, + }, + ) + if combine_format.is_quantized: + # Quantized UBLK reuses the hidden-contiguous STG mappings for 128-bit R2S. + _, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size + ) + else: + sf_store_out_mapping = None + return store_out_mapping, sf_store_out_mapping, fundamental_mapping + + +# (...) -> ((atom_v, 1)) +@cute.jit +def wrap_into_copy_standard_layout(tensor: cute.Tensor): + tensor = cute.coalesce(cute.flatten(tensor)) + tensor = cute.append_ones(tensor, cute.rank(tensor) + 1) + tensor = cute.group_modes(tensor, 0, cute.rank(tensor) - 1) + tensor = cute.group_modes(tensor, 0, cute.rank(tensor)) + return tensor + + +@cute.jit +def fc2_f2fp(*tensors, alpha_val: Optional[cutlass.Float32] = None, **_) -> cute.Tensor: + reorder_dtype = cutlass.BFloat16 + total_size = 0 + for t in tensors: + total_size += cute.size(t) + converted_acc = cute.make_rmem_tensor((total_size,), reorder_dtype) + elems_processed = 0 + for t in tensors: + current_tensor_size = cute.size(t) + dst = cute.make_tensor(converted_acc.iterator + elems_processed, cute.make_layout((current_tensor_size,))) + if cutlass.const_expr(alpha_val is None): + dst.store(t.load().to(reorder_dtype)) + else: + if cutlass.const_expr(current_tensor_size % 2 != 0): + raise ValueError("fc2_f2fp expects even elements for each input tensor.") + scaled = cute.make_rmem_tensor((current_tensor_size,), cutlass.Float32) + for i in cutlass.range_constexpr(0, current_tensor_size, 2): + # scaled[i] = t[i] * alpha_val + s0, s1 = cute.arch.mul_packed_f32x2((t[i], t[i + 1]), (alpha_val, alpha_val)) + scaled[i] = s0 + scaled[i + 1] = s1 + dst.store(scaled.load().to(reorder_dtype)) + elems_processed += current_tensor_size + return converted_acc + + +@cute.jit +def post_f2fp_reorder_identity(*, casted: cute.Tensor, **_): + # UBLK: the f2fp output is already in the pre-store distribution (each lane + # owns one hidden element across the 64 subtile tokens); no reorder needed. + return casted + + +@cute.jit +def fc2_stg_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, **_): + atom_ld16x256 = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x8), cutlass.Float32) + ptr = tmem_subtile_tensor.iterator + half_lane_offset = 16 * TmemTranspose32x64B16Movm._tmem_row_stride + top_view = cute.make_tensor(ptr, TmemTranspose32x64B16Movm._tmem_layout(16, 64)) + bottom_view = cute.make_tensor(ptr + half_lane_offset, TmemTranspose32x64B16Movm._tmem_layout(16, 64)) + top = cute.make_rmem_tensor((32,), cutlass.Float32) + bottom = cute.make_rmem_tensor((32,), cutlass.Float32) + cute.copy(atom_ld16x256, top_view, TmemTranspose32x64B16Movm._rmem_copy_view(top, 32)) + cute.copy(atom_ld16x256, bottom_view, TmemTranspose32x64B16Movm._rmem_copy_view(bottom, 32)) + return top, bottom + + +@cute.jit +def fc2_ublk_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, epi, **_): + # UBLK consumes a warp-local 32-hidden x 64-token slice. The caller passes + # the CTA-level 128-hidden x 64-token subtile view, so select this epi + # warp's hidden block before issuing LDTM.x64. + tmem_subtile_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, epi.warp_idx), None] + raw_regs = cute.make_rmem_tensor((64,), cutlass.Float32) + atom_ld32x32_x64 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x64), cutlass.Float32) + cute.copy( + atom_ld32x32_x64, + wrap_into_copy_standard_layout(tmem_subtile_per_warp), + wrap_into_copy_standard_layout(raw_regs), + ) + return (raw_regs,) + + +@cute.jit +def fc2_stg_post_f2fp_reorder( + *, + casted: cute.Tensor, # (subtile_cnt,) + tmem_subtile_view: cute.Tensor, # (epi_tile_m, epi_tile_n) + **_, +): + if cutlass.const_expr(cute.size(casted) != 64): + raise NotImplementedError("fc2 stg pass expects 64 BF16 registers before store reorder.") + return TmemTranspose32x64B16Movm(tmem_subtile_view.iterator, casted)() + + +@cute.jit +def fc2_redg_post_f2fp_reorder(*, casted: cute.Tensor, tmem_subtile_view: cute.Tensor, **_): + # (epi_tid, elem_idx) -> (token_64, hidden_128), each thread hold token_2 x hidden_32 + natural = fc2_stg_post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_view) + core_matrix_reorder_sttm_atom = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + core_matrix_reorder_ldtm_atom = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), cutlass.Float32) + # ((16, 2), token_32_group) + token_groups = cute.logical_divide(cute.zipped_divide(natural, (32,)), (16, None)) + out = cute.make_rmem_tensor(token_groups.shape, casted.dtype) + out_as_i32 = cute.recast_tensor(out, cutlass.Float32) + # (32, 64) + tmem_warp = cute.flat_divide(tmem_subtile_view, (32, cute.size(tmem_subtile_view, 1)))[None, None, 0, 0] + # (16, 16, 16dp_group, token_32_groups). Note, this tmem can provide 2x cols since the original is bf16. + tmem_groups = cute.flat_divide(tmem_warp, (16, 16)) + for group_idx in cutlass.range_constexpr(cute.size(token_groups, 1)): + sttm_source = cute.recast_tensor(token_groups[None, group_idx], cutlass.Float32) + sttm_destination = tmem_groups[None, None, None, group_idx] + cute.copy( + core_matrix_reorder_sttm_atom, + wrap_into_copy_standard_layout(sttm_source), + wrap_into_copy_standard_layout(sttm_destination), + ) + cute.copy( + core_matrix_reorder_ldtm_atom, + wrap_into_copy_standard_layout(tmem_groups[None, None, 0, group_idx]), + wrap_into_copy_standard_layout(out_as_i32[(None, 0), group_idx]), + ) + cute.copy( + core_matrix_reorder_ldtm_atom, + wrap_into_copy_standard_layout(tmem_groups[None, None, 1, group_idx]), + wrap_into_copy_standard_layout(out_as_i32[(None, 1), group_idx]), + ) + return cute.coalesce(out) + + +@cute.jit +def fc2_stg_store_function( + *, + epi, + subtile: cute.Tensor, # Always BF16 before quantization + subtile_idx: cutlass.Int32, + fc2_output_router: Fc2OutputRouter, + **_, +): + if cutlass.const_expr(epi.combine_format.is_quantized): + data_subtile, sf_regs = QuantImpl(epi.combine_format, "regs_in_thread")(subtile) + else: + data_subtile = subtile + sf_regs = None + stg_width_elems: cutlass.Constexpr[int] = min(32, 256 // data_subtile.element_type.width) + stg_bits: cutlass.Constexpr[int] = stg_width_elems * data_subtile.element_type.width + copy_atom_vec = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=stg_bits) + elems_per_thread: cutlass.Constexpr[int] = cute.size(data_subtile) + if cutlass.const_expr(elems_per_thread % stg_width_elems != 0): + raise ValueError( + "fc2 STG store requires pre-store elems per thread to be divisible " + f"by STG issue width, got {elems_per_thread} and {stg_width_elems}." + ) + + if cutlass.const_expr(sf_regs is not None): + sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block + token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 + for token_group in cutlass.range_constexpr(token_groups_per_subtile): + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group + sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) + if sf_pred != cutlass.Int32(0): + sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) + for scale_idx in cutlass.range_constexpr(sf_scales_per_stg): + sf_dst[scale_idx] = sf_regs[token_group * sf_scales_per_stg + scale_idx] + + iters_per_subtile: cutlass.Constexpr[int] = elems_per_thread // stg_width_elems + copy_src = cute.zipped_divide(data_subtile, (stg_width_elems,)) + single_copy_layout = cute.make_layout(((stg_width_elems, 1),), stride=((1, 0),)) + subtile_iter_base = cutlass.Int32(subtile_idx) * cutlass.Int32(iters_per_subtile) + for local_iter in cutlass.range_constexpr(iters_per_subtile): + global_iter = subtile_iter_base + cutlass.Int32(local_iter) + dst_ptr, pred = fc2_output_router.get_data_dst(global_iter) + if pred != cutlass.Int32(0): + src_i = cute.make_tensor(copy_src[None, local_iter].iterator, single_copy_layout) + dst_i = cute.make_tensor(dst_ptr, single_copy_layout) + cute.copy(copy_atom_vec, cute.recast_tensor(src_i, cutlass.Int32), cute.recast_tensor(dst_i, cutlass.Int32)) + + +@cute.jit +def fc2_tma_store_function( + *, + epi, + subtile: cute.Tensor, # Always BF16 before quantization + work_tile_info: SwapAbFc12WorkTileInfo, + epilogue_iter_idx: cutlass.Int32, + subtile_idx: cutlass.Int32, + fc2_output_router: Fc2OutputRouter, + **_, +): + if cutlass.const_expr(epi.smem_tensor is None or epi.tma_atom_fc2_output is None or epi.fc2_tma_output is None): + raise ValueError("FC2 TMA store requires staged SMEM, a TMA atom, and a token-major output tensor.") + + if cutlass.const_expr(epi.combine_format.is_quantized): + data_subtile, sf_regs = QuantImpl(epi.combine_format, "regs_in_thread")(subtile) + else: + data_subtile = subtile + sf_regs = None + + if cutlass.const_expr(sf_regs is not None): + sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block + token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 + for token_group in cutlass.range_constexpr(token_groups_per_subtile): + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group + sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) + if sf_pred != cutlass.Int32(0): + sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) + for scale_idx in cutlass.range_constexpr(sf_scales_per_stg): + sf_dst[scale_idx] = sf_regs[token_group * sf_scales_per_stg + scale_idx] + + vector_elements: cutlass.Constexpr[int] = 128 // data_subtile.element_type.width + if cutlass.const_expr( + cute.rank(data_subtile) != 1 + or cute.size(data_subtile) != 2 * 32 + or data_subtile.stride[0] != 1 + or 32 % vector_elements != 0 + ): + raise ValueError("FC2 TMA store requires a contiguous two-token by 32-hidden register fragment.") + + stage_idx = epilogue_iter_idx % cutlass.Int32(epi.fc2_tma_stages) + smem_stage = epi.smem_tensor[None, None, stage_idx] + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128) + tiler_mn = (64, 128) + layout_copy_tv = cute.make_layout(((32, 4), (32, 2)), stride=((1, 2048), (64, 32))) + tiled_store = cute.make_tiled_copy(store_atom, layout_copy_tv, tiler_mn) + thread_store = tiled_store.get_slice(epi.tidx) + smem_partition = thread_store.partition_D(smem_stage) + register_partition = cute.composition(data_subtile, cute.make_layout(smem_partition.shape)) + cute.copy(tiled_store, register_partition, smem_partition) + + token_tile_idx = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_n_idx * cutlass.Int32(epi.cta_tile_n) + + subtile_idx * cutlass.Int32(epi._EpilogueTokenTileSize) + ) // cutlass.Int32(epi._EpilogueTokenTileSize) + tiled_output = cute.flat_divide(epi.fc2_tma_output, (epi._EpilogueTokenTileSize, epi._EpilogueFc2HiddenTileSize)) + gmem_subtile = tiled_output[None, None, token_tile_idx, work_tile_info.tile_m_idx, 0] + tma_smem_src, tma_gmem_dst = cpasync.tma_partition( + epi.tma_atom_fc2_output, + 0, + cute.make_layout(1), + cute.group_modes(smem_stage, 0, 2), + cute.group_modes(gmem_subtile, 0, 2), + ) + + epilogue_barrier = SwapABGatedActEpilogue.epilogue_sync_barrier() + cute.arch.fence_proxy("async.shared", space="cta") + epilogue_barrier.arrive_and_wait() + if epi.warp_idx == cutlass.Int32(0): + cute.copy(epi.tma_atom_fc2_output, tma_smem_src, tma_gmem_dst) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(epi.fc2_tma_stages - 1, read=True) + epilogue_barrier.arrive_and_wait() + + +@cute.jit +def fc2_ublk_store_function_impl( + *, + epi, + subtile: cute.Tensor, # Always bf16 pre-quant tensor + epilogue_iter_idx: cutlass.Int32, + subtile_idx: cutlass.Int32, + fc2_output_router: Fc2OutputRouter, + **_, +): + smem_tensor = epi.smem_tensor + if cutlass.const_expr(smem_tensor is None): + raise ValueError("fc2 UBLK store requires epi.smem_tensor.") + quantized: cutlass.Constexpr[bool] = epi.combine_format.is_quantized + if cutlass.const_expr(quantized): + # The transposed fragment gives each thread complete hidden-axis scale blocks. + data_subtile, sf_regs = QuantImpl(epi.combine_format, "regs_in_thread")(subtile) + sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block + token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 + for token_group in cutlass.range_constexpr(token_groups_per_subtile): + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group + sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) + if sf_pred != cutlass.Int32(0): + sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) + for scale_idx in cutlass.range_constexpr(sf_scales_per_stg): + sf_dst[scale_idx] = sf_regs[token_group * sf_scales_per_stg + scale_idx] + else: + data_subtile = subtile + + smem_read_write_bar = SwapABGatedActEpilogue.epilogue_sync_barrier() + warp_idx = epi.warp_idx + lane_idx = epi.lane_idx + stage_cnt: cutlass.Constexpr[int] = epi.fc2_tma_stages + stage_idx = epilogue_iter_idx % cutlass.Int32(stage_cnt) + smem_stage = smem_tensor[None, None, stage_idx] + # Reuse waits here so the previous copy overlaps the next subtile; run() drains the final commit. + cute.arch.cp_async_bulk_wait_group(stage_cnt - 1, read=True) + smem_read_write_bar.arrive_and_wait() + + if cutlass.const_expr(quantized): + vector_elements: cutlass.Constexpr[int] = 128 // data_subtile.element_type.width + if cutlass.const_expr( + cute.rank(data_subtile) != 1 + or cute.size(data_subtile) != 2 * 32 + or data_subtile.stride[0] != 1 + or 32 % vector_elements != 0 + ): + raise ValueError("FC2 UBLK store requires a contiguous two-token by 32-hidden register fragment.") + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128) + tiler_mn = (64, 128) + layout_copy_tv = cute.make_layout(((32, 4), (32, 2)), stride=((1, 2048), (64, 32))) + tiled_store = cute.make_tiled_copy(store_atom, layout_copy_tv, tiler_mn) + thread_store = tiled_store.get_slice(epi.tidx) + smem_partition = thread_store.partition_D(smem_stage) + register_partition = cute.composition(data_subtile, cute.make_layout(smem_partition.shape)) + cute.copy(tiled_store, register_partition, smem_partition) + else: + if cutlass.const_expr(cute.size(data_subtile) != epi._EpilogueTokenTileSize): + raise ValueError("BF16 UBLK staging requires one register per token in the 64-token subtile.") + warp_hidden_base = cutlass.Int32(warp_idx * 32) + for token_idx in cutlass.range_constexpr(epi._EpilogueTokenTileSize): + smem_stage[token_idx, warp_hidden_base + lane_idx] = data_subtile[token_idx] + + cute.arch.fence_proxy("async.shared", space="cta") + smem_read_write_bar.arrive_and_wait() + + copy_elems = cutlass.Int32(epi._EpilogueFc2HiddenTileSize) + if cutlass.const_expr(epi.fc2_hidden_needs_predicate): + copy_elems = cutlass.Int32(fc2_output_router.valid_hidden_this_cta_tile) + copy_bytes = copy_elems * epi.combine_format.act_dtype.width // 8 + scratch_row = warp_idx * cutlass.Int32(16) + lane_idx % cutlass.Int32(16) + copy_iters: cutlass.Constexpr[int] = fc2_output_router.data_mapping.domain.axis_size("iter_idx") + for ublk_iter_idx in cutlass.range_constexpr(copy_iters): + owned_subtile = cutlass.Int32(ublk_iter_idx * 2) + lane_idx // cutlass.Int32(16) + if owned_subtile == subtile_idx: + dst_ptr, pred = fc2_output_router.get_data_dst(ublk_iter_idx) + if pred != cutlass.Int32(0): + src_row = cute.slice_(smem_stage, (scratch_row, None)) + if cutlass.const_expr(epi.reduce_topk_in_epilogue): + cp_reduce_async_bulk_add_bf16_s2g(dst_ptr, src_row.iterator, copy_bytes) + else: + cp_async_bulk_s2g(dst_ptr, src_row.iterator, copy_bytes) + + cute.arch.cp_async_bulk_commit_group() + + +@cute.jit +def fc2_redg_store_function( + *, + epi, + subtile: cute.Tensor, # Always bf16; in-kernel reduce never quantizes + subtile_idx: cutlass.Int32, + fc2_output_router: Fc2OutputRouter, + **_, +): + redg_width_elems: cutlass.Constexpr[int] = 4 + elems_per_thread: cutlass.Constexpr[int] = cute.size(subtile) + if cutlass.const_expr(elems_per_thread % redg_width_elems != 0): + raise ValueError( + "fc2 REDG store requires pre-store elems per thread to be divisible " + f"by REDG issue width, got {elems_per_thread} and {redg_width_elems}." + ) + iters_per_subtile: cutlass.Constexpr[int] = elems_per_thread // redg_width_elems + subtile_iter_base = cutlass.Int32(subtile_idx) * cutlass.Int32(iters_per_subtile) + subtile_by_redg_issue = cute.zipped_divide(subtile, (redg_width_elems,)) + + for local_iter in cutlass.range_constexpr(iters_per_subtile): + global_iter = subtile_iter_base + cutlass.Int32(local_iter) + dst_ptr, pred = fc2_output_router.get_data_dst(global_iter) + if pred != cutlass.Int32(0): + bf16x4 = subtile_by_redg_issue[None, local_iter] + packed_bf16x2 = cute.recast_tensor(bf16x4, cutlass.Float32) + red_add_relaxed_sys_v2_bf16x2(dst_ptr, cutlass.Float32(packed_bf16x2[0]), cutlass.Float32(packed_bf16x2[1])) + + +def make_fc2_stg_process_pipeline( + *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int +) -> Fc2ProcessPipeline: + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size + ) + return Fc2ProcessPipeline( + tmem_acc_load=fc2_stg_tmem_acc_load, + f2fp=fc2_f2fp, + post_f2fp_reorder=fc2_stg_post_f2fp_reorder, + store_function=fc2_stg_store_function, + fc2_cta_tile_mapping=fundamental_mapping, + store_out_mapping=store_out_mapping, + sf_store_out_mapping=sf_store_out_mapping, + require_tmem_trans=True, + ) + + +def make_fc2_tma_process_pipeline( + *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int +) -> Fc2ProcessPipeline: + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size + ) + return Fc2ProcessPipeline( + tmem_acc_load=fc2_stg_tmem_acc_load, + f2fp=fc2_f2fp, + post_f2fp_reorder=fc2_stg_post_f2fp_reorder, + store_function=fc2_tma_store_function, + fc2_cta_tile_mapping=fundamental_mapping, + store_out_mapping=store_out_mapping, + sf_store_out_mapping=sf_store_out_mapping, + require_tmem_trans=True, + ) + + +def make_fc2_redg_process_pipeline( + *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int +) -> Fc2ProcessPipeline: + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_redg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size + ) + return Fc2ProcessPipeline( + tmem_acc_load=fc2_stg_tmem_acc_load, + f2fp=fc2_f2fp, + post_f2fp_reorder=fc2_redg_post_f2fp_reorder, + store_function=fc2_redg_store_function, + fc2_cta_tile_mapping=fundamental_mapping, + store_out_mapping=store_out_mapping, + sf_store_out_mapping=sf_store_out_mapping, + require_tmem_trans=True, + ) + + +def make_fc2_ublk_process_pipeline( + *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int +) -> Fc2ProcessPipeline: + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_ublk_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size + ) + if combine_format.is_quantized: + if combine_format.act_dtype.width != 8: + raise ValueError("FC2 UBLK quantization requires an FP8 combine payload.") + tmem_acc_load = fc2_stg_tmem_acc_load + post_f2fp_reorder = fc2_stg_post_f2fp_reorder + else: + tmem_acc_load = fc2_ublk_tmem_acc_load + post_f2fp_reorder = post_f2fp_reorder_identity + return Fc2ProcessPipeline( + tmem_acc_load=tmem_acc_load, + f2fp=fc2_f2fp, + post_f2fp_reorder=post_f2fp_reorder, + store_function=fc2_ublk_store_function_impl, + fc2_cta_tile_mapping=fundamental_mapping, + store_out_mapping=store_out_mapping, + sf_store_out_mapping=sf_store_out_mapping, + require_tmem_trans=combine_format.is_quantized, + ) + + +__all__ = [ + "Fc2OutputRouter", + "Fc2ProcessPipeline", + "GatedActEpilogueArgs", + "QuantImpl", + "SwapABGatedActEpilogue", + "make_fc2_redg_process_pipeline", + "make_fc2_stg_process_pipeline", + "make_fc2_tma_process_pipeline", + "make_fc2_ublk_process_pipeline", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py new file mode 100644 index 000000000..fc47fe0e3 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py @@ -0,0 +1,164 @@ +"""Block-scaled SwapAb adapter between FC12 scheduling and kernel tensor views.""" + +import dataclasses +from typing import ClassVar, List, Literal, Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass.cute.typing import Pointer +from cutlass.cutlass_dsl import Int32, extract_mlir_values, new_from_mlir_values +from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF + +from .....helpers.dsl_helpers import spin_peek, spin_wait +from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo, peek_ready_bit + + +TensorRole = Literal["a", "b", "sfa", "sfb", "c", "sfc", "topk"] + + +@cute.jit +def _rewrite_tensor_shape(tensor: cute.Tensor, new_shape: Tuple) -> cute.Tensor: + return cute.make_tensor(tensor.iterator, cute.make_layout(new_shape, stride=tensor.stride)) + + +@dataclasses.dataclass(frozen=True) +class BlockScaledSwapAbFc12Extension: + """Kernel-owned work-tile preparation and GMEM view adapter.""" + + work_tile_type: ClassVar[type] = SwapAbFc12WorkTileInfo + + sf_vec_size: int + fc1_done_counter_pointer: Pointer + fc2_spin_threshold: Int32 + fc1_ready_counter_pointer: Optional[Pointer] = None + + def __post_init__(self) -> None: + if self.sf_vec_size <= 0: + raise ValueError(f"sf_vec_size must be positive, got {self.sf_vec_size}.") + object.__setattr__(self, "fc2_spin_threshold", Int32(self.fc2_spin_threshold)) + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + values.extend(extract_mlir_values(self.fc1_done_counter_pointer)) + values.extend(extract_mlir_values(self.fc2_spin_threshold)) + if self.fc1_ready_counter_pointer is not None: + values.extend(extract_mlir_values(self.fc1_ready_counter_pointer)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "BlockScaledSwapAbFc12Extension": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + fc1_done_counter_pointer = rebuild(self.fc1_done_counter_pointer) + fc2_spin_threshold = rebuild(self.fc2_spin_threshold) + fc1_ready_counter_pointer = ( + rebuild(self.fc1_ready_counter_pointer) if self.fc1_ready_counter_pointer is not None else None + ) + if value_index != len(values): + raise ValueError( + f"BlockScaledSwapAbFc12Extension MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return type(self)( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter_pointer, + fc2_spin_threshold=fc2_spin_threshold, + fc1_ready_counter_pointer=fc1_ready_counter_pointer, + ) + + @cute.jit + def prepare_work_tile(self, work_tile: SwapAbFc12WorkTileInfo) -> SwapAbFc12WorkTileInfo: + """Pack kernel readiness observations into the published tile flags.""" + phase_and_flags = work_tile.phase_and_flags + if work_tile.is_valid_tile: + counter_slot = work_tile.cumulative_token_block_count + work_tile.tile_n_idx + is_fc1 = work_tile.phase == Int32(BlockPhase.Linear1) + is_fc2 = work_tile.phase == Int32(BlockPhase.Linear2) + + if cutlass.const_expr(self.fc1_ready_counter_pointer is not None): + if is_fc1: + counter_pointer = self.fc1_ready_counter_pointer + counter_slot + peek_flag = Int32(0) + if spin_peek(counter_pointer, lambda value: value >= work_tile.valid_tokens_in_cta_tile): + peek_flag = Int32(peek_ready_bit) + phase_and_flags = work_tile.phase_and_flags | peek_flag + + if is_fc2: + counter_pointer = self.fc1_done_counter_pointer + counter_slot + peek_flag = Int32(0) + if spin_peek(counter_pointer, lambda value: value >= self.fc2_spin_threshold): + peek_flag = Int32(peek_ready_bit) + phase_and_flags = work_tile.phase_and_flags | peek_flag + + return SwapAbFc12WorkTileInfo( + expert_idx=work_tile.expert_idx, + tile_m_idx=work_tile.tile_m_idx, + tile_n_idx=work_tile.tile_n_idx, + cumulative_data_physical_row=(work_tile.cumulative_data_physical_row), + cumulative_sf_physical_row=(work_tile.cumulative_sf_physical_row), + cumulative_token_block_count=(work_tile.cumulative_token_block_count), + valid_tokens_in_cta_tile=(work_tile.valid_tokens_in_cta_tile), + phase_and_flags=phase_and_flags, + ) + + @cute.jit + def wait_for_input(self, work_tile: SwapAbFc12WorkTileInfo) -> None: + """Wait until the current FC1 input tile is ready.""" + if cutlass.const_expr(self.fc1_ready_counter_pointer is not None): + counter_slot = work_tile.cumulative_token_block_count + work_tile.tile_n_idx + counter_pointer = self.fc1_ready_counter_pointer + counter_slot + spin_wait( + counter_pointer, + lambda value: value >= work_tile.valid_tokens_in_cta_tile, + peek_status=work_tile.peek_ready, + ) + + @cute.jit + def get_gmem_tensor( + self, tensor_name: TensorRole, tensor: cute.Tensor, work_tile: SwapAbFc12WorkTileInfo + ) -> Tuple[cute.Tensor, Optional[Pointer]]: + """Resolve one kernel tensor to its current expert/task-tile view.""" + expert_idx = work_tile.expert_idx + data_token_offset = work_tile.cumulative_data_physical_row + sf_token_offset = work_tile.cumulative_sf_physical_row + shape = tensor.shape + stride = tensor.stride + singleton = Int32(1) + + if cutlass.const_expr(tensor_name == "a"): + result = cute.domain_offset((0, 0, expert_idx), tensor) + return (_rewrite_tensor_shape(result, (shape[0], shape[1], singleton)), None) + + if cutlass.const_expr(tensor_name == "b"): + result = cute.domain_offset((data_token_offset, 0, 0), tensor) + return (_rewrite_tensor_shape(result, (shape[0], shape[1], singleton)), None) + + if cutlass.const_expr(tensor_name == "sfa"): + result = cute.domain_offset((0, 0, expert_idx), tensor) + per_expert_shape = (shape[0], shape[1], singleton) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, self.sf_vec_size) + return (cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), None) + + if cutlass.const_expr(tensor_name in ("sfb", "sfc")): + result = cute.domain_offset((sf_token_offset, 0, 0), tensor) + per_expert_shape = (shape[0], shape[1], singleton) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, self.sf_vec_size) + return (cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), None) + + if cutlass.const_expr(tensor_name == "c"): + result = cute.domain_offset((data_token_offset, 0, 0), tensor) + return (_rewrite_tensor_shape(result, (shape[0], shape[1], singleton)), None) + + if cutlass.const_expr(tensor_name == "topk"): + return (cute.domain_offset((data_token_offset,), tensor), None) + + raise ValueError(f"Unknown tensor_name: {tensor_name!r}.") + + +__all__ = ["BlockScaledSwapAbFc12Extension", "TensorRole"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/topk_reduce.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/topk_reduce.py new file mode 100644 index 000000000..582c4cdc4 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/topk_reduce.py @@ -0,0 +1,484 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Device combine reduce: collapse per-(token, topk) fc2 cells into one row. + +The combine step writes one fc2 output per ``(token, topk)`` cell; this reduces +over the topk axis into the token-centric ``(token, hidden)`` output. The wire +format is described by ``CombineFormat``: + + bf16 -- no staging: bf16 terms reduced directly. + 32e4m3xe8m0 -- MXFP8: fp8 e4m3 data + per-32 e8m0 (power-of-2) scale. + 16e2m1xbf16 -- fp4 e2m1 data + per-16 bf16 amax (one level, no global); + dequant per element x = fp4 * (amax * (1 / 6)). + +Task partition: each worker owns one ``(token, hidden_tile)`` and loops topk; the +flat worker index decodes into ``(token_idx, hidden_tile_idx)`` via a constant +divide by ``hidden_tiles``. The per-block scale is broadcast to a logical +per-hidden view (stride 0) so it tiles by the same worker index as the data. The +activation load stays in the topk loop (too large to hoist); the small scale and +score loads are hoisted ahead of the loop when topk is small. +""" + +import os +from typing import ClassVar, Dict, Optional + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Float32, Int32, T +from cutlass._mlir.dialects import llvm + +from .....quant_def import CombineFormat +from .....helpers.constants import Nvfp4E2M1RcpLimit +from .....helpers.dsl_helpers import mark_alignment + + +# --------------------------------------------------------------------------- +# fp4 (e2m1) -> fp32 register decode. +# +# Blackwell has no e2m1->f32 upconvert: the framework's ``term.load().to(f32)`` +# lowers to an ALU subnormal-normalization path (~60% DRAM SOL). Both helpers +# below force a table-driven decode instead; ``e2m1_reg`` (N e2m1 codes, N % 8 +# == 0) is read as packed b32 words and N fp32 values are written into +# ``fp32_reg`` in code order. The 16 e2m1 values are exact in fp32, so the two +# decoders are bit-for-bit identical (cross-check the optimal one against the +# cvt one over all 16 codes). +# --------------------------------------------------------------------------- + + +@cute.jit +def cvt_e2m1_to_fp32_cvt_ptx(e2m1_reg: cute.Tensor, fp32_reg: cute.Tensor) -> None: + """Decode via the e2m1->f16 HW cvt (``cvt.rn.f16x2.e2m1x2``) then widen f16->f32. + + Safe baseline: the per-pair ``cvt`` instruction is itself a HW PRMT+F2FP, so + the e2m1->f16 step already avoids ALU normalization; f16->f32 is one cheap + ``cvt.f32.f16`` per element. + """ + src_words = cute.recast_tensor(e2m1_reg, Int32) # (N,) e2m1 -> (N/8,) b32 + for w in cutlass.range_constexpr(cute.size(src_words)): + res = llvm.inline_asm( + llvm.StructType.get_literal([T.f32()] * 8), + [src_words[w].ir_value()], + "{\n" + " .reg .b8 b0, b1, b2, b3;\n" + " .reg .b32 p0, p1, p2, p3;\n" + " .reg .b16 c0, d0, c1, d1, c2, d2, c3, d3;\n" + " mov.b32 {b0, b1, b2, b3}, $8;\n" + " cvt.rn.f16x2.e2m1x2 p0, b0;\n" + " cvt.rn.f16x2.e2m1x2 p1, b1;\n" + " cvt.rn.f16x2.e2m1x2 p2, b2;\n" + " cvt.rn.f16x2.e2m1x2 p3, b3;\n" + " mov.b32 {c0, d0}, p0;\n" + " mov.b32 {c1, d1}, p1;\n" + " mov.b32 {c2, d2}, p2;\n" + " mov.b32 {c3, d3}, p3;\n" + " cvt.f32.f16 $0, c0;\n" + " cvt.f32.f16 $1, d0;\n" + " cvt.f32.f16 $2, c1;\n" + " cvt.f32.f16 $3, d1;\n" + " cvt.f32.f16 $4, c2;\n" + " cvt.f32.f16 $5, d2;\n" + " cvt.f32.f16 $6, c3;\n" + " cvt.f32.f16 $7, d3;\n" + "}", + "=f,=f,=f,=f,=f,=f,=f,=f,r", + has_side_effects=False, + ) + for i in cutlass.range_constexpr(8): + fp32_reg[w * 8 + i] = Float32(llvm.extractvalue(T.f32(), res, [i])) + + +@cute.jit +def cvt_e2m1_to_fp32_optimal_ptx(e2m1_reg: cute.Tensor, fp32_reg: cute.Tensor) -> None: + """Decode via a register-resident bf16 LUT + PRMT, landing fp32 directly. + + The 8 e2m1 magnitudes ``{0,.5,1,1.5,2,3,4,6}`` are exact bf16, so a PRMT + byte-gather of the per-magnitude hi/lo bytes builds the bf16; ``fp32 = + bf16 << 16`` makes the widen free. Sign (nibble bit3) is spread to the 4 + output-byte MSBs with one ``prmt`` of ``{word<<4, word}`` (selector 0x5140), + avoiding the non-uniform shift the 4-bit-vs-8-bit stride would otherwise need. + """ + src_words = cute.recast_tensor(e2m1_reg, Int32) # (N,) e2m1 -> (N/8,) b32 + dst_words = cute.recast_tensor(fp32_reg, Int32) # write fp32 bit patterns + for w in cutlass.range_constexpr(cute.size(src_words)): + res = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 8), + [src_words[w].ir_value()], + "{\n" + " .reg .b32 ha, hb, la, lb, inh, wl, ih, il, hl, ll, hh, lh, sl, sh, p0, p1, p2, p3;\n" + " mov.b32 ha, 0x3F3F3F00;\n" # hi byte LUT, magnitudes 0..3 + " mov.b32 hb, 0x40404040;\n" # hi byte LUT, magnitudes 4..7 + " mov.b32 la, 0xC0800000;\n" # lo byte LUT, magnitudes 0..3 + " mov.b32 lb, 0xC0804000;\n" # lo byte LUT, magnitudes 4..7 + " shr.b32 inh, $8, 16;\n" # high 4 elements -> low 16 bits + " and.b32 il, $8, 0x00007777;\n" # low 4 magnitude indices (clear sign) + " and.b32 ih, inh, 0x00007777;\n" # high 4 magnitude indices + " prmt.b32 hl, ha, hb, il;\n" # hi bytes for e0..e3 + " prmt.b32 ll, la, lb, il;\n" # lo bytes for e0..e3 + " prmt.b32 hh, ha, hb, ih;\n" # hi bytes for e4..e7 + " prmt.b32 lh, la, lb, ih;\n" # lo bytes for e4..e7 + " shl.b32 wl, $8, 4;\n" + " prmt.b32 sl, wl, $8, 0x5140;\n" # gather s0..s3 to byte MSBs + " and.b32 sl, sl, 0x80808080;\n" + " or.b32 hl, hl, sl;\n" + " shl.b32 wl, inh, 4;\n" + " prmt.b32 sh, wl, inh, 0x5140;\n" # gather s4..s7 to byte MSBs + " and.b32 sh, sh, 0x80808080;\n" + " or.b32 hh, hh, sh;\n" + " prmt.b32 p0, ll, hl, 0x5140;\n" # {bf16(e0), bf16(e1)} + " prmt.b32 p1, ll, hl, 0x7362;\n" # {bf16(e2), bf16(e3)} + " prmt.b32 p2, lh, hh, 0x5140;\n" # {bf16(e4), bf16(e5)} + " prmt.b32 p3, lh, hh, 0x7362;\n" # {bf16(e6), bf16(e7)} + " shl.b32 $0, p0, 16;\n" + " and.b32 $1, p0, 0xFFFF0000;\n" + " shl.b32 $2, p1, 16;\n" + " and.b32 $3, p1, 0xFFFF0000;\n" + " shl.b32 $4, p2, 16;\n" + " and.b32 $5, p2, 0xFFFF0000;\n" + " shl.b32 $6, p3, 16;\n" + " and.b32 $7, p3, 0xFFFF0000;\n" + "}", + "=r,=r,=r,=r,=r,=r,=r,=r,r", + has_side_effects=False, + ) + for i in cutlass.range_constexpr(8): + dst_words[w * 8 + i] = Int32(llvm.extractvalue(T.i32(), res, [i])) + + +class TopkReduce: + """Combine reduce for a fixed ``(hidden, num_topk, combine_format)``. + + ``__init__`` pins the static shape and format (and the derived launch + geometry); ``__call__`` (a ``@cute.jit`` launcher) sizes a 1D grid from the + runtime token count and dispatches the format's kernel. The caller owns the + torch->cute conversion and the ``cute.compile`` / ``aot_compile``. + """ + + _threads: ClassVar[int] = 128 + # combine_format.name -> hidden elements per worker (one LDG of data: + # bf16 8*2B=16B, e4m3 16*1B=16B, e2m1 16*0.5B=8B). For quantized formats this + # stays <= the scale block, so each worker reads exactly one scale entry. + _hidden_per_thread: ClassVar[Dict[str, int]] = {"bf16": 8, "32e4m3xe8m0": 16, "16e2m1xbf16": 16} + # topk count at/below which the scale + score loads are hoisted ahead of the + # topk loop (small enough to not bloat registers; a CTA-broadcast read). + _prefetch_limit: ClassVar[int] = 16 + + def __init__(self, hidden: int, num_topk: int, combine_format: CombineFormat) -> None: + self.hidden = int(hidden) + self.num_topk = int(num_topk) + self.combine_format = combine_format + self.hidden_per_thread = self._hidden_per_thread[combine_format.name] + # hidden must tile cleanly both into worker slices and into scale blocks. + align = max(combine_format.scale_block or self.hidden_per_thread, self.hidden_per_thread) + if self.hidden % align != 0: + raise ValueError( + f"hidden ({self.hidden}) must be divisible by max(scale_block, " + f"hidden_per_thread) = {align} for combine_format {combine_format}." + ) + self.hidden_tiles = self.hidden // self.hidden_per_thread + # tail guard only needed when the worker count per token is not a whole + # number of CTAs; prefetch only when topk is small enough to hoist. + self.require_predicate = self.hidden_tiles % self._threads != 0 + self.prefetch = self.num_topk <= self._prefetch_limit + + # -- launcher ------------------------------------------------------------- + + @cute.jit + def __call__( + self, + combine_quant: cute.Tensor, # (token, topk, hidden) + combine_sf: Optional[cute.Tensor], # (token, topk, hidden) + reduced_output: cute.Tensor, # (token, hidden) + topk_score: Optional[cute.Tensor], # (token, topk) + stream: cuda.CUstream, + ): + threads = self._threads + total_workers = reduced_output.shape[0] * self.hidden_tiles + grid = [(total_workers + threads - 1) // threads, 1, 1] + block = [threads, 1, 1] + + combine_quant = cute.make_tensor( + combine_quant.iterator, + cute.make_layout((combine_quant.shape[0], self.num_topk, self.hidden), stride=combine_quant.stride), + ) + reduced_output = cute.make_tensor( + reduced_output.iterator, + cute.make_layout((reduced_output.shape[0], self.hidden), stride=reduced_output.stride), + ) + if cutlass.const_expr(topk_score is not None): + topk_score = cute.make_tensor( + topk_score.iterator, cute.make_layout((topk_score.shape[0], self.num_topk), stride=topk_score.stride) + ) + + if cutlass.const_expr(not self.combine_format.is_quantized): + self._reduce_bf16(combine_quant, topk_score, reduced_output).launch(grid=grid, block=block, stream=stream) + return + + # The mega kernel hands sf in already as the depth-2 broadcast layout; a + # plain (torch) sf is depth-1 and gets its hidden mode split into + # (sf_vec, hidden/sf_vec):(0, s_h) so logical hidden h reads block h//sf_vec. + sf_vec = self.combine_format.scale_block + if cutlass.const_expr(cute.depth(combine_sf.layout) >= 2): + sf = cute.make_tensor( + combine_sf.iterator, + cute.make_layout( + (combine_sf.shape[0], self.num_topk, (sf_vec, self.hidden // sf_vec)), stride=combine_sf.stride + ), + ) + else: + sf = cute.make_tensor( + combine_sf.iterator, + cute.make_layout( + (combine_sf.shape[0], self.num_topk, (sf_vec, self.hidden // sf_vec)), + stride=(combine_sf.stride[0], combine_sf.stride[1], (0, combine_sf.stride[2])), + ), + ) + + if cutlass.const_expr(self.combine_format.act_dtype is cutlass.Float8E4M3FN): + self._reduce_mxfp8(combine_quant, sf, topk_score, reduced_output).launch( + grid=grid, block=block, stream=stream + ) + else: + self._reduce_fp4(combine_quant, sf, topk_score, reduced_output).launch( + grid=grid, block=block, stream=stream + ) + + # -- kernels -------------------------------------------------------------- + + @cute.kernel + def _reduce_bf16(self, combine_output: cute.Tensor, topk_score: Optional[cute.Tensor], reduced_output: cute.Tensor): + threads = self._threads + hidden_per_thread = self.hidden_per_thread + hidden_tiles = self.hidden_tiles + num_topk: cutlass.Constexpr[int] = self.num_topk + needs_guard = self.require_predicate + prefetch = self.prefetch + out_dtype = reduced_output.element_type + + worker_idx = cute.arch.block_idx()[0] * Int32(threads) + cute.arch.thread_idx()[0] + token_idx = worker_idx // hidden_tiles + hidden_tile_idx = worker_idx % hidden_tiles + + score_dtype = topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 + score_reg = cute.make_rmem_tensor((num_topk,), score_dtype) + + if (not needs_guard) or token_idx < reduced_output.shape[0]: + # (token, topk, hidden) -> (topk, hidden_per_thread) + terms = cute.zipped_divide(combine_output[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] + + if cutlass.const_expr(topk_score is not None): + if cutlass.const_expr(prefetch): + cute.autovec_copy(topk_score[token_idx, None], score_reg) + else: + for k in cutlass.range_constexpr(num_topk): + score_reg[k] = score_dtype(1) + + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=128) + acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) + + for k in cutlass.range_constexpr(0, num_topk, 1): + term = cute.make_rmem_tensor((hidden_per_thread,), cutlass.BFloat16) + cute.copy( + load_atom, mark_alignment(terms[k, None], hidden_per_thread * cutlass.BFloat16.width // 8), term + ) + if cutlass.const_expr(topk_score is not None and not prefetch): + score_reg[k] = topk_score[token_idx, Int32(k)] + score_pair = (Float32(score_reg[k]), Float32(score_reg[k])) + + for i in cutlass.range_constexpr(0, hidden_per_thread, 2): + value_pair = (Float32(term[i]), Float32(term[i + 1])) + if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(value_pair, score_pair, (acc[i], acc[i + 1])) + else: + if cutlass.const_expr(topk_score is not None): + acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2(value_pair, score_pair) + else: + acc[i] = value_pair[0] + acc[i + 1] = value_pair[1] + + out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) + out.store(acc.load().to(out_dtype)) + cute.copy( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), out_dtype, num_bits_per_copy=128), + out, + mark_alignment(dst, hidden_per_thread * out_dtype.width // 8), + ) + + @cute.kernel + def _reduce_mxfp8( + self, + combine_quant: cute.Tensor, + combine_sf: cute.Tensor, # depth-2 broadcast view: logical (token, topk, hidden) e8m0 + topk_score: Optional[cute.Tensor], + reduced_output: cute.Tensor, + ): + threads = self._threads + hidden_per_thread = self.hidden_per_thread + hidden_tiles = self.hidden_tiles + num_topk: cutlass.Constexpr[int] = self.num_topk + needs_guard = self.require_predicate + prefetch = self.prefetch + out_dtype = reduced_output.element_type + + worker_idx = cute.arch.block_idx()[0] * Int32(threads) + cute.arch.thread_idx()[0] + token_idx = worker_idx // hidden_tiles + hidden_tile_idx = worker_idx % hidden_tiles + + score_dtype = topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 + score_reg = cute.make_rmem_tensor((num_topk,), score_dtype) + scale_reg = cute.make_rmem_tensor((num_topk,), cutlass.Float8E8M0FNU) + + if (not needs_guard) or token_idx < reduced_output.shape[0]: + # (token, topk, hidden) -> (topk, hidden_per_thread) + codes = cute.zipped_divide(combine_quant[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] + # (token, topk, hidden) -> (topk, hidden_per_thread) + sf = cute.zipped_divide(combine_sf[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] + + if cutlass.const_expr(topk_score is not None): + if cutlass.const_expr(prefetch): + cute.autovec_copy(topk_score[token_idx, None], score_reg) + else: + for k in cutlass.range_constexpr(num_topk): + score_reg[k] = score_dtype(1) + if cutlass.const_expr(prefetch): + cute.autovec_copy(sf[None, 0], scale_reg) # one scale per topk slot (stride-0 broadcast) + + fp8_dtype = self.combine_format.act_dtype + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=128) + acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) + + for k in cutlass.range_constexpr(0, num_topk, 1): + term = cute.make_rmem_tensor((hidden_per_thread,), fp8_dtype) + cute.copy(load_atom, mark_alignment(codes[k, None], hidden_per_thread * fp8_dtype.width // 8), term) + value = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) + value.store(term.load().to(cutlass.Float32)) + + if cutlass.const_expr(not prefetch): + scale_reg[k] = sf[k, 0] + if cutlass.const_expr(topk_score is not None): + score_reg[k] = topk_score[token_idx, Int32(k)] + + scale = Float32(scale_reg[k]) # e8m0 -> f32 + scale_pair = (scale, scale) + score_pair = (Float32(score_reg[k]), Float32(score_reg[k])) + + for i in cutlass.range_constexpr(0, hidden_per_thread, 2): + dequant_pair = cute.arch.mul_packed_f32x2((value[i], value[i + 1]), scale_pair) + if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(dequant_pair, score_pair, (acc[i], acc[i + 1])) + else: + if cutlass.const_expr(topk_score is not None): + acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2(dequant_pair, score_pair) + else: + acc[i] = dequant_pair[0] + acc[i + 1] = dequant_pair[1] + + out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) + out.store(acc.load().to(out_dtype)) + cute.copy( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), out_dtype, num_bits_per_copy=256), + out, + mark_alignment(dst, hidden_per_thread * out_dtype.width // 8), + ) + + @cute.kernel + def _reduce_fp4( + self, + combine_quant: cute.Tensor, # (token, topk, hidden) e2m1 (logical) + combine_sf: cute.Tensor, # depth-2 broadcast view: logical (token, topk, hidden) bf16 amax + topk_score: Optional[cute.Tensor], + reduced_output: cute.Tensor, + ): + threads = self._threads + hidden_per_thread = self.hidden_per_thread + hidden_tiles = self.hidden_tiles + num_topk: cutlass.Constexpr[int] = self.num_topk + needs_guard = self.require_predicate + prefetch = self.prefetch + out_dtype = reduced_output.element_type + + worker_idx = cute.arch.block_idx()[0] * Int32(threads) + cute.arch.thread_idx()[0] + token_idx = worker_idx // hidden_tiles + hidden_tile_idx = worker_idx % hidden_tiles + + score_dtype = topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 + score_reg = cute.make_rmem_tensor((num_topk,), score_dtype) + scale_reg = cute.make_rmem_tensor((num_topk,), cutlass.BFloat16) + + if (not needs_guard) or token_idx < reduced_output.shape[0]: + # (token, topk, hidden) -> (topk, hidden_per_thread) + codes = cute.zipped_divide(combine_quant[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] + # (token, topk, hidden) -> (topk, hidden_per_thread) + sf = cute.zipped_divide(combine_sf[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] + + if cutlass.const_expr(topk_score is not None): + if cutlass.const_expr(prefetch): + cute.autovec_copy(topk_score[token_idx, None], score_reg) + else: + for k in cutlass.range_constexpr(num_topk): + score_reg[k] = score_dtype(1) + if cutlass.const_expr(prefetch): + cute.autovec_copy(sf[None, 0], scale_reg) + + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Float4E2M1FN, num_bits_per_copy=64) + acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) + + for k in cutlass.range_constexpr(0, num_topk, 1): + term = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float4E2M1FN) + cute.copy( + load_atom, mark_alignment(codes[k, None], hidden_per_thread * cutlass.Float4E2M1FN.width // 8), term + ) + value = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) + # Dev-only knob (MEGA_F4CVT_USE_MANUAL): manual LUT/PRMT decode vs + # the HW cvt path, so both SASS forms can be compared on device; + # one is kept once chosen. Read inline on purpose -- never a + # customer-facing option. Both decoders are bit-exact. + if cutlass.const_expr(os.environ.get("MEGA_F4CVT_USE_MANUAL", "0") == "1"): + cvt_e2m1_to_fp32_optimal_ptx(term, value) + else: + cvt_e2m1_to_fp32_cvt_ptx(term, value) + + if cutlass.const_expr(not prefetch): + scale_reg[k] = sf[k, 0] + if cutlass.const_expr(topk_score is not None): + score_reg[k] = topk_score[token_idx, Int32(k)] + + # amax (bf16) -> per-element scale; (1/6) folds the fp4 grid max. + scale = Float32(scale_reg[k]) * Float32(Nvfp4E2M1RcpLimit) + scale_pair = (scale, scale) + score_pair = (Float32(score_reg[k]), Float32(score_reg[k])) + + for i in cutlass.range_constexpr(0, hidden_per_thread, 2): + dequant_pair = cute.arch.mul_packed_f32x2((value[i], value[i + 1]), scale_pair) + if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(dequant_pair, score_pair, (acc[i], acc[i + 1])) + elif cutlass.const_expr(topk_score is not None): + acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2(dequant_pair, score_pair) + else: + acc[i] = dequant_pair[0] + acc[i + 1] = dequant_pair[1] + + out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) + out.store(acc.load().to(out_dtype)) + cute.copy( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), out_dtype, num_bits_per_copy=256), + out, + mark_alignment(dst, hidden_per_thread * out_dtype.width // 8), + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/function_mapping.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/function_mapping.py new file mode 100644 index 000000000..83cacc303 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/function_mapping.py @@ -0,0 +1,171 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Pure-Python finite coordinate mappings used by kernel code generation.""" + +import inspect +from dataclasses import dataclass +from math import prod +from typing import Callable, Mapping, Sequence + + +class FunctionMappingError(ValueError): + """Raised when a finite function mapping is malformed.""" + + +@dataclass(frozen=True) +class CoordinateSpace: + """A finite named coordinate space with axis 0 linearized fastest.""" + + names: tuple[str, ...] + sizes: tuple[int, ...] + + def __post_init__(self) -> None: + names = tuple(self.names) + sizes = tuple(self.sizes) + if not names or len(names) != len(sizes): + raise FunctionMappingError( + f"Coordinate-space rank mismatch: {names!r}, {sizes!r}." + ) + if any(not isinstance(name, str) or not name for name in names): + raise FunctionMappingError("Axis names must be non-empty strings.") + if len(set(names)) != len(names): + raise FunctionMappingError(f"Axis names must be unique: {names!r}.") + if any(not isinstance(size, int) or size <= 0 for size in sizes): + raise FunctionMappingError( + f"Axis sizes must be positive Python ints: {sizes!r}." + ) + object.__setattr__(self, "names", names) + object.__setattr__(self, "sizes", sizes) + + @property + def size(self) -> int: + return prod(self.sizes) + + def axis_size(self, name: str) -> int: + try: + return self.sizes[self.names.index(name)] + except ValueError as error: + raise KeyError(f"Unknown coordinate axis {name!r}.") from error + + def delinearize(self, linear_index: int) -> tuple[int, ...]: + if linear_index < 0 or linear_index >= self.size: + raise FunctionMappingError( + f"Linear index {linear_index} is outside [0, {self.size})." + ) + remaining = linear_index + coordinate = [] + for size in self.sizes: + coordinate.append(remaining % size) + remaining //= size + return tuple(coordinate) + + def coordinates(self) -> tuple[tuple[int, ...], ...]: + return tuple(self.delinearize(index) for index in range(self.size)) + + +MappingResult = int | Sequence[int] | Mapping[str, int] + + +@dataclass(frozen=True) +class FunctionMapping: + """A finite coordinate mapping generated by a deterministic Python function.""" + + domain: CoordinateSpace + codomain: CoordinateSpace + function: Callable[..., MappingResult] + + def __post_init__(self) -> None: + if not callable(self.function): + raise FunctionMappingError("FunctionMapping.function must be callable.") + self._validate_signature() + for domain_coordinate in self.domain.coordinates(): + arguments = dict(zip(self.domain.names, domain_coordinate)) + result = self.function(**arguments) + self._normalize_result(result, validate_static=True) + + def _validate_signature(self) -> None: + signature = inspect.signature(self.function) + parameters = signature.parameters + unsupported = [ + name + for name, parameter in parameters.items() + if parameter.kind + in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ) + or parameter.default is not inspect.Parameter.empty + ] + if unsupported: + raise FunctionMappingError( + f"Unsupported mapping parameters {unsupported!r}." + ) + if set(parameters) != set(self.domain.names): + raise FunctionMappingError( + f"Mapping parameters {tuple(parameters)!r} must match " + f"domain axes {self.domain.names!r}." + ) + + def _normalize_result( + self, + result, + *, + validate_static: bool, + ) -> dict[str, object]: + if isinstance(result, Mapping): + if set(result) != set(self.codomain.names): + raise FunctionMappingError( + f"Mapping result keys {tuple(result)!r} must match " + f"codomain axes {self.codomain.names!r}." + ) + coordinate = { + name: result[name] for name in self.codomain.names + } + elif isinstance(result, Sequence) and not isinstance( + result, + (str, bytes), + ): + if len(result) != len(self.codomain.names): + raise FunctionMappingError( + f"Mapping result rank {len(result)} must equal " + f"{len(self.codomain.names)}." + ) + coordinate = dict(zip(self.codomain.names, result)) + elif len(self.codomain.names) == 1: + coordinate = {self.codomain.names[0]: result} + else: + raise FunctionMappingError( + "A multi-axis mapping must return a sequence or named mapping." + ) + + if validate_static: + for name, value in coordinate.items(): + if not isinstance(value, int): + raise FunctionMappingError( + f"Static result {name!r} must be int, got {type(value)}." + ) + size = self.codomain.axis_size(name) + if value < 0 or value >= size: + raise FunctionMappingError( + f"Static result {name!r}={value} is outside [0, {size})." + ) + return coordinate + + def evaluate(self, **domain_coordinate) -> dict[str, object]: + if set(domain_coordinate) != set(self.domain.names): + raise FunctionMappingError( + f"Mapping arguments {tuple(domain_coordinate)!r} must match " + f"domain axes {self.domain.names!r}." + ) + return self._normalize_result( + self.function(**domain_coordinate), + validate_static=False, + ) + + +__all__ = [ + "CoordinateSpace", + "FunctionMapping", + "FunctionMappingError", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/__init__.py new file mode 100644 index 000000000..6b1dde777 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/__init__.py @@ -0,0 +1,29 @@ +"""Scheduler protocols and implementations.""" + +from .base import SchedulerBase, SchedulerConsumer, SchedulerWorkTileBase, WorkIdAcquisitionMode +from .fc12_mapping import ( + BlockPhase, + Fc12WorkTileState, + NonSwapAbFc12WorkTileInfo, + SwapAbFc12WorkTileInfo, + peek_ready_bit, +) +from .fc12_scheduler import BlackwellFusedFc12Scheduler, PhaseInterleavedFc12Scheduler +from .non_clc_mixed_cga import NonClcMixedCgaConfig, NonClcMixedCgaSchedulerWorker + + +__all__ = [ + "BlackwellFusedFc12Scheduler", + "BlockPhase", + "Fc12WorkTileState", + "NonSwapAbFc12WorkTileInfo", + "NonClcMixedCgaConfig", + "NonClcMixedCgaSchedulerWorker", + "PhaseInterleavedFc12Scheduler", + "SchedulerBase", + "SchedulerConsumer", + "SchedulerWorkTileBase", + "SwapAbFc12WorkTileInfo", + "WorkIdAcquisitionMode", + "peek_ready_bit", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/base.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/base.py new file mode 100644 index 000000000..f4a1f40ef --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/base.py @@ -0,0 +1,206 @@ +"""Scheduler façade and architecture-independent work-tile transport.""" + +from abc import ABC, abstractmethod +from typing import Any, ClassVar, Literal, Optional, Type + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass.cutlass_dsl import extract_mlir_values, new_from_mlir_values + +from ...api import ImplDesc, KernelComponent, ProblemDesc +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.smem_workspace import SmemWorkspace + + +WorkIdAcquisitionMode = Literal["grid_stride", "atomic_counter", "cluster_launch_control"] + + +class SchedulerWorkTileBase(ABC): + """Register ABI for one work tile transported through scheduler SMEM.""" + + storage_dtype: ClassVar[type] = cutlass.Int32 + storage_field_count: ClassVar[int] + + @property + @abstractmethod + def is_valid_tile(self): + """Return whether this tile names executable work.""" + ... + + @abstractmethod + def to_rmem(self) -> cute.Tensor: + """Serialize this tile into its one-dimensional register ABI.""" + ... + + @classmethod + @abstractmethod + def from_rmem(cls, registers: cute.Tensor) -> "SchedulerWorkTileBase": + """Deserialize one tile from its register ABI.""" + ... + + +class SchedulerConsumer: + """Per-consumer state for the common scheduler SMEM transport.""" + + def __init__( + self, + scheduler_pipeline: pipeline.PipelineAsync, + smem_buffer: cute.Tensor, + num_stages: int, + work_tile_type: Type[SchedulerWorkTileBase], + ) -> None: + self._pipeline = scheduler_pipeline + self._smem_buffer = smem_buffer + self._work_tile_type = work_tile_type + self._consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, num_stages) + + def __extract_mlir_values__(self) -> list: + return extract_mlir_values(self._consumer_state) + + def __new_from_mlir_values__(self, values: list) -> "SchedulerConsumer": + expected_value_count = len(extract_mlir_values(self._consumer_state)) + if len(values) != expected_value_count: + raise ValueError( + f"SchedulerConsumer MLIR value count mismatch: expected {expected_value_count}, got {len(values)}." + ) + result = type(self).__new__(type(self)) + result._pipeline = self._pipeline + result._smem_buffer = self._smem_buffer + result._work_tile_type = self._work_tile_type + result._consumer_state = new_from_mlir_values(self._consumer_state, values) + return result + + @cute.jit + def consume_work(self) -> SchedulerWorkTileBase: + """Block until the next work tile is available.""" + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self._work_tile_type.storage_dtype, num_bits_per_copy=128 + ) + self._pipeline.consumer_wait(self._consumer_state) + registers = cute.make_rmem_tensor( + (self._work_tile_type.storage_field_count,), self._work_tile_type.storage_dtype + ) + cute.copy(copy_atom, self._smem_buffer[(None, self._consumer_state.index)], registers) + work_tile = self._work_tile_type.from_rmem(registers) + cute.arch.fence_acq_rel_cta() + self._pipeline.consumer_release(self._consumer_state) + self._consumer_state.advance() + return work_tile + + +class SchedulerBase(KernelComponent): + """Common work-tile transport and façade protocol for schedulers.""" + + pipeline_mbarriers_region = "scheduler.pipeline_mbarriers" + work_tiles_region = "scheduler.work_tiles" + num_scheduler_stages = 2 + + @classmethod + def problem_desc_require(cls) -> dict: + return {} + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return {"num_scheduler_consumer_threads": int} + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + self.num_scheduler_consumer_threads = impl_desc["num_scheduler_consumer_threads"] + if self.num_scheduler_consumer_threads <= 0: + raise ValueError("num_scheduler_consumer_threads must be positive.") + + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: + """Register the common work-tile transport regions.""" + if not hasattr(self, "work_tile_type"): + raise AttributeError( + f"{type(self).__name__} must bind work_tile_type before registering scheduler SMEM regions." + ) + work_tile_type = self.work_tile_type + work_tile_field_count = work_tile_type.storage_field_count + smem_workspace.register_mbarrier(self.pipeline_mbarriers_region, self.num_scheduler_stages * 2) + smem_workspace.register_tensor( + self.work_tiles_region, + work_tile_type.storage_dtype, + (work_tile_field_count, self.num_scheduler_stages), + stride=(1, work_tile_field_count), + byte_alignment=16, + ) + + def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: + """Register scheduler-specific GMEM regions when needed.""" + pass + + @cute.jit + def create_scheduler_pipelines(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Create CTA-lifetime scheduler pipelines and transport state.""" + self._pipeline = pipeline.PipelineAsync.create( + num_stages=self.num_scheduler_stages, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_scheduler_consumer_threads), + barrier_storage=smem_workspace.ptr(self.pipeline_mbarriers_region, smem_base), + defer_sync=True, + ) + self._smem_buffer = smem_workspace.tensor(self.work_tiles_region, smem_base) + self._producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_scheduler_stages + ) + + def make_consumer(self) -> SchedulerConsumer: + """Create a consumer with an independent pipeline state.""" + return SchedulerConsumer( + scheduler_pipeline=self._pipeline, + smem_buffer=self._smem_buffer, + num_stages=self.num_scheduler_stages, + work_tile_type=self.work_tile_type, + ) + + @cute.jit + def publish_work(self, work_tile: SchedulerWorkTileBase) -> None: + """Publish one work tile through the common transport pipeline.""" + copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), work_tile.storage_dtype, num_bits_per_copy=128) + self._pipeline.producer_acquire(self._producer_state) + cute.copy(copy_atom, work_tile.to_rmem(), self._smem_buffer[(None, self._producer_state.index)]) + cute.arch.fence_proxy("async.shared", space="cta") + self._pipeline.producer_commit(self._producer_state) + self._producer_state.advance() + + @cute.jit + def produce_tail(self) -> None: + """Wait until every published work tile has been consumed.""" + self._pipeline.producer_tail(self._producer_state) + + def __extract_mlir_values__(self) -> list: + return extract_mlir_values(self._producer_state) + + def __new_from_mlir_values__(self, values: list) -> "SchedulerBase": + expected_value_count = len(extract_mlir_values(self._producer_state)) + if len(values) != expected_value_count: + raise ValueError( + f"SchedulerBase MLIR value count mismatch: expected {expected_value_count}, got {len(values)}." + ) + result = type(self).__new__(type(self)) + result.num_scheduler_consumer_threads = self.num_scheduler_consumer_threads + result.work_tile_type = self.work_tile_type + result._pipeline = self._pipeline + result._smem_buffer = self._smem_buffer + result._producer_state = new_from_mlir_values(self._producer_state, values) + return result + + @abstractmethod + def get_grid_shape(self, *, max_active_clusters: Optional[int] = None, problem_desc: Any = None): + """Return the launch grid from static scheduler policy.""" + ... + + @abstractmethod + def assign_device_members(self, *args, **kwargs) -> None: + """Initialize device members whose ownership spans one CTA lifetime.""" + ... + + @abstractmethod + def gen_next_work(self) -> SchedulerWorkTileBase: + """Claim, map, and return the next work tile.""" + ... + + +__all__ = ["SchedulerBase", "SchedulerConsumer", "SchedulerWorkTileBase", "WorkIdAcquisitionMode"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py new file mode 100644 index 000000000..2afe2ba83 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py @@ -0,0 +1,1175 @@ +"""FC12 work-tile ABI and grouped or phase-interleaved task mapping.""" + +import dataclasses +from enum import IntEnum +from typing import List, Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass.cutlass_dsl import Boolean, Int32, extract_mlir_values, new_from_mlir_values + +from ...helpers.iket_compat import iket +from .base import SchedulerWorkTileBase + + +phase_bits = 16 +phase_mask = (1 << phase_bits) - 1 +peek_ready_bit = 1 << phase_bits + + +class Fc12WorkTileState(IntEnum): + """Sentinel values carried in the expert index field.""" + + Done = -1 + + +class BlockPhase(IntEnum): + """FC1/FC2 phase encoded in one fused work tile.""" + + None_ = 0 + Linear1 = 1 + Linear2 = 2 + + +@dataclasses.dataclass(frozen=True) +class SwapAbFc12WorkTileInfo(SchedulerWorkTileBase): + """Eight-field work tile for the swap-AB FC12 orientation.""" + + storage_field_count = 8 + + expert_idx: Int32 + tile_m_idx: Int32 + tile_n_idx: Int32 + cumulative_data_physical_row: Int32 + cumulative_sf_physical_row: Int32 + cumulative_token_block_count: Int32 + valid_tokens_in_cta_tile: Int32 + phase_and_flags: Int32 + + @property + def is_valid_tile(self): + return self.expert_idx >= Int32(0) + + @property + def phase(self) -> Int32: + return self.phase_and_flags & Int32(phase_mask) + + @property + def peek_ready(self): + return self.phase_and_flags >= Int32(peek_ready_bit) + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in ( + self.expert_idx, + self.tile_m_idx, + self.tile_n_idx, + self.cumulative_data_physical_row, + self.cumulative_sf_physical_row, + self.cumulative_token_block_count, + self.valid_tokens_in_cta_tile, + self.phase_and_flags, + ): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "SwapAbFc12WorkTileInfo": + if len(values) != self.storage_field_count: + raise ValueError( + f"SwapAbFc12WorkTileInfo expects {self.storage_field_count} MLIR values, got {len(values)}." + ) + fields = ( + self.expert_idx, + self.tile_m_idx, + self.tile_n_idx, + self.cumulative_data_physical_row, + self.cumulative_sf_physical_row, + self.cumulative_token_block_count, + self.valid_tokens_in_cta_tile, + self.phase_and_flags, + ) + rebuilt = [new_from_mlir_values(field, [value]) for field, value in zip(fields, values)] + return type(self)(*rebuilt) + + def to_rmem(self) -> cute.Tensor: + registers = cute.make_rmem_tensor((self.storage_field_count,), cutlass.Int32) + registers[0] = self.expert_idx + registers[1] = self.tile_m_idx + registers[2] = self.tile_n_idx + registers[3] = self.cumulative_data_physical_row + registers[4] = self.cumulative_sf_physical_row + registers[5] = self.cumulative_token_block_count + registers[6] = self.valid_tokens_in_cta_tile + registers[7] = self.phase_and_flags + return registers + + @classmethod + def from_rmem(cls, registers: cute.Tensor) -> "SwapAbFc12WorkTileInfo": + return cls( + expert_idx=registers[0], + tile_m_idx=registers[1], + tile_n_idx=registers[2], + cumulative_data_physical_row=registers[3], + cumulative_sf_physical_row=registers[4], + cumulative_token_block_count=registers[5], + valid_tokens_in_cta_tile=registers[6], + phase_and_flags=registers[7], + ) + + +@dataclasses.dataclass(frozen=True) +class NonSwapAbFc12WorkTileInfo(SchedulerWorkTileBase): + """Eight-field work tile for the non-swap-AB FC12 orientation.""" + + storage_field_count = 8 + + expert_idx: Int32 + tile_m_idx: Int32 + tile_n_idx: Int32 + cumulative_data_physical_row: Int32 + cumulative_sf_physical_row: Int32 + cumulative_token_block_count: Int32 + valid_tokens_in_cta_cluster_tile: Int32 + phase_and_flags: Int32 + + @property + def is_valid_tile(self): + return self.expert_idx >= Int32(0) + + @property + def phase(self) -> Int32: + return self.phase_and_flags & Int32(phase_mask) + + @property + def peek_ready(self): + return self.phase_and_flags >= Int32(peek_ready_bit) + + @property + def valid_tokens_in_cta_tile(self) -> Int32: + return self.valid_tokens_in_cta_cluster_tile >> Int32(16) + + @property + def valid_tokens_in_cluster_tile(self) -> Int32: + return self.valid_tokens_in_cta_cluster_tile & Int32(0xFFFF) + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in ( + self.expert_idx, + self.tile_m_idx, + self.tile_n_idx, + self.cumulative_data_physical_row, + self.cumulative_sf_physical_row, + self.cumulative_token_block_count, + self.valid_tokens_in_cta_cluster_tile, + self.phase_and_flags, + ): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "NonSwapAbFc12WorkTileInfo": + if len(values) != self.storage_field_count: + raise ValueError( + f"NonSwapAbFc12WorkTileInfo expects {self.storage_field_count} MLIR values, got {len(values)}." + ) + fields = ( + self.expert_idx, + self.tile_m_idx, + self.tile_n_idx, + self.cumulative_data_physical_row, + self.cumulative_sf_physical_row, + self.cumulative_token_block_count, + self.valid_tokens_in_cta_cluster_tile, + self.phase_and_flags, + ) + rebuilt = [new_from_mlir_values(field, [value]) for field, value in zip(fields, values)] + return type(self)(*rebuilt) + + def to_rmem(self) -> cute.Tensor: + registers = cute.make_rmem_tensor((self.storage_field_count,), cutlass.Int32) + registers[0] = self.expert_idx + registers[1] = self.tile_m_idx + registers[2] = self.tile_n_idx + registers[3] = self.cumulative_data_physical_row + registers[4] = self.cumulative_sf_physical_row + registers[5] = self.cumulative_token_block_count + registers[6] = self.valid_tokens_in_cta_cluster_tile + registers[7] = self.phase_and_flags + return registers + + @classmethod + def from_rmem(cls, registers: cute.Tensor) -> "NonSwapAbFc12WorkTileInfo": + return cls( + expert_idx=registers[0], + tile_m_idx=registers[1], + tile_n_idx=registers[2], + cumulative_data_physical_row=registers[3], + cumulative_sf_physical_row=registers[4], + cumulative_token_block_count=registers[5], + valid_tokens_in_cta_cluster_tile=registers[6], + phase_and_flags=registers[7], + ) + + +class _Fc12TaskCursorState: + """Register-resident cursor for the FC12 group/phase/expert state machine.""" + + def __init__( + self, + current_group_first_expert: Int32, + current_group_last_expert_exclusive: Int32, + current_phase: Int32, + current_expert_idx: Int32, + current_expert_tile_start: Int32, + current_expert_tile_end: Int32, + current_group_fc1_subphase_end: Int32, + current_group_end: Int32, + cumulative_fc1_tiles_at_group_end: Int32, + cumulative_fc2_tiles_at_group_end: Int32, + current_data_cumulative: Int32, + current_sf_cumulative: Int32, + current_token_block_cumulative: Int32, + group_start_data_cumulative: Int32, + group_start_sf_cumulative: Int32, + group_start_token_block_cumulative: Int32, + current_token_block_count: Int32, + current_expert_token_count: Int32, + ) -> None: + self.current_group_first_expert = current_group_first_expert + self.current_group_last_expert_exclusive = current_group_last_expert_exclusive + self.current_phase = current_phase + self.current_expert_idx = current_expert_idx + self.current_expert_tile_start = current_expert_tile_start + self.current_expert_tile_end = current_expert_tile_end + self.current_group_fc1_subphase_end = current_group_fc1_subphase_end + self.current_group_end = current_group_end + self.cumulative_fc1_tiles_at_group_end = cumulative_fc1_tiles_at_group_end + self.cumulative_fc2_tiles_at_group_end = cumulative_fc2_tiles_at_group_end + self.current_data_cumulative = current_data_cumulative + self.current_sf_cumulative = current_sf_cumulative + self.current_token_block_cumulative = current_token_block_cumulative + self.group_start_data_cumulative = group_start_data_cumulative + self.group_start_sf_cumulative = group_start_sf_cumulative + self.group_start_token_block_cumulative = group_start_token_block_cumulative + self.current_token_block_count = current_token_block_count + self.current_expert_token_count = current_expert_token_count + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in self._fields(): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "_Fc12TaskCursorState": + value_index = 0 + rebuilt = [] + for field in self._fields(): + field_value_count = len(extract_mlir_values(field)) + rebuilt.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) + value_index += field_value_count + if value_index != len(values): + raise ValueError( + f"_Fc12TaskCursorState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return type(self)(*rebuilt) + + def _fields(self) -> Tuple: + return ( + self.current_group_first_expert, + self.current_group_last_expert_exclusive, + self.current_phase, + self.current_expert_idx, + self.current_expert_tile_start, + self.current_expert_tile_end, + self.current_group_fc1_subphase_end, + self.current_group_end, + self.cumulative_fc1_tiles_at_group_end, + self.cumulative_fc2_tiles_at_group_end, + self.current_data_cumulative, + self.current_sf_cumulative, + self.current_token_block_cumulative, + self.group_start_data_cumulative, + self.group_start_sf_cumulative, + self.group_start_token_block_cumulative, + self.current_token_block_count, + self.current_expert_token_count, + ) + + +class Fc12TaskMappingState: + """Runtime inputs and cursor for monotonic FC12 linear-ID mapping.""" + + def __init__( + self, + expert_count, + mapping_cta_tile_shape_mnk: Tuple[int, int, int], + mapping_cluster_shape_mn: Tuple[int, int], + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + is_swap_ab: bool, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], + cursor_state: _Fc12TaskCursorState, + num_fc1_intermediate_blocks, + num_fc2_hidden_blocks, + ) -> None: + self.expert_count = expert_count + self.mapping_cta_tile_shape_mnk = mapping_cta_tile_shape_mnk + self.mapping_cluster_shape_mn = mapping_cluster_shape_mn + self.group_hint = group_hint + self.token_padding_block = token_padding_block + self.sf_padding_block = sf_padding_block + self.is_swap_ab = is_swap_ab + self.expert_token_sizes = expert_token_sizes + self.expert_token_prefix_sum = expert_token_prefix_sum + self.cursor_state = cursor_state + self.num_fc1_intermediate_blocks = num_fc1_intermediate_blocks + self.num_fc2_hidden_blocks = num_fc2_hidden_blocks + + @property + def mapping_cluster_tile_m(self) -> int: + return self.mapping_cta_tile_shape_mnk[0] * self.mapping_cluster_shape_mn[0] + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + if isinstance(self.expert_count, Int32): + values.extend(extract_mlir_values(self.expert_count)) + token_counts = self.expert_token_sizes if self.expert_token_sizes is not None else self.expert_token_prefix_sum + values.extend(extract_mlir_values(token_counts)) + values.extend(extract_mlir_values(self.cursor_state)) + if isinstance(self.num_fc1_intermediate_blocks, Int32): + values.extend(extract_mlir_values(self.num_fc1_intermediate_blocks)) + if isinstance(self.num_fc2_hidden_blocks, Int32): + values.extend(extract_mlir_values(self.num_fc2_hidden_blocks)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "Fc12TaskMappingState": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + expert_count = rebuild(self.expert_count) if isinstance(self.expert_count, Int32) else self.expert_count + if self.expert_token_sizes is not None: + expert_token_sizes = rebuild(self.expert_token_sizes) + expert_token_prefix_sum = None + else: + expert_token_sizes = None + expert_token_prefix_sum = rebuild(self.expert_token_prefix_sum) + result = type(self)( + expert_count=expert_count, + mapping_cta_tile_shape_mnk=self.mapping_cta_tile_shape_mnk, + mapping_cluster_shape_mn=self.mapping_cluster_shape_mn, + group_hint=self.group_hint, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + is_swap_ab=self.is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + cursor_state=rebuild(self.cursor_state), + num_fc1_intermediate_blocks=( + rebuild(self.num_fc1_intermediate_blocks) + if isinstance(self.num_fc1_intermediate_blocks, Int32) + else self.num_fc1_intermediate_blocks + ), + num_fc2_hidden_blocks=( + rebuild(self.num_fc2_hidden_blocks) + if isinstance(self.num_fc2_hidden_blocks, Int32) + else self.num_fc2_hidden_blocks + ), + ) + if value_index != len(values): + raise ValueError( + f"Fc12TaskMappingState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return result + + +@cute.jit +def create_fc12_task_mapping_state( + *, + expert_count, + intermediate_gateup_size, + hidden_size, + mapping_cta_tile_shape_mnk: Tuple[int, int, int], + mapping_cluster_shape_mn: Tuple[int, int], + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + is_swap_ab: bool, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], +) -> Fc12TaskMappingState: + """Create the register-resident state for one CTA's FC12 mapper.""" + cursor_state = _Fc12TaskCursorState( + current_group_first_expert=Int32(0), + current_group_last_expert_exclusive=Int32(0), + current_phase=Int32(BlockPhase.Linear1), + current_expert_idx=Int32(-1), + current_expert_tile_start=Int32(0), + current_expert_tile_end=Int32(0), + current_group_fc1_subphase_end=Int32(0), + current_group_end=Int32(0), + cumulative_fc1_tiles_at_group_end=Int32(0), + cumulative_fc2_tiles_at_group_end=Int32(0), + current_data_cumulative=Int32(0), + current_sf_cumulative=Int32(0), + current_token_block_cumulative=Int32(0), + group_start_data_cumulative=Int32(0), + group_start_sf_cumulative=Int32(0), + group_start_token_block_cumulative=Int32(0), + current_token_block_count=Int32(0), + current_expert_token_count=Int32(0), + ) + mapping_cluster_tile_n = mapping_cluster_shape_mn[1] * mapping_cta_tile_shape_mnk[1] + num_fc1_intermediate_blocks = (intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + num_fc2_hidden_blocks = (hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + return Fc12TaskMappingState( + expert_count=expert_count, + mapping_cta_tile_shape_mnk=mapping_cta_tile_shape_mnk, + mapping_cluster_shape_mn=mapping_cluster_shape_mn, + group_hint=group_hint, + token_padding_block=token_padding_block, + sf_padding_block=sf_padding_block, + is_swap_ab=is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + cursor_state=cursor_state, + num_fc1_intermediate_blocks=num_fc1_intermediate_blocks, + num_fc2_hidden_blocks=num_fc2_hidden_blocks, + ) + + +@cute.jit +def _warp_inclusive_sum(value: Int32, lane_idx: Int32) -> Int32: + inclusive = value + for step_log in cutlass.range_constexpr(5): + step = Int32(1 << step_log) + previous = Int32(cute.arch.shuffle_sync(inclusive, lane_idx - step)) + if lane_idx >= step: + inclusive = inclusive + previous + return inclusive + + +@cute.jit +def _first_matching_lane(predicate) -> Int32: + mask = Int32(cute.arch.vote_ballot_sync(predicate)) + first_lane = Int32(-1) + if mask != Int32(0): + lowbit = mask & (-mask) + first_lane = Int32(cute.arch.popc(lowbit - Int32(1))) + return first_lane + + +@cute.jit +def _load_expert_batch_metrics( + mapping_state: Fc12TaskMappingState, batch_base: Int32, active_begin: Int32, active_end: Int32, lane_idx: Int32 +) -> Tuple[Int32, Int32, Int32, Int32, Int32, Int32, Int32]: + expert_idx = batch_base + lane_idx + token_count = Int32(0) + if cutlass.const_expr(mapping_state.expert_token_sizes is not None): + if expert_idx < mapping_state.expert_count: + token_count = mapping_state.expert_token_sizes[expert_idx] + else: + prefix_end = Int32(0) + if expert_idx < mapping_state.expert_count: + prefix_end = mapping_state.expert_token_prefix_sum[expert_idx] + prefix_begin = Int32(cute.arch.shuffle_sync(prefix_end, lane_idx - Int32(1))) + if lane_idx == Int32(0): + prefix_begin = Int32(0) + if batch_base > Int32(0): + prefix_begin = mapping_state.expert_token_prefix_sum[batch_base - Int32(1)] + token_count = prefix_end - prefix_begin + if (expert_idx < active_begin) | (expert_idx >= active_end): + token_count = Int32(0) + token_blocks = (token_count + Int32(mapping_state.mapping_cluster_tile_m - 1)) // Int32( + mapping_state.mapping_cluster_tile_m + ) + data_rows = ( + (token_count + Int32(mapping_state.token_padding_block - 1)) // Int32(mapping_state.token_padding_block) + ) * Int32(mapping_state.token_padding_block) + sf_rows = ( + (token_count + Int32(mapping_state.sf_padding_block - 1)) // Int32(mapping_state.sf_padding_block) + ) * Int32(mapping_state.sf_padding_block) + fc1_tiles = token_blocks * mapping_state.num_fc1_intermediate_blocks + fc2_tiles = token_blocks * mapping_state.num_fc2_hidden_blocks + return expert_idx, token_count, token_blocks, data_rows, sf_rows, fc1_tiles, fc2_tiles + + +def make_fc12_done_tile(is_swap_ab: bool) -> SchedulerWorkTileBase: + """Build the terminal work tile for either FC12 orientation.""" + if cutlass.const_expr(is_swap_ab): + return SwapAbFc12WorkTileInfo( + expert_idx=Int32(Fc12WorkTileState.Done), + tile_m_idx=Int32(0), + tile_n_idx=Int32(0), + cumulative_data_physical_row=Int32(0), + cumulative_sf_physical_row=Int32(0), + cumulative_token_block_count=Int32(0), + valid_tokens_in_cta_tile=Int32(0), + phase_and_flags=Int32(BlockPhase.None_), + ) + return NonSwapAbFc12WorkTileInfo( + expert_idx=Int32(Fc12WorkTileState.Done), + tile_m_idx=Int32(0), + tile_n_idx=Int32(0), + cumulative_data_physical_row=Int32(0), + cumulative_sf_physical_row=Int32(0), + cumulative_token_block_count=Int32(0), + valid_tokens_in_cta_cluster_tile=Int32(0), + phase_and_flags=Int32(BlockPhase.None_), + ) + + +@cute.jit +def _switch_to_fc2(mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: + cursor = mapping_state.cursor_state + cursor.current_phase = Int32(BlockPhase.Linear2) + cursor.current_expert_idx = cursor.current_group_first_expert - Int32(1) + cursor.current_expert_tile_start = cursor.current_group_fc1_subphase_end + cursor.current_expert_tile_end = cursor.current_group_fc1_subphase_end + cursor.current_expert_token_count = Int32(0) + cursor.current_token_block_count = Int32(0) + cursor.current_data_cumulative = cursor.group_start_data_cumulative + cursor.current_sf_cumulative = cursor.group_start_sf_cumulative + cursor.current_token_block_cumulative = cursor.group_start_token_block_cumulative + mapping_state.cursor_state = cursor + return cursor + + +@cute.jit +def _sum_expert_range( + mapping_state: Fc12TaskMappingState, expert_begin: Int32, expert_end: Int32 +) -> Tuple[Int32, Int32, Int32]: + lane_idx = Int32(cute.arch.lane_idx()) + data_rows = Int32(0) + sf_rows = Int32(0) + token_blocks = Int32(0) + batch_base = (expert_begin // Int32(32)) * Int32(32) + while batch_base < expert_end: + (_, _, lane_token_blocks, lane_data_rows, lane_sf_rows, _, _) = _load_expert_batch_metrics( + mapping_state, batch_base, expert_begin, expert_end, lane_idx + ) + data_rows = data_rows + Int32(cute.arch.warp_redux_sync(lane_data_rows, "add")) + sf_rows = sf_rows + Int32(cute.arch.warp_redux_sync(lane_sf_rows, "add")) + token_blocks = token_blocks + Int32(cute.arch.warp_redux_sync(lane_token_blocks, "add")) + batch_base = batch_base + Int32(32) + return data_rows, sf_rows, token_blocks + + +@cute.jit +def _build_group_range( + mapping_state: Fc12TaskMappingState, group_first_expert: Int32, base_fc1_tiles: Int32, base_fc2_tiles: Int32 +) -> Tuple[Int32, Int32, Int32]: + lane_idx = Int32(cute.arch.lane_idx()) + group_threshold = base_fc1_tiles + Int32(mapping_state.group_hint) + group_last_expert = group_first_expert + cumulative_fc1_tiles = base_fc1_tiles + cumulative_fc2_tiles = base_fc2_tiles + batch_base = (group_first_expert // Int32(32)) * Int32(32) + + while batch_base < mapping_state.expert_count and cumulative_fc1_tiles < group_threshold: + (lane_expert_idx, _, _, _, _, lane_fc1_tiles, lane_fc2_tiles) = _load_expert_batch_metrics( + mapping_state, batch_base, group_first_expert, Int32(mapping_state.expert_count), lane_idx + ) + fc1_prefix = _warp_inclusive_sum(lane_fc1_tiles, lane_idx) + reaches_threshold = ( + (lane_expert_idx >= group_first_expert) & (lane_expert_idx < mapping_state.expert_count) + ) & (cumulative_fc1_tiles + fc1_prefix >= group_threshold) + selected_lane = _first_matching_lane(reaches_threshold) + included_fc2_tiles = lane_fc2_tiles + if selected_lane >= Int32(0): + if lane_idx > selected_lane: + included_fc2_tiles = Int32(0) + fc2_batch_tiles = Int32(cute.arch.warp_redux_sync(included_fc2_tiles, "add")) + if selected_lane >= Int32(0): + cumulative_fc1_tiles = cumulative_fc1_tiles + Int32(cute.arch.shuffle_sync(fc1_prefix, selected_lane)) + group_last_expert = batch_base + selected_lane + Int32(1) + else: + cumulative_fc1_tiles = cumulative_fc1_tiles + Int32(cute.arch.shuffle_sync(fc1_prefix, Int32(31))) + batch_base = batch_base + Int32(32) + group_last_expert = cutlass.min(batch_base, Int32(mapping_state.expert_count)) + cumulative_fc2_tiles = cumulative_fc2_tiles + fc2_batch_tiles + + return group_last_expert, cumulative_fc1_tiles, cumulative_fc2_tiles + + +@cute.jit +def _advance_group(mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: + cursor = mapping_state.cursor_state + + residual_begin = cutlass.max(cursor.current_expert_idx, cursor.current_group_first_expert) + residual_data_rows = Int32(0) + residual_sf_rows = Int32(0) + residual_token_blocks = Int32(0) + if residual_begin < cursor.current_group_last_expert_exclusive: + iket.range_push("scheduler.residual_scan") + residual_data_rows, residual_sf_rows, residual_token_blocks = _sum_expert_range( + mapping_state, residual_begin, cursor.current_group_last_expert_exclusive + ) + iket.range_pop() + cursor.current_data_cumulative = cursor.current_data_cumulative + residual_data_rows + cursor.current_sf_cumulative = cursor.current_sf_cumulative + residual_sf_rows + cursor.current_token_block_cumulative = cursor.current_token_block_cumulative + residual_token_blocks + + cursor.group_start_data_cumulative = cursor.current_data_cumulative + cursor.group_start_sf_cumulative = cursor.current_sf_cumulative + cursor.group_start_token_block_cumulative = cursor.current_token_block_cumulative + + base_fc1_tiles = cursor.cumulative_fc1_tiles_at_group_end + base_fc2_tiles = cursor.cumulative_fc2_tiles_at_group_end + cursor.current_group_first_expert = cursor.current_group_last_expert_exclusive + + iket.range_push("scheduler.group_scan") + (cursor.current_group_last_expert_exclusive, cumulative_fc1_tiles, cumulative_fc2_tiles) = _build_group_range( + mapping_state, cursor.current_group_first_expert, base_fc1_tiles, base_fc2_tiles + ) + iket.range_pop() + cursor.cumulative_fc1_tiles_at_group_end = cumulative_fc1_tiles + cursor.cumulative_fc2_tiles_at_group_end = cumulative_fc2_tiles + group_start_tile = cursor.current_group_end + cursor.current_group_fc1_subphase_end = group_start_tile + cumulative_fc1_tiles - base_fc1_tiles + cursor.current_group_end = cursor.current_group_fc1_subphase_end + cumulative_fc2_tiles - base_fc2_tiles + + cursor.current_phase = Int32(BlockPhase.Linear1) + cursor.current_expert_idx = cursor.current_group_first_expert - Int32(1) + cursor.current_expert_tile_start = group_start_tile + cursor.current_expert_tile_end = group_start_tile + cursor.current_expert_token_count = Int32(0) + cursor.current_token_block_count = Int32(0) + mapping_state.cursor_state = cursor + return cursor + + +@cute.jit +def _seek_expert_for_work_id(linear_work_id: Int32, mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: + cursor = mapping_state.cursor_state + + base_tile_end = cursor.current_expert_tile_end + base_data_cumulative = cursor.current_data_cumulative + base_sf_cumulative = cursor.current_sf_cumulative + base_token_block_cumulative = cursor.current_token_block_cumulative + if cursor.current_expert_idx >= cursor.current_group_first_expert: + current_token_count = cursor.current_expert_token_count + base_data_cumulative = base_data_cumulative + ( + (current_token_count + Int32(mapping_state.token_padding_block - 1)) + // Int32(mapping_state.token_padding_block) + ) * Int32(mapping_state.token_padding_block) + base_sf_cumulative = base_sf_cumulative + ( + (current_token_count + Int32(mapping_state.sf_padding_block - 1)) // Int32(mapping_state.sf_padding_block) + ) * Int32(mapping_state.sf_padding_block) + base_token_block_cumulative = base_token_block_cumulative + cursor.current_token_block_count + + search_begin = cutlass.max(cursor.current_expert_idx + Int32(1), cursor.current_group_first_expert) + batch_base = (search_begin // Int32(32)) * Int32(32) + selected_expert = Int32(-1) + selected_token_count = Int32(0) + selected_token_blocks = Int32(0) + selected_tile_start = Int32(0) + selected_tile_end = Int32(0) + selected_data_cumulative = Int32(0) + selected_sf_cumulative = Int32(0) + selected_token_block_cumulative = Int32(0) + lane_idx = Int32(cute.arch.lane_idx()) + + iket.range_push("scheduler.expert_scan") + while selected_expert < Int32(0) and batch_base < cursor.current_group_last_expert_exclusive: + ( + lane_expert_idx, + lane_token_count, + lane_token_blocks, + lane_data_rows, + lane_sf_rows, + lane_fc1_tiles, + lane_fc2_tiles, + ) = _load_expert_batch_metrics( + mapping_state, batch_base, search_begin, cursor.current_group_last_expert_exclusive, lane_idx + ) + lane_phase_tiles = lane_fc1_tiles + if cursor.current_phase == Int32(BlockPhase.Linear2): + lane_phase_tiles = lane_fc2_tiles + tile_prefix = _warp_inclusive_sum(lane_phase_tiles, lane_idx) + candidate_tile_end = base_tile_end + tile_prefix + contains_work = ( + (lane_expert_idx >= search_begin) + & (lane_expert_idx < cursor.current_group_last_expert_exclusive) + & (linear_work_id < candidate_tile_end) + ) + selected_lane = _first_matching_lane(contains_work) + + included_data_rows = lane_data_rows + included_sf_rows = lane_sf_rows + included_token_blocks = lane_token_blocks + if selected_lane >= Int32(0): + if lane_idx > selected_lane: + included_data_rows = Int32(0) + included_sf_rows = Int32(0) + included_token_blocks = Int32(0) + batch_data_rows = Int32(cute.arch.warp_redux_sync(included_data_rows, "add")) + batch_sf_rows = Int32(cute.arch.warp_redux_sync(included_sf_rows, "add")) + batch_token_blocks = Int32(cute.arch.warp_redux_sync(included_token_blocks, "add")) + if selected_lane >= Int32(0): + selected_expert = batch_base + selected_lane + selected_token_count = Int32(cute.arch.shuffle_sync(lane_token_count, selected_lane)) + selected_token_blocks = Int32(cute.arch.shuffle_sync(lane_token_blocks, selected_lane)) + selected_phase_tiles = Int32(cute.arch.shuffle_sync(lane_phase_tiles, selected_lane)) + selected_tile_end = base_tile_end + Int32(cute.arch.shuffle_sync(tile_prefix, selected_lane)) + selected_tile_start = selected_tile_end - selected_phase_tiles + selected_data_rows = Int32(cute.arch.shuffle_sync(lane_data_rows, selected_lane)) + selected_sf_rows = Int32(cute.arch.shuffle_sync(lane_sf_rows, selected_lane)) + selected_data_cumulative = base_data_cumulative + batch_data_rows - selected_data_rows + selected_sf_cumulative = base_sf_cumulative + batch_sf_rows - selected_sf_rows + selected_token_block_cumulative = base_token_block_cumulative + batch_token_blocks - selected_token_blocks + else: + base_tile_end = base_tile_end + Int32(cute.arch.shuffle_sync(tile_prefix, Int32(31))) + base_data_cumulative = base_data_cumulative + batch_data_rows + base_sf_cumulative = base_sf_cumulative + batch_sf_rows + base_token_block_cumulative = base_token_block_cumulative + batch_token_blocks + batch_base = batch_base + Int32(32) + search_begin = batch_base + iket.range_pop() + + cursor.current_expert_idx = selected_expert + cursor.current_expert_token_count = selected_token_count + cursor.current_token_block_count = selected_token_blocks + cursor.current_expert_tile_start = selected_tile_start + cursor.current_expert_tile_end = selected_tile_end + cursor.current_data_cumulative = selected_data_cumulative + cursor.current_sf_cumulative = selected_sf_cumulative + cursor.current_token_block_cumulative = selected_token_block_cumulative + return cursor + + +@cute.jit +def _decode_inside_expert( + linear_work_id: Int32, cta_id_in_mapping_cluster: cute.Coord, mapping_state: Fc12TaskMappingState +) -> SchedulerWorkTileBase: + cursor = mapping_state.cursor_state + cta_tile_m = mapping_state.mapping_cta_tile_shape_mnk[0] + local_work_id = linear_work_id - cursor.current_expert_tile_start + + cluster_token_block_idx = Int32(0) + cluster_output_block_idx = Int32(0) + if cursor.current_phase == Int32(BlockPhase.Linear1): + cluster_token_block_idx = local_work_id // mapping_state.num_fc1_intermediate_blocks + cluster_output_block_idx = local_work_id - cluster_token_block_idx * mapping_state.num_fc1_intermediate_blocks + else: + cluster_token_block_idx = local_work_id // mapping_state.num_fc2_hidden_blocks + cluster_output_block_idx = local_work_id - cluster_token_block_idx * mapping_state.num_fc2_hidden_blocks + + cta_token_block_idx = ( + cluster_token_block_idx * mapping_state.mapping_cluster_shape_mn[0] + cta_id_in_mapping_cluster[0] + ) + cta_output_block_idx = ( + cluster_output_block_idx * mapping_state.mapping_cluster_shape_mn[1] + cta_id_in_mapping_cluster[1] + ) + token_start = cta_token_block_idx * Int32(cta_tile_m) + remaining_tokens = cutlass.max(cursor.current_expert_token_count - token_start, Int32(0)) + valid_tokens_in_cta_tile = cutlass.min(remaining_tokens, Int32(cta_tile_m)) + + if cutlass.const_expr(mapping_state.is_swap_ab): + return SwapAbFc12WorkTileInfo( + expert_idx=cursor.current_expert_idx, + tile_m_idx=cta_output_block_idx, + tile_n_idx=cta_token_block_idx, + cumulative_data_physical_row=cursor.current_data_cumulative, + cumulative_sf_physical_row=cursor.current_sf_cumulative, + cumulative_token_block_count=(cursor.current_token_block_cumulative), + valid_tokens_in_cta_tile=valid_tokens_in_cta_tile, + phase_and_flags=cursor.current_phase, + ) + + cluster_tile_m = mapping_state.mapping_cluster_shape_mn[0] * cta_tile_m + cluster_token_start = cluster_token_block_idx * Int32(cluster_tile_m) + remaining_cluster_tokens = cutlass.max(cursor.current_expert_token_count - cluster_token_start, Int32(0)) + valid_tokens_in_cluster_tile = cutlass.min(remaining_cluster_tokens, Int32(cluster_tile_m)) + valid_tokens_in_cta_cluster_tile = (valid_tokens_in_cta_tile << Int32(16)) | valid_tokens_in_cluster_tile + return NonSwapAbFc12WorkTileInfo( + expert_idx=cursor.current_expert_idx, + tile_m_idx=cta_token_block_idx, + tile_n_idx=cta_output_block_idx, + cumulative_data_physical_row=cursor.current_data_cumulative, + cumulative_sf_physical_row=cursor.current_sf_cumulative, + cumulative_token_block_count=(cursor.current_token_block_cumulative), + valid_tokens_in_cta_cluster_tile=(valid_tokens_in_cta_cluster_tile), + phase_and_flags=cursor.current_phase, + ) + + +@cute.jit +def map_fc12_linear_work_id( + linear_work_id: Int32, cta_id_in_mapping_cluster: cute.Coord, mapping_state: Fc12TaskMappingState +) -> Tuple[SchedulerWorkTileBase, Fc12TaskMappingState]: + """Map one monotonically increasing scalar ID to an FC12 work tile.""" + cursor = mapping_state.cursor_state + work_tile = make_fc12_done_tile(mapping_state.is_swap_ab) + + outer_group_end = cursor.current_group_end + outer_expert_end = cursor.current_group_last_expert_exclusive + while linear_work_id >= outer_group_end and outer_expert_end < mapping_state.expert_count: + mapping_state.cursor_state = _advance_group(mapping_state) + cursor = mapping_state.cursor_state + outer_group_end = cursor.current_group_end + outer_expert_end = cursor.current_group_last_expert_exclusive + cursor = mapping_state.cursor_state + + if linear_work_id < cursor.current_group_end: + if ( + cursor.current_phase == Int32(BlockPhase.Linear1) + and linear_work_id >= cursor.current_group_fc1_subphase_end + ): + mapping_state.cursor_state = _switch_to_fc2(mapping_state) + else: + mapping_state.cursor_state = mapping_state.cursor_state + cursor = mapping_state.cursor_state + + if linear_work_id >= cursor.current_expert_tile_end: + mapping_state.cursor_state = _seek_expert_for_work_id(linear_work_id, mapping_state) + else: + mapping_state.cursor_state = mapping_state.cursor_state + cursor = mapping_state.cursor_state + work_tile = _decode_inside_expert(linear_work_id, cta_id_in_mapping_cluster, mapping_state) + else: + mapping_state.cursor_state = mapping_state.cursor_state + return work_tile, mapping_state + + +class _PhaseFc12CursorState: + """Monotonic expert cursor for one phase-local FC12 work-ID stream.""" + + def __init__( + self, + expert_idx: Int32, + expert_tile_start: Int32, + expert_tile_end: Int32, + current_expert_token_count: Int32, + current_token_block_count: Int32, + data_cumulative: Int32, + sf_cumulative: Int32, + token_block_cumulative: Int32, + blocks_per_token_block: int, + ) -> None: + self.expert_idx = expert_idx + self.expert_tile_start = expert_tile_start + self.expert_tile_end = expert_tile_end + self.current_expert_token_count = current_expert_token_count + self.current_token_block_count = current_token_block_count + self.data_cumulative = data_cumulative + self.sf_cumulative = sf_cumulative + self.token_block_cumulative = token_block_cumulative + self.blocks_per_token_block = blocks_per_token_block + + def _runtime_fields(self) -> Tuple: + return ( + self.expert_idx, + self.expert_tile_start, + self.expert_tile_end, + self.current_expert_token_count, + self.current_token_block_count, + self.data_cumulative, + self.sf_cumulative, + self.token_block_cumulative, + ) + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in self._runtime_fields(): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "_PhaseFc12CursorState": + value_index = 0 + rebuilt_fields = [] + for field in self._runtime_fields(): + field_value_count = len(extract_mlir_values(field)) + rebuilt_fields.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) + value_index += field_value_count + if value_index != len(values): + raise ValueError( + f"_PhaseFc12CursorState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return type(self)(*rebuilt_fields, blocks_per_token_block=self.blocks_per_token_block) + + +class PhaseInterleavedFc12MappingState: + """Runtime inputs and independent FC1/FC2 cursors for phase-local IDs.""" + + def __init__( + self, + expert_count: int, + mapping_cta_tile_shape_mnk: Tuple[int, int, int], + mapping_cluster_shape_mn: Tuple[int, int], + token_padding_block: int, + sf_padding_block: int, + is_swap_ab: bool, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], + fc1_cursor: _PhaseFc12CursorState, + fc2_cursor: _PhaseFc12CursorState, + num_fc1_intermediate_blocks: int, + num_fc2_hidden_blocks: int, + ) -> None: + self.expert_count = expert_count + self.mapping_cta_tile_shape_mnk = mapping_cta_tile_shape_mnk + self.mapping_cluster_shape_mn = mapping_cluster_shape_mn + self.token_padding_block = token_padding_block + self.sf_padding_block = sf_padding_block + self.is_swap_ab = is_swap_ab + self.expert_token_sizes = expert_token_sizes + self.expert_token_prefix_sum = expert_token_prefix_sum + self.fc1_cursor = fc1_cursor + self.fc2_cursor = fc2_cursor + self.num_fc1_intermediate_blocks = num_fc1_intermediate_blocks + self.num_fc2_hidden_blocks = num_fc2_hidden_blocks + + @property + def mapping_cluster_tile_m(self) -> int: + return self.mapping_cta_tile_shape_mnk[0] * self.mapping_cluster_shape_mn[0] + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + token_counts = self.expert_token_sizes if self.expert_token_sizes is not None else self.expert_token_prefix_sum + values.extend(extract_mlir_values(token_counts)) + for field in (self.fc1_cursor, self.fc2_cursor): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "PhaseInterleavedFc12MappingState": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + if self.expert_token_sizes is not None: + expert_token_sizes = rebuild(self.expert_token_sizes) + expert_token_prefix_sum = None + else: + expert_token_sizes = None + expert_token_prefix_sum = rebuild(self.expert_token_prefix_sum) + result = type(self)( + expert_count=self.expert_count, + mapping_cta_tile_shape_mnk=self.mapping_cta_tile_shape_mnk, + mapping_cluster_shape_mn=self.mapping_cluster_shape_mn, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + is_swap_ab=self.is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + fc1_cursor=rebuild(self.fc1_cursor), + fc2_cursor=rebuild(self.fc2_cursor), + num_fc1_intermediate_blocks=self.num_fc1_intermediate_blocks, + num_fc2_hidden_blocks=self.num_fc2_hidden_blocks, + ) + if value_index != len(values): + raise ValueError( + f"PhaseInterleavedFc12MappingState MLIR value count mismatch: " + f"consumed {value_index}, got {len(values)}." + ) + return result + + +def _make_phase_cursor(blocks_per_token_block: int) -> _PhaseFc12CursorState: + return _PhaseFc12CursorState( + expert_idx=Int32(-1), + expert_tile_start=Int32(0), + expert_tile_end=Int32(0), + current_expert_token_count=Int32(0), + current_token_block_count=Int32(0), + data_cumulative=Int32(0), + sf_cumulative=Int32(0), + token_block_cumulative=Int32(0), + blocks_per_token_block=blocks_per_token_block, + ) + + +@cute.jit +def create_phase_interleaved_fc12_mapping_state( + *, + expert_count: int, + intermediate_gateup_size: int, + hidden_size: int, + mapping_cta_tile_shape_mnk: Tuple[int, int, int], + mapping_cluster_shape_mn: Tuple[int, int], + token_padding_block: int, + sf_padding_block: int, + is_swap_ab: bool, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], +) -> PhaseInterleavedFc12MappingState: + """Create independent monotonic mapping cursors for the FC1 and FC2 streams.""" + mapping_cluster_tile_n = mapping_cluster_shape_mn[1] * mapping_cta_tile_shape_mnk[1] + num_fc1_intermediate_blocks = (intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + num_fc2_hidden_blocks = (hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + return PhaseInterleavedFc12MappingState( + expert_count=expert_count, + mapping_cta_tile_shape_mnk=mapping_cta_tile_shape_mnk, + mapping_cluster_shape_mn=mapping_cluster_shape_mn, + token_padding_block=token_padding_block, + sf_padding_block=sf_padding_block, + is_swap_ab=is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + fc1_cursor=_make_phase_cursor(num_fc1_intermediate_blocks), + fc2_cursor=_make_phase_cursor(num_fc2_hidden_blocks), + num_fc1_intermediate_blocks=num_fc1_intermediate_blocks, + num_fc2_hidden_blocks=num_fc2_hidden_blocks, + ) + + +@cute.jit +def _advance_phase_cursor( + cursor: _PhaseFc12CursorState, mapping_state: PhaseInterleavedFc12MappingState +) -> _PhaseFc12CursorState: + previous_token_count = cursor.current_expert_token_count + cursor.data_cumulative = cursor.data_cumulative + ( + (previous_token_count + Int32(mapping_state.token_padding_block - 1)) + // Int32(mapping_state.token_padding_block) + ) * Int32(mapping_state.token_padding_block) + cursor.sf_cumulative = cursor.sf_cumulative + ( + (previous_token_count + Int32(mapping_state.sf_padding_block - 1)) // Int32(mapping_state.sf_padding_block) + ) * Int32(mapping_state.sf_padding_block) + cursor.token_block_cumulative = cursor.token_block_cumulative + cursor.current_token_block_count + + cursor.expert_idx = cursor.expert_idx + Int32(1) + token_count = Int32(0) + if cutlass.const_expr(mapping_state.expert_token_sizes is not None): + token_count = mapping_state.expert_token_sizes[cursor.expert_idx] + else: + prefix_end = mapping_state.expert_token_prefix_sum[cursor.expert_idx] + prefix_begin = Int32(0) + if cursor.expert_idx > Int32(0): + prefix_begin = mapping_state.expert_token_prefix_sum[cursor.expert_idx - Int32(1)] + token_count = prefix_end - prefix_begin + + cursor.current_expert_token_count = token_count + cursor.current_token_block_count = (token_count + Int32(mapping_state.mapping_cluster_tile_m - 1)) // Int32( + mapping_state.mapping_cluster_tile_m + ) + cursor.expert_tile_start = cursor.expert_tile_end + cursor.expert_tile_end = cursor.expert_tile_start + cursor.current_token_block_count * Int32( + cursor.blocks_per_token_block + ) + return cursor + + +@cute.jit +def _seek_phase_cursor( + linear_work_id: Int32, cursor: _PhaseFc12CursorState, mapping_state: PhaseInterleavedFc12MappingState +) -> _PhaseFc12CursorState: + expert_tile_end = cursor.expert_tile_end + next_expert_idx = cursor.expert_idx + Int32(1) + while linear_work_id >= expert_tile_end and next_expert_idx < Int32(mapping_state.expert_count): + cursor = _advance_phase_cursor(cursor, mapping_state) + expert_tile_end = cursor.expert_tile_end + next_expert_idx = cursor.expert_idx + Int32(1) + return cursor + + +@cute.jit +def _decode_phase_work_id( + linear_work_id: Int32, + phase: Int32, + cta_id_in_mapping_cluster: cute.Coord, + cursor: _PhaseFc12CursorState, + mapping_state: PhaseInterleavedFc12MappingState, +) -> SchedulerWorkTileBase: + local_work_id = linear_work_id - cursor.expert_tile_start + cluster_token_block_idx = local_work_id // Int32(cursor.blocks_per_token_block) + cluster_output_block_idx = local_work_id - cluster_token_block_idx * Int32(cursor.blocks_per_token_block) + cta_token_block_idx = ( + cluster_token_block_idx * Int32(mapping_state.mapping_cluster_shape_mn[0]) + cta_id_in_mapping_cluster[0] + ) + cta_output_block_idx = ( + cluster_output_block_idx * Int32(mapping_state.mapping_cluster_shape_mn[1]) + cta_id_in_mapping_cluster[1] + ) + + cta_tile_m = mapping_state.mapping_cta_tile_shape_mnk[0] + token_start = cta_token_block_idx * Int32(cta_tile_m) + remaining_tokens = cutlass.max(cursor.current_expert_token_count - token_start, Int32(0)) + valid_tokens_in_cta_tile = cutlass.min(remaining_tokens, Int32(cta_tile_m)) + + if cutlass.const_expr(mapping_state.is_swap_ab): + return SwapAbFc12WorkTileInfo( + expert_idx=cursor.expert_idx, + tile_m_idx=cta_output_block_idx, + tile_n_idx=cta_token_block_idx, + cumulative_data_physical_row=cursor.data_cumulative, + cumulative_sf_physical_row=cursor.sf_cumulative, + cumulative_token_block_count=cursor.token_block_cumulative, + valid_tokens_in_cta_tile=valid_tokens_in_cta_tile, + phase_and_flags=phase, + ) + + cluster_tile_m = mapping_state.mapping_cluster_shape_mn[0] * cta_tile_m + cluster_token_start = cluster_token_block_idx * Int32(cluster_tile_m) + remaining_cluster_tokens = cutlass.max(cursor.current_expert_token_count - cluster_token_start, Int32(0)) + valid_tokens_in_cluster_tile = cutlass.min(remaining_cluster_tokens, Int32(cluster_tile_m)) + return NonSwapAbFc12WorkTileInfo( + expert_idx=cursor.expert_idx, + tile_m_idx=cta_token_block_idx, + tile_n_idx=cta_output_block_idx, + cumulative_data_physical_row=cursor.data_cumulative, + cumulative_sf_physical_row=cursor.sf_cumulative, + cumulative_token_block_count=cursor.token_block_cumulative, + valid_tokens_in_cta_cluster_tile=(valid_tokens_in_cta_tile << Int32(16)) | valid_tokens_in_cluster_tile, + phase_and_flags=phase, + ) + + +@cute.jit +def map_phase_interleaved_fc12_work_id( + linear_work_id: Int32, + phase: Int32, + cta_id_in_mapping_cluster: cute.Coord, + mapping_state: PhaseInterleavedFc12MappingState, +) -> Tuple[SchedulerWorkTileBase, Boolean, PhaseInterleavedFc12MappingState]: + """Map one phase-local ID and report whether the selected stream contains it.""" + work_tile = make_fc12_done_tile(mapping_state.is_swap_ab) + stream_has_work = Boolean(False) + fc1_cursor = mapping_state.fc1_cursor + fc2_cursor = mapping_state.fc2_cursor + + if phase == Int32(BlockPhase.Linear1): + fc1_cursor = _seek_phase_cursor(linear_work_id, fc1_cursor, mapping_state) + if linear_work_id < fc1_cursor.expert_tile_end: + work_tile = _decode_phase_work_id( + linear_work_id, phase, cta_id_in_mapping_cluster, fc1_cursor, mapping_state + ) + stream_has_work = Boolean(True) + else: + fc2_cursor = _seek_phase_cursor(linear_work_id, fc2_cursor, mapping_state) + if linear_work_id < fc2_cursor.expert_tile_end: + work_tile = _decode_phase_work_id( + linear_work_id, phase, cta_id_in_mapping_cluster, fc2_cursor, mapping_state + ) + stream_has_work = Boolean(True) + + mapping_state.fc1_cursor = fc1_cursor + mapping_state.fc2_cursor = fc2_cursor + return work_tile, stream_has_work, mapping_state + + +__all__ = [ + "BlockPhase", + "Fc12TaskMappingState", + "Fc12WorkTileState", + "NonSwapAbFc12WorkTileInfo", + "PhaseInterleavedFc12MappingState", + "SwapAbFc12WorkTileInfo", + "create_fc12_task_mapping_state", + "create_phase_interleaved_fc12_mapping_state", + "make_fc12_done_tile", + "map_fc12_linear_work_id", + "map_phase_interleaved_fc12_work_id", + "peek_ready_bit", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_scheduler.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_scheduler.py new file mode 100644 index 000000000..ffedfb5c0 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_scheduler.py @@ -0,0 +1,692 @@ +"""Composable grouped and phase-interleaved FC12 schedulers.""" + +import math +from typing import Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Boolean, Int32, Integer, extract_mlir_values, new_from_mlir_values + +from ...api import ImplDesc, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.smem_workspace import SmemWorkspace +from ...helpers.utils import ceil_div +from .base import SchedulerBase, SchedulerWorkTileBase, WorkIdAcquisitionMode +from .fc12_mapping import ( + BlockPhase, + NonSwapAbFc12WorkTileInfo, + SwapAbFc12WorkTileInfo, + create_fc12_task_mapping_state, + create_phase_interleaved_fc12_mapping_state, + make_fc12_done_tile, + map_fc12_linear_work_id, + map_phase_interleaved_fc12_work_id, +) +from .non_clc_mixed_cga import NonClcMixedCgaConfig, NonClcMixedCgaSchedulerWorker + + +def _make_non_clc_mixed_cga_config(impl_desc: ImplDesc) -> NonClcMixedCgaConfig: + return NonClcMixedCgaConfig( + preferred_cluster_shape=impl_desc["cluster_shape_mn"], + fallback_cluster_shape=impl_desc.get("fallback_cluster_shape_mn"), + launch_cluster_count=impl_desc.get("launch_cluster_count"), + preferred_cluster_count=impl_desc.get("preferred_cluster_count"), + fallback_cluster_count=impl_desc.get("fallback_cluster_count"), + ) + + +def _mixed_cga_impl_requirements() -> dict: + return { + "fallback_cluster_shape_mn": OptionalRequirement(Optional[tuple]), + "launch_cluster_count": OptionalRequirement(Optional[int]), + "preferred_cluster_count": OptionalRequirement(Optional[int]), + "fallback_cluster_count": OptionalRequirement(Optional[int]), + } + + +def minimum_phase_interleave_hint( + *, blocks_fc1: int, blocks_fc2: int, launch_cluster_cnt_merge_as_preferred: int +) -> int: + """Return the per-cluster FC1 prologue covering one canonical FC2 claim wave.""" + effective_wave_width = launch_cluster_cnt_merge_as_preferred + max_dependent_token_blocks = ceil_div(effective_wave_width + blocks_fc2 - 1, blocks_fc2) + required_fc1_work = max_dependent_token_blocks * blocks_fc1 + return max(1, ceil_div(required_fc1_work, launch_cluster_cnt_merge_as_preferred)) + + +def _to_fc12_mapping_cta_coord(cta_coord_in_preferred_cluster: cute.Coord, is_swap_ab: bool) -> cute.Coord: + if cutlass.const_expr(not is_swap_ab): + return cta_coord_in_preferred_cluster + return (cta_coord_in_preferred_cluster[1], cta_coord_in_preferred_cluster[0], cta_coord_in_preferred_cluster[2]) + + +class BlackwellFusedFc12Scheduler(SchedulerBase): + """Compose work-ID claim, FC12 mapping, and SMEM tile transport.""" + + pipeline_mbarriers_region = "blackwell.fc12.scheduler.pipeline_mbarriers" + work_tiles_region = "blackwell.fc12.scheduler.work_tiles" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "expert_count": StaticOrRuntimeIntegerType, + "intermediate_gateup_size": StaticOrRuntimeIntegerType, + "hidden_size": StaticOrRuntimeIntegerType, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + **super().impl_desc_require(), + "mma_tiler_mnk": tuple, + "cluster_shape_mn": tuple, + "use_2cta_instrs": bool, + "hint": Optional[int], + "token_padding_block": int, + "sf_padding_block": int, + "work_id_mode": str, + "is_swap_ab": bool, + **_mixed_cga_impl_requirements(), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + super().__init__(problem_desc, impl_desc) + + self.expert_count = problem_desc["expert_count"] + self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] + self.hidden_size = problem_desc["hidden_size"] + self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + self.cluster_shape_mn = impl_desc["cluster_shape_mn"] + self.use_2cta_instrs = impl_desc["use_2cta_instrs"] + self.hint = impl_desc["hint"] + self.group_hint = self.hint + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.work_id_mode: WorkIdAcquisitionMode = impl_desc["work_id_mode"] + self.is_swap_ab = impl_desc["is_swap_ab"] + self.non_clc_mixed_cga_config = _make_non_clc_mixed_cga_config(impl_desc) + self.launch_cluster_cnt_merge_as_preferred = self.non_clc_mixed_cga_config.launch_cluster_cnt_merge_as_preferred + self.work_tile_type = SwapAbFc12WorkTileInfo if self.is_swap_ab else NonSwapAbFc12WorkTileInfo + if self.hint is None: + self.hint = self.launch_cluster_cnt_merge_as_preferred + self.group_hint = self.hint + + self._validate_configuration() + mma_cta_count = 2 if self.use_2cta_instrs else 1 + launch_cta_tile_shape_mnk = ( + self.mma_tiler_mnk[0] // mma_cta_count, + self.mma_tiler_mnk[1], + self.mma_tiler_mnk[2], + ) + if self.is_swap_ab: + self.mapping_cta_tile_shape_mnk = ( + launch_cta_tile_shape_mnk[1], + launch_cta_tile_shape_mnk[0], + launch_cta_tile_shape_mnk[2], + ) + self.mapping_cluster_shape_mn = (self.cluster_shape_mn[1], self.cluster_shape_mn[0]) + else: + self.mapping_cta_tile_shape_mnk = launch_cta_tile_shape_mnk + self.mapping_cluster_shape_mn = self.cluster_shape_mn + + if self.work_id_mode == "cluster_launch_control": + raise NotImplementedError("cluster_launch_control is not implemented for FC12.") + self._work_id_worker = NonClcMixedCgaSchedulerWorker( + config=self.non_clc_mixed_cga_config, work_id_mode=self.work_id_mode, stream_count=1 + ) + + def _validate_configuration(self) -> None: + if len(self.mma_tiler_mnk) != 3: + raise ValueError("mma_tiler_mnk must contain three dimensions.") + if len(self.cluster_shape_mn) != 2: + raise ValueError("cluster_shape_mn must contain two dimensions.") + if any(dimension <= 0 for dimension in self.mma_tiler_mnk): + raise ValueError("mma_tiler_mnk dimensions must be positive.") + if any(dimension <= 0 for dimension in self.cluster_shape_mn): + raise ValueError("cluster_shape_mn dimensions must be positive.") + mma_cta_count = 2 if self.use_2cta_instrs else 1 + if self.mma_tiler_mnk[0] % mma_cta_count != 0: + raise ValueError("mma_tiler M must be divisible by the MMA CTA count.") + if self.group_hint is not None and self.group_hint <= 0: + raise ValueError("group_hint must be positive.") + if self.token_padding_block <= 0: + raise ValueError("token_padding_block must be positive.") + if self.sf_padding_block <= 0: + raise ValueError("sf_padding_block must be positive.") + if self.launch_cluster_cnt_merge_as_preferred <= 0: + raise ValueError("launch_cluster_cnt_merge_as_preferred must be positive.") + if self.work_id_mode not in ("grid_stride", "atomic_counter", "cluster_launch_control"): + raise ValueError( + "work_id_mode must be 'grid_stride', 'atomic_counter', or " + f"'cluster_launch_control', got {self.work_id_mode!r}." + ) + cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] + if self.work_id_mode == "atomic_counter" and cluster_size > 32: + raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per cluster.") + fallback_cluster_shape = self.non_clc_mixed_cga_config.fallback_cluster_shape + if fallback_cluster_shape is not None: + fallback_cluster_size = fallback_cluster_shape[0] * fallback_cluster_shape[1] + if self.work_id_mode == "atomic_counter" and fallback_cluster_size > 32: + raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per fallback cluster.") + for field_name in ("expert_count", "intermediate_gateup_size", "hidden_size"): + value = getattr(self, field_name) + if isinstance(value, int) and value <= 0: + raise ValueError(f"{field_name} must be positive.") + static_dimensions = ( + isinstance(self.expert_count, int), + isinstance(self.intermediate_gateup_size, int), + isinstance(self.hidden_size, int), + ) + if any(static_dimensions) and not all(static_dimensions): + raise ValueError("FC12 expert dimensions must be either all static or all runtime.") + + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: + """Register scheduler-owned SMEM transport and claim regions.""" + super().register_smem_regions(smem_workspace) + self._work_id_worker.register_smem_regions(smem_workspace) + + def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: + """Register work-ID counters and optional fixed-group fallback state.""" + self._work_id_worker.register_device_workspace(device_workspace) + + @cute.jit + def initialize_fallback_group(self) -> None: + """Register this physical fallback cluster with its fixed logical group.""" + self._work_id_worker.initialize_fallback_group() + + def get_grid_shape(self, *, max_active_clusters: Optional[int] = None, problem_desc=None) -> Tuple[int, int, int]: + """Return the persistent launch grid in GEMM-domain orientation.""" + if ( + not self.non_clc_mixed_cga_config.is_mixed + and max_active_clusters is not None + and max_active_clusters < self.launch_cluster_cnt_merge_as_preferred + ): + raise ValueError( + f"max_active_clusters ({max_active_clusters}) must be at least " + "launch_cluster_cnt_merge_as_preferred " + f"({self.launch_cluster_cnt_merge_as_preferred})." + ) + return (self.cluster_shape_mn[0], self.cluster_shape_mn[1], self.launch_cluster_cnt_merge_as_preferred) + + @cute.jit + def assign_device_members( + self, + *, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], + actual_expert_shape: Optional[Tuple], + block_idx: Tuple[Integer, Integer, Integer], + smem_workspace: SmemWorkspace, + smem_base: cute.Pointer, + device_workspace: DeviceWorkspace, + is_fallback_cluster: Optional[Boolean] = None, + ) -> None: + """Initialize all FC12 scheduler state rooted for one CTA lifetime.""" + if cutlass.const_expr((expert_token_sizes is None) == (expert_token_prefix_sum is None)): + raise ValueError("Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided.") + needs_actual_shape = not all( + isinstance(dimension, int) + for dimension in (self.expert_count, self.intermediate_gateup_size, self.hidden_size) + ) + if cutlass.const_expr(needs_actual_shape and actual_expert_shape is None): + raise ValueError("actual_expert_shape is required for runtime dimensions.") + if cutlass.const_expr(isinstance(self.expert_count, int)): + expert_count = self.expert_count + else: + expert_count = actual_expert_shape[0] + if cutlass.const_expr(isinstance(self.intermediate_gateup_size, int)): + intermediate_gateup_size = self.intermediate_gateup_size + else: + intermediate_gateup_size = actual_expert_shape[1] + if cutlass.const_expr(isinstance(self.hidden_size, int)): + hidden_size = self.hidden_size + else: + hidden_size = actual_expert_shape[2] + + self.create_scheduler_pipelines(smem_workspace, smem_base) + self._work_id_worker.assign_device_members( + is_fallback_cluster=is_fallback_cluster, + block_idx=block_idx, + smem_workspace=smem_workspace, + smem_base=smem_base, + device_workspace=device_workspace, + ) + + task_mapping_state = create_fc12_task_mapping_state( + expert_count=expert_count, + intermediate_gateup_size=intermediate_gateup_size, + hidden_size=hidden_size, + mapping_cta_tile_shape_mnk=(self.mapping_cta_tile_shape_mnk), + mapping_cluster_shape_mn=self.mapping_cluster_shape_mn, + group_hint=self.group_hint, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + is_swap_ab=self.is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + ) + + self._task_mapping_state = task_mapping_state + + @cute.jit + def gen_next_work(self) -> SchedulerWorkTileBase: + """Claim and map one work tile without first-tile prefetch.""" + work_id = self._work_id_worker.claim_next_work() + cta_id_in_mapping_cluster = _to_fc12_mapping_cta_coord( + self._work_id_worker.cta_coord_in_preferred_cluster, self.is_swap_ab + ) + work_tile, self._task_mapping_state = map_fc12_linear_work_id( + work_id, cta_id_in_mapping_cluster, self._task_mapping_state + ) + return work_tile + + def __extract_mlir_values__(self) -> list: + values = super().__extract_mlir_values__() + values.extend(extract_mlir_values(self._work_id_worker)) + values.extend(extract_mlir_values(self._task_mapping_state)) + return values + + def __new_from_mlir_values__(self, values: list) -> "BlackwellFusedFc12Scheduler": + base_value_count = len(super().__extract_mlir_values__()) + if len(values) < base_value_count: + raise ValueError( + "BlackwellFusedFc12Scheduler MLIR value count is smaller than " + f"its base state: expected at least {base_value_count}, got {len(values)}." + ) + result = super().__new_from_mlir_values__(values[:base_value_count]) + value_index = base_value_count + + def rebuild(state): + nonlocal value_index + state_value_count = len(extract_mlir_values(state)) + rebuilt_state = new_from_mlir_values(state, values[value_index : value_index + state_value_count]) + value_index += state_value_count + return rebuilt_state + + result._work_id_worker = rebuild(self._work_id_worker) + result._task_mapping_state = rebuild(self._task_mapping_state) + if value_index != len(values): + raise ValueError( + f"BlackwellFusedFc12Scheduler MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + + for field_name in ( + "expert_count", + "intermediate_gateup_size", + "hidden_size", + "mma_tiler_mnk", + "cluster_shape_mn", + "use_2cta_instrs", + "hint", + "group_hint", + "token_padding_block", + "sf_padding_block", + "work_id_mode", + "is_swap_ab", + "launch_cluster_cnt_merge_as_preferred", + "non_clc_mixed_cga_config", + "mapping_cta_tile_shape_mnk", + "mapping_cluster_shape_mn", + ): + setattr(result, field_name, getattr(self, field_name)) + return result + + +class _PhaseInterleaveControlState: + """Per-cluster phase cadence and stream exhaustion state.""" + + def __init__( + self, prologue_remaining: Int32, cycle_position: Int32, fc1_exhausted: Boolean, fc2_exhausted: Boolean + ) -> None: + self.prologue_remaining = prologue_remaining + self.cycle_position = cycle_position + self.fc1_exhausted = fc1_exhausted + self.fc2_exhausted = fc2_exhausted + + def _fields(self) -> Tuple: + return (self.prologue_remaining, self.cycle_position, self.fc1_exhausted, self.fc2_exhausted) + + def __extract_mlir_values__(self) -> list: + values = [] + for field in self._fields(): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: list) -> "_PhaseInterleaveControlState": + value_index = 0 + rebuilt_fields = [] + for field in self._fields(): + field_value_count = len(extract_mlir_values(field)) + rebuilt_fields.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) + value_index += field_value_count + if value_index != len(values): + raise ValueError( + f"_PhaseInterleaveControlState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return type(self)(*rebuilt_fields) + + +class PhaseInterleavedFc12Scheduler(SchedulerBase): + """Schedule independent FC1 and FC2 streams with per-phase atomic counters.""" + + pipeline_mbarriers_region = "fc12.phase_interleaved.scheduler.pipeline_mbarriers" + work_tiles_region = "fc12.phase_interleaved.scheduler.work_tiles" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return {"expert_count": int, "intermediate_gateup_size": int, "hidden_size": int} + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + **super().impl_desc_require(), + "mma_tiler_mnk": tuple, + "cluster_shape_mn": tuple, + "use_2cta_instrs": bool, + "hint": int, + "token_padding_block": int, + "sf_padding_block": int, + "work_id_mode": str, + "is_swap_ab": bool, + **_mixed_cga_impl_requirements(), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + super().__init__(problem_desc, impl_desc) + + self.expert_count = problem_desc["expert_count"] + self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] + self.hidden_size = problem_desc["hidden_size"] + self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + self.cluster_shape_mn = impl_desc["cluster_shape_mn"] + self.use_2cta_instrs = impl_desc["use_2cta_instrs"] + self.hint = impl_desc["hint"] + self.fc1_prologue_tiles = self.hint + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.work_id_mode: WorkIdAcquisitionMode = impl_desc["work_id_mode"] + self.is_swap_ab = impl_desc["is_swap_ab"] + self.non_clc_mixed_cga_config = _make_non_clc_mixed_cga_config(impl_desc) + self.launch_cluster_cnt_merge_as_preferred = self.non_clc_mixed_cga_config.launch_cluster_cnt_merge_as_preferred + self.work_tile_type = SwapAbFc12WorkTileInfo if self.is_swap_ab else NonSwapAbFc12WorkTileInfo + + self._validate_configuration() + mma_cta_count = 2 if self.use_2cta_instrs else 1 + launch_cta_tile_shape_mnk = ( + self.mma_tiler_mnk[0] // mma_cta_count, + self.mma_tiler_mnk[1], + self.mma_tiler_mnk[2], + ) + if self.is_swap_ab: + self.mapping_cta_tile_shape_mnk = ( + launch_cta_tile_shape_mnk[1], + launch_cta_tile_shape_mnk[0], + launch_cta_tile_shape_mnk[2], + ) + self.mapping_cluster_shape_mn = (self.cluster_shape_mn[1], self.cluster_shape_mn[0]) + else: + self.mapping_cta_tile_shape_mnk = launch_cta_tile_shape_mnk + self.mapping_cluster_shape_mn = self.cluster_shape_mn + self._work_id_worker = NonClcMixedCgaSchedulerWorker( + config=self.non_clc_mixed_cga_config, work_id_mode=self.work_id_mode, stream_count=2 + ) + + mapping_cluster_tile_n = self.mapping_cta_tile_shape_mnk[1] * self.mapping_cluster_shape_mn[1] + self.blocks_fc1 = (self.intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + self.blocks_fc2 = (self.hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + interleave_gcd = math.gcd(self.blocks_fc1, self.blocks_fc2) + self.interleave_fc2_slots = self.blocks_fc2 // interleave_gcd + self.interleave_cycle_length = (self.blocks_fc1 + self.blocks_fc2) // interleave_gcd + minimum_hint = minimum_phase_interleave_hint( + blocks_fc1=self.blocks_fc1, + blocks_fc2=self.blocks_fc2, + launch_cluster_cnt_merge_as_preferred=self.launch_cluster_cnt_merge_as_preferred, + ) + if self.fc1_prologue_tiles < minimum_hint: + raise ValueError( + f"phase_interleave hint {self.fc1_prologue_tiles} cannot cover a " + f"{self.launch_cluster_cnt_merge_as_preferred}-cluster FC2 claim wave; " + f"raise hint to at least {minimum_hint}." + ) + + def _validate_configuration(self) -> None: + if len(self.mma_tiler_mnk) != 3: + raise ValueError("mma_tiler_mnk must contain three dimensions.") + if len(self.cluster_shape_mn) != 2: + raise ValueError("cluster_shape_mn must contain two dimensions.") + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in self.mma_tiler_mnk): + raise TypeError("mma_tiler_mnk dimensions must be Python ints.") + if not all( + isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in self.cluster_shape_mn + ): + raise TypeError("cluster_shape_mn dimensions must be Python ints.") + if any(dimension <= 0 for dimension in self.mma_tiler_mnk): + raise ValueError("mma_tiler_mnk dimensions must be positive.") + if any(dimension <= 0 for dimension in self.cluster_shape_mn): + raise ValueError("cluster_shape_mn dimensions must be positive.") + mma_cta_count = 2 if self.use_2cta_instrs else 1 + if self.mma_tiler_mnk[0] % mma_cta_count != 0: + raise ValueError("mma_tiler M must be divisible by the MMA CTA count.") + if ( + isinstance(self.fc1_prologue_tiles, bool) + or not isinstance(self.fc1_prologue_tiles, int) + or self.fc1_prologue_tiles <= 0 + ): + raise ValueError("fc1_prologue_tiles must be a positive Python int resolved by the kernel frontend.") + if self.token_padding_block <= 0: + raise ValueError("token_padding_block must be positive.") + if self.sf_padding_block <= 0: + raise ValueError("sf_padding_block must be positive.") + if self.launch_cluster_cnt_merge_as_preferred <= 0: + raise ValueError("launch_cluster_cnt_merge_as_preferred must be positive.") + if self.work_id_mode != "atomic_counter": + raise ValueError("Phase-interleaved FC12 scheduling currently requires work_id_mode='atomic_counter'.") + cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] + if cluster_size > 32: + raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per cluster.") + fallback_cluster_shape = self.non_clc_mixed_cga_config.fallback_cluster_shape + if fallback_cluster_shape is not None: + fallback_cluster_size = fallback_cluster_shape[0] * fallback_cluster_shape[1] + if fallback_cluster_size > 32: + raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per fallback cluster.") + maximum_int32 = (1 << 31) - 1 + for field_name in ("expert_count", "intermediate_gateup_size", "hidden_size"): + value = getattr(self, field_name) + if isinstance(value, bool) or value <= 0: + raise ValueError(f"{field_name} must be a positive Python int.") + if value > maximum_int32: + raise ValueError(f"{field_name} must fit in a signed Int32, got {value}.") + + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: + """Register work transport and the phase-counter broadcast channel.""" + super().register_smem_regions(smem_workspace) + self._work_id_worker.register_smem_regions(smem_workspace) + + def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: + """Register independently reset FC1 and FC2 work-ID streams.""" + self._work_id_worker.register_device_workspace(device_workspace) + + @cute.jit + def initialize_fallback_group(self) -> None: + """Register this physical fallback cluster with its fixed logical group.""" + self._work_id_worker.initialize_fallback_group() + + def get_grid_shape(self, *, max_active_clusters: Optional[int] = None, problem_desc=None) -> Tuple[int, int, int]: + """Return the statically configured persistent launch grid.""" + if ( + not self.non_clc_mixed_cga_config.is_mixed + and max_active_clusters is not None + and max_active_clusters < self.launch_cluster_cnt_merge_as_preferred + ): + raise ValueError( + f"max_active_clusters ({max_active_clusters}) must be at least " + "launch_cluster_cnt_merge_as_preferred " + f"({self.launch_cluster_cnt_merge_as_preferred})." + ) + return (self.cluster_shape_mn[0], self.cluster_shape_mn[1], self.launch_cluster_cnt_merge_as_preferred) + + @cute.jit + def assign_device_members( + self, + *, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], + actual_expert_shape: Optional[Tuple], + block_idx: Tuple[Integer, Integer, Integer], + smem_workspace: SmemWorkspace, + smem_base: cute.Pointer, + device_workspace: DeviceWorkspace, + is_fallback_cluster: Optional[Boolean] = None, + ) -> None: + """Initialize phase-local mapping, cadence, and counter state.""" + if cutlass.const_expr((expert_token_sizes is None) == (expert_token_prefix_sum is None)): + raise ValueError("Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided.") + self.create_scheduler_pipelines(smem_workspace, smem_base) + self._work_id_worker.assign_device_members( + is_fallback_cluster=is_fallback_cluster, + block_idx=block_idx, + smem_workspace=smem_workspace, + smem_base=smem_base, + device_workspace=device_workspace, + ) + self._task_mapping_state = create_phase_interleaved_fc12_mapping_state( + expert_count=self.expert_count, + intermediate_gateup_size=self.intermediate_gateup_size, + hidden_size=self.hidden_size, + mapping_cta_tile_shape_mnk=self.mapping_cta_tile_shape_mnk, + mapping_cluster_shape_mn=self.mapping_cluster_shape_mn, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + is_swap_ab=self.is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + ) + self._control_state = _PhaseInterleaveControlState( + prologue_remaining=Int32(self.fc1_prologue_tiles), + cycle_position=Int32(0), + fc1_exhausted=Boolean(False), + fc2_exhausted=Boolean(False), + ) + + @cute.jit + def gen_next_work(self) -> SchedulerWorkTileBase: + """Claim until one stream yields valid work or both streams terminate.""" + work_tile = make_fc12_done_tile(self.is_swap_ab) + work_id_worker = self._work_id_worker + task_mapping_state = self._task_mapping_state + control_state = self._control_state + prologue_remaining = control_state.prologue_remaining + cycle_position = control_state.cycle_position + fc1_exhausted = control_state.fc1_exhausted + fc2_exhausted = control_state.fc2_exhausted + resolved = Boolean(False) + + while not resolved: + if fc1_exhausted and fc2_exhausted: + work_tile = make_fc12_done_tile(self.is_swap_ab) + resolved = Boolean(True) + else: + want_fc1 = Boolean(True) + if prologue_remaining <= Int32(0): + is_fc2_slot = (cycle_position * Int32(self.interleave_fc2_slots)) % Int32( + self.interleave_cycle_length + ) < Int32(self.interleave_fc2_slots) + want_fc1 = not is_fc2_slot + if want_fc1 and fc1_exhausted: + want_fc1 = Boolean(False) + if (not want_fc1) and fc2_exhausted: + want_fc1 = Boolean(True) + + atomic_counter_index = Int32(1) + if want_fc1: + atomic_counter_index = Int32(0) + linear_work_id = work_id_worker.claim_next_work(atomic_counter_index) + want_fc1 = work_id_worker.claimed_stream_index == Int32(0) + phase = Int32(BlockPhase.Linear2) + if want_fc1: + phase = Int32(BlockPhase.Linear1) + cta_id_in_mapping_cluster = _to_fc12_mapping_cta_coord( + work_id_worker.cta_coord_in_preferred_cluster, self.is_swap_ab + ) + work_tile, stream_has_work, task_mapping_state = map_phase_interleaved_fc12_work_id( + linear_work_id, phase, cta_id_in_mapping_cluster, task_mapping_state + ) + if stream_has_work: + if prologue_remaining > Int32(0): + prologue_remaining = prologue_remaining - Int32(1) + else: + cycle_position = (cycle_position + Int32(1)) % Int32(self.interleave_cycle_length) + resolved = Boolean(True) + else: + if want_fc1: + fc1_exhausted = Boolean(True) + else: + fc2_exhausted = Boolean(True) + + control_state.prologue_remaining = prologue_remaining + control_state.cycle_position = cycle_position + control_state.fc1_exhausted = fc1_exhausted + control_state.fc2_exhausted = fc2_exhausted + self._work_id_worker = work_id_worker + self._task_mapping_state = task_mapping_state + self._control_state = control_state + return work_tile + + def __extract_mlir_values__(self) -> list: + values = super().__extract_mlir_values__() + for state in (self._work_id_worker, self._task_mapping_state, self._control_state): + values.extend(extract_mlir_values(state)) + return values + + def __new_from_mlir_values__(self, values: list) -> "PhaseInterleavedFc12Scheduler": + base_value_count = len(super().__extract_mlir_values__()) + if len(values) < base_value_count: + raise ValueError( + "PhaseInterleavedFc12Scheduler MLIR value count is smaller than " + f"its base state: expected at least {base_value_count}, got {len(values)}." + ) + result = super().__new_from_mlir_values__(values[:base_value_count]) + value_index = base_value_count + + def rebuild(state): + nonlocal value_index + state_value_count = len(extract_mlir_values(state)) + rebuilt_state = new_from_mlir_values(state, values[value_index : value_index + state_value_count]) + value_index += state_value_count + return rebuilt_state + + result._work_id_worker = rebuild(self._work_id_worker) + result._task_mapping_state = rebuild(self._task_mapping_state) + result._control_state = rebuild(self._control_state) + if value_index != len(values): + raise ValueError( + f"PhaseInterleavedFc12Scheduler MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + + for field_name in ( + "expert_count", + "intermediate_gateup_size", + "hidden_size", + "mma_tiler_mnk", + "cluster_shape_mn", + "use_2cta_instrs", + "hint", + "fc1_prologue_tiles", + "token_padding_block", + "sf_padding_block", + "work_id_mode", + "is_swap_ab", + "launch_cluster_cnt_merge_as_preferred", + "non_clc_mixed_cga_config", + "mapping_cta_tile_shape_mnk", + "mapping_cluster_shape_mn", + "blocks_fc1", + "blocks_fc2", + "interleave_fc2_slots", + "interleave_cycle_length", + ): + setattr(result, field_name, getattr(self, field_name)) + return result + + +__all__ = ["BlackwellFusedFc12Scheduler", "PhaseInterleavedFc12Scheduler", "minimum_phase_interleave_hint"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/non_clc_mixed_cga.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/non_clc_mixed_cga.py new file mode 100644 index 000000000..c0f0fa8a5 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/non_clc_mixed_cga.py @@ -0,0 +1,339 @@ +"""Reusable preferred/fallback cluster scheduling without hardware CLC.""" + +import dataclasses +from typing import List, Optional, Tuple + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass._mlir import ir +from cutlass.cutlass_dsl import Boolean, Int32, Int64, extract_mlir_values, new_from_mlir_values + +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.smem_workspace import SmemWorkspace +from .base import WorkIdAcquisitionMode +from .work_id_claim import ( + AtomicCounterWorkIdState, + FixedGroupMixedCgaAtomicCounterWorkIdState, + GridStrideWorkIdState, + claim_work_id, + initialize_fixed_group_mixed_cga_work_id_state, +) + + +@dataclasses.dataclass(frozen=True) +class NonClcMixedCgaConfig: + """Static launch geometry shared by non-CLC mixed-CGA schedulers.""" + + preferred_cluster_shape: Tuple[int, int] + fallback_cluster_shape: Optional[Tuple[int, int]] + launch_cluster_count: Optional[int] + preferred_cluster_count: Optional[int] + fallback_cluster_count: Optional[int] + mn_split_factors: Tuple[int, int] = dataclasses.field(init=False) + split_factor: int = dataclasses.field(init=False) + launch_cluster_cnt_merge_as_preferred: int = dataclasses.field(init=False) + total_cta_cnt: int = dataclasses.field(init=False) + is_mixed: bool = dataclasses.field(init=False) + + def __post_init__(self) -> None: + self._validate_shape(self.preferred_cluster_shape, "preferred_cluster_shape") + preferred_cluster_size = self._shape_size(self.preferred_cluster_shape) + if self.fallback_cluster_shape is not None: + self._validate_shape(self.fallback_cluster_shape, "fallback_cluster_shape") + has_no_fallback_clusters = ( + isinstance(self.fallback_cluster_count, int) + and not isinstance(self.fallback_cluster_count, bool) + and self.fallback_cluster_count == 0 + ) + if self.fallback_cluster_shape == self.preferred_cluster_shape or has_no_fallback_clusters: + object.__setattr__(self, "fallback_cluster_shape", None) + object.__setattr__(self, "preferred_cluster_count", None) + object.__setattr__(self, "fallback_cluster_count", None) + + if self.fallback_cluster_shape is None: + if not self._is_positive_int(self.launch_cluster_count): + raise ValueError("launch_cluster_count must be positive when fallback_cluster_shape is absent.") + mn_split_factors = (1, 1) + split_factor = 1 + launch_cluster_cnt_merge_as_preferred = self.launch_cluster_count + is_mixed = False + else: + if len(self.fallback_cluster_shape) != len(self.preferred_cluster_shape): + raise ValueError("preferred and fallback cluster shapes must have equal rank.") + if not self._is_positive_int(self.preferred_cluster_count): + raise ValueError("preferred_cluster_count must be positive when fallback_cluster_shape is present.") + if ( + isinstance(self.fallback_cluster_count, bool) + or not isinstance(self.fallback_cluster_count, int) + or self.fallback_cluster_count < 0 + ): + raise ValueError( + "fallback_cluster_count must be a non-negative Python int when fallback_cluster_shape is present." + ) + + mn_split_factors = ( + self.preferred_cluster_shape[0] // self.fallback_cluster_shape[0], + self.preferred_cluster_shape[1] // self.fallback_cluster_shape[1], + ) + if any( + preferred_dimension % fallback_dimension != 0 + for preferred_dimension, fallback_dimension in zip( + self.preferred_cluster_shape, self.fallback_cluster_shape + ) + ): + raise ValueError("Every preferred cluster dimension must be divisible by its fallback dimension.") + split_factor = self._shape_size(mn_split_factors) + is_mixed = self.fallback_cluster_shape != self.preferred_cluster_shape + if is_mixed: + if split_factor > 16 or split_factor & (split_factor - 1): + raise ValueError("The preferred/fallback cluster split factor must be a power of two at most 16.") + if self.fallback_cluster_count % split_factor != 0: + raise ValueError( + "fallback_cluster_count must be divisible by the preferred/fallback cluster split factor." + ) + launch_cluster_cnt_merge_as_preferred = ( + self.preferred_cluster_count + self.fallback_cluster_count // split_factor + ) + else: + launch_cluster_cnt_merge_as_preferred = self.preferred_cluster_count + self.fallback_cluster_count + + if launch_cluster_cnt_merge_as_preferred <= 0: + raise ValueError("The resolved launch must contain at least one cluster merged as preferred.") + object.__setattr__(self, "mn_split_factors", mn_split_factors) + object.__setattr__(self, "split_factor", split_factor) + object.__setattr__(self, "launch_cluster_cnt_merge_as_preferred", launch_cluster_cnt_merge_as_preferred) + object.__setattr__(self, "total_cta_cnt", launch_cluster_cnt_merge_as_preferred * preferred_cluster_size) + object.__setattr__(self, "is_mixed", is_mixed) + + @staticmethod + def _shape_size(shape: Tuple[int, int]) -> int: + result = 1 + for dimension in shape: + result *= dimension + return result + + @staticmethod + def _is_positive_int(value) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + @classmethod + def _validate_shape(cls, shape: Tuple[int, int], field_name: str) -> None: + if not isinstance(shape, tuple) or len(shape) != 2: + raise TypeError(f"{field_name} must be a two-dimensional tuple.") + if not all(cls._is_positive_int(dimension) for dimension in shape): + raise ValueError(f"{field_name} dimensions must be positive Python ints.") + + +class NonClcMixedCgaSchedulerWorker: + """Compose canonical work-ID acquisition with preferred/fallback cluster splitting.""" + + cluster_pipeline_mbarriers_region = "non_clc_mixed_cga.cluster_pipeline_mbarriers" + cluster_broadcast_region = "non_clc_mixed_cga.cluster_broadcast" + work_id_counter_region = "non_clc_mixed_cga.work_id_counters" + fallback_registration_counter_region = "non_clc_mixed_cga.fallback_registration_counter" + fallback_group_token_region = "non_clc_mixed_cga.fallback_group_tokens" + + def __init__(self, *, config: NonClcMixedCgaConfig, work_id_mode: WorkIdAcquisitionMode, stream_count: int) -> None: + if work_id_mode not in ("grid_stride", "atomic_counter"): + raise ValueError("Non-CLC mixed-CGA scheduling supports grid_stride or atomic_counter work IDs.") + if isinstance(stream_count, bool) or not isinstance(stream_count, int) or stream_count <= 0: + raise ValueError("stream_count must be a positive Python int.") + self.config = config + self.work_id_mode = work_id_mode + self.stream_count = stream_count + + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: + """Register the cluster-local claim broadcast channel.""" + if self.work_id_mode == "atomic_counter": + smem_workspace.register_mbarrier(self.cluster_pipeline_mbarriers_region, 2) + smem_workspace.register_tensor(self.cluster_broadcast_region, cutlass.Int32, (1,)) + + def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: + """Register atomic counters and optional fixed-group fallback handoff state.""" + if self.work_id_mode == "atomic_counter": + device_workspace.register( + self.work_id_counter_region, + cutlass.Int32, + (self.stream_count,), + buffer_space="local", + reset="tail_reset", + ) + if self.config.is_mixed: + device_workspace.register( + self.fallback_registration_counter_region, + cutlass.Int32, + (1,), + buffer_space="local", + reset="tail_reset", + ) + device_workspace.register( + self.fallback_group_token_region, + cutlass.Int64, + (self.config.fallback_cluster_count,), + buffer_space="local", + reset="tail_reset", + ) + + @cute.jit + def assign_device_members( + self, + *, + is_fallback_cluster: Optional[Boolean], + block_idx: Tuple, + smem_workspace: SmemWorkspace, + smem_base: cute.Pointer, + device_workspace: DeviceWorkspace, + ) -> None: + """Bind one active cluster to the configured non-CLC claim backend.""" + if cutlass.const_expr(self.config.is_mixed and is_fallback_cluster is None): + raise ValueError("is_fallback_cluster is required for a true mixed-CGA launch.") + + active_cluster_m = self.config.preferred_cluster_shape[0] + active_cluster_n = self.config.preferred_cluster_shape[1] + if cutlass.const_expr(self.config.is_mixed): + if is_fallback_cluster: + active_cluster_m = Int32(self.config.fallback_cluster_shape[0]) + active_cluster_n = Int32(self.config.fallback_cluster_shape[1]) + + cta_coord_in_active_cluster = ( + Int32(block_idx[0]) % Int32(active_cluster_m), + Int32(block_idx[1]) % Int32(active_cluster_n), + Int32(0), + ) + cta_coord_in_preferred_cluster = cta_coord_in_active_cluster + + if cutlass.const_expr(self.config.is_mixed and self.work_id_mode == "grid_stride"): + if is_fallback_cluster: + flattened_index = Int32(0) + dimension_stride = 1 + for dimension_idx in cutlass.range_constexpr(len(self.config.mn_split_factors)): + preferred_dimension = self.config.preferred_cluster_shape[dimension_idx] + fallback_dimension = self.config.fallback_cluster_shape[dimension_idx] + inner_coordinate = (Int32(block_idx[dimension_idx]) % Int32(preferred_dimension)) // Int32( + fallback_dimension + ) + flattened_index = flattened_index + inner_coordinate * Int32(dimension_stride) + dimension_stride *= self.config.mn_split_factors[dimension_idx] + cta_coord_in_preferred_cluster = self._preferred_cluster_cta_coord( + cta_coord_in_active_cluster, flattened_index + ) + self.cta_coord_in_preferred_cluster = cta_coord_in_preferred_cluster + + if cutlass.const_expr(self.work_id_mode == "atomic_counter"): + active_cluster_size = active_cluster_m * active_cluster_n + cluster_pipeline = pipeline.PipelineAsync.create( + num_stages=1, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * active_cluster_size), + barrier_storage=smem_workspace.ptr(self.cluster_pipeline_mbarriers_region, smem_base), + defer_sync=True, + ) + atomic_counter_state = AtomicCounterWorkIdState( + counter_pointer=device_workspace.ptr(self.work_id_counter_region), + counter_count=self.stream_count, + broadcast_pointer=smem_workspace.ptr(self.cluster_broadcast_region, smem_base), + is_leader_cta=(cta_coord_in_active_cluster[0] + cta_coord_in_active_cluster[1]) == Int32(0), + cluster_pipeline=cluster_pipeline, + producer_state=pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1), + consumer_state=pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1), + cluster_size=active_cluster_size, + ) + if cutlass.const_expr(self.config.is_mixed): + self._cta_coord_in_active_cluster = cta_coord_in_active_cluster + self._work_id_state = FixedGroupMixedCgaAtomicCounterWorkIdState( + atomic_counter_state=atomic_counter_state, + registration_counter_pointer=device_workspace.ptr(self.fallback_registration_counter_region), + group_token_pointer=device_workspace.ptr(self.fallback_group_token_region), + split_factor=self.config.split_factor, + fallback_cluster_count=self.config.fallback_cluster_count, + is_fallback_cluster=is_fallback_cluster, + fallback_group_idx=Int32(0), + in_group_idx=Int32(0), + previous_token=Int64(0), + next_generation=Int32(1), + claimed_counter_index=Int32(0), + ) + else: + self._work_id_state = atomic_counter_state + else: + self._work_id_state = GridStrideWorkIdState( + next_work_id=Int32(block_idx[2]), + work_id_stride=Int32(self.config.launch_cluster_cnt_merge_as_preferred), + ) + self.claimed_stream_index = Int32(0) + + @cute.jit + def initialize_fallback_group(self) -> None: + """Register one physical fallback cluster with its fixed logical group.""" + if cutlass.const_expr(self.config.is_mixed and self.work_id_mode == "atomic_counter"): + self._work_id_state = initialize_fixed_group_mixed_cga_work_id_state(self._work_id_state) + if self._work_id_state.is_fallback_cluster: + self.cta_coord_in_preferred_cluster = self._preferred_cluster_cta_coord( + self._cta_coord_in_active_cluster, self._work_id_state.in_group_idx + ) + + @cute.jit + def _preferred_cluster_cta_coord( + self, cta_coord_in_active_cluster: cute.Coord, inner_cluster_idx: Int32 + ) -> cute.Coord: + inner_cluster_m = inner_cluster_idx % Int32(self.config.mn_split_factors[0]) + inner_cluster_n = (inner_cluster_idx // Int32(self.config.mn_split_factors[0])) % Int32( + self.config.mn_split_factors[1] + ) + return ( + cta_coord_in_active_cluster[0] + inner_cluster_m * Int32(self.config.fallback_cluster_shape[0]), + cta_coord_in_active_cluster[1] + inner_cluster_n * Int32(self.config.fallback_cluster_shape[1]), + Int32(0), + ) + + @cute.jit + def claim_next_work(self, stream_index=0) -> Int32: + """Claim the next canonical work ID and update the preferred CTA coordinate.""" + claimed_work_id, self._work_id_state = claim_work_id(self._work_id_state, atomic_counter_index=stream_index) + canonical_work_id = claimed_work_id + if cutlass.const_expr(self.config.is_mixed and self.work_id_mode == "atomic_counter"): + cta_coord_in_preferred_cluster = self._cta_coord_in_active_cluster + if self._work_id_state.is_fallback_cluster: + cta_coord_in_preferred_cluster = self._preferred_cluster_cta_coord( + self._cta_coord_in_active_cluster, self._work_id_state.in_group_idx + ) + self.cta_coord_in_preferred_cluster = cta_coord_in_preferred_cluster + self.claimed_stream_index = self._work_id_state.claimed_counter_index + else: + self.claimed_stream_index = Int32(stream_index) + return canonical_work_id + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + values.extend(extract_mlir_values(self._work_id_state)) + values.extend(extract_mlir_values(self.cta_coord_in_preferred_cluster)) + values.extend(extract_mlir_values(self.claimed_stream_index)) + if self.config.is_mixed and self.work_id_mode == "atomic_counter": + values.extend(extract_mlir_values(self._cta_coord_in_active_cluster)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "NonClcMixedCgaSchedulerWorker": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + result = type(self)(config=self.config, work_id_mode=self.work_id_mode, stream_count=self.stream_count) + result._work_id_state = rebuild(self._work_id_state) + result.cta_coord_in_preferred_cluster = rebuild(self.cta_coord_in_preferred_cluster) + result.claimed_stream_index = rebuild(self.claimed_stream_index) + if self.config.is_mixed and self.work_id_mode == "atomic_counter": + result._cta_coord_in_active_cluster = rebuild(self._cta_coord_in_active_cluster) + if value_index != len(values): + raise ValueError( + f"NonClcMixedCgaSchedulerWorker MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return result + + +__all__ = ["NonClcMixedCgaConfig", "NonClcMixedCgaSchedulerWorker"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/work_id_claim.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/work_id_claim.py new file mode 100644 index 000000000..93a9be875 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/work_id_claim.py @@ -0,0 +1,566 @@ +"""Blackwell persistent work-ID claim backends.""" + +import dataclasses +from typing import List, Tuple + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass._mlir import ir +from cutlass.cutlass_dsl import Boolean, Int32, Int64, extract_mlir_values, new_from_mlir_values + +from ...helpers.ptx_helpers import mbarrier_arrive_expect_tx_on_peer, nanosleep, store_i32_to_peer_cluster_smem_async + + +class GridStrideWorkIdState: + """Register state for one monotonic grid-stride work-ID stream.""" + + def __init__(self, next_work_id: Int32, work_id_stride: Int32) -> None: + self.next_work_id = next_work_id + self.work_id_stride = work_id_stride + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + values.extend(extract_mlir_values(self.next_work_id)) + values.extend(extract_mlir_values(self.work_id_stride)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GridStrideWorkIdState": + next_work_id_value_count = len(extract_mlir_values(self.next_work_id)) + stride_value_count = len(extract_mlir_values(self.work_id_stride)) + expected_value_count = next_work_id_value_count + stride_value_count + if len(values) != expected_value_count: + raise ValueError( + f"GridStrideWorkIdState MLIR value count mismatch: expected {expected_value_count}, got {len(values)}." + ) + return type(self)( + next_work_id=new_from_mlir_values(self.next_work_id, values[:next_work_id_value_count]), + work_id_stride=new_from_mlir_values(self.work_id_stride, values[next_work_id_value_count:]), + ) + + +class AtomicCounterWorkIdState: + """Cluster-wide state for one of several contiguous atomic work-ID streams.""" + + def __init__( + self, + counter_pointer: cute.Pointer, + counter_count: int, + broadcast_pointer: cute.Pointer, + is_leader_cta: Boolean, + cluster_pipeline: pipeline.PipelineAsync, + producer_state, + consumer_state, + cluster_size: int | Int32, + ) -> None: + if isinstance(counter_count, bool) or not isinstance(counter_count, int) or counter_count <= 0: + raise ValueError("counter_count must be a positive Python int.") + self.counter_pointer = counter_pointer + self.counter_count = counter_count + self.broadcast_pointer = broadcast_pointer + self.is_leader_cta = is_leader_cta + self.cluster_pipeline = cluster_pipeline + self.producer_state = producer_state + self.consumer_state = consumer_state + self.cluster_size = cluster_size + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in ( + self.counter_pointer, + self.broadcast_pointer, + self.is_leader_cta, + self.producer_state, + self.consumer_state, + ): + values.extend(extract_mlir_values(field)) + if isinstance(self.cluster_size, Int32): + values.extend(extract_mlir_values(self.cluster_size)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "AtomicCounterWorkIdState": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + result = type(self)( + counter_pointer=rebuild(self.counter_pointer), + counter_count=self.counter_count, + broadcast_pointer=rebuild(self.broadcast_pointer), + is_leader_cta=rebuild(self.is_leader_cta), + cluster_pipeline=self.cluster_pipeline, + producer_state=rebuild(self.producer_state), + consumer_state=rebuild(self.consumer_state), + cluster_size=rebuild(self.cluster_size) if isinstance(self.cluster_size, Int32) else self.cluster_size, + ) + if value_index != len(values): + raise ValueError( + f"AtomicCounterWorkIdState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return result + + +class FixedGroupMixedCgaAtomicCounterWorkIdState: + """Atomic-counter state for fixed groups of physical fallback clusters.""" + + def __init__( + self, + atomic_counter_state: AtomicCounterWorkIdState, + registration_counter_pointer: cute.Pointer, + group_token_pointer: cute.Pointer, + split_factor: int, + fallback_cluster_count: int, + is_fallback_cluster: Boolean, + fallback_group_idx: Int32, + in_group_idx: Int32, + previous_token: Int64, + next_generation: Int32, + claimed_counter_index: Int32, + ) -> None: + if isinstance(split_factor, bool) or not isinstance(split_factor, int) or split_factor <= 1: + raise ValueError("split_factor must be a Python int greater than one.") + if ( + isinstance(fallback_cluster_count, bool) + or not isinstance(fallback_cluster_count, int) + or fallback_cluster_count <= 0 + ): + raise ValueError("fallback_cluster_count must be a positive Python int.") + if fallback_cluster_count % split_factor != 0: + raise ValueError("fallback_cluster_count must be divisible by split_factor.") + if atomic_counter_state.counter_count > 2: + raise ValueError("Fixed fallback groups support at most two work-ID streams.") + self.atomic_counter_state = atomic_counter_state + self.registration_counter_pointer = registration_counter_pointer + self.group_token_pointer = group_token_pointer + self.split_factor = split_factor + self.fallback_cluster_count = fallback_cluster_count + self.is_fallback_cluster = is_fallback_cluster + self.is_preferred_cluster = is_fallback_cluster == Boolean(False) + self.fallback_group_idx = fallback_group_idx + self.in_group_idx = in_group_idx + self.previous_token = previous_token + self.next_generation = next_generation + self.claimed_counter_index = claimed_counter_index + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in ( + self.atomic_counter_state, + self.registration_counter_pointer, + self.group_token_pointer, + self.is_fallback_cluster, + self.fallback_group_idx, + self.in_group_idx, + self.previous_token, + self.next_generation, + self.claimed_counter_index, + ): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "FixedGroupMixedCgaAtomicCounterWorkIdState": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + result = type(self)( + atomic_counter_state=rebuild(self.atomic_counter_state), + registration_counter_pointer=rebuild(self.registration_counter_pointer), + group_token_pointer=rebuild(self.group_token_pointer), + split_factor=self.split_factor, + fallback_cluster_count=self.fallback_cluster_count, + is_fallback_cluster=rebuild(self.is_fallback_cluster), + fallback_group_idx=rebuild(self.fallback_group_idx), + in_group_idx=rebuild(self.in_group_idx), + previous_token=rebuild(self.previous_token), + next_generation=rebuild(self.next_generation), + claimed_counter_index=rebuild(self.claimed_counter_index), + ) + if value_index != len(values): + raise ValueError( + f"FixedGroupMixedCgaAtomicCounterWorkIdState MLIR value count mismatch: " + f"consumed {value_index}, got {len(values)}." + ) + return result + + +@dataclasses.dataclass(frozen=True) +class GridWorkId: + """One CTA-specific coordinate claimed from a three-dimensional grid.""" + + grid_m: Int32 + grid_n: Int32 + grid_l: Int32 + is_valid: Boolean + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in (self.grid_m, self.grid_n, self.grid_l, self.is_valid): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GridWorkId": + if len(values) != 4: + raise ValueError(f"GridWorkId expects four MLIR values, got {len(values)}.") + fields = (self.grid_m, self.grid_n, self.grid_l, self.is_valid) + return type(self)(*(new_from_mlir_values(field, [value]) for field, value in zip(fields, values))) + + +class ClusterLaunchControlWorkIdState: + """Cluster-wide state for hardware-assisted grid-coordinate claims.""" + + def __init__( + self, + response_pending: Boolean, + grid_m: Int32, + grid_n: Int32, + grid_l: Int32, + response_is_valid: Boolean, + cta_coord_in_cluster: cute.Coord, + cluster_pipeline: pipeline.PipelineClcFetchAsync, + producer_state, + consumer_state, + is_leader_cta: Boolean, + response_pointer: cute.Pointer, + ) -> None: + self.response_pending = response_pending + self.grid_m = grid_m + self.grid_n = grid_n + self.grid_l = grid_l + self.response_is_valid = response_is_valid + self.cta_coord_in_cluster = cta_coord_in_cluster + self.cluster_pipeline = cluster_pipeline + self.producer_state = producer_state + self.consumer_state = consumer_state + self.is_leader_cta = is_leader_cta + self.response_pointer = response_pointer + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in ( + self.response_pending, + self.grid_m, + self.grid_n, + self.grid_l, + self.response_is_valid, + self.cta_coord_in_cluster, + self.producer_state, + self.consumer_state, + self.is_leader_cta, + self.response_pointer, + ): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "ClusterLaunchControlWorkIdState": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + rebuilt_field = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return rebuilt_field + + result = type(self)( + response_pending=rebuild(self.response_pending), + grid_m=rebuild(self.grid_m), + grid_n=rebuild(self.grid_n), + grid_l=rebuild(self.grid_l), + response_is_valid=rebuild(self.response_is_valid), + cta_coord_in_cluster=rebuild(self.cta_coord_in_cluster), + cluster_pipeline=self.cluster_pipeline, + producer_state=rebuild(self.producer_state), + consumer_state=rebuild(self.consumer_state), + is_leader_cta=rebuild(self.is_leader_cta), + response_pointer=rebuild(self.response_pointer), + ) + if value_index != len(values): + raise ValueError( + f"ClusterLaunchControlWorkIdState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return result + + +@cute.jit +def _claim_grid_stride_work_id(work_id_state: GridStrideWorkIdState) -> Tuple[Int32, GridStrideWorkIdState]: + """Claim the next ID from one monotonic grid-stride stream.""" + linear_work_id = work_id_state.next_work_id + work_id_state.next_work_id = linear_work_id + work_id_state.work_id_stride + return linear_work_id, work_id_state + + +@cute.jit +def _claim_atomic_counter_work_id( + work_id_state: AtomicCounterWorkIdState, atomic_counter_index=0 +) -> Tuple[Int32, AtomicCounterWorkIdState]: + """Claim from one selected counter and broadcast the ID within the cluster.""" + invalid_static_index = isinstance(atomic_counter_index, int) and ( + atomic_counter_index < 0 or atomic_counter_index >= work_id_state.counter_count + ) + if cutlass.const_expr(invalid_static_index): + raise ValueError( + f"atomic_counter_index must be in [0, {work_id_state.counter_count}), got {atomic_counter_index}." + ) + broadcast_tensor = cute.make_tensor(work_id_state.broadcast_pointer, cute.make_layout((1,))) + cluster_pipeline = work_id_state.cluster_pipeline + selected_counter_pointer = work_id_state.counter_pointer + Int32(atomic_counter_index) + + if work_id_state.is_leader_cta: + cluster_pipeline.producer_acquire(work_id_state.producer_state) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier(work_id_state.producer_state.index) + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + atomic_work_id = Int32(0) + if lane_idx == Int32(0): + atomic_work_id = cute.arch.atomic_add(selected_counter_pointer, Int32(1)) + atomic_work_id = cute.arch.shuffle_sync(atomic_work_id, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) + if lane_idx < Int32(work_id_state.cluster_size): + store_i32_to_peer_cluster_smem_async( + work_id_state.broadcast_pointer, atomic_work_id, full_barrier_pointer, lane_idx + ) + mbarrier_arrive_expect_tx_on_peer(full_barrier_pointer, Int32(4), lane_idx) + work_id_state.producer_state.advance() + + cluster_pipeline.consumer_wait(work_id_state.consumer_state) + linear_work_id = broadcast_tensor[0] + cute.arch.fence_acq_rel_cta() + cluster_pipeline.sync_object_empty.arrive(work_id_state.consumer_state.index, Int32(0)) + work_id_state.consumer_state.advance() + return linear_work_id, work_id_state + + +@cute.jit +def initialize_fixed_group_mixed_cga_work_id_state( + work_id_state: FixedGroupMixedCgaAtomicCounterWorkIdState, +) -> FixedGroupMixedCgaAtomicCounterWorkIdState: + """Register one physical fallback cluster and broadcast its fixed group coordinates.""" + atomic_counter_state = work_id_state.atomic_counter_state + if work_id_state.is_fallback_cluster: + broadcast_tensor = cute.make_tensor(atomic_counter_state.broadcast_pointer, cute.make_layout((1,))) + cluster_pipeline = atomic_counter_state.cluster_pipeline + if atomic_counter_state.is_leader_cta: + cluster_pipeline.producer_acquire(atomic_counter_state.producer_state) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier( + atomic_counter_state.producer_state.index + ) + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + fallback_ordinal = Int32(0) + if lane_idx == Int32(0): + fallback_ordinal = cute.arch.atomic_add( + work_id_state.registration_counter_pointer, Int32(1), sem="relaxed", scope="gpu" + ) + fallback_ordinal = Int32( + cute.arch.shuffle_sync(fallback_ordinal, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) + ) + if lane_idx < Int32(atomic_counter_state.cluster_size): + store_i32_to_peer_cluster_smem_async( + atomic_counter_state.broadcast_pointer, fallback_ordinal, full_barrier_pointer, lane_idx + ) + mbarrier_arrive_expect_tx_on_peer(full_barrier_pointer, Int32(4), lane_idx) + atomic_counter_state.producer_state.advance() + + cluster_pipeline.consumer_wait(atomic_counter_state.consumer_state) + fallback_ordinal = broadcast_tensor[0] + cute.arch.fence_acq_rel_cta() + cluster_pipeline.sync_object_empty.arrive(atomic_counter_state.consumer_state.index, Int32(0)) + atomic_counter_state.consumer_state.advance() + + fallback_group_idx = fallback_ordinal // Int32(work_id_state.split_factor) + work_id_state.fallback_group_idx = fallback_group_idx + work_id_state.in_group_idx = fallback_ordinal - fallback_group_idx * Int32(work_id_state.split_factor) + work_id_state.atomic_counter_state = atomic_counter_state + return work_id_state + + +@cute.jit +def _claim_fixed_group_fallback_work_id( + work_id_state: FixedGroupMixedCgaAtomicCounterWorkIdState, atomic_counter_index=0 +) -> Tuple[Int32, FixedGroupMixedCgaAtomicCounterWorkIdState]: + """Claim one canonical ID and hand it to every member of a fixed fallback group.""" + atomic_counter_state = work_id_state.atomic_counter_state + invalid_static_index = isinstance(atomic_counter_index, int) and ( + atomic_counter_index < 0 or atomic_counter_index >= atomic_counter_state.counter_count + ) + if cutlass.const_expr(invalid_static_index): + raise ValueError( + f"atomic_counter_index must be in [0, {atomic_counter_state.counter_count}), got {atomic_counter_index}." + ) + + broadcast_tensor = cute.make_tensor(atomic_counter_state.broadcast_pointer, cute.make_layout((1,))) + cluster_pipeline = atomic_counter_state.cluster_pipeline + selected_counter_pointer = atomic_counter_state.counter_pointer + Int32(atomic_counter_index) + + if atomic_counter_state.is_leader_cta: + cluster_pipeline.producer_acquire(atomic_counter_state.producer_state) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier(atomic_counter_state.producer_state.index) + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + group_base_offset = work_id_state.fallback_group_idx * Int32(work_id_state.split_factor) + group_token_pointer = work_id_state.group_token_pointer + group_base_offset + claimed_payload = Int32(0) + + if work_id_state.in_group_idx == Int32(0): + all_members_consumed = Boolean(False) + while not all_members_consumed: + observed_token = work_id_state.previous_token + if lane_idx < Int32(work_id_state.split_factor): + observed_token = cute.arch.load(group_token_pointer + lane_idx, Int64, sem="acquire", scope="gpu") + lane_is_ready = (lane_idx >= Int32(work_id_state.split_factor)) | ( + observed_token == work_id_state.previous_token + ) + ready_mask = Int32(cute.arch.vote_ballot_sync(lane_is_ready)) + all_members_consumed = ready_mask == Int32(-1) + if not all_members_consumed: + nanosleep(500) + + claimed_work_id = Int32(0) + if lane_idx == Int32(0): + claimed_work_id = cute.arch.atomic_add(selected_counter_pointer, Int32(1), sem="relaxed", scope="gpu") + claimed_work_id = Int32( + cute.arch.shuffle_sync(claimed_work_id, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) + ) + claimed_payload = claimed_work_id | (Int32(atomic_counter_index) << Int32(31)) + token = (Int64(work_id_state.next_generation) << Int64(32)) | (Int64(claimed_payload) & Int64(0xFFFFFFFF)) + if lane_idx == Int32(0): + cute.arch.store(group_token_pointer, token, sem="relaxed", scope="gpu") + work_id_state.previous_token = token + work_id_state.next_generation = work_id_state.next_generation + Int32(1) + else: + token = work_id_state.previous_token + while token == work_id_state.previous_token: + token_high = Int32(0) + token_low = Int32(0) + if lane_idx == Int32(0): + observed_token = cute.arch.load(group_token_pointer, Int64, sem="relaxed", scope="gpu") + token_high = Int32(observed_token >> Int64(32)) + token_low = Int32(observed_token & Int64(0xFFFFFFFF)) + token_high = Int32(cute.arch.shuffle_sync(token_high, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31)) + token_low = Int32(cute.arch.shuffle_sync(token_low, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31)) + token = (Int64(token_high) << Int64(32)) | (Int64(token_low) & Int64(0xFFFFFFFF)) + if token == work_id_state.previous_token: + nanosleep(500) + claimed_payload = Int32(token & Int64(0xFFFFFFFF)) + if lane_idx == Int32(0): + cute.arch.store(group_token_pointer + work_id_state.in_group_idx, token, sem="relaxed", scope="gpu") + work_id_state.previous_token = token + + if lane_idx < Int32(atomic_counter_state.cluster_size): + store_i32_to_peer_cluster_smem_async( + atomic_counter_state.broadcast_pointer, claimed_payload, full_barrier_pointer, lane_idx + ) + mbarrier_arrive_expect_tx_on_peer(full_barrier_pointer, Int32(4), lane_idx) + atomic_counter_state.producer_state.advance() + + cluster_pipeline.consumer_wait(atomic_counter_state.consumer_state) + claimed_payload = broadcast_tensor[0] + cute.arch.fence_acq_rel_cta() + cluster_pipeline.sync_object_empty.arrive(atomic_counter_state.consumer_state.index, Int32(0)) + atomic_counter_state.consumer_state.advance() + work_id_state.atomic_counter_state = atomic_counter_state + work_id_state.claimed_counter_index = (claimed_payload >> Int32(31)) & Int32(1) + linear_work_id = claimed_payload & Int32(0x7FFFFFFF) + return linear_work_id, work_id_state + + +@cute.jit +def _claim_fixed_group_mixed_cga_work_id( + work_id_state: FixedGroupMixedCgaAtomicCounterWorkIdState, atomic_counter_index=0 +) -> Tuple[Int32, FixedGroupMixedCgaAtomicCounterWorkIdState]: + """Claim directly for a preferred cluster or through its fixed fallback group.""" + linear_work_id = Int32(0) + if work_id_state.is_preferred_cluster: + linear_work_id, atomic_counter_state = _claim_atomic_counter_work_id( + work_id_state.atomic_counter_state, atomic_counter_index + ) + work_id_state.atomic_counter_state = atomic_counter_state + work_id_state.claimed_counter_index = Int32(atomic_counter_index) + else: + linear_work_id, work_id_state = _claim_fixed_group_fallback_work_id(work_id_state, atomic_counter_index) + return linear_work_id, work_id_state + + +@cute.jit +def _claim_cluster_launch_control_work_id( + work_id_state: ClusterLaunchControlWorkIdState, +) -> Tuple[GridWorkId, ClusterLaunchControlWorkIdState]: + """Claim the next canceled cluster and return this CTA's grid coordinate.""" + use_bootstrap = work_id_state.response_pending + state_before_bootstrap = work_id_state + if use_bootstrap: + work_id_state.response_pending = Boolean(False) + else: + work_id_state = state_before_bootstrap + + state_before_query = work_id_state + if not use_bootstrap: + state_before_leader_query = work_id_state + if work_id_state.is_leader_cta: + work_id_state.cluster_pipeline.producer_acquire(work_id_state.producer_state) + response_barrier = work_id_state.cluster_pipeline.producer_get_barrier(work_id_state.producer_state) + with cute.arch.elect_one(): + cute.arch.issue_clc_query(response_barrier, work_id_state.response_pointer) + else: + work_id_state = state_before_leader_query + work_id_state.producer_state.advance() + + work_id_state.cluster_pipeline.consumer_wait(work_id_state.consumer_state) + (cluster_origin_m, cluster_origin_n, grid_l, response_is_valid) = cute.arch.clc_response( + work_id_state.response_pointer + ) + cute.arch.fence_acq_rel_cta() + work_id_state.cluster_pipeline.consumer_release(work_id_state.consumer_state) + work_id_state.consumer_state.advance() + + work_id_state.grid_m = cluster_origin_m + work_id_state.cta_coord_in_cluster[0] + work_id_state.grid_n = cluster_origin_n + work_id_state.cta_coord_in_cluster[1] + work_id_state.grid_l = grid_l + work_id_state.response_is_valid = response_is_valid != Int32(0) + else: + work_id_state = state_before_query + + return ( + GridWorkId( + grid_m=work_id_state.grid_m, + grid_n=work_id_state.grid_n, + grid_l=work_id_state.grid_l, + is_valid=work_id_state.response_is_valid, + ), + work_id_state, + ) + + +@cute.jit +def claim_work_id(work_id_state, atomic_counter_index=0): + """Claim the next work ID using the backend encoded by the state type.""" + if cutlass.const_expr(isinstance(work_id_state, GridStrideWorkIdState)): + return _claim_grid_stride_work_id(work_id_state) + if cutlass.const_expr(isinstance(work_id_state, AtomicCounterWorkIdState)): + return _claim_atomic_counter_work_id(work_id_state, atomic_counter_index) + if cutlass.const_expr(isinstance(work_id_state, FixedGroupMixedCgaAtomicCounterWorkIdState)): + return _claim_fixed_group_mixed_cga_work_id(work_id_state, atomic_counter_index) + if cutlass.const_expr(isinstance(work_id_state, ClusterLaunchControlWorkIdState)): + return _claim_cluster_launch_control_work_id(work_id_state) + raise TypeError(f"Unsupported work-ID state: {type(work_id_state).__name__}.") + + +__all__ = [ + "AtomicCounterWorkIdState", + "ClusterLaunchControlWorkIdState", + "FixedGroupMixedCgaAtomicCounterWorkIdState", + "GridStrideWorkIdState", + "GridWorkId", + "claim_work_id", + "initialize_fixed_group_mixed_cga_work_id_state", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/quant_def.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/quant_def.py new file mode 100644 index 000000000..7ff5730d7 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/quant_def.py @@ -0,0 +1,249 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Canonical quantization kinds and payload formats.""" + +import dataclasses +import enum +from typing import ClassVar, Dict, Literal, Optional, Tuple, Type + +import cutlass + + +class QuantKind(str, enum.Enum): + """One admissible block-scaled quantization mode. + + A member is defined by its (weight, activation) element pair -- under swap-AB the weight is + the MMA A operand and the activation is the B operand. The selected K-throughput mode remains + explicit because one quantization kind can use multiple instruction K extents. + """ + + nvfp4 = "nvfp4" + mxfp4 = "mxfp4" + mxfp8_e4m3 = "mxfp8_e4m3" + mxfp8_e5m2 = "mxfp8_e5m2" + mxfp4_mxfp8 = "mxfp4_mxfp8" + + # Enum's __str__/__format__ for mixed-in types changed across Python 3.10/3.11/3.12, and the + # kernels fold the kind into their compiled-kernel cache key. Pin both to the member value so + # the key cannot silently become "QuantKind.nvfp4" on an interpreter upgrade. + __str__ = str.__str__ + __format__ = str.__format__ + + @property + def weight_dtype(self) -> Type[cutlass.Numeric]: + return _element_pair[self][0] + + @property + def activation_dtype(self) -> Type[cutlass.Numeric]: + """Also the dtype the FC1 epilogue must emit: FC2 consumes it as its activation.""" + return _element_pair[self][1] + + @property + def sf_vec_size(self) -> int: + return 16 if self is QuantKind.nvfp4 else 32 + + @property + def sf_dtype(self) -> Type[cutlass.Numeric]: + # Hardware allows a UE8M0 scale at vec 16 too, but nvfp4 is the only vec-16 mode we build. + return cutlass.Float8E4M3FN if self.sf_vec_size == 16 else cutlass.Float8E8M0FNU + + @property + def umma_kind(self) -> str: + """The ``tcgen05.mma.kind::`` qualifier. + + Mirrors the dispatch in ``blackwell_helpers._make_blockscaled_trivial_tiled_mma_impl``: + only an fp4 pair reaches the fp4-specific kinds, everything else -- including any mixed + pair -- falls back to mxf8f6f4. Keeping the two in sync matters because we build the tiled + MMA through that helper but emit the instruction ourselves. + """ + both_fp4 = self.weight_dtype is cutlass.Float4E2M1FN and self.activation_dtype is cutlass.Float4E2M1FN + if not both_fp4: + return "mxf8f6f4" + return "mxf4nvf4" if self.sf_vec_size == 16 else "mxf4" + + @property + def umma_scale_vec_suffix(self) -> str: + """PTX modifier after ``.block_scale``; mxf8f6f4 takes none (its scale vector is 32).""" + if self.umma_kind == "mxf8f6f4": + return "" + return ".block16" if self.sf_vec_size == 16 else ".block32" + + def instruction_k(self, mma_k_mode: Literal["1x", "2x"]) -> int: + instruction_k_1x = 32 if self.umma_kind == "mxf8f6f4" else 64 + if mma_k_mode == "1x": + return instruction_k_1x + if mma_k_mode == "2x": + return instruction_k_1x * 2 + raise ValueError(f"Invalid MMA K mode {mma_k_mode!r}; expected '1x' or '2x'.") + + @property + def weight_format_code(self) -> int: + """Instruction-descriptor ``a_format_`` under swap-AB.""" + return _instruction_descriptor_format_code(self.umma_kind, self.weight_dtype) + + @property + def activation_format_code(self) -> int: + """Instruction-descriptor ``b_format_`` under swap-AB.""" + return _instruction_descriptor_format_code(self.umma_kind, self.activation_dtype) + + @property + def scale_format_code(self) -> int: + """Instruction-descriptor ``scale_format_``: 0 = UE4M3, 1 = UE8M0.""" + return 0 if self.sf_dtype is cutlass.Float8E4M3FN else 1 + + def needs_unpack_tma(self, architecture: str) -> bool: + """Whether the narrow operand must reach SMEM in 1-byte containers (U4_UNPACK_U8 TMA). + + Blackwell mixed-width MMA consumes the narrow operand through UNPACK TMA. Rubin consumes + mixed FP4 directly from its native packed SMEM image. + """ + normalized_architecture = architecture.lower().replace("_", "") + if normalized_architecture.startswith("sm"): + normalized_architecture = normalized_architecture[2:] + if normalized_architecture not in ("100", "103", "107"): + raise ValueError(f"Unsupported TCGen05 architecture {architecture!r}.") + return normalized_architecture in ("100", "103") and self.weight_dtype.width != self.activation_dtype.width + + @property + def uses_global_scale(self) -> bool: + """Whether the caller supplies per-expert alpha / norm_const. + + Only nvfp4: an e8m0 scale is a pure power of two and already carries the whole rescale. + """ + return self is QuantKind.nvfp4 + + +_element_pair: Dict[QuantKind, Tuple[Type[cutlass.Numeric], Type[cutlass.Numeric]]] = { + QuantKind.nvfp4: (cutlass.Float4E2M1FN, cutlass.Float4E2M1FN), + QuantKind.mxfp4: (cutlass.Float4E2M1FN, cutlass.Float4E2M1FN), + QuantKind.mxfp8_e4m3: (cutlass.Float8E4M3FN, cutlass.Float8E4M3FN), + QuantKind.mxfp8_e5m2: (cutlass.Float8E5M2, cutlass.Float8E5M2), + QuantKind.mxfp4_mxfp8: (cutlass.Float4E2M1FN, cutlass.Float8E4M3FN), +} + + +# The three ``a_format_`` / ``b_format_`` bits are read against one of two disjoint enums, and the +# MMA kind picks which (CUTLASS ``UMMA::MXF4Format`` vs ``UMMA::MXF8F6F4Format`` in +# cute/arch/mma_sm100_desc.hpp). E2M1 is 1 under the first and 5 under the second. Getting it wrong +# still compiles and still runs -- it just computes garbage -- so these live in exactly one place. +_mxf4_format_code: Dict[Type[cutlass.Numeric], int] = {cutlass.Float4E2M1FN: 1} +_mxf8f6f4_format_code: Dict[Type[cutlass.Numeric], int] = { + cutlass.Float8E4M3FN: 0, + cutlass.Float8E5M2: 1, + cutlass.Float4E2M1FN: 5, +} + + +def _instruction_descriptor_format_code(umma_kind: str, dtype: Type[cutlass.Numeric]) -> int: + codes = _mxf8f6f4_format_code if umma_kind == "mxf8f6f4" else _mxf4_format_code + try: + return codes[dtype] + except KeyError as error: + raise ValueError(f"{dtype} has no instruction-descriptor format code under kind::{umma_kind}.") from error + + +@dataclasses.dataclass(frozen=True) +class CombineFormat: + """Data and scale representation of one cross-rank FC2 payload.""" + + _act_by_tag: ClassVar[Dict[str, type]] = {"e2m1": cutlass.Float4E2M1FN, "e4m3": cutlass.Float8E4M3FN} + _scale_by_tag: ClassVar[Dict[str, type]] = {"bf16": cutlass.BFloat16, "e8m0": cutlass.Float8E8M0FNU} + _rejection_reason: ClassVar[Dict[str, str]] = { + "32e5m2xe8m0": ( + "e5m2 costs the same 8.25 bits per element as e4m3 and trades a mantissa bit (6 dB of " + "SNR) for exponent range the e8m0 block scale already provides. Use 32e4m3xe8m0." + ) + } + + act_dtype: type + scale_dtype: Optional[type] + scale_block: Optional[int] + + def __post_init__(self) -> None: + allowed_act = {cutlass.BFloat16, *self._act_by_tag.values()} + allowed_scale = {None, *self._scale_by_tag.values()} + if self.act_dtype not in allowed_act: + raise ValueError(f"Unsupported combine data dtype {self.act_dtype}.") + if self.scale_dtype not in allowed_scale: + raise ValueError(f"Unsupported combine scale dtype {self.scale_dtype}.") + if self.scale_dtype is None: + if self.act_dtype is not cutlass.BFloat16 or self.scale_block is not None: + raise ValueError("The unquantized format must be bf16 without a scale block.") + return + if self.act_dtype is cutlass.BFloat16: + raise ValueError("A quantized format cannot use bf16 data.") + if self.scale_dtype is cutlass.BFloat16 and self.scale_block != 16: + raise ValueError("A bf16 amax scale requires a 16-element block.") + if self.scale_dtype is cutlass.Float8E8M0FNU and self.scale_block != 32: + raise ValueError("An e8m0 scale requires a 32-element block.") + + @property + def is_quantized(self) -> bool: + return self.scale_dtype is not None + + @property + def name(self) -> str: + if not self.is_quantized: + return "bf16" + act_tag = next(tag for tag, dtype in self._act_by_tag.items() if dtype is self.act_dtype) + scale_tag = next(tag for tag, dtype in self._scale_by_tag.items() if dtype is self.scale_dtype) + return f"{self.scale_block}{act_tag}x{scale_tag}" + + def __str__(self) -> str: + return self.name + + @classmethod + def parse(cls, text: str) -> "CombineFormat": + specs = { + "bf16": (cutlass.BFloat16, None, None), + "16e2m1xbf16": (cutlass.Float4E2M1FN, cutlass.BFloat16, 16), + "32e4m3xe8m0": (cutlass.Float8E4M3FN, cutlass.Float8E8M0FNU, 32), + } + token = text.strip().lower() + if token in cls._rejection_reason: + raise ValueError(f"Combine format {token!r} is deliberately unsupported: {cls._rejection_reason[token]}") + if token not in specs: + raise ValueError(f"Invalid combine format {text!r}; expected one of {tuple(specs)}.") + act_dtype, scale_dtype, scale_block = specs[token] + return cls(act_dtype, scale_dtype, scale_block) + + +# Every Blackwell 1x PTX hardware encoding is restated independently of the derivations above, so +# that a typo in a property fails at import rather than at the first wrong numerical result. +# Ordered as (umma_kind, scale_vec_suffix, instruction_k_1x, a_format_, b_format_, scale_format_). +_pinned_hardware_encoding: Dict[QuantKind, Tuple[str, str, int, int, int, int]] = { + QuantKind.nvfp4: ("mxf4nvf4", ".block16", 64, 1, 1, 0), + QuantKind.mxfp4: ("mxf4", ".block32", 64, 1, 1, 1), + QuantKind.mxfp8_e4m3: ("mxf8f6f4", "", 32, 0, 0, 1), + QuantKind.mxfp8_e5m2: ("mxf8f6f4", "", 32, 1, 1, 1), + QuantKind.mxfp4_mxfp8: ("mxf8f6f4", "", 32, 5, 0, 1), +} + + +def _verify_pinned_hardware_encoding() -> None: + unpinned = sorted(kind.name for kind in QuantKind if kind not in _pinned_hardware_encoding) + if unpinned: + raise AssertionError(f"QuantKind members without a pinned hardware encoding: {unpinned}.") + for kind, expected in _pinned_hardware_encoding.items(): + derived = ( + kind.umma_kind, + kind.umma_scale_vec_suffix, + kind.instruction_k("1x"), + kind.weight_format_code, + kind.activation_format_code, + kind.scale_format_code, + ) + if derived != expected: + raise AssertionError(f"QuantKind.{kind.name} derives {derived}, pinned encoding is {expected}.") + expected_instruction_k_2x = expected[2] * 2 + instruction_k_2x = kind.instruction_k("2x") + if instruction_k_2x != expected_instruction_k_2x: + raise AssertionError( + f"QuantKind.{kind.name} derives 2x instruction K {instruction_k_2x}, expected {expected_instruction_k_2x}." + ) + + +_verify_pinned_hardware_encoding() + + +__all__ = ["CombineFormat", "QuantKind"] From efd1695e5c3e0e4177ba65fb4da2f5ec6600b34c Mon Sep 17 00:00:00 2001 From: zhibinz Date: Mon, 24 Aug 2026 16:41:24 -0700 Subject: [PATCH 03/31] feat: add Rubin MXFP8 MegaMoE kernels Vendor the Rubin forward GLU and backward dGLU training kernels needed for native expert-parallel execution on SM107 devices. --- .../cutedsl_src/kernel_src/rubin/__init__.py | 1 + .../kernel_src/rubin/training/__init__.py | 1 + .../rubin/training/mega/__init__.py | 21 + .../rubin/training/mega/bwd_dglu/__init__.py | 22 + .../mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py | 1480 ++++++++++ .../bwd_dglu/dglu_mxfp8_fc12_extension.py | 73 + .../mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py | 2356 ++++++++++++++++ .../bwd_dglu/dglu_mxfp8_mega_moe_kernel.py | 939 +++++++ .../rubin/training/mega/fwd_glu/__init__.py | 18 + .../mega/fwd_glu/glu_mxfp8_col_requant.py | 1276 +++++++++ .../mega/fwd_glu/glu_mxfp8_fc12_epilogue.py | 1647 ++++++++++++ .../mega/fwd_glu/glu_mxfp8_fc12_extension.py | 232 ++ .../mega/fwd_glu/glu_mxfp8_fc12_kernel.py | 2372 +++++++++++++++++ .../mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py | 865 ++++++ .../rubin/training/mega/helpers/__init__.py | 18 + .../rubin/training/mega/helpers/constants.py | 19 + .../rubin/training/mega/helpers/utils.py | 322 +++ .../rubin/training/mega/tmem_transpose.py | 18 + .../rubin/training/mega/topk_reduce.py | 15 + 19 files changed, 11695 insertions(+) create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_extension.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_kernel.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/constants.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/utils.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/tmem_transpose.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/topk_reduce.py diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/__init__.py new file mode 100644 index 000000000..b84f4d01b --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/__init__.py @@ -0,0 +1 @@ +"""Rubin kernel implementations.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/__init__.py new file mode 100644 index 000000000..4e1908109 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/__init__.py @@ -0,0 +1 @@ +"""Rubin SM107 training kernel package.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/__init__.py new file mode 100644 index 000000000..14b0bb10d --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Rubin training MegaMoE (mxfp8) kernels. + +Organised into two subpackages: + +* ``fwd_glu`` -- the forward fused FC1+SwiGLU+FC2 MoE kernel. +* ``bwd_dglu`` -- the backward fused dfc2+dswiglu+dfc1 MoE kernel. + +The forward symbols are re-exported here so existing ``rubin.training.mega`` +importers keep resolving after the fwd_glu/bwd_dglu reorg. +""" + +from .fwd_glu import ( + Fc2OutputDest, + GluMxFp8Fc12SchedExtension, + GluMxfp8Epilogue, + Sm107MegaMoEMxfp8GluKernel, + Sm107Mxfp8GluFc12Kernel, + TensorRole, +) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/__init__.py new file mode 100644 index 000000000..81bd8e3bd --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Rubin training MegaMoE (mxfp8 dGLU backward) kernel components. + +The backward fused dfc2 + dswiglu + dfc1 MoE kernel used by the training path. Computes +grad_x (activation gradient -> output_activation) and dprob (routing-weight gradient, +pool region). Built on top of the forward GLU package (subclasses the shared +``Fc2OutputDest`` peer-store) and mirrors the forward mega structure. +""" + +from .dglu_mxfp8_fc12_epilogue import DgluMxfp8Epilogue +from .dglu_mxfp8_fc12_extension import DgluMxFp8Fc12SchedExtension +from .dglu_mxfp8_fc12_kernel import Sm107Mxfp8DgluDfc21Kernel +from .dglu_mxfp8_mega_moe_kernel import Sm107MegaMoEMxfp8DgluKernel + + +__all__ = [ + "DgluMxfp8Epilogue", + "DgluMxFp8Fc12SchedExtension", + "Sm107Mxfp8DgluDfc21Kernel", + "Sm107MegaMoEMxfp8DgluKernel", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py new file mode 100644 index 000000000..46b57451d --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py @@ -0,0 +1,1480 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +from typing import Optional, Tuple, Type, Union + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import AddressSpace +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils + +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith as _arith +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import dsl_user_op, Int32 as _epi_Int32, Int64 +from cutlass.cute.typing import Float32 + +from ......helpers.iket_compat import iket +from ......helpers.flag_batch import GpuReleaseFlagBatchTracker +from ......helpers.ptx_helpers import ( + red_add_relaxed_sys_f32 as _red_add_relaxed_sys_f32, + red_add_relaxed_sys_v2_bf16x2 as _red_add_relaxed_sys_v2_bf16x2, + stg_e8m0_from_f32, + stg_e8m0x8_from_f32, +) +from ..helpers.utils import swiglu_act, dswiglu_act, quant_sfd_row, quant_sfd_col +from ......quant_def import CombineFormat +from .....schedulers import BlockPhase +from ..tmem_transpose import _TmemTranspose16x32Core +from ..fwd_glu.glu_mxfp8_fc12_epilogue import Fc2OutputDest + +Fc1GateUpInterleave = 32 +EpilogueTileN = 32 +Fc1EpilogueOutputTileM = 128 +Fc1EpilogueOutputTileN = 128 +WarpThreadCount = 32 +EpiWarpCount = 4 + + +@cute.jit +def dprob_reduce_gmem( + real_dprob: cute.Tensor, + dprob_val: cutlass.Float32, + is_valid: bool, + expert_local_token_idx, + system_scope: bool = False, +) -> None: + """Atomically reduce per-tile dprob accumulator into GMEM.""" + if is_valid: + if cutlass.const_expr(system_scope): + _red_add_relaxed_sys_f32( + real_dprob.iterator + expert_local_token_idx, + cutlass.Float32(dprob_val), + ) + else: + cute.arch.atomic_add( + real_dprob.iterator + expert_local_token_idx, + cutlass.Float32(dprob_val), + sem="relaxed", + scope="gpu", + ) + +# ============================================================================= +# DgluMxfp8Epilogue +# ============================================================================= + +class DgluMxfp8Epilogue: + + _SubtileBarIdBase = 4 + + def __init__( + self, + *, + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + use_2cta_instrs: bool, + sf_vec_size: int, + fc1_output_dtype: Type[cutlass.Numeric], + fc1_output_layout: utils.LayoutEnum, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + sf_dtype: Type[cutlass.Numeric] = cutlass.Float8E8M0FNU, + epilog_sync_bar_id: int = 1, + epilogue_warp_ids: Tuple[int, ...] = (0, 1, 2, 3), + static_expert_shape: Optional[Tuple[int, int, int]] = None, + token_back_by_dispatch: bool = False, + epi_flag_batch: Optional[Tuple[int, int]] = (1, 1), + dfc2_recompute: bool = False, + dfc2_col_output: bool = False, + fc2_in_kernel_topk_reduce: bool = False, + combine_format: Optional[CombineFormat] = None, + combine_hidden: Optional[int] = None, + act_func: str = "swiglu", + gate_up_clamp: Optional[float] = None, + ) -> None: + self._act_func = act_func + self._gate_up_clamp = ( + cutlass.Float32(gate_up_clamp) if gate_up_clamp is not None else None + ) + self.fc1_output_dtype = fc1_output_dtype + self.fc1_output_layout = fc1_output_layout + self.acc_dtype = acc_dtype + self.sf_dtype = sf_dtype + self._sf_vec_size = sf_vec_size + self._epilog_sync_bar_id = epilog_sync_bar_id + self._epilogue_warp_ids = epilogue_warp_ids + self._use_2cta_instrs = use_2cta_instrs + + self._atom_thr_size = 2 if use_2cta_instrs else 1 + self._cta_tile_m = mma_tiler_mnk[0] // self._atom_thr_size + self._cta_tile_n = mma_tiler_mnk[1] + self._mma_tiler_k = mma_tiler_mnk[2] + self._mma_tiler = tuple(mma_tiler_mnk) # for partition_C in the C-load + self._cta_tile_n_sfb = ((mma_tiler_mnk[1] + 127) // 128) * 128 + self._static_expert_shape = static_expert_shape + if ( + static_expert_shape is not None + and static_expert_shape[2] % (self._cta_tile_m * cluster_shape_mn[0]) == 0 + ): + self._fc2_stg_needs_predicate: bool = False + else: + self._fc2_stg_needs_predicate: bool = True + + self._epi_tile = (EpilogueTileN, Fc1EpilogueOutputTileM) + self._subtile_cnt = self._cta_tile_n // 2 // EpilogueTileN + + self._num_acc_stage = 2 + self._num_acc_pipeline_stages = self._num_acc_stage + + k = self._mma_tiler_k + self._num_sfa_tmem_cols = self._cta_tile_m * k // sf_vec_size * 4 // 4 // 128 + self._num_sfb_tmem_cols = ( + self._cta_tile_n_sfb * k // sf_vec_size * 4 // 4 // 128 + ) + self._num_sf_tmem_cols = 32 # self._num_sfa_tmem_cols + self._num_sfb_tmem_cols + + self._num_accumulator_tmem_cols = self._cta_tile_n * self._num_acc_stage + + self._token_back_by_dispatch = token_back_by_dispatch + # In-kernel top-k reduce + self._reduce_topk_in_epilogue = ( + fc2_in_kernel_topk_reduce and not token_back_by_dispatch + ) + _fc1_batch, _fc2_batch = (1, 1) if epi_flag_batch is None else epi_flag_batch + self._epi_fc1_batch = max(1, min(32, int(_fc1_batch))) + self._epi_fc2_batch = max(1, min(32, int(_fc2_batch))) + + self._dfc2_recompute = dfc2_recompute + self._dfc2_col_output = dfc2_col_output + + # combine_format determines the dfc1 (final grad_x) combine encoding. + if combine_format is None: + combine_format = CombineFormat.parse("bf16") + self._combine_format = combine_format + self._combine_mxfp8 = combine_format.is_quantized + # sf_block_pad for the dfc1 MXFP8 combine + if self._combine_mxfp8 and combine_hidden is not None: + _hidden_dfc1 = combine_hidden + _sf_blocks_dfc1 = _hidden_dfc1 // EpilogueTileN + self._dfc1_sf_block_pad = ((_sf_blocks_dfc1 + 15) // 16) * 16 + self._hidden_dfc1 = _hidden_dfc1 + else: + self._dfc1_sf_block_pad = 0 + self._hidden_dfc1 = 0 + # batching stg.64 SF + self._dfc1_sf_batch8 = ( + self._combine_mxfp8 + and self._hidden_dfc1 > 0 + and (self._hidden_dfc1 % self._cta_tile_n == 0) + and (self._cta_tile_n // EpilogueTileN == 8) + ) + + pass + + # -- Codegen-time queries -- + + @property + def epi_tile(self) -> Tuple[int, int]: + return self._epi_tile + + @property + def num_acc_pipeline_stages(self) -> int: + return self._num_acc_pipeline_stages + + @property + def num_acc_stage(self) -> int: + return self._num_acc_stage + + @property + def subtile_cnt(self) -> int: + return self._subtile_cnt + + @property + def cta_tile_n(self) -> int: + return self._cta_tile_n + + @property + def num_sf_tmem_cols(self) -> int: + return self._num_sf_tmem_cols + + @property + def num_sfa_tmem_cols(self) -> int: + return self._num_sfa_tmem_cols + + @property + def num_sfb_tmem_cols(self) -> int: + return self._num_sfb_tmem_cols + + @property + def num_accumulator_tmem_cols(self) -> int: + return self._num_accumulator_tmem_cols + + def staged_smem_layout( + self, + n_stages: int, + ) -> Union[cute.Layout, cute.ComposedLayout]: + return sm100_utils.make_smem_layout_epi( + self.fc1_output_dtype, + self.fc1_output_layout, + self._epi_tile, + n_stages, + ) + + @property + def smem_layout_one_stage(self) -> Union[cute.Layout, cute.ComposedLayout]: + staged = self.staged_smem_layout(1) + return cute.select(staged, mode=[0, 1]) + + @property + def bytes_per_stage(self) -> int: + return cute.size_in_bytes(self.fc1_output_dtype, self.smem_layout_one_stage) + + # -- grad_y1 (dfc2 output) sD staging: reference-style shared (128×32) box -- + @property + def d_epi_tile(self) -> Tuple[int, int]: + return (self._cta_tile_m, EpilogueTileN) + + def d_staged_smem_layout( + self, + n_stages: int, + ) -> Union[cute.Layout, cute.ComposedLayout]: + return sm100_utils.make_smem_layout_epi( + self.fc1_output_dtype, + self.fc1_output_layout, + self.d_epi_tile, + n_stages, + ) + + @property + def d_smem_layout_one_stage(self) -> Union[cute.Layout, cute.ComposedLayout]: + staged = self.d_staged_smem_layout(1) + return cute.select(staged, mode=[0, 1]) + + @property + def d_bytes_per_stage(self) -> int: + return cute.size_in_bytes(self.fc1_output_dtype, self.d_smem_layout_one_stage) + + # forward pre-activation (dswiglu C input) staging + @property + def preact_epi_tile(self) -> Tuple[int, int]: + return (self._cta_tile_m, EpilogueTileN) + + def preact_staged_smem_layout( + self, n_stages: int + ) -> Union[cute.Layout, cute.ComposedLayout]: + return sm100_utils.make_smem_layout_epi( + cutlass.BFloat16, + self.fc1_output_layout, + self.preact_epi_tile, + n_stages, + ) + + @property + def preact_smem_layout_one_stage(self) -> Union[cute.Layout, cute.ComposedLayout]: + staged = self.preact_staged_smem_layout(1) + return cute.select(staged, mode=[0, 1]) + + @property + def preact_bytes_per_stage(self) -> int: + return cute.size_in_bytes(cutlass.BFloat16, self.preact_smem_layout_one_stage) + + @cute.jit + def _run_dfc2_task_tile( + self, + work_tile_info, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + acc_consumer_state, + sched_ext, + gmem_fc1_output: cute.Tensor, + gmem_fc1_output_sf: cute.Tensor, + gmem_fc1_recompute: cute.Tensor, + gmem_fc1_recompute_sf: cute.Tensor, + gmem_fc1_col_output: cute.Tensor, + gmem_fc1_col_output_sf: cute.Tensor, + c_pipeline, + smem_preact_buffer: cute.Tensor, + c_consumer_state, + smem_d_buffer: cute.Tensor, + tma_atom_grad_y1: cute.CopyAtom, + warp_idx: int, + tidx, + norm_const, + gmem_topk_scores: cute.Tensor, + gmem_beta: cute.Tensor, + gmem_dprob: cute.Tensor, + d_pipeline, + d_num_stage, + token_comm_args=None, + ): + """dfc2 task-tile — c_pipeline CONSUMER. """ + real_fc1_output, _ = sched_ext.get_gmem_tensor("d", gmem_fc1_output, work_tile_info) + real_fc1_output_sf, _ = sched_ext.get_gmem_tensor("sfd", gmem_fc1_output_sf, work_tile_info) + if cutlass.const_expr(token_comm_args is None): + real_dprob, _ = sched_ext.get_gmem_tensor("topk", gmem_dprob, work_tile_info) + else: + real_dprob = None + if cutlass.const_expr(self._dfc2_recompute): + real_fc1_recompute, _ = sched_ext.get_gmem_tensor("recompute", gmem_fc1_recompute, work_tile_info) + real_fc1_recompute_sf, _ = sched_ext.get_gmem_tensor("sfrecompute", gmem_fc1_recompute_sf, work_tile_info) + else: + real_fc1_recompute = None + real_fc1_recompute_sf = None + if cutlass.const_expr(self._dfc2_col_output): + real_fc1_col_output, _ = sched_ext.get_gmem_tensor("col_output", gmem_fc1_col_output, work_tile_info) + real_fc1_col_output_sf, _ = sched_ext.get_gmem_tensor("sfcol_output", gmem_fc1_col_output_sf, work_tile_info) + else: + real_fc1_col_output = None + real_fc1_col_output_sf = None + + acc_pipeline.consumer_wait(acc_consumer_state) + iket.range_push("mxfp8_dfc2_epi_tile") + + subtile_cnt = self._cta_tile_n // EpilogueTileN # 8 (256 / 32) + start_subtile = 0 + tmem_t = self._subtile_dfc12_tmem_tensor( + tmem_acc_tensor, cutlass.Int32(start_subtile), warp_idx, + ) + tmem_forward_cols = EpilogueTileN + + rmem_sf = cute.make_rmem_tensor( + cute.make_layout(2 * (self._cta_tile_n // EpilogueTileN)).shape, self.acc_dtype, + ) + # fc1_recompute SF accumulator: ONE SF per subtile (recompute N = half of dfc2). + rmem_sf_recompute = cute.make_rmem_tensor( + cute.make_layout(self._cta_tile_n // EpilogueTileN).shape, self.acc_dtype, + ) + # fc1_col_output SF accumulator + rmem_sf_col_output = cute.make_rmem_tensor( + cute.make_layout(2 * (self._cta_tile_n // EpilogueTileN)).shape, self.acc_dtype, + ) + thread_in_warp = tidx % WarpThreadCount + token_row_in_cta = cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + expert_local_token_idx = ( + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + token_row_in_cta + ) + + # beta / prob / dprob setup + beta_val = gmem_beta[work_tile_info.expert_idx] + # mProb: load from topk_scores for valid tokens; default 1.0 (unused) for invalid. + rmem_prob = cute.make_rmem_tensor(cute.make_layout(1).shape, self.acc_dtype) + rmem_prob[0] = cutlass.Float32(1.0) + if token_row_in_cta < valid_tokens: + real_topk, _ = sched_ext.get_gmem_tensor("topk", gmem_topk_scores, work_tile_info) + rmem_prob[0] = real_topk[expert_local_token_idx] + + # Per-tile dprob accumulator (single scalar). + dprob = cutlass.Float32(0.0) + + # Output col-strips per N-tile. + n_col_strips_per_tile = (self._cta_tile_n * 2) // EpilogueTileN + base_token_tile = work_tile_info.tile_m_idx # 128-row tile index + + _epilog_sync = pipeline.NamedBarrier( + barrier_id=self._epilog_sync_bar_id, + num_threads=WarpThreadCount * len(self._epilogue_warp_ids), + ) + + # Build tiled copies for SMEM↔REG (reference pattern). + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self._mma_tiler, + self.fc1_output_layout, + self.fc1_output_dtype, + self.acc_dtype, + self.d_epi_tile, + self._use_2cta_instrs, + ) + tAcc_epi = cute.flat_divide( + tmem_acc_tensor[((None, None), 0, 0)], + self.d_epi_tile, + ) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_rAcc_full = thr_copy_t2r.partition_D(tAcc_epi) + copy_atom_s2r = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16) + tiled_copy_s2r = cute.make_tiled_copy_D(copy_atom_s2r, tiled_copy_t2r) + thr_copy_s2r = tiled_copy_s2r.get_slice(tidx) + tRS_sPre = thr_copy_s2r.partition_D(smem_preact_buffer) + + r_layout = cute.make_layout(((1, EpilogueTileN,), 1, 1,), stride=((0, 1,), 0, 0,)) + r_gate_bf = cute.make_rmem_tensor(r_layout, cutlass.BFloat16) + r_up_bf = cute.make_rmem_tensor(r_layout, cutlass.BFloat16) + copy_atom_r2s = sm100_utils.get_smem_store_op( + self.fc1_output_layout, self.fc1_output_dtype, self.acc_dtype, tiled_copy_t2r + ) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + + for i in cutlass.range(0, subtile_cnt, 1): + subtile_idx = cutlass.Int32(i) + c_consumer_state, subtile_dprob = self._run_dfc2_subtile( + subtile_idx=subtile_idx, + subtile_i=i, + t_subtile=tmem_t, + smem_d=smem_d_buffer, + tiled_copy_r2s=tiled_copy_r2s, + tiled_copy_s2r=tiled_copy_s2r, + tRS_sPre=tRS_sPre, + c_pipeline=c_pipeline, + c_consumer_state=c_consumer_state, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + r_gate_bf=r_gate_bf, + r_up_bf=r_up_bf, + work_tile_info=work_tile_info, + warp_idx=warp_idx, + tidx=tidx, + norm_const=norm_const, + rmem_sf=rmem_sf, + rmem_sf_recompute=rmem_sf_recompute, + real_fc1_recompute=real_fc1_recompute, + rmem_sf_col_output=rmem_sf_col_output, + real_fc1_col_output=real_fc1_col_output, + beta=beta_val, + prob=rmem_prob[0], + epilog_sync=_epilog_sync, + d_pipeline=d_pipeline, + d_num_stage=d_num_stage, + ) + dprob = dprob + subtile_dprob + + tmem_t = self._advance_fc2_tmem_tensor(tmem_t, tmem_forward_cols) + + # BARRIER: fence_proxy makes R2S (stmatrix) writes visible to TMA + cute.arch.fence_proxy("async.shared", space="cta") + _epilog_sync.arrive_and_wait() + + # Compute GMEM tile pointers for this subtile. + gate_col_idx = ( + work_tile_info.tile_n_idx * cutlass.Int32(n_col_strips_per_tile) + + subtile_idx * cutlass.Int32(2) + ) + g_gate = cute.local_tile( + real_fc1_output, + (self._cta_tile_m, EpilogueTileN, 1), + (base_token_tile, gate_col_idx, cutlass.Int32(0)), + )[(None, None, 0)] + g_up = cute.local_tile( + real_fc1_output, + (self._cta_tile_m, EpilogueTileN, 1), + (base_token_tile, gate_col_idx + cutlass.Int32(1), cutlass.Int32(0)), + )[(None, None, 0)] + # TMA issue (warp 0 only). + d_n_slots = cutlass.const_expr(d_num_stage // 2) + d_slot = cutlass.Int32(2) * (cutlass.Int32(i) % cutlass.Int32(d_n_slots)) + if warp_idx == self._epilogue_warp_ids[0]: + self.tma_store_dfc2_output( + smem_d_buffer, + tma_atom_grad_y1, + g_gate, + g_up, + valid_tokens, + d_pipeline, + d_slot, + ) + + self._acc_pipeline_consumer_release(acc_pipeline, acc_consumer_state, True) + + valid_inter = real_fc1_output.shape[1] + self._stg_sf_dfc2(rmem_sf, real_fc1_output_sf, work_tile_info, tidx, valid_inter) + + # fc1_recompute SFs + if cutlass.const_expr(self._dfc2_recompute): + valid_inter_recompute = real_fc1_recompute.shape[1] + self._stg_sf_recompute( + rmem_sf_recompute, real_fc1_recompute_sf, + work_tile_info, tidx, valid_inter_recompute, valid_tokens, + ) + # fc1_col_output SFs + if cutlass.const_expr(self._dfc2_col_output): + valid_inter_col_output = real_fc1_col_output.shape[1] + self._stg_sf_col_output( + rmem_sf_col_output, real_fc1_col_output_sf, + work_tile_info, tidx, valid_inter_col_output, valid_tokens, + ) + # MegaMoE maps the receiver-pool row back to the source rank's combine slot. + if cutlass.const_expr(token_comm_args is not None): + if token_row_in_cta < valid_tokens: + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + metadata_u32 = cute.recast_tensor( + token_comm_args.token_src_metadata, cutlass.Uint32, + ) + dprob_output_dest = Fc2OutputDest( + tensor=token_comm_args.dprob_output, + metadata=metadata_u32, + peer_rank_ptr_mapper=token_comm_args.peer_rank_ptr_mapper, + ) + dest_row = dprob_output_dest.resolve_token_row(pool_token_global) + dprob_reduce_gmem( + dest_row, + dprob, + True, + cutlass.Int32(0), + system_scope=True, + ) + else: + dprob_reduce_gmem( + real_dprob, + dprob, + token_row_in_cta < valid_tokens, + expert_local_token_idx, + ) + iket.range_pop() + return c_consumer_state + + @cute.jit + def _dfc2_load_c( + self, + c_pipeline, + c_consumer_state, + tiled_copy_s2r, + tRS_sPre: cute.Tensor, + r_gate_bf, + r_up_bf, + ): + """Consume gate then up c_pipeline stages into register tiles.""" + c_pipeline.consumer_wait(c_consumer_state) + c_slot = 2 * c_consumer_state.index + cute.copy( + tiled_copy_s2r, + tRS_sPre[(None, None, None, c_slot)], + r_gate_bf, + ) + cute.copy( + tiled_copy_s2r, + tRS_sPre[(None, None, None, c_slot + 1)], + r_up_bf, + ) + cute.arch.fence_proxy("async.shared", space="cta") + c_pipeline.consumer_release(c_consumer_state) + c_consumer_state.advance() + + return c_consumer_state + + @cute.jit + def _run_dfc2_subtile( + self, + subtile_idx, + subtile_i, + t_subtile: cute.Tensor, + smem_d: cute.Tensor, + tiled_copy_r2s, + tiled_copy_s2r, + tRS_sPre: cute.Tensor, + c_pipeline, + c_consumer_state, + acc_pipeline, + acc_consumer_state, + r_gate_bf: cute.Tensor, + r_up_bf: cute.Tensor, + work_tile_info, + warp_idx: int, + tidx, + norm_const, + rmem_sf: cute.Tensor, + rmem_sf_recompute: cute.Tensor, + real_fc1_recompute: cute.Tensor, + rmem_sf_col_output: cute.Tensor, + real_fc1_col_output: cute.Tensor, + beta: cutlass.Float32, + prob: cutlass.Float32, + epilog_sync, + d_pipeline, + d_num_stage, + ): + iket.range_push("mxfp8_dfc2_epilogue_subtile") + EN = EpilogueTileN + + r_layout = cute.make_layout((((EN,), 1),), stride=(((1,), 0),)) + atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), self.acc_dtype, + ) + r_acc = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + cute.copy(atom_t2r, t_subtile, r_acc) + + thread_in_warp = tidx % WarpThreadCount + token_row_in_cta = cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + subtile_dprob = cutlass.Float32(0.0) + + # load c from shared memory to registers + c_consumer_state = self._dfc2_load_c( + c_pipeline, + c_consumer_state, + tiled_copy_s2r, + tRS_sPre, + r_gate_bf, + r_up_bf, + ) + + # c_gate / c_up declared outside the validity guard: stmatrix (tiled_copy_r2s) + # is warp-cooperative and all threads must call it regardless of token validity. + c_shape = cute.make_layout(((1, EN,), 1, 1), stride=((0, 1,), 0, 0)).shape + c_gate = cute.make_rmem_tensor(c_shape, self.fc1_output_dtype) + c_up = cute.make_rmem_tensor(c_shape, self.fc1_output_dtype) + # c_recompute: flat MXFP8 row for per-thread STG + c_recompute = cute.make_rmem_tensor(cute.make_layout(EN).shape, self.fc1_output_dtype) + + is_valid_row = token_row_in_cta < valid_tokens + + r_gate = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + r_up = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + for j in cutlass.range_constexpr(EN): + r_gate[j] = r_gate_bf[j].to(self.acc_dtype) + r_up[j] = r_up_bf[j].to(self.acc_dtype) + # Zero invalid rows' inputs (their r_acc/r_gate/r_up are padding garbage) so the + # warp-wide column amax is not polluted and no NaN propagates into the reduction. + if token_row_in_cta >= valid_tokens: + for j in cutlass.range_constexpr(EN): + r_acc[j] = self.acc_dtype(0.0) + r_gate[j] = self.acc_dtype(0.0) + r_up[j] = self.acc_dtype(0.0) + + # dswiglu backward: acc(grad_h) x (gate, up) -> (d_gate, d_up) ---- + d_gate = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + d_up = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + if cutlass.const_expr(self._act_func == "swiglu"): + subtile_dprob = dswiglu_act( + d_gate, d_up, r_acc, r_gate, r_up, beta, prob, self._gate_up_clamp + ) + + if cutlass.const_expr(self._dfc2_col_output): + # Snapshot d_gate / d_up BEFORE quant_sfd_row mutates them in place; the col + # path col-quants these copies (quant_sfd_col mutates its input). + d_gate_col = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + d_up_col = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + for j in cutlass.range_constexpr(EN): + d_gate_col[j] = d_gate[j] + d_up_col[j] = d_up[j] + c_gate_col = cute.make_rmem_tensor(cute.make_layout(EN).shape, self.fc1_output_dtype) + c_up_col = cute.make_rmem_tensor(cute.make_layout(EN).shape, self.fc1_output_dtype) + qg_col = quant_sfd_col( + d_gate_col, c_gate_col, norm_const, + self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype, + ) + qu_col = quant_sfd_col( + d_up_col, c_up_col, norm_const, + self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype, + ) + for _k in cutlass.range_constexpr(self._cta_tile_n // EN): + if subtile_idx == cutlass.Int32(_k): + rmem_sf_col_output[2 * _k] = qg_col + rmem_sf_col_output[2 * _k + 1] = qu_col + + # STG gate + up MXFP8 cols -- only valid rows and in-range N strips. + if is_valid_row: + n_col_strips = 2 * (self._cta_tile_n // EN) + gate_strip_idx = ( + work_tile_info.tile_n_idx * cutlass.Int32(n_col_strips) + + subtile_idx * cutlass.Int32(2) + ) + up_strip_idx = gate_strip_idx + cutlass.Int32(1) + token_idx = ( + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + token_row_in_cta + ) + for strip_idx, r_data in ((gate_strip_idx, c_gate_col), (up_strip_idx, c_up_col)): + if strip_idx * cutlass.Int32(EN) < real_fc1_col_output.shape[1]: + strip_base = cute.local_tile( + real_fc1_col_output, (1, EN, 1), + (token_idx, strip_idx, cutlass.Int32(0)), + ) + strip_ptr = cute.make_ptr( + self.fc1_output_dtype, + strip_base.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=EN, + ) + gmem_strip = cute.make_tensor(strip_ptr, cute.make_layout(EN)) + cute.autovec_copy(r_data, gmem_strip) + + # quantize each half to MXFP8 + E8M0 row SF (per-thread, no warp reduction) ---- + qg = quant_sfd_row( + d_gate, c_gate, norm_const, self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype, + ) + qu = quant_sfd_row( + d_up, c_up, norm_const, self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype, + ) + + # accumulate the 2 E8M0 row SFs into rmem + for _k in cutlass.range_constexpr(self._cta_tile_n // EN): + if subtile_idx == cutlass.Int32(_k): + rmem_sf[2 * _k] = qg + rmem_sf[2 * _k + 1] = qu + + # dfc2_recompute: forward swiglu + col quant + per-thread STG + if cutlass.const_expr(self._dfc2_recompute): + c_recompute_f32 = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + if cutlass.const_expr(self._act_func == "swiglu"): + swiglu_act( + c_recompute_f32, + r_up, + r_gate, + prob, + self._gate_up_clamp, + ) + qc = quant_sfd_col( + c_recompute_f32, c_recompute, norm_const, + self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype, + ) + for _k in cutlass.range_constexpr(self._cta_tile_n // EN): + if subtile_idx == cutlass.Int32(_k): + rmem_sf_recompute[_k] = qc + + # STG c_recompute -- only valid rows and in-range N strips. + if is_valid_row: + c_col_idx = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n // EN) + + subtile_idx + ) + expert_local_token_idx = ( + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + token_row_in_cta + ) + if c_col_idx * cutlass.Int32(EN) < real_fc1_recompute.shape[1]: + c_base = cute.local_tile( + real_fc1_recompute, (1, EN, 1), + (expert_local_token_idx, c_col_idx, cutlass.Int32(0)), + ) + c_ptr = cute.make_ptr( + self.fc1_output_dtype, + c_base.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=EN, + ) + gmem_c = cute.make_tensor(c_ptr, cute.make_layout(EN)) + cute.autovec_copy(c_recompute, gmem_c) + + # BARRIER: drain PREVIOUS subtile's TMA BEFORE R2S. + if warp_idx == self._epilogue_warp_ids[0]: + d_pipeline.producer_acquire() + epilog_sync.arrive_and_wait() + + # Write d to smem. + d_n_slots = cutlass.const_expr(d_num_stage // 2) + d_slot = cutlass.Int32(2) * (subtile_i % cutlass.Int32(d_n_slots)) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + sd = thr_copy_r2s.partition_D(smem_d) + cute.copy(tiled_copy_r2s, c_gate, sd[(None, None, None, d_slot)]) + cute.copy(tiled_copy_r2s, c_up, sd[(None, None, None, d_slot + cutlass.Int32(1))]) + + iket.range_pop() + return c_consumer_state, subtile_dprob + + @cute.jit + def _stg_sf_dfc2( + self, + rmem_sf_f32: cute.Tensor, + real_fc1_output_sf: cute.Tensor, + work_tile_info, + tidx, + valid_inter, + ) -> None: + """Store the dfc2 grad_y1 E8M0 row SFs, 4 blocks per 128-col region.""" + if tidx < work_tile_info.valid_tokens_in_cta_tile: + token_idx = ( + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + tidx + ) + n_regions = (self._cta_tile_n * 2) // Fc1EpilogueOutputTileN + region_col = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n * 2) + ) + for r in cutlass.range_constexpr(n_regions): + sf_base = cute.local_tile( + real_fc1_output_sf, (1, 1, 1), + (token_idx, region_col, cutlass.Int32(0)), + ) + r_sf4_f32 = cute.make_rmem_tensor(cute.make_layout(4).shape, self.acc_dtype) + for idx in cutlass.range_constexpr(4): + r_sf4_f32[idx] = rmem_sf_f32[r * 4 + idx] + if region_col < valid_inter: + sf_ptr = cute.make_ptr( + self.sf_dtype, + sf_base.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=4, + ) + gmem_sf4 = cute.make_tensor(sf_ptr, cute.make_layout(4)) + r_sf4 = cute.make_rmem_tensor(cute.make_layout(4).shape, self.sf_dtype) + r_sf4.store(r_sf4_f32.load().to(self.sf_dtype)) + cute.autovec_copy(r_sf4, gmem_sf4) + region_col += cutlass.Int32(Fc1EpilogueOutputTileN) + + @cute.jit + def _stg_col_sf_atom_value( + self, + real_sf: cute.Tensor, + row_block, + col, + hidden_atoms, + sf_value, + ) -> None: + """Store one col-SF in the MN-major 128-column × 4-token-block atom.""" + token_atom = row_block // cutlass.Int32(4) + token_bank = row_block % cutlass.Int32(4) + hidden_atom = col // cutlass.Int32(128) + hidden_bank = (col // cutlass.Int32(32)) % cutlass.Int32(4) + hidden_lane = col % cutlass.Int32(32) + atom_idx = Int64(token_atom) * Int64(hidden_atoms) + Int64(hidden_atom) + byte_offset = ( + atom_idx * Int64(512) + + Int64(hidden_lane) * Int64(16) + + Int64(hidden_bank) * Int64(4) + + Int64(token_bank) + ) + sf_ptr = cute.make_ptr( + self.sf_dtype, + real_sf.iterator.toint() + byte_offset, + cute.AddressSpace.gmem, + assumed_align=1, + ) + gmem_sf1 = cute.make_tensor(sf_ptr, cute.make_layout(1)) + r_sf1 = cute.make_rmem_tensor(cute.make_layout(1).shape, self.sf_dtype) + r_sf1[0] = sf_value.to(self.sf_dtype) + cute.autovec_copy(r_sf1, gmem_sf1) + + @cute.jit + def _stg_sf_recompute( + self, + rmem_sf_f32: cute.Tensor, + real_fc1_recompute_sf: cute.Tensor, + work_tile_info, + tidx, + valid_inter, + valid_tokens, + ) -> None: + """Store fc1_recompute SFs in MN-major 128×4 atoms.""" + EN = EpilogueTileN # 32 + sf_vec_size = self._sf_vec_size + warp_lane_idx = tidx % cutlass.Int32(32) + warp_idx_local = tidx // cutlass.Int32(32) + hidden_atoms = (valid_inter + cutlass.Int32(127)) // cutlass.Int32(128) + + # Row-block within the M-tile: 4 warps × 1 row-block each (128 / 32). + row_blocks_per_m_tile = self._cta_tile_m // sf_vec_size + row_block = ( + work_tile_info.tile_m_idx * cutlass.Int32(row_blocks_per_m_tile) + + warp_idx_local + ) + col_base = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + ) + # Row predicate: a warp stores its SF only if its 32-row block overlaps + # the CTA tile's valid rows. + warp_rows_valid = (warp_idx_local * cutlass.Int32(sf_vec_size)) < valid_tokens + for s in cutlass.range_constexpr(self._cta_tile_n // EN): + col = col_base + cutlass.Int32(s * EN) + warp_lane_idx + if warp_rows_valid and col < valid_inter: + self._stg_col_sf_atom_value( + real_fc1_recompute_sf, + row_block, + col, + hidden_atoms, + rmem_sf_f32[s], + ) + + @cute.jit + def _stg_sf_col_output( + self, + rmem_sf_f32: cute.Tensor, + real_fc1_col_output_sf: cute.Tensor, + work_tile_info, + tidx, + valid_inter, + valid_tokens, + ) -> None: + """Store fc1_col_output SFs in MN-major 128×4 atoms.""" + EN = EpilogueTileN # 32 + sf_vec_size = self._sf_vec_size + warp_lane_idx = tidx % cutlass.Int32(32) + warp_idx_local = tidx // cutlass.Int32(32) + hidden_atoms = (valid_inter + cutlass.Int32(127)) // cutlass.Int32(128) + + row_blocks_per_m_tile = self._cta_tile_m // sf_vec_size + row_block = ( + work_tile_info.tile_m_idx * cutlass.Int32(row_blocks_per_m_tile) + + warp_idx_local + ) + # Doubled N: cta_tile_n * 2 cols per fc1 N-tile. + col_base = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n * 2) + ) + warp_rows_valid = (warp_idx_local * cutlass.Int32(sf_vec_size)) < valid_tokens + for s in cutlass.range_constexpr(self._cta_tile_n // EN): + for gu in cutlass.range_constexpr(2): + col = col_base + cutlass.Int32((2 * s + gu) * EN) + warp_lane_idx + if warp_rows_valid and col < valid_inter: + self._stg_col_sf_atom_value( + real_fc1_col_output_sf, + row_block, + col, + hidden_atoms, + rmem_sf_f32[2 * s + gu], + ) + + @cute.jit + def tma_store_dfc2_output( + self, + smem_d_buffer: cute.Tensor, + tma_atom: cute.CopyAtom, + g_gate_2d: cute.Tensor, + g_up_2d: cute.Tensor, + valid_tokens, + d_pipeline, + d_slot, + ) -> None: + """TMA-store one subtile's gate+up from shared sD stages to grad_y1 GMEM.""" + sD_gate = cute.slice_(smem_d_buffer, (None, None, d_slot)) + sD_up = cute.slice_(smem_d_buffer, (None, None, d_slot + cutlass.Int32(1))) + bSG_sD_gate, bSG_g_gate = cpasync.tma_partition( + tma_atom, 0, cute.make_layout(1), + cute.group_modes(sD_gate, 0, 2), + cute.group_modes(g_gate_2d, 0, 2), + ) + bSG_sD_up, bSG_g_up = cpasync.tma_partition( + tma_atom, 0, cute.make_layout(1), + cute.group_modes(sD_up, 0, 2), + cute.group_modes(g_up_2d, 0, 2), + ) + tile_is_valid = valid_tokens > cutlass.Int32(0) + if tile_is_valid: + cute.copy(tma_atom, bSG_sD_gate, bSG_g_gate) + cute.copy(tma_atom, bSG_sD_up, bSG_g_up) + d_pipeline.producer_commit() + + + @cute.jit + def _subtile_dfc12_tmem_tensor( + self, + tmem_acc_tensor: cute.Tensor, + subtile_idx, + warp_idx, + ) -> cute.Tensor: + """ + Per-warp TMEM view for one fc2 subtile (EpilogueTileN=32 cols). + """ + base = tmem_acc_tensor.iterator + warp_lane_off = warp_idx * WarpThreadCount + subtile_col_off = subtile_idx * EpilogueTileN + total = (warp_lane_off << 16) + subtile_col_off + subtile_ptr = base + cute.assume(total, divby=16) + return cute.make_tensor( + subtile_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ) + + @cute.jit + def _advance_fc2_tmem_tensor( + self, + tmem_tensor: cute.Tensor, + col_offset: int, + ) -> cute.Tensor: + new_ptr = tmem_tensor.iterator + cute.assume(col_offset, divby=16) + return cute.make_tensor( + new_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ) + + @cute.jit + def _acc_pipeline_consumer_release( + self, + acc_pipeline, + acc_consumer_state, + is_release: bool, + ) -> None: + """Release the acc pipeline consumer.""" + if is_release: + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + @cute.jit + def _run_dfc1_subtile( + self, + subtile_idx, + subtile_i, + t_subtile: cute.Tensor, + real_fc2_output: cute.Tensor, + work_tile_info, + valid_hidden, + warp_idx: int, + tidx, + acc_pipeline, + acc_consumer_state, + token_comm_args=None, + rmem_sf_dfc1=None, + *, + preload_acc=None, + ) -> None: + """fc2 subtile: LDTM + fp32->bf16 + STG.""" + iket.range_push("mxfp8_fc2_epilogue_subtile") + dfc1_subtile_cnt = self._cta_tile_n // EpilogueTileN # = 8 + r_acc_layout = cute.make_layout((((EpilogueTileN,), 1),), stride=(((1,), 0),)) + atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), self.acc_dtype, + ) + r_acc = cute.make_rmem_tensor(r_acc_layout.shape, self.acc_dtype) + cute.copy(atom_t2r, t_subtile, r_acc) + + hidden_group = ( + work_tile_info.tile_n_idx * cutlass.Int32(dfc1_subtile_cnt) + subtile_idx + ) + hidden_col_start = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + + subtile_idx * cutlass.Int32(EpilogueTileN) + ) + r_bf16 = cute.make_rmem_tensor(r_acc_layout.shape, cutlass.BFloat16) + r_bf16.store(r_acc.load().to(cutlass.BFloat16)) + thread_in_warp = tidx % WarpThreadCount + token_row_in_cta = cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + is_valid = token_row_in_cta < valid_tokens and hidden_col_start < valid_hidden + + if cutlass.const_expr( + token_comm_args is not None + and not self._token_back_by_dispatch + and self._combine_mxfp8 + ): + fp8_dtype = self._combine_format.act_dtype + r_fp8 = cute.make_rmem_tensor(r_acc_layout.shape, fp8_dtype) + qpvscale = quant_sfd_row( + r_acc, r_fp8, 1.0, EpilogueTileN, + cutlass.Float8E8M0FNU, fp8_dtype, + ) + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + metadata_u32 = cute.recast_tensor( + token_comm_args.token_src_metadata, cutlass.Uint32, + ) + fc2_output_dest = Fc2OutputDest( + tensor=token_comm_args.combine_output, + metadata=metadata_u32, + peer_rank_ptr_mapper=token_comm_args.peer_rank_ptr_mapper, + ) + dest_row = fc2_output_dest.resolve_token_row(pool_token_global) + r_fp8_flat = cute.make_tensor(r_fp8.iterator, cute.make_layout(32)) + stg_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=256, + ) + dest_fp8_ptr = cute.make_ptr( + fp8_dtype, + dest_row.iterator.toint() + Int64(hidden_col_start), + cute.AddressSpace.gmem, + assumed_align=32, + ) + if is_valid: + cute.copy( + stg_fp8_atom, r_fp8_flat, + cute.make_tensor(dest_fp8_ptr, cute.make_layout(32)), + ) + self._write_sf_dfc1_buffer(rmem_sf_dfc1, subtile_idx, qpvscale) + elif cutlass.const_expr( + self._token_back_by_dispatch and self._combine_mxfp8 + ): + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + fp8_dtype = self._combine_format.act_dtype + r_fp8 = cute.make_rmem_tensor(r_acc_layout.shape, fp8_dtype) + qpvscale = quant_sfd_row( + r_acc, r_fp8, 1.0, EpilogueTileN, + cutlass.Float8E8M0FNU, fp8_dtype, + ) + fp8_byte_addr = ( + token_comm_args.fc2_output_workspace.iterator.toint() + + Int64(pool_token_global) * Int64(self._hidden_dfc1) + + Int64(hidden_col_start) + ) + stg_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=256, + ) + aligned_fp8_iter = cute.make_ptr( + fp8_dtype, + fp8_byte_addr, + cute.AddressSpace.gmem, + assumed_align=32, + ) + r_fp8_flat = cute.make_tensor(r_fp8.iterator, cute.make_layout(EpilogueTileN)) + if is_valid: + cute.copy( + stg_fp8_atom, r_fp8_flat, + cute.make_tensor(aligned_fp8_iter, cute.make_layout(EpilogueTileN)), + ) + self._write_sf_dfc1_buffer(rmem_sf_dfc1, subtile_idx, qpvscale) + else: + # BF16 path (default): fp32->bf16, two 256-bit STGs. + stg_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=256, + ) + if cutlass.const_expr( + token_comm_args is not None + and not self._token_back_by_dispatch + ): + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + metadata_u32 = cute.recast_tensor( + token_comm_args.token_src_metadata, cutlass.Uint32, + ) + fc2_output_dest = Fc2OutputDest( + tensor=token_comm_args.combine_output, + metadata=metadata_u32, + peer_rank_ptr_mapper=token_comm_args.peer_rank_ptr_mapper, + # Collapse topk -> src_topk=0 so every contribution of a + # source token resolves to the SAME combine row (red-added + # below). No-op when reduce is off. + reduce_topk_in_kernel=self._reduce_topk_in_epilogue, + ) + dest_row = fc2_output_dest.resolve_token_row(pool_token_global) + for stg_half in cutlass.range(EpilogueTileN // 16, unroll_full=True): + reg_view = cute.make_tensor( + r_bf16.iterator + stg_half * 16, + cute.make_layout(16), + ) + if cutlass.const_expr( + token_comm_args is not None + and not self._token_back_by_dispatch + ): + # epi_warps: peer-write grad_x directly to combine_output + hidden_off = hidden_col_start + cutlass.Int32(stg_half * 16) + if cutlass.const_expr(self._reduce_topk_in_epilogue): + if is_valid: + reg_u32 = cute.recast_tensor(reg_view, cutlass.Uint32) + for redg_i in cutlass.range_constexpr(16 // 4): + chunk_ptr = cute.make_ptr( + cutlass.BFloat16, + dest_row.iterator.toint() + + (hidden_off + cutlass.Int32(redg_i * 4)) + * cutlass.Int64(2), + cute.AddressSpace.gmem, + assumed_align=8, + ) + _red_add_relaxed_sys_v2_bf16x2( + chunk_ptr, + cutlass.Uint32(reg_u32[2 * redg_i]), + cutlass.Uint32(reg_u32[2 * redg_i + 1]), + ) + else: + dest_ptr = cute.make_ptr( + cutlass.BFloat16, + dest_row.iterator.toint() + hidden_off * cutlass.Int64(2), + cute.AddressSpace.gmem, + assumed_align=32, + ) + if is_valid: + cute.copy( + stg_atom, reg_view, + cute.make_tensor(dest_ptr, cute.make_layout(16)), + ) + else: + # Lean path (token_comm_args is None) OR dispatch-push + # (token_back_by_dispatch) + g_fc2_output_tile = cute.local_tile( + real_fc2_output, + (self._cta_tile_m, EpilogueTileN, 1), + (work_tile_info.tile_m_idx, hidden_group, 0), + ) + g_fc2_slice = cute.slice_(g_fc2_output_tile, (None, None, 0)) + g_thread_row = cute.local_tile( + g_fc2_slice, (1, 16), (token_row_in_cta, stg_half), + ) + g_flat = cute.coalesce(g_thread_row) + aligned_iter = cute.make_ptr( + cutlass.BFloat16, + g_flat.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=32, + ) + if is_valid: + cute.copy(stg_atom, reg_view, cute.make_tensor(aligned_iter, g_flat.layout)) + + iket.range_pop() + + @cute.jit + def _write_sf_dfc1_buffer(self, rmem_sf_dfc1, subtile_idx, qpvscale) -> None: + """Scatter one subtile's E8M0 scale into the per-tile SF buffer.""" + for j in cutlass.range_constexpr(self._cta_tile_n // EpilogueTileN): + if subtile_idx == cutlass.Int32(j): + rmem_sf_dfc1[j] = qpvscale + + @cute.jit + def _stg_sf_dfc1( + self, + rmem_sf_dfc1: cute.Tensor, + token_comm_args, + work_tile_info, + valid_hidden, + warp_idx: int, + tidx, + ) -> None: + """Flush a task tile's dfc1 E8M0 scales to local fc2_output_sf.""" + dfc1_subtile_cnt = self._cta_tile_n // EpilogueTileN + thread_in_warp = tidx % WarpThreadCount + token_row_in_cta = cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + if token_row_in_cta < work_tile_info.valid_tokens_in_cta_tile: + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + hidden_group_base = ( + work_tile_info.tile_n_idx * cutlass.Int32(dfc1_subtile_cnt) + ) + sf_byte_addr = ( + token_comm_args.fc2_output_sf.iterator.toint() + + Int64(pool_token_global) * Int64(self._dfc1_sf_block_pad) + + Int64(hidden_group_base) + ) + if cutlass.const_expr(self._dfc1_sf_batch8): + stg_e8m0x8_from_f32( + sf_byte_addr, + rmem_sf_dfc1[0], rmem_sf_dfc1[1], rmem_sf_dfc1[2], rmem_sf_dfc1[3], + rmem_sf_dfc1[4], rmem_sf_dfc1[5], rmem_sf_dfc1[6], rmem_sf_dfc1[7], + ) + else: + for j in cutlass.range_constexpr(dfc1_subtile_cnt): + block_hidden_start = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + + cutlass.Int32(j * EpilogueTileN) + ) + if block_hidden_start < valid_hidden: + stg_e8m0_from_f32(sf_byte_addr + Int64(j), rmem_sf_dfc1[j]) + + + @cute.jit + def _run_dfc1_task_tile( + self, + work_tile_info, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + acc_consumer_state, + sched_ext, + gmem_fc2_output: cute.Tensor, + valid_hidden, + warp_idx: int, + tidx, + token_comm_args=None, + ) -> None: + """fc2 (Linear2) task-tile body following fc1 pattern exactly.""" + real_fc2_output, _ = sched_ext.get_gmem_tensor( + "d", gmem_fc2_output, work_tile_info, + ) + acc_pipeline.consumer_wait(acc_consumer_state) + iket.range_push("mxfp8_dfc1_epi_tile") + + dfc1_subtile_cnt = self._cta_tile_n // EpilogueTileN # = 8 + + # Start subtile mirrors fc1: last for odd turn, first for even. + start_subtile = 0 + tmem_t = self._subtile_dfc12_tmem_tensor( + tmem_acc_tensor, cutlass.Int32(start_subtile), warp_idx, + ) + tmem_forward_cols = EpilogueTileN + + # Quantized combine: buffer the per-subtile E8M0 scales and flush them + # in one stg.64 after the loop (see _stg_sf_dfc1). + if cutlass.const_expr(self._combine_mxfp8 and token_comm_args is not None): + layout_sf_dfc1 = cute.make_layout(dfc1_subtile_cnt) + rmem_sf_dfc1 = cute.make_rmem_tensor(layout_sf_dfc1.shape, self.acc_dtype) + else: + rmem_sf_dfc1 = None + + for i in cutlass.range(0, dfc1_subtile_cnt, 1, unroll=1): + self._run_dfc1_subtile( + subtile_idx=cutlass.Int32(i), + subtile_i=i, + t_subtile=tmem_t, + real_fc2_output=real_fc2_output, + work_tile_info=work_tile_info, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + token_comm_args=token_comm_args, + rmem_sf_dfc1=rmem_sf_dfc1, + ) + + tmem_t = self._advance_fc2_tmem_tensor(tmem_t, tmem_forward_cols) + + # Release AFTER all subtile reads (never early-release for FC2). + self._acc_pipeline_consumer_release(acc_pipeline, acc_consumer_state, True) + + # Flush the buffered E8M0 scales (one stg.64 per thread when aligned). + if cutlass.const_expr(self._combine_mxfp8 and token_comm_args is not None): + self._stg_sf_dfc1( + rmem_sf_dfc1=rmem_sf_dfc1, + token_comm_args=token_comm_args, + work_tile_info=work_tile_info, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + ) + + iket.range_pop() + + + @cute.jit + def run( + self, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + sched_consumer, + sched_ext, + gmem_fc1_output: cute.Tensor, + gmem_fc1_output_sf: cute.Tensor, + gmem_fc1_recompute: Optional[cute.Tensor], + gmem_fc1_recompute_sf: Optional[cute.Tensor], + gmem_fc1_col_output: Optional[cute.Tensor], + gmem_fc1_col_output_sf: Optional[cute.Tensor], + smem_preact_buffer: cute.Tensor, + c_pipeline, + c_num_stage, + smem_d_buffer: cute.Tensor, + d_pipeline, + d_num_stage, + tma_atom_grad_y1: cute.CopyAtom, + gmem_topk_scores: cute.Tensor, + gmem_fc2_output: cute.Tensor, + gmem_fc1_done_counter: cute.Tensor, + warp_idx: int, + tidx, + alpha, + norm_const, + gmem_beta: cute.Tensor, + gmem_dprob: cute.Tensor, + token_comm_args=None, + ) -> None: + """ + Run the full MXFP8 dfc2+dfc1-fused (backward) epilogue task-tile loop. + """ + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._num_acc_pipeline_stages + ) + task_tile_boundary_bar = pipeline.NamedBarrier( + barrier_id=self._epilog_sync_bar_id, + num_threads=32 * len(self._epilogue_warp_ids), + ) + + valid_hidden = cutlass.Int32(gmem_fc2_output.shape[1]) + + c_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, c_num_stage + ) + + bidx, bidy, bidz = cute.arch.block_idx() + work_tile_info = sched_consumer.consume_work() + + flag_tracker = GpuReleaseFlagBatchTracker( + flag_address=Int64(0), + accumulated_flags=cutlass.Int32(0), + phase=cutlass.Int32(work_tile_info.phase), + thread_idx=tidx % (len(self._epilogue_warp_ids) * WarpThreadCount), + ) + + while work_tile_info.is_valid_tile: + acc_stage_index = acc_consumer_state.index + tmem_acc_stage_tesnor = tmem_acc_tensor[(None, None, None, acc_stage_index)] + + if work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1): + c_consumer_state = self._run_dfc2_task_tile( + work_tile_info=work_tile_info, + tmem_acc_tensor=tmem_acc_stage_tesnor, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + sched_ext=sched_ext, + gmem_fc1_output=gmem_fc1_output, + gmem_fc1_output_sf=gmem_fc1_output_sf, + gmem_fc1_recompute=gmem_fc1_recompute, + gmem_fc1_recompute_sf=gmem_fc1_recompute_sf, + gmem_fc1_col_output=gmem_fc1_col_output, + gmem_fc1_col_output_sf=gmem_fc1_col_output_sf, + c_pipeline=c_pipeline, + smem_preact_buffer=smem_preact_buffer, + c_consumer_state=c_consumer_state, + smem_d_buffer=smem_d_buffer, + tma_atom_grad_y1=tma_atom_grad_y1, + warp_idx=warp_idx, + tidx=tidx, + norm_const=norm_const, + gmem_topk_scores=gmem_topk_scores, + gmem_beta=gmem_beta, + gmem_dprob=gmem_dprob, + d_pipeline=d_pipeline, + d_num_stage=d_num_stage, + token_comm_args=token_comm_args, + ) + else: + self._run_dfc1_task_tile( + work_tile_info=work_tile_info, + tmem_acc_tensor=tmem_acc_stage_tesnor, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + sched_ext=sched_ext, + gmem_fc2_output=gmem_fc2_output, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + token_comm_args=token_comm_args, + ) + + acc_consumer_state.advance() + + cur_was_linear1 = work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + cur_fc1_counter_slot = ( + work_tile_info.cumulative_token_block_count + + work_tile_info.tile_m_idx // cutlass.Int32(self._atom_thr_size) + ) + cur_fc2_expert_idx = work_tile_info.expert_idx + + work_tile_info = sched_consumer.consume_work() + + if cur_was_linear1: + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.fence_proxy("async") + cute.arch.fence_acq_rel_gpu() + + task_tile_boundary_bar.arrive_and_wait() + + if cur_was_linear1: + flag_tracker = flag_tracker.accumulate( + work_tile_info.phase, + self._epi_fc1_batch, + (gmem_fc1_done_counter.iterator + cur_fc1_counter_slot).toint(), + ) + else: + if cutlass.const_expr( + self._token_back_by_dispatch or self._combine_mxfp8 + ): + # Fence before (deferred) counter release: make the fc2 + # pool-output STG writes device-visible. + cute.arch.fence_acq_rel_gpu() + fc2_flag_addr = ( + token_comm_args.fc2_done_counter.iterator + cur_fc2_expert_idx + ).toint() + else: + fc2_flag_addr = Int64(0) + no_fire: cutlass.Constexpr = not ( + self._token_back_by_dispatch or self._combine_mxfp8 + ) + flag_tracker = flag_tracker.accumulate( + work_tile_info.phase, + self._epi_fc2_batch, + fc2_flag_addr, + no_fire, + ) + + flag_tracker.fire() + + d_pipeline.producer_tail() diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py new file mode 100644 index 000000000..11aafe176 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py @@ -0,0 +1,73 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Sched extension for the fused fc1+fc2 dGLU-backward MXFP8 kernel.""" + +from typing import Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import Pointer + +from ..fwd_glu.glu_mxfp8_fc12_extension import GluMxFp8Fc12SchedExtension +from .....schedulers.fc12_mapping import NonSwapAbFc12WorkTileInfo + + +def _rewrite_tensor_shape(tensor: cute.Tensor, new_shape: Tuple) -> cute.Tensor: + return cute.make_tensor(tensor.iterator, cute.make_layout(new_shape, stride=tensor.stride)) + + +class DgluMxFp8Fc12SchedExtension(GluMxFp8Fc12SchedExtension): + """ + Sched extension for the fused fc1+fc2 dGLU-backward MXFP8 kernel. + """ + + @cute.jit + def get_gmem_tensor( + self, + tensor_name: str, + gmem_tensor_in_moe_view: cute.Tensor, + work_tile_info: NonSwapAbFc12WorkTileInfo, + ) -> Tuple[cute.Tensor, Optional[Pointer]]: + """dGLU-backward operand views; every other name delegates to the base.""" + data_token_offset = work_tile_info.cumulative_data_physical_row + sf_token_offset = work_tile_info.cumulative_sf_physical_row + + shape = gmem_tensor_in_moe_view.shape + c1 = cutlass.Int32(1) + sf_vec_size = self.sf_vec_size + + if cutlass.const_expr(tensor_name == "recompute"): + # Forward-swiglu recompute data tensor. + real = cute.domain_offset( + (data_token_offset, 0, 0), gmem_tensor_in_moe_view + ) + real = _rewrite_tensor_shape(real, (shape[0], shape[1], c1)) # type: ignore[index] + return (real, None) + + elif cutlass.const_expr(tensor_name == "sfrecompute"): + # Per-expert base for atom-packed col-SF of the forward-swiglu recompute. + real = cute.domain_offset( + (sf_token_offset // sf_vec_size, 0, 0), gmem_tensor_in_moe_view + ) + real = _rewrite_tensor_shape(real, (shape[0], shape[1], c1)) # type: ignore[index] + return (real, None) + + elif cutlass.const_expr(tensor_name == "col_output"): + # Col-quant grad_y1 data tensor (alongside row-quant "d"). + real = cute.domain_offset( + (data_token_offset, 0, 0), gmem_tensor_in_moe_view + ) + real = _rewrite_tensor_shape(real, (shape[0], shape[1], c1)) # type: ignore[index] + return (real, None) + + elif cutlass.const_expr(tensor_name == "sfcol_output"): + # Per-expert base for atom-packed col-SF of the grad_y1 col output. + real = cute.domain_offset( + (sf_token_offset // sf_vec_size, 0, 0), gmem_tensor_in_moe_view + ) + real = _rewrite_tensor_shape(real, (shape[0], shape[1], c1)) # type: ignore[index] + return (real, None) + + return GluMxFp8Fc12SchedExtension.get_gmem_tensor( + self, tensor_name, gmem_tensor_in_moe_view, work_tile_info + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py new file mode 100644 index 000000000..c1f51a930 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py @@ -0,0 +1,2356 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +""" +Fused fc1+fc2 GLU MXFP8 MegaMoE kernel for SM100. +""" + +import dataclasses +from typing import Any, Literal, Optional, Tuple, Type, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute + +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +import cutlass.utils.rubin_helpers as sm107_utils +from cutlass.cute.nvgpu.tcgen05 import CollectorOp + +from ..tmem_transpose import _TmemTranspose16x32Core +from .dglu_mxfp8_fc12_epilogue import DgluMxfp8Epilogue +from .....schedulers import BlockPhase +from .....schedulers.base import WorkIdAcquisitionMode +from .....schedulers.fc12_scheduler import BlackwellFusedFc12Scheduler +from .dglu_mxfp8_fc12_extension import DgluMxFp8Fc12SchedExtension +from ......api import ImplDesc, KernelClass, ProblemDesc, StaticOrRuntimeIntegerType +from ..helpers.constants import ( + SupportedMmaTileM, + SupportedMmaTileN, +) +from ......helpers.iket_compat import iket +from ......helpers.device_workspace import DeviceWorkspace +from ......helpers.smem_workspace import SmemWorkspace +from ......helpers.dsl_helpers import spin_wait +from ......quant_def import CombineFormat, QuantKind +from ......communication.nvlink_domain.token_comm import TokenCommArgs + + +@dataclasses.dataclass(frozen=True) +class _EpilogueCommView: + """Fields the dGLU epilogue reads for cross-rank dfc1/dprob routing.""" + token_src_metadata: Any + combine_output: Any + dprob_output: Any + peer_rank_ptr_mapper: Any + fc2_output_sf: Any = None + fc2_done_counter: Any = None + fc2_output_workspace: Any = None + + +# ============================================================================= +# Sm107Mxfp8DgluDfc21Kernel +# ============================================================================= + +class Sm107Mxfp8DgluDfc21Kernel: + + # SMEM budget for buffers like mbarriers, sched, work-tile buffer, TMEM allocator state + _SmemMiscBudget = 1024 + + # Supported (ab_dtype, sf_vec_size) pairings. + # MXFP8 → Float8E4M3FN / Float8E5M2 + sf_vec_size=32 (FP8-E8M0 scales, MmaMXF8Op) + VALID_AB_DTYPE_SF_SIZE: dict = { + 32: (cutlass.Float8E4M3FN, cutlass.Float8E5M2,), + } + + # Interleave granularity for gate and up in SwiGLU / GeGlu + GateUpInterleave: int = 32 + + def __init__( + self, + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mnk: Tuple[int, int, int], + use_2cta_instrs: bool, + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + load_balance_mode: Literal["static", "atomic_counter"] = "static", + static_expert_shape: Optional[Tuple[int, int, int]] = None, + force_static_sched: bool = True, + clc_bundle_size: Optional[int] = None, + num_sched_stages: Optional[int] = None, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + sf_vec_size: int = 32, + ab_dtype: Type[cutlass.Numeric] = cutlass.Float4E2M1FN, + epi_flag_batch: Optional[Tuple[int, int]] = (1, 1), + dfc2_recompute: bool = False, + dfc2_col_output: bool = False, + fc2_in_kernel_topk_reduce: bool = False, + act_func: str = "swiglu", + gate_up_clamp: Optional[float] = None, + ) -> None: + if not force_static_sched: + raise NotImplementedError( + "v1 only implements force_static_sched=True (lean 7-warp). " + "Dynamic CLC (force_static_sched=False) is future work." + ) + + # Validate (ab_dtype, sf_vec_size) pairing. + if sf_vec_size in self.VALID_AB_DTYPE_SF_SIZE: + valid_ab = self.VALID_AB_DTYPE_SF_SIZE[sf_vec_size] + if ab_dtype not in valid_ab: + raise ValueError( + f"ab_dtype={ab_dtype.__name__} is not valid for " + f"sf_vec_size={sf_vec_size}. " + f"Expected one of: {[t.__name__ for t in valid_ab]}." + ) + else: + valid_sf_vec_sizes = tuple(self.VALID_AB_DTYPE_SF_SIZE) + raise NotImplementedError( + f"sf_vec_size must be one of {valid_sf_vec_sizes} (MXFP8); got {sf_vec_size}." + ) + + + if load_balance_mode not in ("static", "atomic_counter"): + raise ValueError( + f"load_balance_mode must be 'static' or 'atomic_counter'; " + f"got {load_balance_mode!r}." + ) + if act_func not in ("swiglu", "geglu"): + raise ValueError( + f"act_func must be 'swiglu' or 'geglu'; got {act_func!r}." + ) + if act_func != "swiglu": + raise NotImplementedError( + f"act_func={act_func!r} is not yet implemented; only " + "'swiglu' is currently supported (geglu support is planned)." + ) + + # Store ab_dtype so workspace-size helpers can use it without tensors. + self.ab_dtype = ab_dtype + self.act_func = act_func + + self.acc_dtype = acc_dtype + self.mma_tiler_mnk = mma_tiler_mnk + self.cluster_shape_mn = (cluster_shape_mnk[0], cluster_shape_mnk[1]) + self.use_2cta_instrs = use_2cta_instrs + self.force_static_sched = force_static_sched + # static_expert_shape / clc_bundle_size / num_sched_stages + self.static_expert_shape = static_expert_shape + self.clc_bundle_size = clc_bundle_size + self.num_sched_stages = num_sched_stages + + # Fused fc12 sched-side knobs + self.group_hint = group_hint + self.token_padding_block = token_padding_block + self.sf_padding_block = sf_padding_block + self.load_balance_mode = load_balance_mode + + self.sf_vec_size = sf_vec_size + self.arch = "sm_107" + self.epi_flag_batch = epi_flag_batch + self.dfc2_recompute = dfc2_recompute + self.dfc2_col_output = dfc2_col_output + self.fc2_in_kernel_topk_reduce = fc2_in_kernel_topk_reduce + self.gate_up_clamp = abs(gate_up_clamp) if gate_up_clamp is not None else None + + self._validate_mma_tiler_and_cluster_shape() + self.mma_tiler = mma_tiler_mnk + + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + # Warp specialization (9-warp / 288 thread: + dedicated preact-C load warp) + self.occupancy = 1 + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_a_warp_id = 5 + self.tma_b_warp_id = 6 + self.sched_warp_id = 7 + # Dedicated TMA-load warp for the forward pre-activation (dswiglu C), + self.c_load_warp_id = 8 + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_a_warp_id, + self.tma_b_warp_id, + self.sched_warp_id, + self.c_load_warp_id, + *self.epilogue_warp_id, + ) + ) + + # NamedBarriers. + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + self.epi_subtile_bar_ids = (4, 5, 6, 7) + + self.smem_capacity = utils.get_smem_capacity_in_bytes() + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols( + self.arch + ) + + # Warp-specialized register split. + self.epi_reg_cnt = 256 + self.task_reg_cnt = 72 + + # Token-comm (MegaMoE) + self.enable_token_comm: bool = False + self.dispatch_warp_id: Optional[Tuple[int, int, int, int]] = None + self.token_back_by_dispatch: bool = False + self.token_back_standalone: bool = False + self.token_back_warp_id: Optional[Tuple[int, int, int, int]] = None + + def _validate_mma_tiler_and_cluster_shape(self) -> None: + """Validate user-provided geometry against v1 fused-fc12 constraints.""" + m, n, k = self.mma_tiler_mnk + cm, cn = self.cluster_shape_mn + + if m not in SupportedMmaTileM: + raise ValueError( + f"mma_tiler M ({m}) must be one of {SupportedMmaTileM}" + ) + + per_cta_m = m // (2 if self.use_2cta_instrs else 1) + if per_cta_m != 128: + raise ValueError( + f"per-CTA mma_tiler M must be 128, got {per_cta_m} " + f"(mma_tiler_m={m}, use_2cta_instrs={self.use_2cta_instrs})" + ) + + for _name, _blk in ( + ("token_padding_block", self.token_padding_block), + ("sf_padding_block", self.sf_padding_block), + ): + if _blk <= 0 or _blk % self.sf_vec_size != 0: + raise ValueError( + f"{_name} ({_blk}) must be a positive multiple of " + f"sf_vec_size ({self.sf_vec_size}); the col-quant epilogue " + f"turns a per-expert row offset into a col-SF row-block " + f"index by an exact '// sf_vec_size' division." + ) + + if n not in SupportedMmaTileN: + raise ValueError( + f"mma_tiler N ({n}) must be one of {SupportedMmaTileN} in fused fc12 " + f"(N=64 SFB hack is dropped; swap-AB sched handles short-N " + f"via subtile early-exit)." + ) + + sf_k_granularity = self.sf_vec_size * 4 + if k % sf_k_granularity != 0: + raise ValueError( + f"mma_tiler K ({k}) must be a multiple of " + f"sf_vec_size * 4 = {sf_k_granularity}" + ) + + if cm % (2 if self.use_2cta_instrs else 1) != 0: + raise ValueError( + f"cluster_shape M ({cm}) must be even when use_2cta_instrs=True" + ) + + is_pow2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if cm * cn > 16 or not is_pow2(cm) or not is_pow2(cn) or cm > 4 or cn > 4: + raise ValueError( + f"Invalid cluster_shape ({cm}, {cn}): each dim must be " + f"a power of 2 and <= 4, product must be <= 16" + ) + + # v1 swap-AB requires cluster_n == 1. + if cn != 1: + raise NotImplementedError( + f"v1 fused fc12 requires cluster_n == 1 (got {cn}). " + f"cluster_n > 1 needs sentinel-style acc/ab pipeline release." + ) + + def _create_tiled_mmas(self) -> Tuple[cute.TiledMma, cute.TiledMma]: + """Return (tiled_mma, tiled_mma_sfb).""" + common = ( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + ) + # Rubin: the SM107 blockscaled FP8 MMA op hard-requires instruction + tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma( + *common, self.cta_group, + (*self.mma_inst_shape_mn, 64), + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma( + *common, tcgen05.CtaGroup.ONE, + (*self.mma_inst_shape_mn_sfb, 64), + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + return tiled_mma, tiled_mma_sfb + + def _build_scheduler( + self, *, expert_cnt, intermediate_gateup, hidden_dim, launch_cluster_count + ) -> None: + """Construct FC12 scheduler and its SMEM/device workspaces.""" + work_id_mode = "grid_stride" if self.load_balance_mode == "static" else "atomic_counter" + num_scheduler_consumer_threads = 32 * (len(self.epilogue_warp_id) + 4) + if self.static_expert_shape is not None: + expert_cnt, intermediate_gateup, hidden_dim = self.static_expert_shape + problem_desc = ProblemDesc( + { + "expert_count": expert_cnt, + "intermediate_gateup_size": intermediate_gateup, + "hidden_size": hidden_dim, + } + ) + impl_desc = ImplDesc( + { + "num_scheduler_consumer_threads": num_scheduler_consumer_threads, + "mma_tiler_mnk": self.mma_tiler, + "cluster_shape_mn": self.cluster_shape_mn, + "use_2cta_instrs": self.use_2cta_instrs, + "hint": self.group_hint, + "token_padding_block": self.token_padding_block, + "sf_padding_block": self.sf_padding_block, + "work_id_mode": work_id_mode, + "is_swap_ab": False, + "launch_cluster_count": launch_cluster_count, + } + ) + self.scheduler = BlackwellFusedFc12Scheduler(problem_desc, impl_desc) + + sched_smem_ws = SmemWorkspace() + self.scheduler.register_smem_regions(sched_smem_ws) + sched_smem_ws.finalize(max_bytes=self.smem_capacity) + self.sched_smem_ws = sched_smem_ws + + sched_device_ws = DeviceWorkspace() + self.scheduler.register_device_workspace(sched_device_ws) + sched_device_ws.finalize() + self.sched_device_ws = sched_device_ws + + def _setup_attributes(self) -> None: + """Set up MMA / cluster / tile shapes, SMEM layouts, stage counts. + + The fc12 path shares ``mma_tiler_mnk`` and SMEM layouts across phases. + """ + self.mma_inst_shape_mn = (self.mma_tiler[0], self.mma_tiler[1]) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + tiled_mma, tiled_mma_sfb = self._create_tiled_mmas() + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + assert self.mma_tiler[2] % mma_inst_shape_k == 0, ( + f"mma_tiler K ({self.mma_tiler[2]}) must be a multiple of " + f"MMA instruction K ({mma_inst_shape_k})" + ) + + # SFB-specific tiler: rounded-up MN; same K as main tiler. + self.mma_tiler_sfb = ( + self.mma_inst_shape_mn_sfb[0], + self.mma_inst_shape_mn_sfb[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + # Multicast CTA counts + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + _epi_common = dict( + mma_tiler_mnk=self.mma_tiler, + cluster_shape_mn=self.cluster_shape_mn, + use_2cta_instrs=self.use_2cta_instrs, + sf_vec_size=self.sf_vec_size, + fc1_output_dtype=self.fc1_output_dtype, + fc1_output_layout=self.fc1_output_layout, + acc_dtype=self.acc_dtype, + epilog_sync_bar_id=self.epilog_sync_bar_id, + epilogue_warp_ids=self.epilogue_warp_id, + static_expert_shape=self.static_expert_shape, + epi_flag_batch=self.epi_flag_batch, + token_back_by_dispatch=self.token_back_by_dispatch, + dfc2_recompute=self.dfc2_recompute, + dfc2_col_output=self.dfc2_col_output, + fc2_in_kernel_topk_reduce=self.fc2_in_kernel_topk_reduce, + combine_format=getattr(self, "combine_format", None), + combine_hidden=getattr(self, "hidden", None), + act_func=self.act_func, + gate_up_clamp=self.gate_up_clamp, + ) + self.epilogue = DgluMxfp8Epilogue(**_epi_common) + + if self.num_sched_stages is None: + self.num_sched_stages = 2 + + # Reserve SMEM for the preact (dswiglu C) pipeline staging buffer + self.num_c_stage = 2 + assert self.num_c_stage % 2 == 0, f"num_c_stage must be even, got {self.num_c_stage}" + self.num_c_pipe_stage = self.num_c_stage // 2 + self.num_d_stage = 2 + c_bytes_total = self.num_c_stage * self.epilogue.preact_bytes_per_stage + d_bytes_total = self.num_d_stage * self.epilogue.d_bytes_per_stage + self.c_bytes_total = c_bytes_total + self.d_bytes_total = d_bytes_total + + ( + self.num_acc_stage, + self.num_ab_stage, + self.num_sched_stages, + ) = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.sf_dtype, + self.sf_vec_size, + self.c_bytes_total, + self.d_bytes_total, + self.smem_capacity, + self.occupancy, + self.num_sched_stages, + self._smem_misc_budget_bytes() - self._SmemMiscBudget, + ) + + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + + # Read epilogue's accumulator and scale-factor sizing decisions. + self.num_acc_pipeline_stages = self.epilogue.num_acc_pipeline_stages + self.num_acc_stage = self.epilogue.num_acc_stage + self.num_sfa_tmem_cols = self.epilogue.num_sfa_tmem_cols + self.num_sfb_tmem_cols = self.epilogue.num_sfb_tmem_cols + self.num_accumulator_tmem_cols = self.epilogue.num_accumulator_tmem_cols + + # TMA load bytes per stage (A + B + SFA + SFB). + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + self.atom_thr_size = atom_thr_size # store as Python int for use in @cute.kernel + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = ( + a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + ) * atom_thr_size + + # SMEM usage report (all sizes are per-CTA) + _ab_per_stage = a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + _misc_total = self._smem_misc_budget_bytes() + _fixed = _misc_total + self.c_bytes_total + self.d_bytes_total + _total_used = _fixed + self.num_ab_stage * _ab_per_stage + _per_cta_budget = self.smem_capacity // self.occupancy + _free = _per_cta_budget - _total_used + _extra_misc = _misc_total - self._SmemMiscBudget + print( + f"[smem] capacity={self.smem_capacity}B ({self.smem_capacity//1024}KB)" + f" occupancy={self.occupancy}" + f" per-CTA budget={_per_cta_budget}B ({_per_cta_budget//1024}KB)\n" + f" AB stages: {self.num_ab_stage} × {_ab_per_stage}B ({_ab_per_stage/1024:.1f}KB)" + f" = {self.num_ab_stage * _ab_per_stage}B" + f" [A={a_copy_size}B B={b_copy_size}B" + f" SFA={sfa_copy_size}B SFB={sfb_copy_size}B]\n" + f" fixed: misc={_misc_total}B (base={self._SmemMiscBudget}B" + f" + subclass_extra={_extra_misc}B)" + f" preact(C)={self.num_c_stage}×{self.epilogue.preact_bytes_per_stage}B" + f" sD(D)={self.num_d_stage}×{self.epilogue.d_bytes_per_stage}B" + f" used={_total_used}B ({_total_used/1024:.1f}KB)" + f" free={_free}B ({_free/1024:.1f}KB)\n" + ) + + @staticmethod + def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_bytes_total: int, + d_bytes_total: int, + smem_capacity: int, + occupancy: int, + num_sched_stages: int, + extra_misc_bytes: int = 0, + ) -> Tuple[int, int, int]: + """Compute stage counts for ACC, AB+SF, and scheduler. + """ + num_acc_stage = 2 + + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1, + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 1, + ) + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, mma_tiler_mnk, sf_vec_size, 1, + ) + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, mma_tiler_mnk, sf_vec_size, 1, + ) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + + fixed_overhead = ( + Sm107Mxfp8DgluDfc21Kernel._SmemMiscBudget + extra_misc_bytes + c_bytes_total + d_bytes_total + ) + + num_ab_stage = ( + smem_capacity // occupancy - fixed_overhead + ) // ab_bytes_per_stage + return num_acc_stage, num_ab_stage, num_sched_stages + + def get_workspace_size_in_bytes( + self, + fc1_activation_tensor, + fc1_weight_tensor, + ) -> int: + """Compute opaque workspace size for one fused dfc2+dfc1 launch.""" + sf_padding_block = self.sf_padding_block + sf_vec_size = self.sf_vec_size + + mma_tiler_n = self.mma_tiler_mnk[1] + + data_total_rows, _hidden = fc1_activation_tensor.shape + experts, _hidden_w, dfc2_weight_n = fc1_weight_tensor.shape + # grad_y1 (doubled dswiglu output) width = intermediate = 2 * inter_half. + intermediate_out = dfc2_weight_n * 2 + + # Conservative upper bound for sf_total_rows. + sf_total_rows_upper = data_total_rows + experts * sf_padding_block + + # grad_y1 byte size (MXFP8, 8-bit: 1 element per byte). + fc1_output_bytes = ( + data_total_rows * intermediate_out * self.ab_dtype.width // 8 + ) + + # grad_y1 SF sf_vec_size matches the kernel's sf_vec_size. + fc1_out_sf_vec_size = self.sf_vec_size + sf_block_cols = ( + (intermediate_out // fc1_out_sf_vec_size) + 3 + ) // 4 * 4 + fc1_output_sf_bytes = sf_total_rows_upper * sf_block_cols + + # fc1_recompute (forward-swiglu recompute): N = inter_half = intermediate_out // 2. + fc1_recompute_bytes = ( + data_total_rows * dfc2_weight_n * self.ab_dtype.width // 8 + ) + fc1_recompute_row_blocks_upper = sf_total_rows_upper // fc1_out_sf_vec_size + fc1_recompute_sf_bytes = fc1_recompute_row_blocks_upper * dfc2_weight_n + + # fc1_col_output (col-quant grad_y1): N = intermediate_out (same as + # grad_y1's row-quant fc1_output). Col-SF: row_blocks × intermediate. + fc1_col_output_bytes = fc1_output_bytes + fc1_col_output_sf_bytes = fc1_recompute_row_blocks_upper * intermediate_out + + # fc1_done_counter: one Int32 per CTA-level token block (each cluster block + # has atom_thr_size CTAs, each with its own per-CTA counter slot). + counter_slots_upper = ( + (data_total_rows + mma_tiler_n - 1) // mma_tiler_n + + experts + ) + fc1_done_counter_bytes = counter_slots_upper * 4 + + # load_balance_counter: Int32 scalar. + if self.load_balance_mode == "atomic_counter": + load_balance_counter_bytes = 4 + else: + load_balance_counter_bytes = 0 + + total = ( + fc1_output_bytes + + fc1_output_sf_bytes + + fc1_recompute_bytes + + fc1_recompute_sf_bytes + + fc1_col_output_bytes + + fc1_col_output_sf_bytes + + fc1_done_counter_bytes + + load_balance_counter_bytes + ) + + # 128B align (TMA tensor base address alignment requirement). + alignment = 128 + total = ((total + alignment - 1) // alignment) * alignment + return total + + def mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """SMEM → TMEM tiled copy + partition for SFA / SFB.""" + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t, tCsSF_compact_s2t_ + ) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + # ========================================================================= + # Token-comm hook surface (MegaMoE-only; lean base = no-op stubs) + # + # Mirrors the hook interface in ``moe_mxfp8_glu.kernel_mxfp8_glu_fc12`` + # so that ``Sm107MegaMoEMxfp8DgluKernel`` can override exactly the same + # methods. The mega wrapper realigns dispatch onto warps 8-11 (128-aligned + # for next token_comm) and relocates ``c_load_warp_id`` above the transfer + # block (warp 12 or 16); the lean base keeps c_load at warp 8. + # ========================================================================= + + def _smem_misc_budget_bytes(self) -> int: + """SMEM reserved for non-problem-tensor buffers (mbarriers, sched, TMEM state). + + MegaMoE subclass adds dispatch-warp SMEM on top via:: + + return super()._smem_misc_budget_bytes() + self._dispatch_smem_bytes() + """ + return self._SmemMiscBudget + + def token_comm_extra_smem_storage_class(self) -> type: + """Return a ``@cute.struct`` for dispatch-warp SMEM, or None.""" + return None + + def token_comm_hook_fc1_ready_counter_ptr(self, token_comm_args): + """Return dispatch->fc1 release counter pointer, or None (lean: disabled).""" + return None + + def sched_ext_fc1_peek_threshold(self) -> int: + """Return the fc1 ready-counter peek threshold for DgluMxFp8Fc12SchedExtension.""" + return 0 + + def sched_ext_fc1_counter_cumul_scale(self) -> int: + """Return the scale factor for the fc1 ready-counter slot formula.""" + return 1 + + @cute.jit + def token_comm_hook_sched_warp_pre_init_wait(self, token_comm_args): + """Sched warp: wait for dispatch barrier before reading sizes. No-op base.""" + pass + + @cute.jit + def token_comm_hook_fc1_tma_b_predispatch_spin(self, token_comm_args, work_tile_info): + """TMA-A warp: spin until dispatch-pulled tokens are resident. No-op base.""" + pass + + @cute.jit + def token_comm_hook_dispatch_warp_body( + self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx, + ): + """Body for dispatch warps 8-11 (MegaMoE-only). No-op base.""" + pass + + @cute.jit + def token_comm_hook_token_back_warp_body( + self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx, + ): + """Body for standalone token-back warps 12-15 (MegaMoE-only). No-op base.""" + pass + + @cute.jit + def token_comm_hook_kernel_tail(self, token_comm_args, *, warp_idx, lane_idx, tidx): + """All-warp kernel tail (NVLink release, etc.). No-op base.""" + pass + + @cute.jit + def __call__( + self, + activation: cute.Tensor, # (token_sum_padded, hidden) = grad_out + fc1_weight: cute.Tensor, # (experts, hidden, inter_half) + activation_sf: cute.Tensor, # (token_sum_padded_sf, hidden / sf_vec_size) + fc1_weight_sf: cute.Tensor, # dfc2-weight SF + fc1_output: cute.Tensor, # (token_sum_padded, intermediate) + fc1_output_sf: cute.Tensor, # (token_sum_padded_sf, intermediate / sf_vec_size) + fc1_recompute: Optional[cute.Tensor], # (token_sum_padded, inter_half) + fc1_recompute_sf: Optional[cute.Tensor], # (token_sum_padded_sf, inter_half / sf_vec_size) + fc1_col_output: Optional[cute.Tensor], # (token_sum_padded, intermediate) + fc1_col_output_sf: Optional[cute.Tensor], # (sf_row_blocks, intermediate) col-SF + fc2_weight: cute.Tensor, # (experts, intermediate, hidden) + fc2_weight_sf: cute.Tensor, # dfc1-weight SF + fc2_output: cute.Tensor, # (token_sum_padded, hidden) BFloat16 = grad_x + fc1_preact: cute.Tensor, # (token_sum_padded, intermediate) BFloat16 + topk_scores: cute.Tensor, # (token_sum_padded,) Float32 + beta: cute.Tensor, # (experts,) Float32 + dprob: cute.Tensor, # (token_sum_padded,) Float32 + fc1_done_counter: cute.Tensor, # (max_token_block_per_rank,) Int32 + offs: Optional[cute.Tensor] = None, # (experts,) Int32 cumulative end offsets + max_active_clusters: cutlass.Constexpr = None, + stream: cuda.CUstream = None, + norm_const_tensor: Optional[cute.Tensor] = None, + global_activation_sf: Optional[cute.Tensor] = None, + global_fc1_weight_sf: Optional[cute.Tensor] = None, + load_balance_counter: Optional[cute.Tensor] = None, + expert_token_sizes: Optional[cute.Tensor] = None, + token_comm_args=None, + overflow_flag: cute.Tensor = None, + mega_peer_rank_ptr_mapper=None, + mega_local_rank: Optional[cutlass.Int32] = None, + mega_local_workspace: Optional[cute.Pointer] = None, + mega_shared_workspace: Optional[cute.Pointer] = None, + mega_activation: Optional[cute.Tensor] = None, + mega_activation_sf: Optional[cute.Tensor] = None, + mega_pre_reduced_activation: Optional[cute.Tensor] = None, + mega_pre_reduced_activation_sf: Optional[cute.Tensor] = None, + ) -> None: + """Launch the fused dfc2+dfc1 dGLU MXFP8 (backward) kernel.""" + if cutlass.const_expr(self.static_expert_shape is not None): + ( + experts_static, + intermediate_gateup_static, # inter_half = dfc2 weight N + hidden_static, + ) = self.static_expert_shape + intermediate_out_static = intermediate_gateup_static * 2 # grad_y1 / dfc1-K + + fc1_weight = cute.make_tensor( + fc1_weight.iterator, + cute.make_layout( + (experts_static, hidden_static, intermediate_gateup_static), + stride=fc1_weight.stride, + ), + ) + fc2_weight = cute.make_tensor( + fc2_weight.iterator, + cute.make_layout( + (experts_static, intermediate_out_static, hidden_static), + stride=fc2_weight.stride, + ), + ) + activation = cute.make_tensor( + activation.iterator, + cute.make_layout( + (activation.shape[0], hidden_static), + stride=activation.stride, + ), + ) + fc1_output = cute.make_tensor( + fc1_output.iterator, + cute.make_layout( + (fc1_output.shape[0], intermediate_out_static), + stride=fc1_output.stride, + ), + ) + fc1_recompute = cute.make_tensor( + fc1_recompute.iterator, + cute.make_layout( + (fc1_recompute.shape[0], intermediate_gateup_static), + stride=fc1_recompute.stride, + ), + ) + fc1_col_output = cute.make_tensor( + fc1_col_output.iterator, + cute.make_layout( + (fc1_col_output.shape[0], intermediate_out_static), + stride=fc1_col_output.stride, + ), + ) + fc1_preact = cute.make_tensor( + fc1_preact.iterator, + cute.make_layout( + (fc1_preact.shape[0], intermediate_out_static), + stride=fc1_preact.stride, + ), + ) + if cutlass.const_expr(len(fc2_output.shape) == 3): + fc2_output = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (fc2_output.shape[0], fc2_output.shape[1], hidden_static), + stride=fc2_output.stride, + ), + ) + else: + fc2_output = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (fc2_output.shape[0], hidden_static), + stride=fc2_output.stride, + ), + ) + + # ── GEMM-domain transform for fc1 phase ── + c1 = cutlass.Int32(1) + c0 = cutlass.Int32(0) + + # A_gemm (fc1 activations): (tokens_sum, hidden) -> (M=tokens, K=hidden, L=1). + tokens_sum, hidden = activation.shape + activation_gemm = cute.make_tensor( + activation.iterator, + cute.make_layout( + (tokens_sum, hidden, 1), + stride=(activation.stride[0], activation.stride[1], 0), + ), + ) + + # B_gemm (fc1 weights): (experts, hidden, intermediate_gateup) with hidden stride-1 (K-major) + # -> (N=intermediate_gateup, K=hidden, L=experts). + experts, hidden_b, intermediate_gateup = fc1_weight.shape + fc1_weight_gemm = cute.make_tensor( + fc1_weight.iterator, + cute.make_layout( + (intermediate_gateup, hidden_b, experts), + stride=(fc1_weight.stride[2], fc1_weight.stride[1], fc1_weight.stride[0]), + ), + ) + + intermediate_downproj = fc1_output.shape[1] + fc1_output_gemm = cute.make_tensor( + fc1_output.iterator, + cute.make_layout( + (tokens_sum, intermediate_downproj, 1), + stride=(fc1_output.stride[0], fc1_output.stride[1], 0), + ), + ) + + # fc1_recompute: forward swiglu recomputed from the fc1 c-tensor. + intermediate_downproj_half = fc1_recompute.shape[1] + fc1_recompute_gemm = cute.make_tensor( + fc1_recompute.iterator, + cute.make_layout( + (tokens_sum, intermediate_downproj_half, 1), + stride=(fc1_recompute.stride[0], fc1_recompute.stride[1], 0), + ), + ) + + # fc1_col_output: col-quantized grad_y1 alongside row-quant fc1_output. + fc1_col_output_gemm = cute.make_tensor( + fc1_col_output.iterator, + cute.make_layout( + (tokens_sum, intermediate_downproj, 1), + stride=(fc1_col_output.stride[0], fc1_col_output.stride[1], 0), + ), + ) + + # Forward pre-activation (gate||up) + fc1_preact_gemm = cute.make_tensor( + fc1_preact.iterator, + cute.make_layout( + (tokens_sum, intermediate_downproj, 1), + stride=(fc1_preact.stride[0], fc1_preact.stride[1], 0), + ), + ) + + # SFA / SFB scale tensors (atom-tiled) + tokens_sum_padded = activation_sf.shape[0] + hidden_padded = activation_sf.shape[1] * self.sf_vec_size + activation_sf_gemm = cute.make_tensor( + activation_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (tokens_sum_padded, hidden_padded, 1), self.sf_vec_size + ), + ) + intermediate_gateup_padded_mul_hidden_padded = fc1_weight_sf.shape[1] + intermediate_gateup_padded = ( + intermediate_gateup_padded_mul_hidden_padded * self.sf_vec_size + ) // hidden_padded + fc1_weight_sf_gemm = cute.make_tensor( + fc1_weight_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (intermediate_gateup_padded, hidden_padded, experts), + self.sf_vec_size, + ), + ) + + # GEMM-domain transform for fc2 phase ── + experts2, intermediate_downproj_b2, hidden_b2 = fc2_weight.shape + fc2_weight_gemm = cute.make_tensor( + fc2_weight.iterator, + cute.make_layout( + (hidden_b2, intermediate_downproj_b2, experts2), + stride=(fc2_weight.stride[2], fc2_weight.stride[1], fc2_weight.stride[0]), + ), + ) + + if cutlass.const_expr(len(fc2_output.shape) == 3): + fc2_hidden_out = fc2_output.shape[2] + fc2_output_gemm = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (fc2_output.shape[0], fc2_hidden_out, c1), + stride=(fc2_output.stride[0], fc2_output.stride[2], c0), + ), + ) + else: + fc2_hidden_out = fc2_output.shape[1] + fc2_output_gemm = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (tokens_sum, fc2_hidden_out, c1), + stride=(fc2_output.stride[0], fc2_output.stride[1], c0), + ), + ) + + fc1_out_sf_vec_size = self.sf_vec_size + tokens_sum_padded_sf = fc1_output_sf.shape[0] + intermediate_downproj_padded = fc1_output_sf.shape[1] * fc1_out_sf_vec_size + fc1_output_sf_gemm_for_fc2_load = cute.make_tensor( + fc1_output_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (tokens_sum_padded_sf, intermediate_downproj_padded, 1), + fc1_out_sf_vec_size, + ), + ) + + hidden_padded_fc2_mul_intermediate_downproj_padded = fc2_weight_sf.shape[1] + hidden_padded_fc2 = ( + hidden_padded_fc2_mul_intermediate_downproj_padded * self.sf_vec_size + ) // intermediate_downproj_padded + fc2_weight_sf_gemm = cute.make_tensor( + fc2_weight_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (hidden_padded_fc2, intermediate_downproj_padded, experts2), + self.sf_vec_size, + ), + ) + + expert_cnt = experts + hidden_dim = hidden + + # Infer dtypes and major modes + self.a_dtype: Type[cutlass.Numeric] = activation_gemm.element_type + self.b_dtype: Type[cutlass.Numeric] = fc1_weight_gemm.element_type + self.fc1_output_dtype: Type[cutlass.Numeric] = fc1_output_gemm.element_type + self.sf_dtype: Type[cutlass.Numeric] = activation_sf_gemm.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(activation_gemm).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(fc1_weight_gemm).mma_major_mode() + self.fc1_output_layout = utils.LayoutEnum.from_tensor(fc1_output_gemm) + + self._setup_attributes() + tiled_mma, tiled_mma_sfb = self._create_tiled_mmas() + + # fc1 TMA atoms load A1 + a_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_fc1_activation, tma_tensor_fc1_activation = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + activation_gemm, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load B1 + b_op = sm100_utils.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_fc1_weight, tma_tensor_fc1_weight = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + fc1_weight_gemm, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load SFA1 + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_fc1_activation_sf, tma_tensor_fc1_activation_sf = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + activation_sf_gemm, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64, + ) + + # TMA load SFB1 + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_fc1_weight_sf, tma_tensor_fc1_weight_sf = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + fc1_weight_sf_gemm, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Uint64, + ) + + # Coalesced TMA G2S load of the forward preact (dswiglu C input). + preact_tma_op = cpasync.CopyBulkTensorTileG2SOp() + tma_atom_fc1_preact, tma_tensor_fc1_preact = cpasync.make_tiled_tma_atom( + preact_tma_op, + fc1_preact_gemm, + self.epilogue.preact_smem_layout_one_stage, + self.epilogue.preact_epi_tile, + ) + + # Coalesced TMA S2G store of grad_y1 (dfc2 fp8 output). + grad_y1_tma_op = cpasync.CopyBulkTensorTileS2GOp() + tma_atom_grad_y1, tma_tensor_grad_y1 = cpasync.make_tiled_tma_atom( + grad_y1_tma_op, + fc1_output_gemm, + self.epilogue.d_smem_layout_one_stage, + self.epilogue.d_epi_tile, + ) + + # fc1 SFC GMEM tensor (= fc1_output_sf user view). No TMA atom; it is + # per-thread STG. + fc1_output_sf_gemm = cute.make_tensor( + fc1_output_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (tokens_sum_padded, intermediate_downproj, 1), + self.sf_vec_size, + ), + ) + + # fc1_recompute SFC GMEM storage. The epilogue uses the iterator as the + # base of per-expert MN-major 128-column × 4-token-block atoms. + fc1_recompute_sf_row_blocks = fc1_recompute_sf.shape[0] + fc1_recompute_sf_gemm = cute.make_tensor( + fc1_recompute_sf.iterator, + cute.make_layout( + (fc1_recompute_sf_row_blocks, intermediate_downproj_half, 1), + stride=(fc1_recompute_sf.stride[0], fc1_recompute_sf.stride[1], 0), + ), + ) + + # fc1_col_output SFC GMEM storage, likewise atom-packed by the epilogue. + fc1_col_output_sf_row_blocks = fc1_col_output_sf.shape[0] + fc1_col_output_sf_gemm = cute.make_tensor( + fc1_col_output_sf.iterator, + cute.make_layout( + (fc1_col_output_sf_row_blocks, intermediate_downproj, 1), + stride=(fc1_col_output_sf.stride[0], fc1_col_output_sf.stride[1], 0), + ), + ) + + # ── fc2 TMA atoms: fc1_output → A-side (M=tokens), fc2_weight → B-side (N=hidden) ── + tma_atom_fc2_activation, tma_tensor_fc2_activation = ( + cute.nvgpu.make_tiled_tma_atom_A( + a_op, + fc1_output_gemm, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + ) + tma_atom_fc2_weight, tma_tensor_fc2_weight = ( + cute.nvgpu.make_tiled_tma_atom_B( + b_op, + fc2_weight_gemm, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + ) + tma_atom_fc2_activation_sf, tma_tensor_fc2_activation_sf = ( + cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + fc1_output_sf_gemm_for_fc2_load, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64, + ) + ) + tma_atom_fc2_weight_sf, tma_tensor_fc2_weight_sf = ( + cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + fc2_weight_sf_gemm, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Uint64, + ) + ) + + # ── Scheduler params + grid + launch ── + if cutlass.const_expr(self.load_balance_mode == "atomic_counter"): + if cutlass.const_expr(load_balance_counter is None): + raise ValueError( + "load_balance_counter must be provided when " + "load_balance_mode == 'atomic_counter'" + ) + load_balance_counter_ptr = load_balance_counter.iterator + else: + load_balance_counter_ptr = None + + # On the MegaMoE path the per-expert sizes come from the Router (device-side), so the + # caller supplies neither offs nor expert_token_sizes. + if cutlass.const_expr(not self.enable_token_comm): + if cutlass.const_expr((offs is None) == (expert_token_sizes is None)): + raise ValueError( + "Exactly one of `offs` / `expert_token_sizes` must be " + "non-None. Got offs=" + f"{'set' if offs is not None else 'None'}, " + f"expert_token_sizes=" + f"{'set' if expert_token_sizes is not None else 'None'}." + ) + + self._build_scheduler( + expert_cnt=expert_cnt, + intermediate_gateup=intermediate_gateup, + hidden_dim=hidden_dim, + launch_cluster_count=max_active_clusters, + ) + grid = self.scheduler.get_grid_shape(max_active_clusters=max_active_clusters) + + self.kernel( + tiled_mma, + tiled_mma_sfb, + # fc1 TMA atoms / tensors (A=activations, B=weights) + tma_atom_fc1_activation, + tma_tensor_fc1_activation, + tma_atom_fc1_weight, + tma_tensor_fc1_weight, + tma_atom_fc1_activation_sf, + tma_tensor_fc1_activation_sf, + tma_atom_fc1_weight_sf, + tma_tensor_fc1_weight_sf, + # fc2 TMA atoms / tensors (fc1_output→A, fc2_weight→B) + tma_atom_fc2_activation, + tma_tensor_fc2_activation, + tma_atom_fc2_weight, + tma_tensor_fc2_weight, + tma_atom_fc2_activation_sf, + tma_tensor_fc2_activation_sf, + tma_atom_fc2_weight_sf, + tma_tensor_fc2_weight_sf, + # GEMM-domain tensors (fc1) + activation_gemm, + fc1_weight_gemm, + fc1_output_gemm, + activation_sf_gemm, + fc1_weight_sf_gemm, + fc1_output_sf_gemm, + # GEMM-domain tensors (fc2) + fc2_weight_gemm, + fc2_output_gemm, + fc2_weight_sf_gemm, + fc1_output_sf_gemm_for_fc2_load, + # forward pre-activation (dswiglu input) — TMA G2S into SMEM + tma_atom_fc1_preact, + tma_tensor_fc1_preact, + tma_atom_grad_y1, + tma_tensor_grad_y1, + # fc1_recompute (forward swiglu) — per-thread STG, N = inter_half + fc1_recompute_gemm, + fc1_recompute_sf_gemm, + # fc1_col_output (col-quant grad_y1) — per-thread STG, N = intermediate + fc1_col_output_gemm, + fc1_col_output_sf_gemm, + # topk / beta / dprob + cross-phase sync workspace + topk_scores, + beta, + dprob, + overflow_flag, + fc1_done_counter, + # Scheduling + offs, + expert_token_sizes, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + # SMEM layouts + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + token_comm_args, + # MegaMoE push-model token-comm inputs (None on the lean path) + mega_peer_rank_ptr_mapper, + mega_local_rank, + mega_local_workspace, + mega_shared_workspace, + mega_activation, + mega_activation_sf, + mega_pre_reduced_activation, + mega_pre_reduced_activation_sf, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=self.occupancy, + ) + + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + # fc1 TMA atoms / tensors + tma_atom_fc1_activation_1: cute.CopyAtom, + tma_tensor_fc1_activation_1: cute.Tensor, + tma_atom_weight: cute.CopyAtom, + tma_tensor_weight: cute.Tensor, + tma_atom_fc1_activation_1_sf: cute.CopyAtom, + tma_tensor_fc1_activation_1_sf: cute.Tensor, + tma_atom_fc1_weight_sf: cute.CopyAtom, + tma_tensor_fc1_weight_sf: cute.Tensor, + # fc2 TMA atoms / tensors (fc1_output→A, fc2_weight→B) + tma_atom_fc2_activation: cute.CopyAtom, + tma_tensor_fc2_activation: cute.Tensor, + tma_atom_fc2_weight: cute.CopyAtom, + tma_tensor_fc2_weight: cute.Tensor, + tma_atom_fc2_activation_sf: cute.CopyAtom, + tma_tensor_fc2_activation_sf: cute.Tensor, + tma_atom_fc2_weight_sf: cute.CopyAtom, + tma_tensor_fc2_weight_sf: cute.Tensor, + # GEMM-domain tensors (fc1) + activation_gemm: cute.Tensor, + fc1_weight_gemm: cute.Tensor, + fc1_output_gemm: cute.Tensor, + activation_sf_gemm: cute.Tensor, + fc1_weight_sf_gemm: cute.Tensor, + fc1_output_sf_gemm: cute.Tensor, + # GEMM-domain tensors (fc2) + fc2_weight_gemm: cute.Tensor, + fc2_output_gemm: cute.Tensor, + fc2_weight_sf_gemm: cute.Tensor, + fc1_output_sf_gemm_for_fc2_load: cute.Tensor, + # forward pre-activation (dswiglu input) — TMA G2S into SMEM + tma_atom_fc1_preact: cute.CopyAtom, + tma_tensor_fc1_preact: cute.Tensor, + # grad_y1 (dfc2 output) — TMA S2G store + tma_atom_grad_y1: cute.CopyAtom, + tma_tensor_grad_y1: cute.Tensor, + # fc1_recompute (forward swiglu) — per-thread STG, N = inter_half + fc1_recompute_gemm: cute.Tensor, + fc1_recompute_sf_gemm: cute.Tensor, + # fc1_col_output (col-quant grad_y1) — per-thread STG, N = intermediate + fc1_col_output_gemm: cute.Tensor, + fc1_col_output_sf_gemm: cute.Tensor, + # topk / beta / dprob + cross-phase sync workspace + topk_scores: cute.Tensor, + beta: cute.Tensor, + dprob: cute.Tensor, + overflow_flag: cute.Tensor, + fc1_done_counter: cute.Tensor, + # Scheduling + offs: Optional[cute.Tensor], + expert_token_sizes: Optional[cute.Tensor], + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + # SMEM layouts + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + token_comm_args=None, + # MegaMoE push-model token-comm inputs + mega_peer_rank_ptr_mapper=None, + mega_local_rank: Optional[cutlass.Int32] = None, + mega_local_workspace: Optional[cute.Pointer] = None, + mega_shared_workspace: Optional[cute.Pointer] = None, + mega_activation: Optional[cute.Tensor] = None, + mega_activation_sf: Optional[cute.Tensor] = None, + mega_pre_reduced_activation: Optional[cute.Tensor] = None, + mega_pre_reduced_activation_sf: Optional[cute.Tensor] = None, + ): + """Device kernel for fused fc1+fc2 swap-AB GLU MXFP8 grouped GEMM.""" + a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, None, 0)) + b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, None, 0)) + sfa_smem_layout = cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)) + sfb_smem_layout = cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)) + + # MegaMoE (push model): bind the device workspace here so the fc1_ready counter + # pointer that the scheduler extension spins on (built just below) resolves. + if cutlass.const_expr(self.enable_token_comm): + self._mega_device_workspace.assign_device_members( + mega_local_workspace, mega_shared_workspace + ) + + # fc2 waits for all fc1 intermediate N-tiles in the same token block. + ext_fc2_spin_threshold = ( + fc1_weight_gemm.shape[0] + self.cta_tile_shape_mnk[1] - 1 + ) // self.cta_tile_shape_mnk[1] * self.epilogue._atom_thr_size + + ext = DgluMxFp8Fc12SchedExtension( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter.iterator, + fc2_spin_threshold=ext_fc2_spin_threshold, + # MegaMoE: peek the dispatch->fc1 ready counter (None on the lean + # path). Parity with the forward kernel's SchedExtension wiring. + fc1_ready_counter_pointer=self.token_comm_hook_fc1_ready_counter_ptr( + token_comm_args + ), + # Fold the 2 CTAs of a cluster onto one fc1_ready slot + cluster_m=self.epilogue._atom_thr_size, + ) + + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + tidx, _, _ = cute.arch.thread_idx() + + # MegaMoE (push model): bind token-comm device members (transfer-warp state + the + # NVLink barrier's peer mapper) before any token_in / token_back / size-wait runs. + if cutlass.const_expr(self.enable_token_comm): + _mega_token_comm_args = TokenCommArgs( + mega_activation, + mega_activation_sf, + mega_pre_reduced_activation, + mega_pre_reduced_activation_sf, + mega_peer_rank_ptr_mapper, + ) + _mega_cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] + _, _, _mega_cluster_idx = cute.arch.block_idx() + _mega_linear_cta_idx = cta_rank_in_cluster + _mega_cluster_idx * _mega_cluster_size + self.token_comm.assign_device_members( + device_workspace=self._mega_device_workspace, + token_comm_args=_mega_token_comm_args, + local_rank=mega_local_rank, + linear_cta_idx=_mega_linear_cta_idx, + ) + + # preact (dswiglu C) pipeline + num_c_stage = self.num_c_stage + num_c_pipe_stage = self.num_c_pipe_stage + num_d_stage = self.num_d_stage + + # SharedStorage (mainloop + epilogue SMEM). next's scheduler owns its own + # SMEM workspace, allocated separately below. + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_pipeline_stages * 2 + ] + c_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_c_pipe_stage * 2] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + sPre: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(self.epilogue.preact_staged_smem_layout(num_c_stage).outer), + ], + 1024, + ] + # grad_y1 (dfc2 output) store staging — stage 0 = gate, stage 1 = up. + sD: cute.struct.Align[ + cute.struct.MemRange[ + self.fc1_output_dtype, + cute.cosize(self.epilogue.d_staged_smem_layout(num_d_stage).outer), + ], + 1024, + ] + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # next scheduler SMEM: a self-contained workspace carved from the same + # allocator; its transport regions resolve against ``sched_smem_base``. + sched_storage = smem.allocate(self.sched_smem_ws.storage_class()) + sched_smem_base = sched_storage.buffer.data_ptr() + + # MegaMoE-only dispatch-warp SMEM (pull_buffer, mbarriers, etc.). + # Kept out of ``SharedStorage`` so the lean path never allocates it. + TokenCommStorageCls = self.token_comm_extra_smem_storage_class() + if cutlass.const_expr(TokenCommStorageCls is not None): + token_comm_storage = smem.allocate(TokenCommStorageCls) + else: + token_comm_storage = None + + # ── Pipelines: two TMA producer warps share the AB pipeline. ── + + ab_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, 2 + ) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes // 2, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = ( + len(self.epilogue_warp_id) * 32 * (2 if use_2cta_instrs else 1) + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_pipeline_stages, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # preact (dswiglu C) pipeline + c_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + c_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len(self.epilogue_warp_id) + ) + c_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.c_full_mbar_ptr.data_ptr(), + num_stages=num_c_pipe_stage, + producer_group=c_pipeline_producer_group, + consumer_group=c_pipeline_consumer_group, + tx_count=2 * self.epilogue.preact_bytes_per_stage, + defer_sync=True, + ) + # d pipeline + d_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + d_pipeline = pipeline.PipelineTmaStore.create( + num_stages=num_d_stage // 2, + producer_group=d_producer_group, + ) + + + # TMEM allocator + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, + arch=self.arch, + ) + + # Sched + scheduler = self.scheduler + if cutlass.const_expr(self.enable_token_comm): + _sched_expert_sizes = self.token_comm.local_expert_sizes( + self._mega_device_workspace, mega_local_rank + ) + _sched_prefix_sum = None + # Bind the scheduler's own device workspace + self.sched_device_ws.assign_device_members( + cute.make_ptr( + cutlass.Uint8, + mega_local_workspace.toint() + + self._mega_device_workspace.offset(self.sched_work_id_region), + cute.AddressSpace.gmem, + assumed_align=16, + ), + mega_shared_workspace, + ) + else: + _sched_expert_sizes = expert_token_sizes + _sched_prefix_sum = offs + scheduler.assign_device_members( + expert_token_sizes=_sched_expert_sizes, + expert_token_prefix_sum=_sched_prefix_sum, + actual_expert_shape=None, + block_idx=cute.arch.block_idx(), + smem_workspace=self.sched_smem_ws, + smem_base=sched_smem_base, + device_workspace=self.sched_device_ws, + ) + sched_consumer = scheduler.make_consumer() + + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # SMEM tensors A / B / SFA / SFB (shared by fc1 / fc2) + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + sSFA = smem.allocate_tensor( + element_type=self.sf_dtype, + layout=sfa_smem_layout_staged, + byte_alignment=128, + ) + sSFB = smem.allocate_tensor( + element_type=self.sf_dtype, + layout=sfb_smem_layout_staged, + byte_alignment=128, + ) + + # preact (dswiglu C) staging tensor + preact_smem_layout_staged = self.epilogue.preact_staged_smem_layout( + num_c_stage + ) + sPre = storage.sPre.get_tensor( + preact_smem_layout_staged.outer, + swizzle=preact_smem_layout_staged.inner, + ) + + # grad_y1 (dfc2 output) store staging tensor (stage 0 = gate, stage 1 = up). + d_smem_layout_staged = self.epilogue.d_staged_smem_layout(num_d_stage) + sD = storage.sD.get_tensor( + d_smem_layout_staged.outer, + swizzle=d_smem_layout_staged.inner, + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + + # acc_fake layout: (MMA, MMA_M, MMA_N, STAGE) + acc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # Cluster wait before TMEM alloc. + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + mma_tiler_k = self.mma_tiler[2] + k_tile_cnt_fc1 = (fc1_weight_gemm.shape[1] + mma_tiler_k - 1) // mma_tiler_k + k_tile_cnt_fc2 = (fc2_weight_gemm.shape[1] + mma_tiler_k - 1) // mma_tiler_k + # fc2 spin threshold: number of N-tiles per CTA (per-CTA counter now). + fc2_spin_threshold = ( + (fc1_weight_gemm.shape[0] + self.cta_tile_shape_mnk[1] - 1) + // self.cta_tile_shape_mnk[1] + ) * self.epilogue._atom_thr_size + + # ════════════════════════════════════════════════════════════════════ + # Scheduler warp (warp 7) — lean path + # ════════════════════════════════════════════════════════════════════ + if warp_idx == self.sched_warp_id: + cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) + # MegaMoE: block until the Router has published this rank's per-expert sizes + self.token_comm_hook_sched_warp_pre_init_wait(token_comm_args) + work_tile = scheduler.gen_next_work() + while work_tile.is_valid_tile: + scheduler.publish_work(ext.prepare_work_tile(work_tile)) + work_tile = scheduler.gen_next_work() + # Sentinel publish (the tile is already invalid here). + scheduler.publish_work(work_tile) + scheduler.produce_tail() + + # ════════════════════════════════════════════════════════════════════ + # TMA load warps (warps 5 / 6) + # ════════════════════════════════════════════════════════════════════ + # + # TMA-A loads weights/SFA; TMA-B loads activations/SFB and waits for + # fc1 workspace readiness in fc2 phase. Both feed the same AB pipeline. + + # ── TMA-A warp (warp 5) ───────────────────────────────────────────── + if warp_idx == self.tma_a_warp_id: + cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) + a_full_mcast_mask = None + sfa_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + + b_full_mcast_mask = None + if cutlass.const_expr(self.is_b_mcast or use_2cta_instrs): + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + sfa_cta_layout = a_cta_layout + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + + work_tile_info = sched_consumer.consume_work() + + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + if is_phase_linear1: + iket.range_push("tma_weight_fc1") + # MegaMoE: spin until the dispatch (token_in) warps have pulled + ext.wait_for_input(work_tile_info) + self.token_comm_hook_fc1_tma_b_predispatch_spin( + token_comm_args, work_tile_info, + ) + + k_tile_cnt = k_tile_cnt_fc1 + real_a, desc_ptr_a = ext.get_gmem_tensor( + "fc1_activation", tma_tensor_fc1_activation_1, work_tile_info, + ) + real_sfa, desc_ptr_sfa = ext.get_gmem_tensor( + "fc1_activation_sf", tma_tensor_fc1_activation_1_sf, work_tile_info, + ) + + gA_mkl = cute.local_tile( + real_a, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + gSFA_mkl = cute.local_tile( + real_sfa, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + tCgA = thr_mma.partition_A(gA_mkl) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + + tAsA, tAgA = cpasync.tma_partition( + tma_atom_fc1_activation_1, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_fc1_activation_1_sf, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + mma_tile_m = work_tile_info.tile_m_idx // cute.size( + tiled_mma.thr_id.shape + ) + tAgA_slice = tAgA[(None, mma_tile_m, None, 0)] + tAgSFA_slice = tAgSFA[(None, mma_tile_m, None, 0)] + + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance( + peek_ab_empty_status + ) + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + cute.copy( + tma_atom_fc1_activation_1, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_a, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_fc1_activation_1_sf, + tAgSFA_slice[(None, handle.count)], + tAsSFA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfa, + mcast_mask=sfa_full_mcast_mask, + ) + iket.range_pop() + + else: + # fc2 phase A-side: load fc1_output (M=tokens) + wait for fc1 done + iket.range_push("tma_token_fc2") + counter_slot = ( + work_tile_info.cumulative_token_block_count + + work_tile_info.tile_m_idx // cutlass.Int32(self.epilogue._atom_thr_size) + ) + counter_ptr = fc1_done_counter.iterator + counter_slot + iket.range_push("tma_token_fc2_a_wait") + spin_wait( + counter_ptr, + lambda v: v >= fc2_spin_threshold, + sleep_cycles=20, + ) + iket.range_pop() + k_tile_cnt = k_tile_cnt_fc2 + real_a, desc_ptr_a = ext.get_gmem_tensor( + "fc2_activation", tma_tensor_fc2_activation, work_tile_info, + ) + real_sfa, desc_ptr_sfa = ext.get_gmem_tensor( + "fc2_activation_sf", tma_tensor_fc2_activation_sf, work_tile_info, + ) + + gA_mkl = cute.local_tile( + real_a, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + gSFA_mkl = cute.local_tile( + real_sfa, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + tCgA = thr_mma.partition_A(gA_mkl) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + + tAsA, tAgA = cpasync.tma_partition( + tma_atom_fc2_activation, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_fc2_activation_sf, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + mma_tile_m = work_tile_info.tile_m_idx // cute.size( + tiled_mma.thr_id.shape + ) + tAgA_slice = tAgA[(None, mma_tile_m, None, 0)] + tAgSFA_slice = tAgSFA[(None, mma_tile_m, None, 0)] + + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance( + peek_ab_empty_status + ) + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + cute.copy( + tma_atom_fc2_activation, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_a, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_fc2_activation_sf, + tAgSFA_slice[(None, handle.count)], + tAsSFA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfa, + mcast_mask=sfa_full_mcast_mask, + ) + + iket.range_pop() + work_tile_info = sched_consumer.consume_work() + + ab_producer.tail() + + # ── TMA-B warp (warp 6) ───────────────────────────────────────────── + if warp_idx == self.tma_b_warp_id: + cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) + b_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr(self.is_b_mcast or use_2cta_instrs): + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, + block_in_cluster_coord_sfb_vmnk, + mcast_mode=1, + ) + + # FC1: weight (B) is multicast (like original A) + a_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + sfb_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) + + work_tile_info = sched_consumer.consume_work() + + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + + if is_phase_linear1: + iket.range_push("tma_token_fc1") + + k_tile_cnt = k_tile_cnt_fc1 + real_b, desc_ptr_b = ext.get_gmem_tensor( + "fc1_weight", tma_tensor_weight, work_tile_info, + ) + real_sfb, desc_ptr_sfb = ext.get_gmem_tensor( + "fc1_weight_sf", tma_tensor_fc1_weight_sf, work_tile_info, + ) + + # N-K tiling for N-side weight (N=intermediate, K=hidden). + gB_nkl = cute.local_tile( + real_b, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + gSFB_nkl = cute.local_tile( + real_sfb, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + tCgB = thr_mma.partition_B(gB_nkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + + tBsB, tBgB = cpasync.tma_partition( + tma_atom_weight, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_fc1_weight_sf, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + tBgB_slice = tBgB[(None, work_tile_info.tile_n_idx, None, 0)] + tBgSFB_slice = tBgSFB[(None, work_tile_info.tile_n_idx, None, 0)] + + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance( + peek_ab_empty_status + ) + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + cute.copy( + tma_atom_weight, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_b, + mcast_mask=a_full_mcast_mask, # same as A-loading for weights + ) + cute.copy( + tma_atom_fc1_weight_sf, + tBgSFB_slice[(None, handle.count)], + tBsSFB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfb, + mcast_mask=sfb_full_mcast_mask, + ) + iket.range_pop() + + else: + # fc2 phase B-side: load fc2_weight (N=hidden), no counter wait + iket.range_push("tma_weight_fc2") + k_tile_cnt = k_tile_cnt_fc2 + real_b, desc_ptr_b = ext.get_gmem_tensor( + "fc2_weight", tma_tensor_fc2_weight, work_tile_info, + ) + real_sfb, desc_ptr_sfb = ext.get_gmem_tensor( + "fc2_weight_sf", tma_tensor_fc2_weight_sf, work_tile_info, + ) + + gB_nkl = cute.local_tile( + real_b, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + gSFB_nkl = cute.local_tile( + real_sfb, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + tCgB = thr_mma.partition_B(gB_nkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + + tBsB, tBgB = cpasync.tma_partition( + tma_atom_fc2_weight, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_fc2_weight_sf, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + fc2_b_hidden_tile = work_tile_info.tile_n_idx + tBgB_slice = tBgB[(None, fc2_b_hidden_tile, None, 0)] + tBgSFB_slice = tBgSFB[(None, fc2_b_hidden_tile, None, 0)] + + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance( + peek_ab_empty_status + ) + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + cute.copy( + tma_atom_fc2_weight, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_b, + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_fc2_weight_sf, + tBgSFB_slice[(None, handle.count)], + tBsSFB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfb, + mcast_mask=sfb_full_mcast_mask, + ) + iket.range_pop() + work_tile_info = sched_consumer.consume_work() + + ab_producer.tail() + + # ════════════════════════════════════════════════════════════════════ + # MMA warp (warp 4) + # ════════════════════════════════════════════════════════════════════ + # + # Both phases share tiled_mma and TMEM; only K-tile count differs. + if warp_idx == self.mma_warp_id: + cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + acc_base = cute.make_tensor(acc_tmem_ptr, acc_fake.layout) + + # SFA TMEM tensor (placed after the acc cols). + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + # SFB TMEM tensor (after acc + SFA cols). + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + + ( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t, + tCtSFA_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + ( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t, + tCtSFB_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_pipeline_stages + ) + + work_tile_info = sched_consumer.consume_work() + + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + # Prebind k_tile_cnt due to DSL AST. + k_tile_cnt = cutlass.Int32(0) + if is_phase_linear1: + k_tile_cnt = k_tile_cnt_fc1 + iket.range_push("mma_dfc2") + else: + k_tile_cnt = k_tile_cnt_fc2 + iket.range_push("mma_dfc1") + + acc_stage_index = acc_producer_state.index + + if is_leader_cta: + tCtAcc = acc_base[(None, None, None, acc_stage_index)] + + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if k_tile_cnt > 0: + peek_ab_full_status = ab_consumer.try_wait() + acc_pipeline.producer_acquire(acc_producer_state) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + iket.range_push("mma_ab_wait") + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + iket.range_pop() + + s2t_stage_coord = (None, None, None, None, handle.index) + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t[s2t_stage_coord], + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t[s2t_stage_coord], + tCtSFB_compact_s2t, + ) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + tile_crd = (None, None, None, handle.index) + cute.gemm( + tiled_mma, + tCtAcc, + [tCrA[tile_crd], tCtSFA], + [tCrB[tile_crd], tCtSFB], + tCtAcc, + ) + handle.release() + + if k_tile_cnt > 0: + acc_pipeline.producer_commit(acc_producer_state) + if k_tile_cnt > 0: + acc_producer_state.advance() + + iket.range_pop() + + work_tile_info = sched_consumer.consume_work() + + acc_pipeline.producer_tail(acc_producer_state) + + # ════════════════════════════════════════════════════════════════════ + # Dedicated preact-C TMA-load warp (c_load_warp_id) — c_pipeline PRODUCER + # ════════════════════════════════════════════════════════════════════ + # + # Mirrors the reference's epilog_load_tma warp: consume the same work + # tiles in lockstep, and for each Linear1 (dfc2) tile TMA-load gate + # (epi-tile 2*s) then up (2*s+1) into successive c_pipeline stages. + if warp_idx == self.c_load_warp_id: + cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) + + c_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, num_c_pipe_stage + ) + thr_mma_c = tiled_mma.get_slice(mma_tile_coord_v) + preact_epi_tile = self.epilogue.preact_epi_tile + c_subtile_cnt = self.cta_tile_shape_mnk[1] // 32 # 8 + + work_tile_info = sched_consumer.consume_work() + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + if is_phase_linear1: + real_preact, _ = ext.get_gmem_tensor( + "c", tma_tensor_fc1_preact, work_tile_info + ) + gC_mnl = cute.local_tile( + real_preact, cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + tCgC = thr_mma_c.partition_C(gC_mnl) + gC_epi = cute.flat_divide( + tCgC[((None, None), 0, 0, None, None, None)], preact_epi_tile + ) + bGS_sPre, bGS_gC = cpasync.tma_partition( + tma_atom_fc1_preact, 0, cute.make_layout(1), + cute.group_modes(sPre, 0, 2), + cute.group_modes(gC_epi, 0, 2), + ) + mma_m_coord = work_tile_info.tile_m_idx // cutlass.Int32(self.atom_thr_size) + mma_n_coord = work_tile_info.tile_n_idx * cutlass.Int32(2) + bGS_gC = bGS_gC[(None, None, None, mma_m_coord, mma_n_coord, 0)] + bGS_gC = cute.group_modes(bGS_gC, 1, cute.rank(bGS_gC)) + + for i in cutlass.range(0, c_subtile_cnt, 1, unroll=1): + subtile_idx = cutlass.Int32(i) + # gate (2*subtile) then up (2*subtile+1) + c_pipeline.producer_acquire(c_producer_state) + c_bar = c_pipeline.producer_get_barrier(c_producer_state) + c_slot = 2 * c_producer_state.index + cute.copy( + tma_atom_fc1_preact, + bGS_gC[(None, subtile_idx * cutlass.Int32(2) + cutlass.Int32(0))], + bGS_sPre[(None, c_slot)], + tma_bar_ptr=c_bar + ) + cute.copy( + tma_atom_fc1_preact, + bGS_gC[(None, subtile_idx * cutlass.Int32(2) + cutlass.Int32(1))], + bGS_sPre[(None, c_slot + 1)], + tma_bar_ptr=c_bar, + ) + c_producer_state.advance() + + work_tile_info = sched_consumer.consume_work() + + c_pipeline.producer_tail(c_producer_state) + + # ════════════════════════════════════════════════════════════════════ + # Epilogue warps (warps 0-3) + # ════════════════════════════════════════════════════════════════════ + # + # Fully delegated to ``self.epilogue.run(...)`` -- the epilogue owns + # the entire 2-phase task-tile loop. + if warp_idx < self.mma_warp_id: + cute.arch.warpgroup_reg_alloc(self.epi_reg_cnt) + epi_warp_idx = warp_idx + + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + acc_tensor = cute.make_tensor(acc_tmem_ptr, acc_fake.layout) + + # The epilogue is the preact c_pipeline CONSUMER (the dedicated + # c_load warp is the producer); it reads gate/up from sPre stages. + _run_kwargs = dict( + tmem_acc_tensor=acc_tensor, + acc_pipeline=acc_pipeline, + sched_consumer=sched_consumer, + sched_ext=ext, + gmem_fc1_output=tma_tensor_grad_y1, + gmem_fc1_output_sf=fc1_output_sf_gemm, + gmem_fc1_recompute=fc1_recompute_gemm, + gmem_fc1_recompute_sf=fc1_recompute_sf_gemm, + gmem_fc1_col_output=fc1_col_output_gemm, + gmem_fc1_col_output_sf=fc1_col_output_sf_gemm, + smem_preact_buffer=sPre, + c_pipeline=c_pipeline, + c_num_stage=num_c_pipe_stage, + smem_d_buffer=sD, + d_pipeline=d_pipeline, + d_num_stage=num_d_stage, + tma_atom_grad_y1=tma_atom_grad_y1, + gmem_topk_scores=topk_scores, + gmem_fc2_output=fc2_output_gemm, + gmem_fc1_done_counter=fc1_done_counter, + warp_idx=epi_warp_idx, + tidx=tidx, + alpha=cutlass.Float32(1.0), + norm_const=cutlass.Float32(1.0), + gmem_beta=beta, + gmem_dprob=dprob, + ) + + # MegaMoE: pass token_comm_args only when it is a real bundle (not + # None). Passing Python None explicitly to @cute.jit methods + # triggers a CuteDSL codegen issue; const_expr dispatch avoids any + # None-as-JIT-argument path. + if cutlass.const_expr(self.enable_token_comm): + # MegaMoE (push model): bridge next's TokenComm accessors + peer mapper into + # the dGLU epilogue's Fc2OutputDest peer-store expectations for grad_x combine. + _epi_comm = _EpilogueCommView( + token_src_metadata=self.token_comm.token_src_metadata_tensor( + self._mega_device_workspace + ), + combine_output=mega_pre_reduced_activation, + dprob_output=dprob, + peer_rank_ptr_mapper=mega_peer_rank_ptr_mapper, + fc2_output_sf=self.token_comm.fc2_activation_sf_tensor(self._mega_device_workspace), + fc2_done_counter=self.token_comm.fc2_done_counter_tensor(self._mega_device_workspace), + fc2_output_workspace=self.token_comm.fc2_activation_tensor(self._mega_device_workspace), + ) + self.epilogue.run(**_run_kwargs, token_comm_args=_epi_comm) + elif cutlass.const_expr(token_comm_args is not None): + self.epilogue.run(**_run_kwargs, token_comm_args=token_comm_args) + else: + self.epilogue.run(**_run_kwargs) + + tmem.relinquish_alloc_permit() + tmem.free(acc_tmem_ptr) + if cutlass.const_expr(self.enable_token_comm): + cute.arch.fence_acq_rel_sys() + + # ════════════════════════════════════════════════════════════════════ + # Dispatch / token_back warps hook (warps 8-11 [+ 12-15]; MegaMoE-only) + # ════════════════════════════════════════════════════════════════════ + # + # ``enable_token_comm=False`` → these warps don't exist (lean base has 9 + # warps), so the guard is const_expr-eliminated in the lean path. + # NOTE: c_load now lives ABOVE the transfer block (warp 12 or 16), so the + # gate must be UPPER-bounded at the last transfer warp — otherwise the + # c_load warp (already run above) would re-enter the dispatch body. + if cutlass.const_expr(self.enable_token_comm): + _last_transfer_warp = ( + self.token_back_warp_id[-1] + if self.token_back_standalone + else self.dispatch_warp_id[-1] + ) + if (warp_idx >= self.dispatch_warp_id[0]) & (warp_idx <= _last_transfer_warp): + cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) + lane_idx_for_dispatch = cute.arch.lane_idx() + if cutlass.const_expr(self.token_back_standalone): + if warp_idx < self.token_back_warp_id[0]: + self.token_comm_hook_dispatch_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + else: + self.token_comm_hook_token_back_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + else: + self.token_comm_hook_dispatch_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + + # ════════════════════════════════════════════════════════════════════ + # Kernel tail hook (MegaMoE-only; lean base = no-op) + # ════════════════════════════════════════════════════════════════════ + lane_idx = cute.arch.lane_idx() + self.token_comm_hook_kernel_tail( + token_comm_args, + warp_idx=warp_idx, + lane_idx=lane_idx, + tidx=tidx, + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py new file mode 100644 index 000000000..b7844237a --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py @@ -0,0 +1,939 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Full MegaMoE (multi-rank) mxfp8 dGLU training-backward kernel.""" + +from types import SimpleNamespace +from typing import Any, Literal, Optional, Tuple, Type + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int64 + +from ......api import ImplDesc, KernelClass, ProblemDesc, StaticOrRuntimeIntegerType +from ......helpers.device_workspace import DeviceWorkspace +from ......helpers.smem_workspace import SmemWorkspace +from ......helpers.utils import ceil_div, round_up +from ......quant_def import CombineFormat, QuantKind +from ......communication.nvlink_domain.token_comm_deterministic import TokenCommDeterministic +from ..topk_reduce import TopkReduce +from ..fwd_glu.glu_mxfp8_col_requant import Mxfp8ColRequant +from .dglu_mxfp8_fc12_kernel import Sm107Mxfp8DgluDfc21Kernel + + +_AB_DTYPE_TO_QUANT_KIND = {cutlass.Float8E4M3FN: QuantKind.mxfp8_e4m3, cutlass.Float8E5M2: QuantKind.mxfp8_e5m2} +_QUANT_KIND_TO_AB_DTYPE = {str(k): d for d, k in _AB_DTYPE_TO_QUANT_KIND.items()} + +# TVM-FFI export symbol for the AOT-compiled callable. +_aot_symbol_prefix = "rubin_mega_moe_dglu_mxfp8_aot" + + +class Sm107MegaMoEMxfp8DgluKernel(Sm107Mxfp8DgluDfc21Kernel, KernelClass): + """Multi-rank MegaMoE wrapper around the lean mxfp8 dGLU kernel.""" + + # grad_y1 (dfc2 output), its SF, cross-phase counter, and internal dGLU pools. + fc1_output_region = "rubin.dglu_mxfp8.mega.fc1_output" + fc1_output_sf_region = "rubin.dglu_mxfp8.mega.fc1_output_sf" + fc1_done_counter_region = "rubin.dglu_mxfp8.mega.fc1_done_counter" + load_balance_counter_region = "rubin.dglu_mxfp8.mega.load_balance_counter" + sched_work_id_region = "rubin.dglu_mxfp8.mega.sched_work_id" + # Host-side local mirror of the token_comm shared token_src_metadata, so the + # legacy dfc2_recompute / dfc2_col_output validation (which reads it from the + # LOCAL workspace) works with the next's shared-heap metadata layout. + token_src_metadata_local_region = "rubin.dglu_mxfp8.mega.token_src_metadata_local" + fc1_preact_region = "rubin.dglu_mxfp8.mega.fc1_preact" + grad_y2_sizes_region = "rubin.dglu_mxfp8.mega.grad_y2_expert_token_sizes" + + # Reserved on top of the exact token_comm/sched SMEM to cover smem.allocate + # inter-allocation alignment padding that _compute_stages does not model. + _SMEM_ALLOC_MARGIN = 2048 + + @classmethod + def problem_desc_require(cls): + return { + "expert_count": StaticOrRuntimeIntegerType, + "intermediate_gateup_size": StaticOrRuntimeIntegerType, + "hidden_size": StaticOrRuntimeIntegerType, + "quant_kind": str, + "combine_format": CombineFormat, + "world_size": int, + "local_rank": int, + "topk": int, + "max_tokens_per_rank": int, + "max_recv_size_per_rank": int, + "gate_up_clamp": Optional[float], + } + + @classmethod + def impl_desc_require(cls): + return { + "mma_tiler_mnk": tuple, + "cluster_shape_mnk": tuple, + "use_2cta_instrs": bool, + "group_hint": int, + "token_padding_block": int, + "sf_padding_block": int, + "load_balance_mode": str, + "force_static_sched": bool, + "clc_bundle_size": Optional[int], + "num_sched_stages": Optional[int], + "acc_dtype": type, + "sf_vec_size": int, + "launch_cluster_count": int, + "drop_on_overflow": bool, + "fc2_in_kernel_topk_reduce": bool, + "token_back_mode": str, + "epi_flag_batch": tuple, + "flag_batch": int, + "act_func": str, + "dfc2_recompute": bool, + "dfc2_col_output": bool, + "enable_grad_y2_col_quant": bool, + "num_ctas_grad_y2_col_quant": int, + } + + def name(self) -> str: + return ( + f"sm107_megamoe_dglu_{self.quant_kind}_m{self.mma_tiler_mnk[0]}n{self.mma_tiler_mnk[1]}" + f"k{self.mma_tiler_mnk[2]}_e{self.expert_count}_ep{self.world_size}_topk{self.topk}_" + f"h{self.hidden_size}_i{self.intermediate_gateup_size}_combine{self.combine_format}_" + f"clamp{self.gate_up_clamp}_" + f"tokenback{self.token_back_mode}_hint{self.group_hint}_" + f"epi{self.epi_flag_batch[0]}x{self.epi_flag_batch[1]}_tif{self.flag_batch}_" + f"deterministic_mtpr{self.max_tokens_per_rank}_mrpr{self.max_recv_size_per_rank}_" + f"drop{int(self.drop_on_overflow)}_lc{self.launch_cluster_count}_" + f"recompute{int(self.dfc2_recompute)}x{int(self.dfc2_col_output)}_" + f"redtopk{int(self.reduce_topk_in_kernel)}_preactarg1" + ) + + def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): + """Compile against fake (metadata-only) inputs; ``out_path=None`` returns the in-memory callable.""" + import math + + from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_ptr + from cutlass.cute.typing import AddressSpace, sym_int64 + from cutlass.cutlass_dsl import Int32, Int64 + + from ......communication.nvlink_domain.symmetric_buffer import SymmetricBufferHost + + def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): + extents = tuple( + sym_int64(divisibility=math.gcd(int(extent), 128)) if axis in dynamic_axes else int(extent) + for axis, extent in enumerate(shape) + ) + return make_fake_compact_tensor(dtype, extents, stride_order=stride_order, assumed_align=alignment) + + tokens = self.max_tokens_per_rank + hidden = self.hidden + inter_half = self.intermediate_downproj # dfc2 weight N + gate_up = self.intermediate_gateup # grad_y1 width / dfc1 K = 2 * inter_half + experts = self.num_experts_per_rank + vec = self.sf_vec_size + activation_dtype = self.token_comm.activation_dtype # grad_out fp8 + sf_dtype = self.token_comm.activation_sf_dtype # E8M0 + # Atom-swizzled weight SF extents (to_blocked pads rows->128, cols->4). + fc1_weight_sf_columns = round_up(inter_half, 128) * round_up(hidden // vec, 4) + fc2_weight_sf_columns = round_up(hidden, 128) * round_up(gate_up // vec, 4) + aux_shapes = self.get_aux_output_shapes() + + fake_arguments = dict( + grad_out=fake_tensor(activation_dtype, (tokens, hidden), (1, 0), {0}, 16), + grad_out_sf=fake_tensor(sf_dtype, (tokens, self.token_comm.activation_sf_hidden_padded), (1, 0), {0}, 16), + topk_idx=fake_tensor(cutlass.Int64, (tokens, self.num_topk), (1, 0), {0}, 16), + topk_weights=fake_tensor(cutlass.Float32, (tokens, self.num_topk), (1, 0), {0}, 4), + fc1_weight=fake_tensor(self.ab_dtype, (experts, hidden, inter_half), (2, 0, 1), {0, 2}, 16), + fc1_weight_sf=fake_tensor(sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16), + fc2_weight=fake_tensor(self.ab_dtype, (experts, gate_up, hidden), (2, 0, 1), {0, 2}, 16), + fc2_weight_sf=fake_tensor(sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16), + beta=fake_tensor(cutlass.Float32, (experts,), (0,), {0}, 4), + fc1_preact=fake_tensor(cutlass.BFloat16, self.get_fc1_preact_shape(), (1, 0), set(), 128), + output_activation=fake_tensor(cutlass.BFloat16, (tokens, hidden), (1, 0), {0}, 16), + overflow_flag=fake_tensor(cutlass.Int32, (1,), (0,), set(), 4), + dprob=fake_tensor(cutlass.Float32, aux_shapes["dprob"], (1, 0), {0}, 16), + fc1_recompute=fake_tensor(self.ab_dtype, aux_shapes["fc1_recompute"], (1, 0), set(), 128), + fc1_recompute_sf=fake_tensor(sf_dtype, aux_shapes["fc1_recompute_sf"], (1, 0), set(), 128), + fc1_col_output=fake_tensor(self.ab_dtype, aux_shapes["fc1_col_output"], (1, 0), set(), 128), + fc1_col_output_sf=fake_tensor(sf_dtype, aux_shapes["fc1_col_output_sf"], (1, 0), set(), 128), + local_workspace=make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), + shared_workspace=make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), + peer_rank_ptr_mapper_host=SymmetricBufferHost( + base_address=Int64(0), + offsets=tuple(Int64(0) for _ in range(self.world_size)), + rank=Int32(0), + max_ranks=self.world_size, + ), + stream=make_fake_stream(), + ) + fake_arguments["grad_y2"] = fake_tensor( + self.ab_dtype, aux_shapes["grad_y2"], (1, 0), {0}, 16 + ) + fake_arguments["grad_y2_sf"] = fake_tensor( + cutlass.Uint8, aux_shapes["grad_y2_sf"], (0,), set(), 16 + ) + + compiled = cute.compile[cute.EnableTVMFFI(True)](self, **fake_arguments) + if out_path is None: + return compiled + compiled.export_to_c(out_path, function_name=_aot_symbol_prefix, export_only_tvm_ffi_symbols=True) + return out_path + + @staticmethod + def load_compiled(path: str): + from cutlass.cute.runtime import load_module + + return load_module(path, enable_tvm_ffi=True)[_aot_symbol_prefix] + + @classmethod + def from_kwargs( + cls, + # Base-class (lean dfc2+dfc1) kwargs. + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mnk: Tuple[int, int, int], + use_2cta_instrs: bool, + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + load_balance_mode: str = "static", + static_expert_shape: Optional[Tuple[int, int, int]] = None, + force_static_sched: bool = True, + clc_bundle_size: Optional[int] = None, + num_sched_stages: Optional[int] = None, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + ab_dtype: Type[cutlass.Numeric] = cutlass.Float8E4M3FN, + sf_vec_size: int = 32, + *, + world_size: int, + local_rank: int, + num_topk: int, + max_tokens_per_rank: int, + max_recv_size_per_rank: int, + hidden: int, + launch_cluster_count: int, + drop_on_overflow: bool, + fc2_in_kernel_topk_reduce: bool = False, + token_back_mode: Literal["epi_warps", "standalone_warps", "reuse_dispatch_warps"] = "epi_warps", + epi_flag_batch: Optional[Tuple[int, int]] = (1, 1), + flag_batch: int = 1, + combine_format: Optional[CombineFormat] = None, + act_func: str = "swiglu", + gate_up_clamp: Optional[float] = None, + dfc2_recompute: bool = False, + dfc2_col_output: bool = False, + enable_grad_y2_col_quant: bool = False, + num_ctas_grad_y2_col_quant: int = 2368, + ) -> "Sm107MegaMoEMxfp8DgluKernel": + """Build the ``(ProblemDesc, ImplDesc)`` pair from the legacy flat signature.""" + if static_expert_shape is None: + raise NotImplementedError("Sm107MegaMoEMxfp8DgluKernel requires a static_expert_shape.") + if hidden != static_expert_shape[2]: + raise ValueError(f"hidden ({hidden}) must equal static_expert_shape[2] ({static_expert_shape[2]}).") + if ab_dtype not in _AB_DTYPE_TO_QUANT_KIND: + raise ValueError(f"ab_dtype {ab_dtype} has no mxfp8 QuantKind.") + num_experts_per_rank, intermediate_gateup, _hidden = static_expert_shape + combine_format = CombineFormat.parse("bf16" if combine_format is None else str(combine_format)) + problem_desc = ProblemDesc( + { + "expert_count": world_size * num_experts_per_rank, + "intermediate_gateup_size": intermediate_gateup, + "hidden_size": hidden, + "quant_kind": str(_AB_DTYPE_TO_QUANT_KIND[ab_dtype]), + "combine_format": combine_format, + "world_size": world_size, + "local_rank": local_rank, + "topk": num_topk, + "max_tokens_per_rank": max_tokens_per_rank, + "max_recv_size_per_rank": max_recv_size_per_rank, + "gate_up_clamp": gate_up_clamp, + } + ) + impl_desc = ImplDesc( + { + "mma_tiler_mnk": tuple(mma_tiler_mnk), + "cluster_shape_mnk": tuple(cluster_shape_mnk), + "use_2cta_instrs": use_2cta_instrs, + "group_hint": group_hint, + "token_padding_block": token_padding_block, + "sf_padding_block": sf_padding_block, + "load_balance_mode": load_balance_mode, + "force_static_sched": force_static_sched, + "clc_bundle_size": clc_bundle_size, + "num_sched_stages": num_sched_stages, + "acc_dtype": acc_dtype, + "sf_vec_size": sf_vec_size, + "launch_cluster_count": launch_cluster_count, + "drop_on_overflow": drop_on_overflow, + "fc2_in_kernel_topk_reduce": fc2_in_kernel_topk_reduce, + "token_back_mode": token_back_mode, + "epi_flag_batch": tuple(epi_flag_batch) if epi_flag_batch is not None else (1, 1), + "flag_batch": flag_batch, + "act_func": act_func, + "dfc2_recompute": dfc2_recompute, + "dfc2_col_output": dfc2_col_output, + "enable_grad_y2_col_quant": enable_grad_y2_col_quant, + "num_ctas_grad_y2_col_quant": num_ctas_grad_y2_col_quant, + } + ) + return cls(problem_desc, impl_desc) + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + # -- Extract descriptors into locals matching the legacy param names. -- + world_size = problem_desc["world_size"] + local_rank = problem_desc["local_rank"] + num_topk = problem_desc["topk"] + max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + max_recv_size_per_rank = min( + problem_desc["max_recv_size_per_rank"], world_size * max_tokens_per_rank * num_topk + ) + hidden = problem_desc["hidden_size"] + combine_format = problem_desc["combine_format"] + gate_up_clamp = problem_desc["gate_up_clamp"] + _quant_kind = problem_desc["quant_kind"] + ab_dtype = _QUANT_KIND_TO_AB_DTYPE[_quant_kind] + static_expert_shape = ( + problem_desc["expert_count"] // world_size, + problem_desc["intermediate_gateup_size"], + hidden, + ) + + mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + cluster_shape_mnk = impl_desc["cluster_shape_mnk"] + use_2cta_instrs = impl_desc["use_2cta_instrs"] + group_hint = impl_desc["group_hint"] + token_padding_block = impl_desc["token_padding_block"] + sf_padding_block = impl_desc["sf_padding_block"] + load_balance_mode = impl_desc["load_balance_mode"] + force_static_sched = impl_desc["force_static_sched"] + clc_bundle_size = impl_desc["clc_bundle_size"] + num_sched_stages = impl_desc["num_sched_stages"] + acc_dtype = impl_desc["acc_dtype"] + sf_vec_size = impl_desc["sf_vec_size"] + launch_cluster_count = impl_desc["launch_cluster_count"] + drop_on_overflow = impl_desc["drop_on_overflow"] + fc2_in_kernel_topk_reduce = impl_desc["fc2_in_kernel_topk_reduce"] + token_back_mode = impl_desc["token_back_mode"] + epi_flag_batch = impl_desc["epi_flag_batch"] + flag_batch = impl_desc["flag_batch"] + act_func = impl_desc["act_func"] + dfc2_recompute = impl_desc["dfc2_recompute"] + dfc2_col_output = impl_desc["dfc2_col_output"] + self.enable_grad_y2_col_quant = impl_desc["enable_grad_y2_col_quant"] + self.num_ctas_grad_y2_col_quant = impl_desc["num_ctas_grad_y2_col_quant"] + + if hidden != static_expert_shape[2]: + raise ValueError(f"hidden ({hidden}) must equal static_expert_shape[2] ({static_expert_shape[2]}).") + token_back_by_dispatch = token_back_mode != "epi_warps" + combine_format = CombineFormat.parse("bf16" if combine_format is None else str(combine_format)) + # in-kernel topk reduce only conflicts with a QUANTIZED combine (no per-topk + # reduced-plane accumulation for quantized). It DOES work with the + # standalone/reuse_dispatch token-back modes (the token_comm token_back path + # honours token_back_reduce_topk), so those are allowed (mirrors the legacy). + if fc2_in_kernel_topk_reduce and combine_format.is_quantized: + raise ValueError("fc2_in_kernel_topk_reduce requires a non-quantized (bf16) combine.") + if token_back_mode not in ("epi_warps", "standalone_warps", "reuse_dispatch_warps"): + raise ValueError(f"unsupported token_back_mode={token_back_mode!r}.") + if ab_dtype not in _AB_DTYPE_TO_QUANT_KIND: + raise ValueError(f"ab_dtype {ab_dtype} has no mxfp8 QuantKind.") + + super().__init__( + mma_tiler_mnk=mma_tiler_mnk, + cluster_shape_mnk=cluster_shape_mnk, + use_2cta_instrs=use_2cta_instrs, + group_hint=group_hint, + token_padding_block=token_padding_block, + sf_padding_block=sf_padding_block, + load_balance_mode=load_balance_mode, + static_expert_shape=static_expert_shape, + force_static_sched=force_static_sched, + clc_bundle_size=clc_bundle_size, + num_sched_stages=num_sched_stages, + acc_dtype=acc_dtype, + ab_dtype=ab_dtype, + sf_vec_size=sf_vec_size, + epi_flag_batch=epi_flag_batch, + dfc2_recompute=dfc2_recompute, + dfc2_col_output=dfc2_col_output, + fc2_in_kernel_topk_reduce=fc2_in_kernel_topk_reduce, + act_func=act_func, + gate_up_clamp=gate_up_clamp, + ) + + # --- Warp topology (realigned for next's TokenCommDeterministic). --- + # next derives each transfer warp's transfer index as ``thread_idx % 128`` + # (token_comm.py:1421-1424 / 1911), which HARD-REQUIRES the dispatch warps -- + # and, standalone, the token_back warps -- to begin on a 4-warp / 128-thread + # boundary. So dispatch sits at warps 8-11 (thread 256 -> 256%128=0) and + # standalone token_back at 12-15 (thread 384 -> 384%128=0), mirroring the + # forward GLU. The dGLU c_load warp does NOT use the transfer index, so it + # moves ABOVE the transfer block (warp 16 iff standalone else 12), overriding + # the base kernel's default (warp 8, now occupied by dispatch). + self.enable_token_comm = True + self.dispatch_warp_id = (8, 9, 10, 11) + self.token_back_mode = token_back_mode + # Thread token_back_by_dispatch to the FC12 base (it hardcodes False) so the + # epilogue, built later in _setup_attributes(), fires the fc2_done counter for + # the standalone / reuse_dispatch token-back warps. Without this the dedicated + # token-back warps spin forever on fc2_done < target and the block-wide + # sync_threads() in kernel_tail deadlocks (M09/M10/M14/M15 hang). epi_warps is + # unaffected (it peer-writes grad_x directly and never reads fc2_done). + self.token_back_by_dispatch = token_back_by_dispatch + self.token_back_standalone = token_back_by_dispatch and token_back_mode == "standalone_warps" + self.token_back_warp_id = (12, 13, 14, 15) if self.token_back_standalone else None + num_token_back_warps = len(self.token_back_warp_id) if self.token_back_standalone else 0 + self.c_load_warp_id = 16 if self.token_back_standalone else 12 + + # Register re-balance for the mega warp layout. The base kernel sizes + # ``epi_reg_cnt`` (256) for the lean 9-warp dGLU; mega adds the 4 dispatch + # warps (+4 token-back if standalone) and the dedicated c_load warp, so the + # per-CTA register file can no longer grant 256 regs to all 4 epilogue warps + # -- the epilogue warpgroup then stalls forever inside + # ``warpgroup_reg_alloc`` and the mma/tmem barrier deadlocks. Mirror the + # legacy mega dGLU (megamoe_kernel_mxfp8_dglu.py:181-184). + self.epi_reg_cnt = 168 if self.token_back_standalone else 200 + self.threads_per_cta = 32 * ( + len(self.epilogue_warp_id) # 4 (warps 0-3) + + 1 # mma (warp 4) + + 1 # tma_a (warp 5) + + 1 # tma_b (warp 6) + + 1 # sched (warp 7) + + len(self.dispatch_warp_id) # 4 (warps 8-11) + + num_token_back_warps # 4 iff standalone_warps (warps 12-15) + + 1 # c_load (warp 12 or 16, dGLU-specific) + ) + + # --- MegaMoE constants. --- + self.world_size = world_size + self.local_rank = local_rank + self.num_topk = num_topk + self.max_tokens_per_rank = max_tokens_per_rank + self.max_recv_size_per_rank = max_recv_size_per_rank + self.hidden = hidden + self.launch_cluster_count = launch_cluster_count + self.drop_on_overflow = drop_on_overflow + self.combine_format = combine_format + self.num_experts_per_rank = static_expert_shape[0] + self.intermediate_downproj = static_expert_shape[1] + self.intermediate_gateup = self.intermediate_downproj * 2 + self.num_total_experts = world_size * self.num_experts_per_rank + self.reduce_topk_in_kernel = fc2_in_kernel_topk_reduce + self.token_back_schedule_mode = load_balance_mode if load_balance_mode == "atomic_counter" else "static" + + # --- next Router-push token communication component. --- + # dGLU dispatches raw grad_out tokens + the per-token routing prob (topk score) + # into the pool; the dfc2 epilogue folds the prob into d_gate/d_up. So the router + # ALWAYS pushes scores into the pool -> apply_topk_at_fc1=True. + mma_cta_count = 2 if use_2cta_instrs else 1 + cta_tile_m = mma_tiler_mnk[0] // mma_cta_count + cluster_m, cluster_n = self.cluster_shape_mn + tokens_per_fc1_ready_slot = cta_tile_m * cluster_m + hidden_per_fc2_cluster_tile = cta_tile_m * cluster_m + fc2_done_signals_per_token_tile = ceil_div(hidden, hidden_per_fc2_cluster_tile) * cluster_m * cluster_n + promised_launchable_sm_count = launch_cluster_count * cluster_m * cluster_n + quant_kind = _AB_DTYPE_TO_QUANT_KIND[ab_dtype] + tc_problem_desc = ProblemDesc( + { + "world_size": world_size, + "expert_count": self.num_total_experts, + "topk": num_topk, + "max_tokens_per_rank": max_tokens_per_rank, + "max_recv_size_per_rank": max_recv_size_per_rank, + "hidden_size": hidden, + "quant_kind": str(quant_kind), + "combine_format": combine_format, + "apply_topk_at_fc1": True, + } + ) + tc_impl_desc = ImplDesc( + { + "token_padding_block": token_padding_block, + "sf_padding_block": sf_padding_block, + "tokens_per_fc1_ready_slot": tokens_per_fc1_ready_slot, + "fc2_done_signals_per_token_tile": fc2_done_signals_per_token_tile, + "promised_launchable_sm_count": promised_launchable_sm_count, + "drop_on_overflow": drop_on_overflow, + "token_in_flag_batch": flag_batch, + "token_back_mode": token_back_mode, + "token_back_schedule_mode": self.token_back_schedule_mode, + "reduce_topk_in_kernel": fc2_in_kernel_topk_reduce, + } + ) + self.token_comm = TokenCommDeterministic(tc_problem_desc, tc_impl_desc) + self.pool_token_capacity = self.token_comm.worst_case_token_count + + # --- SMEM sub-buffer for the token_comm transport. --- + tc_smem_ws = SmemWorkspace() + self.token_comm.register_smem_regions(tc_smem_ws) + tc_smem_ws.finalize(max_bytes=self.smem_capacity) + self.tc_smem_ws = tc_smem_ws + self._token_comm_smem_bytes = tc_smem_ws.total_bytes + + # Build the scheduler NOW (launch_cluster_count known at construction) so its + # separately-allocated SMEM is reservable by ``_smem_misc_budget_bytes``. + _ec, _ig, _hd = static_expert_shape + self._build_scheduler( + expert_cnt=_ec, intermediate_gateup=_ig, hidden_dim=_hd, launch_cluster_count=launch_cluster_count + ) + self._sched_smem_bytes = self.sched_smem_ws.total_bytes + + # --- Post-kernel top-k reduction (skipped under in-kernel reduce). --- + self._topk_reduce = None if fc2_in_kernel_topk_reduce else TopkReduce(hidden, num_topk, combine_format) + + # --- Device workspace (next model): dGLU pools + token_comm regions. --- + self._mega_device_workspace = self._build_megamoe_device_workspace() + + # --- Bind every KernelClass schema field under its schema name. --- + self.expert_count = self.num_total_experts + self.intermediate_gateup_size = self.intermediate_downproj + self.hidden_size = hidden + self.quant_kind = _quant_kind + self.topk = num_topk + self.cluster_shape_mnk = tuple(cluster_shape_mnk) + self.mma_tiler_mnk = tuple(mma_tiler_mnk) + self.group_hint = group_hint + self.token_padding_block = token_padding_block + self.sf_padding_block = sf_padding_block + self.load_balance_mode = load_balance_mode + self.force_static_sched = force_static_sched + self.clc_bundle_size = clc_bundle_size + self.num_sched_stages = num_sched_stages + self.acc_dtype = acc_dtype + self.sf_vec_size = sf_vec_size + self.fc2_in_kernel_topk_reduce = fc2_in_kernel_topk_reduce + self.epi_flag_batch = tuple(epi_flag_batch) + self.flag_batch = flag_batch + self.act_func = act_func + self.dfc2_recompute = dfc2_recompute + self.dfc2_col_output = dfc2_col_output + self.use_2cta_instrs = use_2cta_instrs + + # Optional post-kernel token-axis MXFP8 requantization of the routed + # grad_out pool consumed as the dfc2 input. + if self.enable_grad_y2_col_quant: + col_quant_type = "mxfp8_e4m3" if ab_dtype is cutlass.Float8E4M3FN else "mxfp8_e5m2" + self.grad_y2_col_quant = Mxfp8ColRequant( + hidden=self.hidden, + num_experts=self.num_experts_per_rank, + max_total_tokens=( + self.world_size + * self.max_tokens_per_rank + * min(self.num_topk, self.num_experts_per_rank) + ), + quant_type=col_quant_type, + num_persistent_ctas=self.num_ctas_grad_y2_col_quant, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + ) + + def _smem_misc_budget_bytes(self) -> int: + """Reserve the token_comm transport + scheduler SMEM on top of the base misc budget.""" + _sched = getattr(self, "_sched_smem_bytes", 0) + return super()._smem_misc_budget_bytes() + self._token_comm_smem_bytes + _sched + self._SMEM_ALLOC_MARGIN + + def get_aux_output_shapes(self) -> dict: + """Shapes of the fixed-ABI dFC2 auxiliary outputs.""" + data_token_capacity = self.token_comm.worst_case_token_count + sf_token_capacity = self.token_comm.worst_case_sf_token_count + column_sf_row_count = sf_token_capacity // self.sf_vec_size + return { + "dprob": (self.max_tokens_per_rank, self.num_topk), + "fc1_recompute": (data_token_capacity, self.intermediate_downproj), + "fc1_recompute_sf": (column_sf_row_count, self.intermediate_downproj), + "fc1_col_output": (data_token_capacity, self.intermediate_gateup), + "fc1_col_output_sf": (column_sf_row_count, self.intermediate_gateup), + "grad_y2": (data_token_capacity, self.hidden), + "grad_y2_sf": (sf_token_capacity * (self.hidden // self.sf_vec_size),), + } + + def get_fc1_preact_shape(self) -> Tuple[int, int]: + """Shape of the externally supplied, pool-indexed gate||up pre-activations.""" + return (self.token_comm.worst_case_token_count, self.intermediate_gateup) + + @cute.jit + def _validate_fixed_pool_tensor(self, tensor: cute.Tensor, dtype, expected_shape) -> None: + if cutlass.const_expr(tensor.element_type is not dtype): + raise TypeError("pool-domain tensor has an unexpected element type.") + if cutlass.const_expr(cute.rank(tensor.layout) != 2): + raise ValueError("pool-domain tensor must be rank 2.") + if cutlass.const_expr( + not isinstance(tensor.shape[0], int) + or not isinstance(tensor.shape[1], int) + or tensor.shape[0] != expected_shape[0] + or tensor.shape[1] != expected_shape[1] + ): + raise ValueError(f"pool-domain tensor must have static shape {expected_shape}.") + if cutlass.const_expr(tensor.stride[0] != expected_shape[1] or tensor.stride[1] != 1): + raise ValueError("pool-domain tensor must be compact row-major.") + + def _build_megamoe_device_workspace(self) -> DeviceWorkspace: + """Register internal dGLU pools, counters, and token-comm regions.""" + sf_dtype = cutlass.Float8E8M0FNU + sf_vec_size = self.sf_vec_size + data_token_capacity = self.token_comm.worst_case_token_count + sf_token_capacity = self.token_comm.worst_case_sf_token_count + inter_gateup = self.intermediate_gateup # grad_y1 width (DOUBLED) + + # grad_y1 SF: row-quant, DOUBLED N columns. + sf_block_cols_back = round_up(ceil_div(inter_gateup, sf_vec_size), 4) + counter_slot_count = self.token_comm.max_fc1_ready_slot_count + + dw = DeviceWorkspace() + # grad_y1 (dfc2 output), consumed as the dfc1 (fc2) GEMM-B. + dw.register( + self.fc1_output_region, + self.ab_dtype, + (data_token_capacity, inter_gateup), + buffer_space="local", + mem_order=(1, 0), + byte_alignment=128, + ) + dw.register( + self.fc1_output_sf_region, + sf_dtype, + (sf_token_capacity, sf_block_cols_back), + buffer_space="local", + mem_order=(1, 0), + byte_alignment=128, + ) + # cross-phase fc1->fc2 done counter. + dw.register( + self.fc1_done_counter_region, + cutlass.Int32, + (counter_slot_count,), + buffer_space="local", + byte_alignment=16, + reset="tail_reset", + ) + # Dynamic load-balance atomic counter (scheduler claims work by atomic-inc). + # Always registered (cheap 1-int); the base kernel only reads it in + # atomic_counter mode, but it must be zeroed between back-to-back launches. + dw.register( + self.load_balance_counter_region, + cutlass.Int32, + (1,), + buffer_space="local", + byte_alignment=16, + reset="tail_reset", + ) + # Reserve a slot for the scheduler's atomic work-id counter (persistent-grid + # dynamic work distribution in atomic_counter mode). + dw.register( + self.sched_work_id_region, cutlass.Int32, (4,), buffer_space="local", byte_alignment=16, reset="tail_reset" + ) + # Local mirror of the shared token_src_metadata (Int64 per pool slot), filled + # host-side after the launch so the recompute/col-output validation can read it. + dw.register( + self.token_src_metadata_local_region, + cutlass.Int64, + (data_token_capacity,), + buffer_space="local", + byte_alignment=16, + ) + if self.enable_grad_y2_col_quant: + dw.register( + self.grad_y2_sizes_region, + cutlass.Int32, + (self.num_experts_per_rank,), + buffer_space="local", + byte_alignment=16, + ) + self.token_comm.register_device_workspace(dw) + dw.finalize() + return dw + + @property + def _local_offsets(self) -> dict: + """Legacy-name -> byte-offset map for inherited tester pool reads.""" + dw = self._mega_device_workspace + return { + "fc1_output": dw.offset(self.fc1_output_region), + "fc1_output_sf": dw.offset(self.fc1_output_sf_region), + "fc1_done_counter": dw.offset(self.fc1_done_counter_region), + # For the legacy dfc2_recompute / dfc2_col_output validation: + "l1_token_buffer": dw.offset(self.token_comm.fc1_activation_region), + "token_src_metadata": dw.offset(self.token_src_metadata_local_region), + } + + @property + def _shared_metadata_offset(self) -> int: + """Byte offset of the token_comm shared token_src_metadata (for the host mirror copy).""" + return self._mega_device_workspace.offset(self.token_comm._router.token_src_metadata_region) + + @property + def _local_region_by_name(self) -> dict: + """Legacy-name -> object exposing ``.nbytes`` for the inherited tester's pool reads.""" + dw = self._mega_device_workspace + name_to_region = { + "fc1_output": self.fc1_output_region, + "fc1_output_sf": self.fc1_output_sf_region, + "fc1_done_counter": self.fc1_done_counter_region, + "l1_token_buffer": self.token_comm.fc1_activation_region, + "token_src_metadata": self.token_src_metadata_local_region, + } + return {name: SimpleNamespace(nbytes=dw.nbytes(region)) for name, region in name_to_region.items()} + + def get_workspace_sizes(self) -> Tuple[int, int]: + """Return required (local, shared/symmetric) workspace bytes.""" + return self._mega_device_workspace.local_and_shared_bytes + + @property + def require_zero_workspace_leading_bytes(self) -> Tuple[int, int]: + return self._mega_device_workspace.require_zero_workspace_leading_bytes + + # ========================================================================= + # token_comm_hook_* -- filled with next's Router-push TokenCommDeterministic calls. + # ========================================================================= + + def token_comm_extra_smem_storage_class(self) -> type: + return self.tc_smem_ws.storage_class() + + def token_comm_hook_fc1_ready_counter_ptr(self, token_comm_args): + return self.token_comm.fc1_ready_counter_pointer(self._mega_device_workspace) + + def sched_ext_fc1_peek_threshold(self) -> int: # noqa: D401 + return super().sched_ext_fc1_peek_threshold() + + @cute.jit + def token_comm_hook_sched_warp_pre_init_wait(self, token_comm_args): + """The scheduler warp must wait for the Router to publish per-expert sizes.""" + self.token_comm.wait_for_sizes_ready(self._mega_device_workspace) + + @cute.jit + def token_comm_hook_fc1_tma_b_predispatch_spin(self, token_comm_args, work_tile_info): + """No-op: FC1 input readiness is enforced by the scheduler extension's fc1_ready spin.""" + pass + + @cute.jit + def token_comm_hook_dispatch_warp_body(self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx): + """Transfer warps (8-11): pull grad_out from peers into the local FC1 pool.""" + self.token_comm.token_in(self.tc_smem_ws, token_comm_storage.buffer.data_ptr()) + if cutlass.const_expr(self.token_comm.token_back_enabled and not self.token_back_standalone): + self.token_comm.token_back(self.tc_smem_ws, token_comm_storage.buffer.data_ptr()) + + @cute.jit + def token_comm_hook_token_back_warp_body(self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx): + """Standalone token-back warps (12-15): push grad_x back to source ranks.""" + self.token_comm.token_back(self.tc_smem_ws, token_comm_storage.buffer.data_ptr()) + + @cute.jit + def token_comm_hook_kernel_tail(self, token_comm_args, *, warp_idx, lane_idx, tidx): + """Cross-rank drain + workspace tail reset, performed by the transfer warps.""" + if cutlass.const_expr(self.enable_grad_y2_col_quant): + self._snapshot_grad_y2_expert_sizes(tidx) + cute.arch.sync_threads() + if (warp_idx >= self.dispatch_warp_id[0]) & (warp_idx <= self.dispatch_warp_id[-1]): + self.token_comm.reset_tail() + self.token_comm.remove_device_members() + + @cute.jit + def _snapshot_grad_y2_expert_sizes(self, tidx) -> None: + """Preserve local expert counts before token_comm tail reset.""" + from cutlass.cutlass_dsl import Int32 + + dw = self._mega_device_workspace + if self.token_comm._linear_cta_idx == Int32(0): + sizes = self.token_comm.local_expert_sizes(dw, self.token_comm._local_rank) + snapshot = dw.tensor(self.grad_y2_sizes_region) + block_dim_x, _, _ = cute.arch.block_dim() + expert_idx = tidx + while expert_idx < Int32(self.num_experts_per_rank): + snapshot[expert_idx] = Int32(sizes[expert_idx]) + expert_idx = expert_idx + block_dim_x + + # ========================================================================= + # Host launch: Router kernel -> fused MegaMoE backward kernel -> top-k reduction. + # ========================================================================= + + @cute.jit + def __call__( + self, + grad_out: cute.Tensor, # (max_tokens_per_rank, hidden) fp8 + grad_out_sf: cute.Tensor, # (max_tokens_per_rank, hidden // sf_vec_size) E8M0 + topk_idx: cute.Tensor, # (max_tokens_per_rank, num_topk) + topk_weights: cute.Tensor, # (max_tokens_per_rank, num_topk) Float32 (prob) + fc1_weight: cute.Tensor, # W2^T: (experts_per_rank, hidden, inter_downproj) + fc1_weight_sf: cute.Tensor, + fc2_weight: cute.Tensor, # W1^T: (experts_per_rank, intermediate, hidden) + fc2_weight_sf: cute.Tensor, + beta: cute.Tensor, # (experts_per_rank,) Float32 + fc1_preact: cute.Tensor, # (pool_token_capacity, intermediate_gateup) BFloat16 + output_activation: cute.Tensor, # (max_tokens_per_rank, topk, hidden) BF16 + overflow_flag: cute.Tensor, # (1,) Int32, per-rank FC12 overflow output + dprob: cute.Tensor, # (max_tokens_per_rank, topk) Float32; symmetric, pre-zeroed + fc1_recompute: cute.Tensor, # (pool_token_capacity, inter_downproj) + fc1_recompute_sf: cute.Tensor, # (col_sf_rows, inter_downproj) E8M0 + fc1_col_output: cute.Tensor, # (pool_token_capacity, intermediate_gateup) + fc1_col_output_sf: cute.Tensor, # (col_sf_rows, intermediate_gateup) E8M0 + grad_y2: cute.Tensor, # (pool_token_capacity, hidden) token-axis MXFP8 + grad_y2_sf: cute.Tensor, # flat MN-major E8M0 bytes + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, # symmetric (NVLink) heap base + peer_rank_ptr_mapper_host, + stream: cuda.CUstream, + ) -> None: + """Launch the Router, then the fused backward main kernel, then (optionally) the top-k reduce.""" + + dw = self._mega_device_workspace + local_rank = peer_rank_ptr_mapper_host.rank + aux_shapes = self.get_aux_output_shapes() + self._validate_fixed_pool_tensor( + fc1_preact, cutlass.BFloat16, self.get_fc1_preact_shape() + ) + self._validate_fixed_pool_tensor( + fc1_recompute, self.ab_dtype, aux_shapes["fc1_recompute"] + ) + self._validate_fixed_pool_tensor( + fc1_recompute_sf, + self.token_comm.activation_sf_dtype, + aux_shapes["fc1_recompute_sf"], + ) + self._validate_fixed_pool_tensor( + fc1_col_output, self.ab_dtype, aux_shapes["fc1_col_output"] + ) + self._validate_fixed_pool_tensor( + fc1_col_output_sf, + self.token_comm.activation_sf_dtype, + aux_shapes["fc1_col_output_sf"], + ) + self.token_comm.launch_router( + topk_indices=topk_idx, + topk_scores=topk_weights, + local_rank=local_rank, + local_workspace=local_workspace, + shared_workspace=shared_workspace, + peer_rank_ptr_mapper_host=peer_rank_ptr_mapper_host, + device_workspace=dw, + overflow_flag=overflow_flag, + stream=stream, + ) + peer_mapper = peer_rank_ptr_mapper_host.make_device_object() + dw.assign_device_members(local_workspace, shared_workspace) + + activation_pool = self.token_comm.fc1_activation_tensor(dw) + _sf_pool_atom = self.token_comm.fc1_activation_sf_tensor(dw) + activation_sf_pool = cute.make_tensor( + _sf_pool_atom.iterator, + cute.make_layout( + (self.token_comm.worst_case_sf_token_count, self.hidden // self.sf_vec_size), + stride=(self.token_comm.activation_sf_hidden_padded, 1), + ), + ) + fc1_output = dw.tensor(self.fc1_output_region) + fc1_output_sf = dw.tensor(self.fc1_output_sf_region) + fc1_done_counter = dw.tensor(self.fc1_done_counter_region) + load_balance_counter = ( + dw.tensor(self.load_balance_counter_region) if self.token_back_schedule_mode == "atomic_counter" else None + ) + pool_topk_scores = self.token_comm.fc1_topk_scores_tensor(dw) + + if cutlass.const_expr(self.reduce_topk_in_kernel): + # In-kernel top-k reduce (epi_warps + bf16 combine): the epilogue red-adds each + # topk grad_x contribution straight into the (pre-zeroed) 2D output. + pre_reduced = cute.make_tensor( + output_activation.iterator, + cute.make_layout( + (output_activation.shape[0], 1, output_activation.shape[1]), + stride=(output_activation.stride[0], output_activation.stride[0], output_activation.stride[1]), + ), + ) + pre_reduced_sf = None + else: + pre_reduced = self.token_comm.pre_reduced_activation_tensor(dw) + pre_reduced_sf = self.token_comm.pre_reduced_activation_sf_tensor(dw) + + if cutlass.const_expr(self.token_comm.token_back_push_data): + # token_back-by-dispatch: the epilogue writes grad_x to the LOCAL fc2_activation + # pool (same pool token_back reads), in its native (tokens, 1, hidden) shape. + fc2_output = self.token_comm.fc2_activation_tensor(dw) + else: + # epi_warps: the epilogue peer-writes grad_x directly (combine_output = pre_reduced), + _combine_hidden = pre_reduced.shape[2] + fc2_output = cute.make_tensor( + pre_reduced.iterator, + cute.make_layout( + (pre_reduced.shape[0] * pre_reduced.shape[1], _combine_hidden), stride=(_combine_hidden, 1) + ), + ) + + # dprob is a source-domain combine plane. Add a singleton value mode so + # the epilogue can reuse Fc2OutputDest's (token, topk, value) resolver. + dprob_combine = cute.make_tensor( + dprob.iterator, + cute.make_layout( + (dprob.shape[0], dprob.shape[1], 1), + stride=(dprob.stride[0], dprob.stride[1], 0), + ), + ) + + super().__call__( + activation_pool, + fc1_weight, + activation_sf_pool, + fc1_weight_sf, + fc1_output, + fc1_output_sf, + fc1_recompute, + fc1_recompute_sf, + fc1_col_output, + fc1_col_output_sf, + fc2_weight, + fc2_weight_sf, + fc2_output, + fc1_preact, + pool_topk_scores, + beta, + dprob_combine, + fc1_done_counter, + offs=None, + load_balance_counter=load_balance_counter, + max_active_clusters=self.launch_cluster_count, + stream=stream, + overflow_flag=overflow_flag, + mega_peer_rank_ptr_mapper=peer_mapper, + mega_local_rank=local_rank, + mega_local_workspace=local_workspace, + mega_shared_workspace=shared_workspace, + mega_activation=grad_out, + mega_activation_sf=grad_out_sf, + mega_pre_reduced_activation=pre_reduced, + mega_pre_reduced_activation_sf=pre_reduced_sf, + ) + + # Post-kernel top-k reduction: dequant + K-sum into the final output. + if cutlass.const_expr(not self.reduce_topk_in_kernel): + self._topk_reduce(pre_reduced, pre_reduced_sf, output_activation, None, stream) + + # Export the routed dfc2 input in token-axis MXFP8 form. The source + # grad_out pool and its row-wise SF remain resident after reset_tail. + if cutlass.const_expr(self.enable_grad_y2_col_quant): + lw = local_workspace + data_offset = dw.offset(self.token_comm.fc1_activation_region) + sf_offset = dw.offset(self.token_comm.fc1_activation_sf_region) + sizes_offset = dw.offset(self.grad_y2_sizes_region) + sf_pool_bytes = self.token_comm.worst_case_sf_token_count * (self.hidden // self.sf_vec_size) + src_data = cute.make_tensor( + cute.make_ptr( + self.ab_dtype, lw.toint() + Int64(data_offset), AddressSpace.gmem, assumed_align=128 + ), + cute.make_layout( + (self.token_comm.worst_case_token_count, self.hidden), stride=(self.hidden, 1) + ), + ) + src_sf_u8 = cute.make_tensor( + cute.make_ptr(cutlass.Uint8, lw.toint() + Int64(sf_offset), AddressSpace.gmem, assumed_align=16), + cute.make_layout((sf_pool_bytes,)), + ) + expert_sizes = cute.make_tensor( + cute.make_ptr(cutlass.Int32, lw.toint() + Int64(sizes_offset), AddressSpace.gmem, assumed_align=16), + cute.make_layout((self.num_experts_per_rank,)), + ) + self.grad_y2_col_quant( + src_data, + src_sf_u8, + expert_sizes, + grad_y2, + grad_y2_sf, + stream, + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/__init__.py new file mode 100644 index 000000000..4e9e25577 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Rubin training MegaMoE (mxfp8 GLU) kernel components.""" + +from .glu_mxfp8_fc12_epilogue import Fc2OutputDest, GluMxfp8Epilogue +from .glu_mxfp8_fc12_extension import GluMxFp8Fc12SchedExtension, TensorRole +from .glu_mxfp8_fc12_kernel import Sm107Mxfp8GluFc12Kernel +from .glu_mxfp8_mega_moe_kernel import Sm107MegaMoEMxfp8GluKernel + + +__all__ = [ + "Fc2OutputDest", + "GluMxFp8Fc12SchedExtension", + "GluMxfp8Epilogue", + "Sm107MegaMoEMxfp8GluKernel", + "Sm107Mxfp8GluFc12Kernel", + "TensorRole", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py new file mode 100644 index 000000000..bd54c92f3 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py @@ -0,0 +1,1276 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Warp-specialized, persistent, TMA MXFP8 row-scale -> column-scale requant. + +Data uses the padded row-major dispatch pool. Destination SF is concatenated +per expert in ``[hidden_atom][token_atom]`` order. +""" + +from typing import Literal + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.cpasync as cpasync +from cutlass.cutlass_dsl import ( + Float32, + Int32, + Int64, + T, + Uint8, + dsl_user_op, +) +from cutlass._mlir.dialects import arith, llvm +from cutlass.cute.typing import AddressSpace + +from ......helpers.constants import ( + Fp8E4M3FNMax, + Fp8E5M2Max, +) + +# The next/ helpers do not export these two block-size constants; they are fixed +# by the MXFP8 spec / the dispatch pool's 32x4x4 SF atom layout, so define them +# locally. +Mxfp8BlockSize = 32 +SfPaddingBlock = 128 + + +def _lcm(a: int, b: int) -> int: + from math import gcd + + return a * b // gcd(a, b) + + +def _smem_capacity() -> int: + """Max dynamic SMEM per CTA, from CUTLASS's per-arch table.""" + try: + from cutlass.utils import get_smem_capacity_in_bytes + + return int(get_smem_capacity_in_bytes()) + except Exception: + return 227 * 1024 + + +# ptxas 13.2 accepts ``scaled::n1`` only on sm_107a -- not sm_107, sm_107f, nor +# the later sm_110a / sm_120a -- so this is an exact set, not a floor. +_SCALED_CVT_ARCHS = frozenset({(10, 7)}) + + +def _target_arch_tuple() -> "tuple[int, int, str]": + """``(major, minor, suffix)`` of the active cuTeDSL compilation target.""" + from cutlass.cutlass_dsl import CuTeDSL + + arch = CuTeDSL._get_dsl().get_arch_enum() + return int(arch.major), int(arch.minor), (getattr(arch, "suffix", "") or "") + + +def _scaled_cvt_available() -> bool: + """Can this target assemble ``cvt...scaled::n1::ue8m0.e4m3x2.bf16x2``?""" + major, minor, suffix = _target_arch_tuple() + if (major, minor) not in _SCALED_CVT_ARCHS: + return False + if suffix != "a": + raise ValueError( + f"MXFP8 column requant targets sm_{major}{minor}{suffix}, but its " + f"block-scaled requant instruction " + f"'cvt.rn.satfinite.scaled::n1::ue8m0.e4m3x2.bf16x2' is accepted by " + f"ptxas only for the 'a' architecture variant; sm_{major}{minor} and " + f"sm_{major}{minor}f both fail with \"Arguments mismatch for " + f"instruction 'cvt'\". Compile for sm_{major}{minor}a, or pass " + f"scaled_cvt=False to select the portable requant path." + ) + return True + + +_SM_COUNT_CACHE: "list[int | None]" = [None] + + +def _resolve_sm_count(default: int) -> int: + """SM count of the current device.""" + if _SM_COUNT_CACHE[0] is None: + n = 0 + try: + import ctypes + + lib = ctypes.CDLL("libcuda.so.1") + lib.cuInit(0) + dev = ctypes.c_int() + if lib.cuDeviceGet(ctypes.byref(dev), 0) == 0: + val = ctypes.c_int() + # CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT + if lib.cuDeviceGetAttribute(ctypes.byref(val), 16, dev) == 0: + n = int(val.value) + except Exception: + n = 0 + _SM_COUNT_CACHE[0] = n + return _SM_COUNT_CACHE[0] or default + + +def _address_value(pointer_or_address, *, loc=None, ip=None): + if isinstance(pointer_or_address, Int64): + return pointer_or_address.ir_value() + return pointer_or_address.toint(loc=loc, ip=ip).ir_value() + + +@dsl_user_op +def tma_load_1d( + destination_smem, source_gmem, mbarrier_smem, num_bytes, *, loc=None, ip=None, +) -> None: + """Issue a 1D GMEM-to-SMEM bulk copy.""" + llvm.inline_asm( + None, + [ + destination_smem.toint(loc=loc, ip=ip).ir_value(), + _address_value(source_gmem, loc=loc, ip=ip), + num_bytes.ir_value(), + mbarrier_smem.toint(loc=loc, ip=ip).ir_value(), + ], + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1], $2, [$3];", + "r,l,r,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cp_async_bulk_s2g(destination_gmem, source_smem, num_bytes, *, loc=None, ip=None) -> None: + """Issue a 1D SMEM-to-GMEM bulk copy; the caller commits the group.""" + llvm.inline_asm( + None, + [ + _address_value(destination_gmem, loc=loc, ip=ip), + source_smem.toint(loc=loc, ip=ip).ir_value(), + num_bytes.ir_value(), + ], + "cp.async.bulk.global.shared::cta.bulk_group [$0], [$1], $2;", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _fp8x2_mnemonic(fp8_type) -> str: + if fp8_type is cutlass.Float8E4M3FN: + return "e4m3x2" + if fp8_type is cutlass.Float8E5M2: + return "e5m2x2" + raise TypeError(f"unsupported FP8 type {fp8_type}") + + +@dsl_user_op +def cvt_scaled_up_bf16x2(pair_b32, scale_b32, half: int, fp8_type, *, loc=None, ip=None) -> Int32: + """Two FP8 values + their E8M0 scale -> BF16x2, in one SASS instruction.""" + mn = _fp8x2_mnemonic(fp8_type) + asm = ( + "{\n" + " .reg .b16 a0,a1,s0,s1;\n" + " mov.b32 {a0,a1}, $1;\n" + " mov.b32 {s0,s1}, $2;\n" + f" cvt.rn.scaled::n2::ue8m0.bf16x2.{mn} $0, a{half}, s0;\n" + "}" + ) + return Int32( + llvm.inline_asm( + T.i32(), + [Int32(pair_b32).ir_value(loc=loc, ip=ip), Int32(scale_b32).ir_value(loc=loc, ip=ip)], + asm, + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def cvt_scaled_dn_fp8x2(v_bf16x2, raw_b32, fp8_type, *, loc=None, ip=None) -> Int32: + """BF16x2 + one E8M0 scale -> two FP8 bytes, in one SASS instruction.""" + mn = _fp8x2_mnemonic(fp8_type) + asm = ( + "{\n" + " .reg .b16 q, s16, junk;\n" + " .reg .b8 sb;\n" + " mov.b32 {s16, junk}, $2;\n" + " cvt.u8.u16 sb, s16;\n" + f" cvt.rn.satfinite.scaled::n1::ue8m0.{mn}.bf16x2 q, $1, sb;\n" + " cvt.u32.u16 $0, q;\n" + "}" + ) + return Int32( + llvm.inline_asm( + T.i32(), + [Int32(v_bf16x2).ir_value(loc=loc, ip=ip), Int32(raw_b32).ir_value(loc=loc, ip=ip)], + asm, + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def cvt_dn_fp8x2_portable(v_bf16x2, inv_lo, inv_hi, fp8_type, *, loc=None, ip=None) -> Int32: + """Two hidden columns of one token, as BF16x2, plus those two columns' + exact FP32 reciprocal scales -> two FP8 bytes in bits [15:0] (byte 0 from + the low BF16 half, byte 1 from the high half). + + Same rounding as ``cvt_scaled_dn_fp8x2``, not an approximation of it: BF16 + -> FP32 is an exact left shift, the reciprocal is an exact power of two, and + ``cvt.rn.satfinite`` is the RNE-and-saturate the hardware instruction + applies. Taking a scale per half is what lets the caller skip the transpose + that the one-scale-per-pair hardware instruction forces. + """ + mn = _fp8x2_mnemonic(fp8_type) + asm = ( + "{\n" + " .reg .b32 a, b;\n" + " .reg .b16 q;\n" + " shl.b32 a, $1, 16;\n" + " and.b32 b, $1, 0xffff0000;\n" + " mul.f32 a, a, $2;\n" + " mul.f32 b, b, $3;\n" + # ``cvt d, a, b`` yields d[15:8] = cvt(a) and d[7:0] = cvt(b), so the + # HIGH column has to be the first source for byte 0 to be the low one. + f" cvt.rn.satfinite.{mn}.f32 q, b, a;\n" + " cvt.u32.u16 $0, q;\n" + "}" + ) + return Int32( + llvm.inline_asm( + T.i32(), + [ + Int32(v_bf16x2).ir_value(loc=loc, ip=ip), + Float32(inv_lo).ir_value(loc=loc, ip=ip), + Float32(inv_hi).ir_value(loc=loc, ip=ip), + ], + asm, + "=r,r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def max_xorsign_abs_bf16x2(a, b, *, loc=None, ip=None) -> Int32: + """Packed magnitude max of two BF16x2.""" + return Int32( + llvm.inline_asm( + T.i32(), + [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)], + "max.xorsign.abs.bf16x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def e8m0_raw_from_bf16(bf16_bits, limit_exponent: int, *, loc=None, ip=None) -> Int32: + """E8M0 raw byte for a non-negative BF16 magnitude.""" + bits = Int32(bf16_bits) << Int32(16) + biased = (bits + Int32(0x1FFFFF - (limit_exponent << 23))) >> Int32(23) + return Int32( + arith.select( + (bits >= Int32(0x7F800000)).ir_value(loc=loc, ip=ip), + Int32(254).ir_value(loc=loc, ip=ip), + cutlass.max(Int32(0), biased).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def bits_f32(x: Int32, *, loc=None, ip=None) -> Float32: + return Float32(llvm.bitcast(T.f32(), Int32(x).ir_value(loc=loc, ip=ip), loc=loc, ip=ip)) + + +def fp8_limit_exponent(limit: float) -> int: + """Exponent ``k`` of an FP8 limit written as ``1.75 * 2**k``.""" + scaled = limit / 1.75 + exponent = int(scaled).bit_length() - 1 + if float(1 << exponent) != scaled: + raise ValueError( + f"FP8 limit {limit} is not of the form 1.75 * 2**k, which the integer E8M0 encoding assumes." + ) + return exponent + + +class Mxfp8ColRequant: + """Warp-specialized persistent TMA launcher for token-axis MXFP8 requant.""" + + # --- fixed MXFP8 / SF-atom geometry (do not retune) --------------------- + TokensPerBlock: int = Mxfp8BlockSize # 32: one E8M0 scale covers 32 values + SfAtomBytes: int = 512 + SfAtomNonK: int = 128 + SfAtomKBanks: int = 4 + + TokenPaddingBlocks: tuple = (128, 256) + + # --- this kernel's tile ------------------------------------------------ + TILE_TOK: int = 128 # == SfAtomNonK, so a tile is one whole SF atom row + NSTAGE: int = 2 # 2 is what leaves shared memory for 2 CTAs per SM + + # Tile width and per-lane access width, per requant path. Both are tuned: + # the two paths have different shared-memory budgets and different + # instruction mixes, and the pairs below are the measured optima. + TileHidScaled: int = 512 + ColsPerLaneScaled: int = 8 + TileHidPortable: int = 256 + ColsPerLanePortable: int = 4 + + # Sub-tile width of the consumer, in hidden columns. See _sp_tile_body. + # Must divide ColsPerLane and divide 32. + SP_LW: int = 4 + + # Declared upper bound on the block size. ptxas budgets + # 512//ceil(T/128) registers from it, so 576 yields 96 registers per + # thread, which is what two resident CTAs need. + MaxNTidTarget: int = 576 + + # Padding between consecutive SF atoms in shared memory, in bytes; it is + # what keeps the consumer's strided scale accesses off one bank. + SfInPad: int = 16 + SfOutPad: int = 16 + + # Grid depth, in resident waves. The curve is broad and flat above this. + GridWaves: int = 24 + + ProducerWarps: int = 1 + # ConsumerWarps is derived: + # (TILE_TOK / TokensPerBlock) * (TILE_HID / (32 * ColsPerLane)) + + SmCount: int = 148 # only a fallback; _resolve_sm_count queries the driver + # Upper bound on the grid: below this many tiles per CTA the + # O(num_experts) prologue dominates. Stated per 256 experts. + MinTilesPerCta: int = 4 + MinTilesRefExperts: int = 256 + + ConsumerBarrierId: int = 1 + + @classmethod + def _require_token_padding_block(cls, block) -> int: + """Refuse any token padding block this kernel is not built for.""" + if isinstance(block, bool) or not isinstance(block, int): + raise ValueError( + f"token_padding_block must be an int, got {block!r} " + f"({type(block).__name__})." + ) + if block not in cls.TokenPaddingBlocks: + raise ValueError( + f"token_padding_block must be one of " + f"{tuple(cls.TokenPaddingBlocks)}, got {block}. The work tile is " + f"TILE_TOK = SfAtomNonK = {cls.SfAtomNonK} tokens, i.e. one " + f"{cls.SfAtomBytes} B SF atom of {cls.SfAtomNonK}x{cls.SfAtomNonK} " + f"(token x hidden), so an expert's padded extent has to be a whole " + f"number of {cls.SfAtomNonK}-token tiles for a tile to belong to " + f"exactly one expert; {cls.TokenPaddingBlocks} are the only blocks " + f"the dispatch pool emits and the only ones validated." + ) + return int(block) + + def __init__( + self, + hidden: int, + num_experts: int, + max_total_tokens: int, + quant_type: Literal["mxfp8_e4m3", "mxfp8_e5m2"], + num_persistent_ctas: int = -1, + token_padding_block: int = SfPaddingBlock, + sf_padding_block: int = SfPaddingBlock, + *, + scaled_cvt: "bool | None" = None, + ) -> None: + """``scaled_cvt`` selects the requant path: ``None`` asks the + compilation target, ``True`` and ``False`` force the block-scaled and + the portable path so that either can be exercised on one machine. + Everything else is derived from the problem.""" + self.hidden = int(hidden) + self.num_experts = int(num_experts) + self.max_total_tokens = int(max_total_tokens) + self.quant_type = quant_type + self.token_padding_block = int(token_padding_block) + self.sf_padding_block = int(sf_padding_block) + + self._require_token_padding_block(self.token_padding_block) + if self.sf_padding_block != self.SfAtomNonK: + raise ValueError( + f"sf_padding_block must be {self.SfAtomNonK} for the 32x4x4 atom layout, " + f"got {self.sf_padding_block}." + ) + + if scaled_cvt is None: + self.scaled_cvt = _scaled_cvt_available() + elif scaled_cvt: + if not _scaled_cvt_available(): + _major, _minor, _suffix = _target_arch_tuple() + raise ValueError( + f"scaled_cvt=True forces the block-scaled cvt consumer, but " + f"the target is sm_{_major}{_minor}{_suffix} and " + f"'cvt.rn.satfinite.scaled::n1::ue8m0.e4m3x2.bf16x2' is not " + f"available there. Pass scaled_cvt=None to let the target " + f"choose." + ) + self.scaled_cvt = True + else: + self.scaled_cvt = False + + if self.hidden <= 0: + raise ValueError(f"hidden must be positive, got {self.hidden}.") + + # The tile has to be a multiple of lcm(SF atom, 32*C) and divide hidden, + # so a shape the tuned pair cannot tile falls back to a narrower access. + if self.scaled_cvt: + _cols_choices = (self.ColsPerLaneScaled, self.ColsPerLanePortable) + _preferred_tile_hid = self.TileHidScaled + else: + _cols_choices = (self.ColsPerLanePortable,) + _preferred_tile_hid = self.TileHidPortable + _picked = None + for _c in _cols_choices: + _grain = _lcm(self.SfAtomNonK, 32 * _c) + _choices = [ + t + for t in range(_grain, _preferred_tile_hid + 1, _grain) + if self.hidden % t == 0 + ] + if _choices: + _picked = (_c, _choices) + break + if _picked is None: + raise ValueError( + f"hidden={self.hidden} is not supported: it must be a multiple " + f"of {_lcm(self.SfAtomNonK, 32 * _cols_choices[-1])}." + ) + self.ColsPerLane, self._tile_hid_choices = _picked + self.TILE_HID = self._tile_hid_choices[-1] + + # Narrow shapes can end up with fewer columns per lane than the sub-tile + # width, so clamp instead of refusing; the split is then trivial. + self.sp_lw = min(self.SP_LW, self.ColsPerLane) + if self.num_experts <= 0: + raise ValueError(f"num_experts must be positive, got {self.num_experts}.") + if self.max_total_tokens <= 0: + raise ValueError(f"max_total_tokens must be positive, got {self.max_total_tokens}.") + + if quant_type == "mxfp8_e4m3": + self.quant_dtype = cutlass.Float8E4M3FN + self._data_limit_exponent = fp8_limit_exponent(float(Fp8E4M3FNMax)) + elif quant_type == "mxfp8_e5m2": + self.quant_dtype = cutlass.Float8E5M2 + self._data_limit_exponent = fp8_limit_exponent(float(Fp8E5M2Max)) + else: + raise ValueError(f"Unsupported quant_type: {quant_type!r}") + self.sf_dtype = cutlass.Float8E8M0FNU + + self.sf_in_pad = self.SfInPad + self.sf_out_pad = self.SfOutPad + self.smem_capacity = _smem_capacity() + while ( + self._smem_bytes_for(self.TILE_HID, self.NSTAGE) > self.smem_capacity + and len(self._tile_hid_choices) > 1 + ): + self._tile_hid_choices.pop() + self.TILE_HID = self._tile_hid_choices[-1] + + self.NumStages = self.NSTAGE + self.TmaBoxHidU32 = self.TILE_HID // 4 + + self._hidden_atoms = self.hidden // self.SfAtomNonK + self.HidAtomsPerTile = self.TILE_HID // self.SfAtomNonK + self.HidSegs = self.TILE_HID // (32 * self.ColsPerLane) + self.TokBlocks = self.TILE_TOK // self.TokensPerBlock # 4 + self.ConsumerWarps = self.HidSegs * self.TokBlocks + self.SfInStride = self.SfAtomBytes + self.sf_in_pad + self.SfTileBytes = self.HidAtomsPerTile * self.SfInStride + self.SfTileXferBytes = self.HidAtomsPerTile * self.SfAtomBytes + + if (self.ProducerWarps + self.ConsumerWarps) * 32 > 1024: + raise ValueError( + f"hidden={self.hidden} needs " + f"{(self.ProducerWarps + self.ConsumerWarps) * 32} threads per " + f"CTA, over the 1024-thread hardware limit." + ) + self.WarpsPerCta = self.ProducerWarps + self.ConsumerWarps + self.ThreadsPerCta = self.WarpsPerCta * 32 + # ``.maxntid`` is an upper bound, so it can never be below the launch. + self.MaxNTid = max(self.MaxNTidTarget, self.ThreadsPerCta) + + # --- SMEM ------------------------------------------------------------ + self.smem_data_bytes = self.NumStages * self.TILE_TOK * self.TILE_HID + self.smem_sf_in_bytes = self.NumStages * self.SfTileBytes + self.SfOutStride = self.SfAtomBytes + self.sf_out_pad + self.smem_sf_out_bytes = self.HidAtomsPerTile * self.SfOutStride + self.smem_table_bytes = 3 * (self.num_experts + 1) * 4 + self.smem_bytes = ( + self.smem_data_bytes + + self.smem_sf_in_bytes + + self.smem_sf_out_bytes + + self.smem_table_bytes + + 2 * self.NumStages * 8 + + 256 + ) + if self.smem_bytes > self.smem_capacity: + raise ValueError( + f"hidden={self.hidden} needs {self.smem_bytes} B of shared " + f"memory per CTA, over this target's " + f"{self.smem_capacity} B limit." + ) + + self.hidden_groups = self.hidden // self.TILE_HID + + # --- grid ------------------------------------------------------------ + # The grid quantum is RES = (resident CTAs) = CtasPerSm * SM_count: a + # grid that is not a multiple of RES leaves a fractional resident wave, + # which is a large loss. Across multiples of RES the curve is broad and + # flat, so one tuned wave count serves every production size. + self.SmCount = _resolve_sm_count(type(self).SmCount) + # Both gates: SMEM, and the 8-warps-per-scheduler cap (warps go to the + # 4 schedulers round robin, so one CTA occupies ceil(warps/4) slots). + _warp_gate = 8 // -(-self.WarpsPerCta // 4) + self.CtasPerSm = max(1, min(2, self.smem_capacity // self.smem_bytes, _warp_gate)) + self.ResidentCtas = self.CtasPerSm * self.SmCount + # Every CTA pays an O(num_experts) prefix-table prologue, so the grid has + # an upper bound, and that bound rises in proportion to the expert count: + # with too few tiles per CTA the prologue dominates the tile work. + _min_tiles = self.MinTilesPerCta * max( + 1, -(-self.num_experts // self.MinTilesRefExperts) + ) + # The wave count is tuned for 2 resident CTAs. A shape that gets only + # one keeps a single-wave grid rather than extrapolating that tuning + # point outside the regime it was taken in. + _want = self.GridWaves if self.CtasPerSm >= 2 else 1 + _max_tiles = -(-self.max_total_tokens // self.TILE_TOK) * self.hidden_groups + _waves = max(1, min(_want, _max_tiles // (_min_tiles * self.ResidentCtas))) + if num_persistent_ctas > 0: + self.num_persistent_ctas = int(num_persistent_ctas) + else: + self.num_persistent_ctas = _waves * self.ResidentCtas + self.grid = self.num_persistent_ctas + + # Reported by the runner's PASS line; this kernel has a fixed split. + self.HiddenPerCta = self.TILE_HID + self.hidden_tiles_per_work = 1 + + # Binary-search ladder over the valid-token prefix table: the powers of + # two below num_experts, largest first. + steps = [] + span = 1 + while span < self.num_experts: + span <<= 1 + span >>= 1 + while span >= 1: + steps.append(span) + span >>= 1 + self._search_steps = tuple(steps) + self._search_needs_guard = (self.num_experts & (self.num_experts - 1)) != 0 + self._experts_per_lane = (self.num_experts + 31) // 32 + + # ------------------------------------------------------------------ host + @cute.jit + def __call__( + self, + src_data: cute.Tensor, + src_sf_u8: cute.Tensor, + expert_token_sizes: cute.Tensor, + dst_data: cute.Tensor, + dst_sf_u8: cute.Tensor, + cuda_stream: cuda.CUstream, + token_padding_block: cutlass.Constexpr = None, + ) -> None: + TOKPAD = cutlass.const_expr( + self.token_padding_block if token_padding_block is None else token_padding_block + ) + self._require_token_padding_block(TOKPAD) + if cutlass.const_expr(TOKPAD != self.token_padding_block): + raise ValueError( + f"token_padding_block passed to __call__ ({TOKPAD}) disagrees with " + f"the one the Mxfp8ColRequant was constructed with " + f"({self.token_padding_block}); the launch geometry is derived from " + f"the constructor value, so construct a new instance instead." + ) + if cutlass.const_expr(src_data.element_type is not self.quant_dtype): + raise TypeError(f"src_data must use {self.quant_dtype}, got {src_data.element_type}.") + if cutlass.const_expr(dst_data.element_type is not self.quant_dtype): + raise TypeError(f"dst_data must use {self.quant_dtype}, got {dst_data.element_type}.") + + HID_U32 = cutlass.const_expr(self.hidden // 4) + BOX_H = cutlass.const_expr(self.TmaBoxHidU32) + BOX_T = cutlass.const_expr(self.TILE_TOK) + src_u32 = cute.make_tensor( + cute.recast_ptr(src_data.iterator, dtype=cutlass.Uint32), + cute.make_layout((src_data.shape[0], HID_U32), stride=(HID_U32, 1)), + ) + tma_atom, tma_tensor = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + src_u32, + cute.make_layout((BOX_T, BOX_H), stride=(BOX_H, 1)), + (BOX_T, BOX_H), + ) + + dst_u32 = cute.make_tensor( + cute.recast_ptr(dst_data.iterator, dtype=cutlass.Uint32), + cute.make_layout((dst_data.shape[0], HID_U32), stride=(HID_U32, 1)), + ) + tma_atom_st, tma_tensor_st = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + dst_u32, + cute.make_layout((BOX_T, BOX_H), stride=(BOX_H, 1)), + (BOX_T, BOX_H), + ) + k = self.ws_kernel( + src_data, src_sf_u8, expert_token_sizes, dst_data, dst_sf_u8, + tma_atom, tma_tensor, tma_atom_st, tma_tensor_st, TOKPAD, + ) + # Validated in __init__, not here: __call__ is DSL-preprocessed and a + # plain if/raise in a traced body is rejected at trace time. + _ntid = cutlass.const_expr(self.MaxNTid) + k.launch( + grid=[self.grid, 1, 1], + block=[self.ThreadsPerCta, 1, 1], + max_number_threads=[_ntid, 1, 1], + stream=cuda_stream, + ) + + # ------------------------------------------------------- prefix tables + @cute.jit + def _pad_up(self, count, block: int): + if cutlass.const_expr(block > 0 and (block & (block - 1)) == 0): + return (count + Int32(block - 1)) & Int32(-block) + return ((count + Int32(block - 1)) // Int32(block)) * Int32(block) + + @cute.jit + def warp_prefix_sum(self, value, lane_idx): + acc = Int32(value) + for shift in cutlass.range_constexpr(0, 5, 1): + step = 1 << shift + other = Int32(cute.arch.shuffle_sync_up(acc, Int32(step), mask_and_clamp=0)) + if lane_idx >= Int32(step): + acc = acc + other + return acc + + @cute.jit + def build_prefix_tables( + self, expert_token_sizes, tbl_vend, tbl_data, tbl_sf, lane_idx, + token_padding_block: cutlass.Constexpr = None, + ): + """Exclusive prefixes of the padded-data / padded-SF row counts.""" + E = cutlass.const_expr(self.num_experts) + EPL = cutlass.const_expr(self._experts_per_lane) + TOKPAD = cutlass.const_expr( + self.token_padding_block if token_padding_block is None else token_padding_block + ) + self._require_token_padding_block(TOKPAD) + + counts = cute.make_rmem_tensor((EPL,), cutlass.Int32) + local_data = Int32(0) + local_sf = Int32(0) + for slot in cutlass.range_constexpr(0, EPL, 1): + expert = lane_idx * Int32(EPL) + Int32(slot) + count = Int32(0) + if expert < Int32(E): + count = Int32(expert_token_sizes[expert]) + counts[slot] = count + local_data = local_data + self._pad_up(count, TOKPAD) + local_sf = local_sf + self._pad_up(count, self.sf_padding_block) + + base_data = self.warp_prefix_sum(local_data, lane_idx) - local_data + base_sf = self.warp_prefix_sum(local_sf, lane_idx) - local_sf + + for slot in cutlass.range_constexpr(0, EPL, 1): + expert = lane_idx * Int32(EPL) + Int32(slot) + count = Int32(counts[slot]) + if expert <= Int32(E): + tbl_vend[expert] = base_data + count + tbl_data[expert] = base_data + tbl_sf[expert] = base_sf + base_data = base_data + self._pad_up(count, TOKPAD) + base_sf = base_sf + self._pad_up(count, self.sf_padding_block) + + # When 32 * EPL == E no lane owns index E, so it is written separately. + if (lane_idx + Int32(1)) * Int32(EPL) == Int32(E): + tbl_vend[E] = base_data + tbl_data[E] = base_data + tbl_sf[E] = base_sf + + @cute.jit + def find_expert(self, tbl, key): + E = cutlass.const_expr(self.num_experts) + lo = Int32(0) + for step in self._search_steps: + probe = lo + Int32(step) + if cutlass.const_expr(self._search_needs_guard): + if probe < Int32(E) and Int32(tbl[probe]) <= key: + lo = probe + else: + if Int32(tbl[probe]) <= key: + lo = probe + return lo + + # ---------------------------------------------------------------- kernel + @cute.kernel + def ws_kernel( + self, + src_data: cute.Tensor, + src_sf_u8: cute.Tensor, + expert_token_sizes: cute.Tensor, + dst_data: cute.Tensor, + dst_sf_u8: cute.Tensor, + tma_atom=None, + tma_tensor=None, + tma_atom_st=None, + tma_tensor_st=None, + token_padding_block: cutlass.Constexpr = None, + ) -> None: + # Compile-time constant, folded before a single instruction is emitted. + TOKPAD = cutlass.const_expr( + self.token_padding_block if token_padding_block is None else token_padding_block + ) + self._require_token_padding_block(TOKPAD) + TOK = cutlass.const_expr(self.TILE_TOK) + W = cutlass.const_expr(self.TILE_HID) + S = cutlass.const_expr(self.NumStages) + SFB = cutlass.const_expr(self.SfTileBytes) + NCONS = cutlass.const_expr(self.ConsumerWarps) + table_len = cutlass.const_expr(self.num_experts + 1) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + grid_dim_x, _, _ = cute.arch.grid_dim() + warp_idx = tidx // Int32(32) + lane_idx = tidx % Int32(32) + + smem = cutlass.utils.SmemAllocator() + mbar_full = smem.allocate_array(cutlass.Int64, S) + mbar_empty = smem.allocate_array(cutlass.Int64, S) + tbl_vend = smem.allocate_tensor(cutlass.Int32, cute.make_layout((table_len,)), 16) + tbl_data = smem.allocate_tensor(cutlass.Int32, cute.make_layout((table_len,)), 16) + tbl_sf = smem.allocate_tensor(cutlass.Int32, cute.make_layout((table_len,)), 16) + smem_sf_in = smem.allocate_array(self.sf_dtype, S * SFB, byte_alignment=128) + smem_sf_out = smem.allocate_array( + self.sf_dtype, cutlass.const_expr(self.smem_sf_out_bytes), byte_alignment=128 + ) + smem_data = smem.allocate_array( + self.quant_dtype, self.NumStages * TOK * W, byte_alignment=128 + ) + + if tidx == Int32(0): + for s in cutlass.range_constexpr(0, S, 1): + cute.arch.mbarrier_init(mbar_full + s, 1) + cute.arch.mbarrier_init(mbar_empty + s, NCONS) + cute.arch.mbarrier_init_fence() + + if warp_idx == Int32(0): + self.build_prefix_tables( + expert_token_sizes, tbl_vend, tbl_data, tbl_sf, lane_idx, TOKPAD + ) + cute.arch.sync_threads() + + total_tiles = Int32(tbl_data[self.num_experts]) // Int32(TOK) + + smem_data_base = smem_data.toint() + smem_sf_in_base = smem_sf_in.toint() + smem_sf_out_base = smem_sf_out.toint() + src_sf_base = src_sf_u8.iterator.toint() + dst_sf_base = dst_sf_u8.iterator.toint() + + if warp_idx < Int32(self.ProducerWarps): + self.produce( + smem_data_base, smem_sf_in_base, mbar_full, mbar_empty, + tbl_data, tbl_sf, src_sf_base, + bidx, grid_dim_x, total_tiles, lane_idx, + tma_atom, tma_tensor, TOKPAD, + ) + else: + self.consume_scaled( + smem_data_base, smem_sf_in_base, smem_sf_out_base, mbar_full, mbar_empty, + tbl_vend, tbl_data, tbl_sf, dst_sf_base, + bidx, grid_dim_x, total_tiles, + warp_idx - Int32(self.ProducerWarps), lane_idx, + tma_atom_st, tma_tensor_st, TOKPAD, + ) + + # -------------------------------------------------------------- producer + @cute.jit + def produce( + self, smem_data_base, smem_sf_in_base, mbar_full, mbar_empty, + tbl_data, tbl_sf, src_sf_base, + bidx, grid_dim_x, total_tiles, lane_idx, + tma_atom=None, tma_tensor=None, + token_padding_block: cutlass.Constexpr = None, + ): + TOKPAD = cutlass.const_expr( + self.token_padding_block if token_padding_block is None else token_padding_block + ) + self._require_token_padding_block(TOKPAD) + TOK = cutlass.const_expr(self.TILE_TOK) + W = cutlass.const_expr(self.TILE_HID) + S = cutlass.const_expr(self.NumStages) + SFB = cutlass.const_expr(self.SfTileBytes) + # data bytes + SF bytes, both fenced by bar_full[stage] + EXPECT = cutlass.const_expr(TOK * W + self.SfTileXferBytes) + SF_PREFIX_DIFFERS = cutlass.const_expr(TOKPAD != self.sf_padding_block) + + BOX_H = cutlass.const_expr(self.TmaBoxHidU32) + sD = cute.make_tensor( + cute.make_ptr( + cutlass.Uint32, smem_data_base, AddressSpace.smem, assumed_align=128, + ), + cute.make_layout((TOK, BOX_H, S), stride=(BOX_H, 1, TOK * BOX_H)), + ) + gD = cute.group_modes( + cute.local_tile(tma_tensor, (TOK, BOX_H), (None, None)), 0, 2 + ) + tDsD, tDgD = cpasync.tma_partition( + tma_atom, 0, cute.make_layout(1), cute.group_modes(sD, 0, 2), gD, + ) + cpasync.prefetch_descriptor(tma_atom) + + t = Int32(0) + work_idx = Int32(bidx) + total_work = total_tiles * Int32(self.hidden_groups) + while work_idx < total_work: + stage = t % Int32(S) + token_tile = work_idx // Int32(self.hidden_groups) + hid_begin = (work_idx % Int32(self.hidden_groups)) * Int32(W) + + data_row0 = token_tile * Int32(TOK) + sf_row0 = data_row0 + if cutlass.const_expr(SF_PREFIX_DIFFERS): + owner = self.find_expert(tbl_data, data_row0) + sf_row0 = Int32(tbl_sf[owner]) + (data_row0 - Int32(tbl_data[owner])) + sf_row0 = cutlass.min( + sf_row0, Int32(tbl_sf[owner + Int32(1)]) - Int32(self.SfAtomNonK) + ) + + if t >= Int32(S): + cute.arch.mbarrier_wait(mbar_empty + stage, ((t // Int32(S)) - Int32(1)) % Int32(2)) + + if lane_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(mbar_full + stage, Int32(EXPECT)) + cute.arch.sync_warp() + + if lane_idx == Int32(0): + sf_src = ( + Int64(src_sf_base) + + Int64(sf_row0 // Int32(self.SfAtomNonK)) * Int64(self._hidden_atoms * self.SfAtomBytes) + + Int64(hid_begin // Int32(self.SfAtomNonK)) * Int64(self.SfAtomBytes) + ) + # One copy per SF atom: the atoms are padded apart in shared + # memory, so they are not one contiguous run. + for a in cutlass.range_constexpr(0, cutlass.const_expr(self.HidAtomsPerTile), 1): + tma_load_1d( + cute.make_ptr( + self.sf_dtype, + smem_sf_in_base + stage * Int32(SFB) + + Int32(a * self.SfInStride), + AddressSpace.smem, assumed_align=16, + ), + sf_src + Int64(a * self.SfAtomBytes), + mbar_full + stage, + Int32(self.SfAtomBytes), + ) + + cute.copy( + tma_atom, + tDgD[(None, token_tile, hid_begin // Int32(W))], + tDsD[(None, stage)], + tma_bar_ptr=mbar_full + stage, + ) + + t = t + Int32(1) + work_idx = work_idx + grid_dim_x + + # ------------------------------------------------- consumer (scaled cvt) + @cute.jit + def consume_scaled( + self, smem_data_base, smem_sf_in_base, smem_sf_out_base, mbar_full, mbar_empty, + tbl_vend, tbl_data, tbl_sf, dst_sf_base, + bidx, grid_dim_x, total_tiles, cw, lane_idx, + tma_atom_st=None, tma_tensor_st=None, + token_padding_block: cutlass.Constexpr = None, + ): + """The single-pass all-BF16 consumer, shared by every target.""" + TOKPAD = cutlass.const_expr( + self.token_padding_block if token_padding_block is None else token_padding_block + ) + self._require_token_padding_block(TOKPAD) + TOK = cutlass.const_expr(self.TILE_TOK) + W = cutlass.const_expr(self.TILE_HID) + S = cutlass.const_expr(self.NumStages) + SFB = cutlass.const_expr(self.SfTileBytes) + NB = cutlass.const_expr(self.TokensPerBlock) + CONS_THREADS = cutlass.const_expr(self.ConsumerWarps * 32) + HATOMS = cutlass.const_expr(self.HidAtomsPerTile) + SF_PREFIX_DIFFERS = cutlass.const_expr(TOKPAD != self.sf_padding_block) + SFOUT_ONE = cutlass.const_expr(self.HidAtomsPerTile * self.SfOutStride) + C = cutlass.const_expr(self.ColsPerLane) # hidden columns per lane + SEGW = cutlass.const_expr(32 * C) # columns per consumer segment + + tb = cw // Int32(self.HidSegs) + seg = cw % Int32(self.HidSegs) + + LW = cutlass.const_expr(self.sp_lw) + ldsw = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=LW * 8 + ) + data_lane_off = tb * Int32(NB * W) + seg * Int32(SEGW) + lane_idx * Int32(LW) + + hb0 = seg * Int32(C) + (lane_idx * Int32(LW)) // Int32(32) + sf_lane_off = tb * Int32(4) + + BOX_HS = cutlass.const_expr(self.TmaBoxHidU32) + sDo = cute.make_tensor( + cute.make_ptr( + cutlass.Uint32, smem_data_base, AddressSpace.smem, assumed_align=128, + ), + cute.make_layout((TOK, BOX_HS, S), stride=(BOX_HS, 1, TOK * BOX_HS)), + ) + gDo = cute.group_modes( + cute.local_tile(tma_tensor_st, (TOK, BOX_HS), (None, None)), 0, 2 + ) + tDsDo, tDgDo = cpasync.tma_partition( + tma_atom_st, 0, cute.make_layout(1), cute.group_modes(sDo, 0, 2), gDo, + ) + cpasync.prefetch_descriptor(tma_atom_st) + + t = Int32(0) + work_idx = Int32(bidx) + total_work = total_tiles * Int32(self.hidden_groups) + while work_idx < total_work: + stage = t % Int32(S) + token_tile = work_idx // Int32(self.hidden_groups) + hid_begin = (work_idx % Int32(self.hidden_groups)) * Int32(W) + + data_row0 = token_tile * Int32(TOK) + owner = self.find_expert(tbl_data, data_row0) + valid_rows = cutlass.min(Int32(TOK), Int32(tbl_vend[owner]) - data_row0) + + # Destination SF is concatenated by expert, with token atoms + # contiguous inside each hidden atom. + sf_expert_token_atom = Int32(tbl_sf[owner]) // Int32(TOK) + sf_token_atom = ( + data_row0 - Int32(tbl_data[owner]) + ) // Int32(TOK) + sf_token_atoms = ( + Int32(tbl_sf[owner + Int32(1)]) - Int32(tbl_sf[owner]) + ) // Int32(TOK) + sf_live = Int32(1) + if cutlass.const_expr(SF_PREFIX_DIFFERS): + sf_live = cutlass.min(Int32(1), cutlass.max(Int32(0), valid_rows)) + + cute.arch.mbarrier_wait(mbar_full + stage, (t // Int32(S)) % Int32(2)) + + stage_data = smem_data_base + stage * Int32(TOK * W) + data_lane_off + stage_sf = smem_sf_in_base + stage * Int32(SFB) + sf_lane_off + # A tile's sf_out store is drained before the stage is handed + # back, so one buffer is enough. + sfout = smem_sf_out_base + + self._sp_tile_body( + stage_data, stage_sf, sfout, ldsw, hb0, seg, tb, lane_idx, valid_rows, + ) + + # Cross-proxy ordering, and it is NOT optional. The consumer warps + # wrote this tile's SMEM through the generic proxy; the store below + # reads it through the async proxy. Without this fence the store may + # observe stale bytes, and the kernel becomes nondeterministic: the + # same input yields different outputs from run to run. + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.barrier(barrier_id=self.ConsumerBarrierId, number_of_threads=CONS_THREADS) + + # The whole tile goes out, padding rows included; those were + # neutralised to zero in shared memory, which is what the pool + # expects to find there. + if cw == Int32(0): + cute.copy( + tma_atom_st, + tDsDo[(None, stage)], + tDgDo[(None, token_tile, hid_begin // Int32(W))], + ) + if cw == Int32(0) and lane_idx >= Int32(16) and lane_idx < Int32(16) + Int32(HATOMS) * sf_live: + atom = lane_idx - Int32(16) + cp_async_bulk_s2g( + Int64(dst_sf_base) + + ( + Int64(sf_expert_token_atom) * Int64(self._hidden_atoms) + + ( + Int64(hid_begin // Int32(self.SfAtomNonK)) + + Int64(atom) + ) + * Int64(sf_token_atoms) + + Int64(sf_token_atom) + ) + * Int64(self.SfAtomBytes), + cute.make_ptr( + self.sf_dtype, + sfout + atom * Int32(self.SfOutStride), + AddressSpace.smem, + assumed_align=16, + ), + Int32(self.SfAtomBytes), + ) + cute.arch.cp_async_bulk_commit_group() + # Drain every store before the stage goes back to the producer. + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.barrier(barrier_id=self.ConsumerBarrierId, number_of_threads=CONS_THREADS) + if lane_idx == Int32(0): + cute.arch.mbarrier_arrive(mbar_empty + stage) + + t = t + Int32(1) + work_idx = work_idx + grid_dim_x + + @cute.jit + def _sp_tile_body( + self, stage_data, stage_sf_base, sfout, ldsw, hb0, seg, tb, lane_idx, valid_rows, + ): + """The single-pass arithmetic for one lane's share of one tile.""" + TOK = cutlass.const_expr(self.TILE_TOK) + W = cutlass.const_expr(self.TILE_HID) + NB = cutlass.const_expr(self.TokensPerBlock) + C = cutlass.const_expr(self.ColsPerLane) + LW = cutlass.const_expr(self.sp_lw) + NWc = cutlass.const_expr(LW // 4) # 4-byte words per sub-tile row + NPc = cutlass.const_expr(LW // 2) # BF16x2 registers per sub-tile row + NCH = cutlass.const_expr(C // LW) # sub-tiles per lane per tile + SEGW = cutlass.const_expr(32 * C) # columns per consumer segment + QT = self.quant_dtype + + # Dead rows ARE reachable: token_padding_block constrains an expert's + # PADDED extent, not its valid count, so counts like 127,127,127 leave a + # padded tail in every expert's last tile. Neutralising them here needs + # no barrier: a consumer thread only reads the bytes it just wrote. + zeros = cute.make_rmem_tensor((NWc,), cutlass.Int32) + for w in cutlass.range_constexpr(0, NWc, 1): + zeros[w] = Int32(0) + + for ch in cutlass.range_constexpr(0, NCH, 1): + base = stage_data + Int32(ch * 32 * LW) + hbc = hb0 + Int32(ch * LW) + sfb = ( + stage_sf_base + + (hbc // Int32(4)) * Int32(self.SfInStride) + + hbc % Int32(4) + ) + + # The dead row's SCALE has to be neutralised as well as its data. + # ``cvt.rn.scaled::n2::ue8m0.bf16x2`` turns a NaN scale (raw 0xFF) + # into a NaN BF16 even from a zero payload. The amax survives -- + # ``max.xorsign.abs`` returns the non-NaN operand -- so the payload + # stays right and only the PADDING row is written back as 0x7F + # instead of 0x00. The fc1 pool really does leave 0xFF in padding + # scale bytes. Writing 127 is the same value the masked arm + # substitutes, paid once per tile instead of once per read. + if valid_rows < Int32(TOK): + for tt in cutlass.range_constexpr(0, NB, 1): + if tb * Int32(NB) + Int32(tt) >= valid_rows: + self._store_words( + base + Int32(tt * W), ldsw, zeros, NWc, LW + ) + sf_t = cute.make_tensor( + cute.make_ptr( + cutlass.Uint8, sfb + Int32(tt * 16), + AddressSpace.smem, assumed_align=1, + ), + cute.make_layout((1,)), + ) + sf_t[0] = Uint8(127) + + # ---- the one scan: unpack-with-scale, keep BF16, accumulate amax -- + d = [[None] * NPc for _ in range(NB)] + acc = [Int32(0)] * NPc + for tt in cutlass.range_constexpr(0, NB, 1): + # Both loads are issued before either is consumed, so the scale + # load overlaps the data load. + words = self._load_words(base + Int32(tt * W), ldsw, NWc, LW) + raw_sf = self._src_scale_raw(sfb + Int32(tt * 16)) + s16 = raw_sf | (raw_sf << Int32(8)) + for w in cutlass.range_constexpr(0, NWc, 1): + qw = Int32(words[w]) + lo = cvt_scaled_up_bf16x2(qw, s16, 0, QT) + hi = cvt_scaled_up_bf16x2(qw, s16, 1, QT) + # Kept LIVE across the amax -- this is the single pass. + d[tt][2 * w] = lo + d[tt][2 * w + 1] = hi + acc[2 * w] = max_xorsign_abs_bf16x2(acc[2 * w], lo) + acc[2 * w + 1] = max_xorsign_abs_bf16x2(acc[2 * w + 1], hi) + + # max.xorsign.abs leaves junk in every sign bit. + raws = [None] * LW + for k in cutlass.range_constexpr(0, NPc, 1): + a = acc[k] & Int32(0x7FFF7FFF) + raws[2 * k] = e8m0_raw_from_bf16(a & Int32(0xFFFF), self._data_limit_exponent) + raws[2 * k + 1] = e8m0_raw_from_bf16( + (a >> Int32(16)) & Int32(0xFFFF), self._data_limit_exponent + ) + scs = raws + if cutlass.const_expr(self.scaled_cvt): + invs = None + else: + # The portable down-convert scales by an exact FP32 reciprocal + # instead of handing an E8M0 byte to the hardware. + invs = [ + bits_f32( + cutlass.max( + (Int32(254) - raws[j]) << Int32(23), Int32(0x400000) + ) + ) + for j in range(LW) + ] + + col0 = seg * Int32(SEGW) + Int32(ch * 32 * LW) + lane_idx * Int32(LW) + for j in cutlass.range_constexpr(0, LW, 1): + col = col0 + Int32(j) + off = ( + sfout + + (col // Int32(128)) * Int32(self.SfOutStride) + + (col % Int32(32)) * Int32(16) + + ((col % Int32(128)) // Int32(32)) * Int32(4) + + tb + ) + out_t = cute.make_tensor( + cute.make_ptr(cutlass.Uint8, off, AddressSpace.smem, assumed_align=1), + cute.make_layout((1,)), + ) + out_t[0] = Uint8(raws[j]) + + for tt in cutlass.range_constexpr(0, NB, 2): + out0, out1 = self._requant_token_pair( + d[tt], d[tt + 1], scs, invs, NPc, NWc + ) + self._store_words(base + Int32(tt * W), ldsw, out0, NWc, LW) + self._store_words(base + Int32((tt + 1) * W), ldsw, out1, NWc, LW) + + def _requant_token_pair(self, d0, d1, scs, invs, NPc, NWc): + """Requantise one lane's two consecutive tokens. THE ONLY PLACE THIS + KERNEL DEPENDS ON THE TARGET ARCHITECTURE. + + ``d0``/``d1`` are lists of NPc Int32, each a token-major BF16x2: register + ``k`` is hidden columns ``(2k, 2k+1)`` of one token. ``out0``/``out1`` + come back as NWc packed FP8 words per token, in the byte order + ``_load_words`` read, so ``_store_words`` can put them straight back. + ``scs[j]`` is column ``j``'s output E8M0 raw byte and ``invs[j]`` its + exact FP32 reciprocal; each arm reads only the one it needs. + + Deliberately NOT ``@cute.jit``: it has to inline into the caller's trace. + A jitted callee would fail to marshal the Python lists and would change + the emitted IR. + + The transpose belongs inside this function: it is not a layout + preference but a consequence of ``scaled::n1`` taking a single scale for + both halves, which forces the two elements to share a hidden column + while SMEM is token-major. + """ + QT = self.quant_dtype + if cutlass.const_expr(self.scaled_cvt): + # Allocate the rmem tensors AFTER the cvt chain: the allocation + # order reaches the IR, so moving them changes the emitted code. + o = [None] * (2 * NPc) + # token-major -> column-major: o[j] = (col j of t0, t1) + for k in range(0, NPc, 1): + lo = Int32(cute.arch.prmt(d0[k], d1[k], Int32(0x5410))) + hi = Int32(cute.arch.prmt(d0[k], d1[k], Int32(0x7632))) + o[2 * k] = cvt_scaled_dn_fp8x2(lo, scs[2 * k], QT) + o[2 * k + 1] = cvt_scaled_dn_fp8x2(hi, scs[2 * k + 1], QT) + + # column-major -> token-major, one 4-byte word per token/w + out0 = cute.make_rmem_tensor((NWc,), cutlass.Int32) + out1 = cute.make_rmem_tensor((NWc,), cutlass.Int32) + for w in range(0, NWc, 1): + b = 4 * w + m01 = (o[b] & Int32(0xFFFF)) | (o[b + 1] << Int32(16)) + m23 = (o[b + 2] & Int32(0xFFFF)) | (o[b + 3] << Int32(16)) + out0[w] = Int32(cute.arch.prmt(m01, m23, Int32(0x6420))) + out1[w] = Int32(cute.arch.prmt(m01, m23, Int32(0x7531))) + else: + # Word w is columns 4w..4w+3; d[2w] is (4w, 4w+1) and d[2w+1] is + # (4w+2, 4w+3), and each helper returns byte 0 = its low column, so + # the OR below reproduces the byte order _load_words saw. + out0 = cute.make_rmem_tensor((NWc,), cutlass.Int32) + out1 = cute.make_rmem_tensor((NWc,), cutlass.Int32) + for w in range(0, NWc, 1): + k0 = 2 * w + k1 = 2 * w + 1 + p00 = cvt_dn_fp8x2_portable(d0[k0], invs[2 * k0], invs[2 * k0 + 1], QT) + p01 = cvt_dn_fp8x2_portable(d0[k1], invs[2 * k1], invs[2 * k1 + 1], QT) + p10 = cvt_dn_fp8x2_portable(d1[k0], invs[2 * k0], invs[2 * k0 + 1], QT) + p11 = cvt_dn_fp8x2_portable(d1[k1], invs[2 * k1], invs[2 * k1 + 1], QT) + out0[w] = (p00 & Int32(0xFFFF)) | (p01 << Int32(16)) + out1[w] = (p10 & Int32(0xFFFF)) | (p11 << Int32(16)) + return out0, out1 + + @cute.jit + def _load_words(self, byte_addr, atom, NW, C): + """One LDS of C raw bytes (no FP8 decode -- the scaled cvt does that).""" + regs = cute.make_rmem_tensor((NW,), cutlass.Int32) + cute.copy( + atom, + cute.make_tensor( + cute.make_ptr(cutlass.Int32, byte_addr, AddressSpace.smem, assumed_align=C), + cute.make_layout((NW,)), + ), + regs, + ) + return regs + + @cute.jit + def _store_words(self, byte_addr, atom, regs, NW, C): + cute.copy( + atom, + regs, + cute.make_tensor( + cute.make_ptr(cutlass.Int32, byte_addr, AddressSpace.smem, assumed_align=C), + cute.make_layout((NW,)), + ), + ) + + def _smem_bytes_for(self, tile_hid: int, stages: int) -> int: + """SMEM a (tile_hid, stages) pair would need, in bytes.""" + hid_atoms = tile_hid // self.SfAtomNonK + return ( + stages * self.TILE_TOK * tile_hid + + stages * hid_atoms * (self.SfAtomBytes + self.sf_in_pad) + + hid_atoms * (self.SfAtomBytes + self.sf_out_pad) + + 3 * (self.num_experts + 1) * 4 + + 2 * stages * 8 + + 256 + ) + + @cute.jit + def _src_scale_raw(self, byte_addr): + """Raw source E8M0 byte, forced to 0..255. + + The mask is load-bearing: the pointer type says unsigned, but the + emitted load sign-extends, and the caller packs this into both halves + of a word with `raw | (raw << 8)`. Without the mask a scale byte >= + 0x80 poisons the E8M0 pair. Removing it took the mega suite from 28/28 + to 0/28. + """ + return ( + Int32( + cute.make_tensor( + cute.make_ptr(cutlass.Uint8, byte_addr, AddressSpace.smem, assumed_align=1), + cute.make_layout((1,)), + )[0] + ) + & Int32(0xFF) + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py new file mode 100644 index 000000000..552965c7b --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py @@ -0,0 +1,1647 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Autonomous epilogue for the fused fc1+fc2 swap-AB MegaMoE kernel. + +Component boundaries use ``TensorWithContract`` to keep per-thread RMEM layout +semantics explicit. See ``megamoe_design.md`` for the epilogue dataflow. +""" + +from typing import Optional, Tuple, Type, Union, Any +import dataclasses + +import cutlass +import cutlass.cute as cute + +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import AddressSpace +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils + +from cutlass._mlir import ir +from cutlass._mlir.dialects import arith as _arith +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import dsl_user_op, Int32 as _epi_Int32, Int64 + +from cutlass.cute.typing import Float32 + +from ......helpers.iket_compat import iket +from ......helpers.flag_batch import GpuReleaseFlagBatchTracker +from ......helpers.ptx_helpers import ( + cp_async_bulk_s2g, + red_add_relaxed_sys_v2_bf16x2 as _red_add_relaxed_sys_v2_bf16x2, + stg_e8m0_from_f32, + stg_e8m0x8_from_f32, +) +from ..helpers.utils import quant_sfd_row, swiglu_act +from ......quant_def import CombineFormat +from ......communication.token_protocol import TokenSrcMetadata +from .....schedulers import BlockPhase + +# The 16x32 TMEM transpose core is architecture-neutral; reuse it via the local +# source-copy shim rather than re-porting the transpose math or reaching into the +# inference deliverable's directory. +from ..tmem_transpose import _TmemTranspose16x32Core + +Fc1GateUpInterleave = 32 +EpilogueTileN = 32 +Fc1EpilogueOutputTileM = 256 +Fc1EpilogueOutputTileN = 128 +WarpThreadCount = 32 +EpiWarpCount = 4 +Fc1CTMAStages = 1 + + +# ============================================================================= +# Fc2OutputDest (MegaMoE fc2 STG destination resolver, non-swap MXFP8 path) +# ============================================================================= + + +@dataclasses.dataclass(frozen=True) +class Fc2OutputDest: + """fc2 output destination in MoE-domain ``(token_max, topk, hidden)`` layout.""" + + tensor: cute.Tensor + metadata: Optional[cute.Tensor] = None + peer_rank_ptr_mapper: Any = None + reduce_topk_in_kernel: bool = False + + def __post_init__(self) -> None: + if (self.metadata is None) != (self.peer_rank_ptr_mapper is None): + raise ValueError( + "Fc2OutputDest: ``metadata`` and ``peer_rank_ptr_mapper`` must be " + "both None (direct mode) or both non-None (MegaMoE / indirect " + "mode). Got metadata=" + f"{'set' if self.metadata is not None else 'None'}, " + f"peer_rank_ptr_mapper=" + f"{'set' if self.peer_rank_ptr_mapper is not None else 'None'}." + ) + + @cute.jit + def resolve_token_row(self, pool_token_global) -> cute.Tensor: + """Return the ``(hidden,)`` BF16 GMEM row this pool token's STG lands on.""" + if cutlass.const_expr(self.metadata is None): + # Int64 token coord: the (token, topk, hidden) row offset is token*(topk*hidden), + # which overflows int32 once max_tokens*topk*hidden > 2^31 (mirrors inference). + return cute.slice_(self.tensor, (Int64(pool_token_global), 0, None)) + + md = TokenSrcMetadata.load( + self.metadata.iterator.toint() + + Int64(pool_token_global) * Int64(TokenSrcMetadata.nbytes) + ) + src_rank = md.src_rank + src_token = md.src_token + if cutlass.const_expr(self.reduce_topk_in_kernel): + src_topk = cutlass.Int32(0) + else: + src_topk = md.src_topk + # Int64 token coord: src_token*(topk*hidden) overflows int32 once + local_row = cute.slice_(self.tensor, (Int64(src_token), src_topk, None)) + # next's SymmetricBufferDevice exposes ``map_pointer`` + peer_iter = self.peer_rank_ptr_mapper.map_pointer( + local_row.iterator, src_rank, + ) + return cute.make_tensor(peer_iter, local_row.layout) + + +# ============================================================================= +# GluMxfp8Epilogue +# ============================================================================= + +class GluMxfp8Epilogue: + + _SubtileBarIdBase = 4 + # Named barrier for cross-warp sync during raw-C TMA stores. + _CStoreBarId = 10 + + def __init__( + self, + *, + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + use_2cta_instrs: bool, + sf_vec_size: int, + fc1_output_dtype: Type[cutlass.Numeric], + fc1_output_layout: utils.LayoutEnum, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + sf_dtype: Type[cutlass.Numeric] = cutlass.Float8E8M0FNU, + c_dtype: Type[cutlass.Numeric] = cutlass.BFloat16, + glu_clamp: Optional[float] = None, + epilog_sync_bar_id: int = 1, + epilogue_warp_ids: Tuple[int, ...] = (0, 1, 2, 3), + static_expert_shape: Optional[Tuple[int, int, int]] = None, + fc2_in_kernel_topk_reduce: bool = False, + token_back_by_dispatch: bool = False, + epi_flag_batch: Optional[Tuple[int, int]] = (1, 1), + apply_topk_in_fc1: bool = False, + generate_c: bool = False, + use_stg_fc1: bool = False, + combine_format: Optional[Any] = None, + act_func: str = "swiglu", + fc2_use_bulk: bool = False, + fc2_tma_stages: Optional[int] = None, + ) -> None: + self._act_func = act_func + self.fc1_output_dtype = fc1_output_dtype + self.fc1_output_layout = fc1_output_layout + self.acc_dtype = acc_dtype + self.sf_dtype = sf_dtype + self._sf_vec_size = sf_vec_size + self._c_dtype = c_dtype + self._epilog_sync_bar_id = epilog_sync_bar_id + self._epilogue_warp_ids = epilogue_warp_ids + self._use_2cta_instrs = use_2cta_instrs + + self._atom_thr_size = 2 if use_2cta_instrs else 1 + self._cta_tile_m = mma_tiler_mnk[0] // self._atom_thr_size + self._cta_tile_n = mma_tiler_mnk[1] + self._mma_tiler_k = mma_tiler_mnk[2] + self._cta_tile_n_sfb = ((mma_tiler_mnk[1] + 127) // 128) * 128 + self._static_expert_shape = static_expert_shape + if ( + static_expert_shape is not None + and static_expert_shape[2] % (self._cta_tile_m * cluster_shape_mn[0]) == 0 + ): + self._fc2_stg_needs_predicate: bool = False + else: + self._fc2_stg_needs_predicate: bool = True + + # TMA tile is (EpilogueTileN tokens, Fc1EpilogueOutputTileN intermediates) + self._epi_tile = (EpilogueTileN, Fc1EpilogueOutputTileN) + self._subtile_cnt = self._cta_tile_n // 2 // EpilogueTileN + + self._num_acc_stage = 2 + self._num_acc_pipeline_stages = self._num_acc_stage + + k = self._mma_tiler_k + self._num_sfa_tmem_cols = self._cta_tile_m * k // sf_vec_size * 4 // 4 // 128 + self._num_sfb_tmem_cols = ( + self._cta_tile_n_sfb * k // sf_vec_size * 4 // 4 // 128 + ) + + self._num_accumulator_tmem_cols = self._cta_tile_n * self._num_acc_stage + + self._fc2_in_kernel_topk_reduce = fc2_in_kernel_topk_reduce + self._token_back_by_dispatch = token_back_by_dispatch + self._apply_topk_in_fc1 = apply_topk_in_fc1 + self._generate_c = generate_c + self._use_stg_fc1 = use_stg_fc1 + # combine_format determines fc2 output encoding: bf16 (default) or quantized. + if combine_format is None: + combine_format = CombineFormat.parse("bf16") + self._combine_format = combine_format + self._combine_mxfp8 = combine_format.is_quantized + # sf_block_pad for fc2 MXFP8 combine + if self._combine_mxfp8 and static_expert_shape is not None: + _hidden_fc2 = static_expert_shape[2] + _sf_blocks_fc2 = _hidden_fc2 // EpilogueTileN + self._fc2_sf_block_pad = ((_sf_blocks_fc2 + 15) // 16) * 16 + self._hidden_fc2 = _hidden_fc2 + else: + self._fc2_sf_block_pad = 0 + self._hidden_fc2 = 0 + # batching stg.64 SF + self._fc2_sf_batch8 = ( + self._combine_mxfp8 + and self._hidden_fc2 > 0 + and (self._hidden_fc2 % self._cta_tile_n == 0) + and (self._cta_tile_n // EpilogueTileN == 8) + ) + self._epi_tile_c = (self._cta_tile_m, 2 * Fc1GateUpInterleave) + self._epi_fc1_batch = max(1, epi_flag_batch[0]) + self._epi_fc2_batch = max(1, epi_flag_batch[1]) + + self.glu_clamp = ( + cutlass.Float32(glu_clamp) if glu_clamp is not None else None + ) + + self._fc2_subtile_cnt = self._cta_tile_n // EpilogueTileN # = 8 + self._fc2_use_tma = ( + bool(fc2_use_bulk) and token_back_by_dispatch and self._combine_mxfp8 + ) + _ublk_hidden_ok = ( + self._combine_mxfp8 + and static_expert_shape is not None + and self._hidden_fc2 % self._cta_tile_n == 0 + ) + self._fc2_use_ublk = ( + bool(fc2_use_bulk) and (not token_back_by_dispatch) and _ublk_hidden_ok + ) + self._fc2_needs_staging = self._fc2_use_tma or self._fc2_use_ublk + if fc2_tma_stages is not None and not 1 <= fc2_tma_stages <= self._fc2_subtile_cnt: + raise ValueError( + f"fc2_tma_stages must be in [1, {self._fc2_subtile_cnt}], got {fc2_tma_stages}." + ) + if self._fc2_needs_staging: + self._fc2_tma_stages = ( + fc2_tma_stages if fc2_tma_stages is not None else min(2, self._fc2_subtile_cnt) + ) + else: + self._fc2_tma_stages = 0 + self._fc2_wire_dtype = ( + self._combine_format.act_dtype if self._combine_mxfp8 else cutlass.BFloat16 + ) + self._fc2_tma_stage_bytes = ( + self._cta_tile_m * EpilogueTileN * self._fc2_wire_dtype.width // 8 + ) + self._fc2_tma_staging_bytes = self._fc2_tma_stages * self._fc2_tma_stage_bytes + + self._fc2_reduce_coalesce = bool(self._fc2_in_kernel_topk_reduce) + # bf16 staging tile: (cta_tile_m tokens x EpilogueTileN hidden), 2 B/elem. + self._fc2_reduce_staging_bytes = ( + self._cta_tile_m * EpilogueTileN * cutlass.BFloat16.width // 8 + if self._fc2_reduce_coalesce + else 0 + ) + + # -- Codegen-time queries -- + + @property + def epi_tile(self) -> Tuple[int, int]: + return self._epi_tile + + @property + def num_acc_pipeline_stages(self) -> int: + return self._num_acc_pipeline_stages + + @property + def num_acc_stage(self) -> int: + return self._num_acc_stage + + @property + def subtile_cnt(self) -> int: + return self._subtile_cnt + + @property + def cta_tile_n(self) -> int: + return self._cta_tile_n + + @property + def num_sfa_tmem_cols(self) -> int: + return self._num_sfa_tmem_cols + + @property + def num_sfb_tmem_cols(self) -> int: + return self._num_sfb_tmem_cols + + @property + def num_accumulator_tmem_cols(self) -> int: + return self._num_accumulator_tmem_cols + + def staged_smem_layout( + self, + n_stages: int, + ) -> Union[cute.Layout, cute.ComposedLayout]: + return sm100_utils.make_smem_layout_epi( + self.fc1_output_dtype, + self.fc1_output_layout, + self._epi_tile, + n_stages, + ) + + @property + def smem_layout_one_stage(self) -> Union[cute.Layout, cute.ComposedLayout]: + staged = self.staged_smem_layout(1) + return cute.select(staged, mode=[0, 1]) + + @property + def bytes_per_stage(self) -> int: + return cute.size_in_bytes(self.fc1_output_dtype, self.smem_layout_one_stage) + + @property + def epi_tile_c(self) -> Tuple[int, int]: + """TMA tile for raw gate+up output: (cta_tile_m=128, 2*Fc1GateUpInterleave=64) fp32.""" + return self._epi_tile_c + + def staged_c_smem_layout(self, n_stages: int): + """SMEM layout for n_stages of raw gate+up (Float32, row-major, epi_tile_c).""" + return sm100_utils.make_smem_layout_epi( + self._c_dtype, + self.fc1_output_layout, # same N-major direction as fc1 output + self._epi_tile_c, + n_stages, + ) + + @property + def c_smem_layout_one_stage(self): + return cute.select(self.staged_c_smem_layout(1), mode=[0, 1]) + + @property + def c_bytes_per_stage(self) -> int: + return cute.size_in_bytes(self._c_dtype, self.c_smem_layout_one_stage) + + # ── FC2 TMASTG staging ─────────────────────────────────────────────── + @property + def fc2_use_tma(self) -> bool: + return self._fc2_use_tma + + @property + def fc2_use_ublk(self) -> bool: + return self._fc2_use_ublk + + @property + def fc2_needs_staging(self) -> bool: + """True when the FC2 store stages into the SMEM tile: either the dispatch + TMASTG path (``fc2_use_tma``) or the epi_warps UBLK path (``fc2_use_ublk``).""" + return self._fc2_needs_staging + + @property + def fc2_tma_stages(self) -> int: + return self._fc2_tma_stages + + @property + def fc2_tma_staging_bytes(self) -> int: + return self._fc2_tma_staging_bytes + + @property + def fc2_tma_tile(self) -> Tuple[int, int]: + """TMA store tile: (cta_tile_m=128 tokens, EpilogueTileN=32 hidden).""" + return (self._cta_tile_m, EpilogueTileN) + + def fc2_tma_staged_smem_layout(self, n_stages: int): + """Row-major (128 tokens, 32 hidden, n_stages) staging tile for the FC2 + bulk store. No swizzle: each thread owns one token row and writes its 32 + contiguous wire-dtype elements as a single 256-bit STS, and a 32-byte + innermost box is a valid (unswizzled) TMA tile. ``select(mode=[0,1])`` + yields the single-stage tile used to build the TMA atom.""" + cta_tile_m = self._cta_tile_m + stage_stride = cta_tile_m * EpilogueTileN + layout = cute.make_layout( + (cta_tile_m, EpilogueTileN, n_stages), + stride=(EpilogueTileN, 1, stage_stride if n_stages > 1 else 0), + ) + return layout + + @property + def fc2_tma_smem_layout_one_stage(self): + return cute.select(self.fc2_tma_staged_smem_layout(1), mode=[0, 1]) + + # ── FC2 in-kernel reduce coalescing ────────────────────────────────── + @property + def fc2_reduce_coalesce(self) -> bool: + return self._fc2_reduce_coalesce + + @property + def fc2_reduce_staging_bytes(self) -> int: + return self._fc2_reduce_staging_bytes + + def fc2_reduce_smem_layout(self): + """Row-major (cta_tile_m tokens, EpilogueTileN hidden) bf16 transpose tile. + Each epi thread writes its token's contiguous hidden row (64 B STS); the + coalesced-issue re-partition reads 4-elem chunks hidden-major.""" + return cute.make_layout( + (self._cta_tile_m, EpilogueTileN), stride=(EpilogueTileN, 1) + ) + + @staticmethod + @cute.jit + def tma_store_fc1_output( + warp_idx, + sC, + store_idx, + tma_atom_fc1_output: cute.CopyAtom, + g_fc1_output_subtile_view: cute.Tensor, + valid_tokens, + ) -> None: + """Per-warp TMA store for FC1 output.""" + cute.arch.fence_proxy("async.shared", space="cta") + sC_stage = cute.slice_(sC, (None, None, store_idx)) + g_fc1_output_2d = cute.slice_(g_fc1_output_subtile_view, (None, None, 0)) + bSG_sC, bSG_g = cpasync.tma_partition( + tma_atom_fc1_output, + 0, + cute.make_layout(1), + cute.group_modes(sC_stage, 0, 2), + cute.group_modes(g_fc1_output_2d, 0, 2), + ) + + leader_warp = store_idx + tile_has_valid = ( + store_idx * cutlass.Int32(EpilogueTileN) < valid_tokens + ) + + bar_id = store_idx + cutlass.Int32(GluMxfp8Epilogue._SubtileBarIdBase) + bar = pipeline.NamedBarrier( + barrier_id=bar_id, + num_threads=EpiWarpCount * WarpThreadCount, + ) + if warp_idx == leader_warp: + bar.arrive_and_wait() + # TMA bulk-tensor stores are all-or-nothing per issue (no per-element + # `pred` mask, no scalar predicate arg), so guard the whole copy with + # a runtime `if`. Skips fully-padding token-tiles that would alias + # the next expert's region. + if tile_has_valid: + cute.copy(tma_atom_fc1_output, bSG_sC, bSG_g) + else: + bar.arrive() + + @cute.jit + def _store_fc1_c_subtile( + self, + r_gate: cute.Tensor, + r_up: cute.Tensor, + smem_c_buffer: cute.Tensor, + tma_atom_c: cute.CopyAtom, + gmem_c_subtile_view: cute.Tensor, + c_buffer_idx, + work_tile_info, + warp_idx: int, + tidx, + c_pipeline, + ) -> None: + """Store pre-SwiGLU gate/up accumulators to the global C tensor via SMEM staging.""" + r_layout = cute.make_layout((((Fc1GateUpInterleave,), 1),), stride=(((1,), 0),)) + + # Cast acc_dtype (Float32) → c_dtype (e.g. BFloat16) before R2S. + r_gate_c = cute.make_rmem_tensor(r_layout.shape, self._c_dtype) + r_up_c = cute.make_rmem_tensor(r_layout.shape, self._c_dtype) + r_gate_c.store(r_gate.load().to(self._c_dtype)) + r_up_c.store(r_up.load().to(self._c_dtype)) + + r2s_c_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self._c_dtype, num_bits_per_copy=128, + ) + thread_in_warp_c = tidx % cutlass.Int32(WarpThreadCount) + c_row = cutlass.Int32(warp_idx * EpilogueTileN) + thread_in_warp_c + sC_raw_stage = cute.slice_(smem_c_buffer, (None, None, c_buffer_idx)) + c_gate_smem = cute.local_tile( + sC_raw_stage, (1, Fc1GateUpInterleave), (c_row, cutlass.Int32(0)), + ) + cute.copy(r2s_c_atom, cute.coalesce(r_gate_c), cute.coalesce(c_gate_smem)) + c_up_smem = cute.local_tile( + sC_raw_stage, (1, Fc1GateUpInterleave), (c_row, cutlass.Int32(1)), + ) + cute.copy(r2s_c_atom, cute.coalesce(r_up_c), cute.coalesce(c_up_smem)) + + # Fence + barrier: ensure all warps have written before TMA issue. + cute.arch.fence_proxy("async.shared", space="cta") + c_store_bar = pipeline.NamedBarrier( + barrier_id=self._CStoreBarId, + num_threads=EpiWarpCount * WarpThreadCount, + ) + c_store_bar.arrive_and_wait() + + # Warp 0 issues TMA S2G, commits, then pre-acquires the next stage so + # smem is guaranteed free before the next call's R2S writes. + if warp_idx == 0: + if work_tile_info.valid_tokens_in_cta_tile > cutlass.Int32(0): + g_c_2d = cute.slice_(gmem_c_subtile_view, (None, None, 0)) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, 0, cute.make_layout(1), + cute.group_modes(sC_raw_stage, 0, 2), + cute.group_modes(g_c_2d, 0, 2), + ) + cute.copy(tma_atom_c, bSG_sC, bSG_gC) + c_pipeline.producer_commit() + + def _subtile_local_tmem_tensor_pair( + self, + tmem_acc_tensor: cute.Tensor, + subtile_idx, + warp_idx, + ) -> cute.Tensor: + """ + Build a (gate, up) pair of TMEM tensor views for the MXFP8 fc1 epilogue. + """ + base = tmem_acc_tensor.iterator + warp_lane_off = warp_idx * WarpThreadCount + subtile_col_off = subtile_idx * EpilogueTileN * 2 + total = (warp_lane_off << 16) + subtile_col_off + subtile_gate_ptr = base + cute.assume(total, divby=16) + subtile_up_ptr = base + cute.assume(total + Fc1GateUpInterleave, divby=16) + return ( + cute.make_tensor( + subtile_gate_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ), + cute.make_tensor( + subtile_up_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ) + ) + + def _subtile_forward_tmem_tensor( + self, + tmem_gate_tensor: cute.Tensor, + tmem_up_tensor: cute.Tensor, + col_offset: int, + ) -> (cute.Tensor, cute.Tensor): + """Move the tmem tensor to the correct position.""" + tmem_gate_ptr = tmem_gate_tensor.iterator + cute.assume(col_offset, divby=16) + tmem_up_ptr = tmem_up_tensor.iterator + cute.assume(col_offset, divby=16) + return ( + cute.make_tensor( + tmem_gate_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ), + cute.make_tensor( + tmem_up_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ) + ) + + # -- fc1 subtile: SM100 path -- + @cute.jit + def _run_fc1_task_tile( + self, + work_tile_info, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + acc_consumer_state, + smem_fc1_output_buffer: cute.Tensor, + tma_atom_fc1_output: cute.CopyAtom, + sched_ext, + gmem_fc1_output: cute.Tensor, + gmem_fc1_output_sf: cute.Tensor, + gmem_topk_scores: cute.Tensor, + warp_idx: int, + tidx, + alpha, + norm_const, + smem_c_buffer: cute.Tensor, + tma_atom_c: cute.CopyAtom, + gmem_c: cute.Tensor, + c_pipeline, + ) -> None: + """MXFP8 fc1 task-tile: TmemTranspose16x32 TMEM loading + cross-warp E8M0 exchange.""" + real_fc1_output, _ = sched_ext.get_gmem_tensor("d", gmem_fc1_output, work_tile_info) + if cutlass.const_expr(self.fc1_output_dtype.width == 8): + real_fc1_output_sf, _ = sched_ext.get_gmem_tensor("sfd", gmem_fc1_output_sf, work_tile_info) + else: + real_fc1_output_sf = None + + real_topk_scores = gmem_topk_scores + if cutlass.const_expr(self._apply_topk_in_fc1): + real_topk_scores, _ = sched_ext.get_gmem_tensor( + "topk", gmem_topk_scores, work_tile_info + ) + + if cutlass.const_expr(self._generate_c): + real_c, _ = sched_ext.get_gmem_tensor("c", gmem_c, work_tile_info) + c_n_base = work_tile_info.tile_n_idx * cutlass.Int32(self._subtile_cnt) + + acc_pipeline.consumer_wait(acc_consumer_state) + if warp_idx == 0: + iket.range_push("fc1_epi_tile") + + subtile_cnt = self._subtile_cnt + tmem_gate, tmem_up = self._subtile_local_tmem_tensor_pair( + tmem_acc_tensor, 0, warp_idx, + ) + tmem_forward_cols = Fc1GateUpInterleave * 2 + + if cutlass.const_expr(self.fc1_output_dtype.width == 8): + layout_sf = cute.make_layout(4) + rmem_sf = cute.make_rmem_tensor(layout_sf.shape, self.acc_dtype) + else: + rmem_sf = None + + # Set up RMEM→SMEM copy atom (direct CopyUniversalOp) + r2s_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.fc1_output_dtype, + num_bits_per_copy=128, + ) + tRS_sC = None + + for i in cutlass.range(0, subtile_cnt, 1, unroll=1): + subtile_idx = cutlass.Int32(i) + + if cutlass.const_expr(self._generate_c): + c_buffer_idx = cutlass.Int32(i % Fc1CTMAStages) + g_c_subtile = cute.local_tile( + real_c, + (self._cta_tile_m, 2 * Fc1GateUpInterleave, 1), + (work_tile_info.tile_m_idx, c_n_base + subtile_idx, cutlass.Int32(0)), + ) + smem_c_buf_arg = smem_c_buffer + tma_atom_c_arg = tma_atom_c + gmem_c_subtile_arg = g_c_subtile + else: + c_buffer_idx = cutlass.Int32(0) + smem_c_buf_arg = smem_fc1_output_buffer + tma_atom_c_arg = tma_atom_fc1_output + gmem_c_subtile_arg = real_fc1_output + + self._run_fc1_subtile( + subtile_idx=subtile_idx, + tmem_gate_tensor=tmem_gate, + tmem_up_tensor=tmem_up, + real_fc1_output=real_fc1_output, + real_fc1_output_sf=real_fc1_output_sf, + real_topk_scores=real_topk_scores, + work_tile_info=work_tile_info, + smem_fc1_output_buffer=smem_fc1_output_buffer, + tma_atom_fc1_output=tma_atom_fc1_output, + r2s_copy_atom=r2s_copy_atom, + warp_idx=warp_idx, + tidx=tidx, + alpha=alpha, + norm_const=norm_const, + rmem_sf=rmem_sf, + smem_c_buffer=smem_c_buf_arg, + tma_atom_c=tma_atom_c_arg, + gmem_c_subtile_view=gmem_c_subtile_arg, + c_buffer_idx=c_buffer_idx, + c_pipeline=c_pipeline, + ) + + tmem_gate, tmem_up = self._subtile_forward_tmem_tensor(tmem_gate, tmem_up, tmem_forward_cols) + + self._acc_pipeline_consumer_release(acc_pipeline, acc_consumer_state, True) + + if cutlass.const_expr(self.fc1_output_dtype.width == 8): + self._stg_sf_fc1(rmem_sf, real_fc1_output_sf, work_tile_info, tidx) + + # TMA store: 4 stores (one per warp group) after all subtiles + if cutlass.const_expr(not self._use_stg_fc1): + base_token_tile = ( + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m // EpilogueTileN) + ) + for idx in cutlass.range_constexpr(EpiWarpCount): + g_fc1_output_warp_view = cute.local_tile( + real_fc1_output, + (EpilogueTileN, Fc1EpilogueOutputTileN, 1), + (base_token_tile + idx, work_tile_info.tile_n_idx, 0), + ) + GluMxfp8Epilogue.tma_store_fc1_output( + warp_idx, smem_fc1_output_buffer, idx, + tma_atom_fc1_output, g_fc1_output_warp_view, + work_tile_info.valid_tokens_in_cta_tile, + ) + + if warp_idx == 0: + iket.range_pop() + + @cute.jit + def _swiglu_act( + self, + t_swiglu: cute.Tensor, + t_up: cute.Tensor, + t_gate: cute.Tensor, + prob: Optional[Float32] = None, + ) -> None: + """SwiGLU hook consumed by ``_run_fc1_subtile`` """ + swiglu_act(t_swiglu, t_up, t_gate, prob) + + @cute.jit + def _run_fc1_subtile( + self, + subtile_idx, + tmem_gate_tensor: cute.Tensor, + tmem_up_tensor: cute.Tensor, + real_fc1_output: cute.Tensor, + real_fc1_output_sf: cute.Tensor, + real_topk_scores: cute.Tensor, + work_tile_info, + smem_fc1_output_buffer: cute.Tensor, + tma_atom_fc1_output: cute.CopyAtom, + r2s_copy_atom: cute.CopyAtom, + warp_idx: int, + tidx, + alpha, + norm_const, + rmem_sf: cute.Tensor, + smem_c_buffer: cute.Tensor, + tma_atom_c: cute.CopyAtom, + gmem_c_subtile_view: cute.Tensor, + c_buffer_idx, + c_pipeline, + ) -> None: + """MXFP8 fc1 subtile: GLU + E8M0 SF + fp8 R2S.""" + if warp_idx == 0: + iket.range_push("fc1_epilogue_subtile") + + r_layout = cute.make_layout((((Fc1GateUpInterleave,), 1),), stride=(((1,), 0),)) + r_gate = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + r_up = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + + atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), self.acc_dtype, + ) + cute.copy(atom_t2r, tmem_gate_tensor, r_gate) + cute.copy(atom_t2r, tmem_up_tensor, r_up) + + # ── generate_c: store raw gate+up to GMEM C tensor via SMEM staging ── + if cutlass.const_expr(self._generate_c): + self._store_fc1_c_subtile( + r_gate=r_gate, + r_up=r_up, + smem_c_buffer=smem_c_buffer, + tma_atom_c=tma_atom_c, + gmem_c_subtile_view=gmem_c_subtile_view, + c_buffer_idx=c_buffer_idx, + work_tile_info=work_tile_info, + warp_idx=warp_idx, + tidx=tidx, + c_pipeline=c_pipeline, + ) + + if cutlass.const_expr(self.glu_clamp is not None): + for i in cutlass.range_constexpr(cute.size(r_up)): + r_gate[i] = cute.arch.fmin(r_gate[i], self.glu_clamp) + r_up[i] = cute.arch.fmin(r_up[i], self.glu_clamp) + r_up[i] = cute.arch.fmax(r_up[i], -self.glu_clamp) + + topk = None + if cutlass.const_expr(self._apply_topk_in_fc1): + thread_in_warp = tidx % cutlass.Int32(WarpThreadCount) + token_in_tile = ( + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + ) + topk = Float32(real_topk_scores[token_in_tile]) + + swiglu = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) + if cutlass.const_expr(self._act_func == "swiglu"): + self._swiglu_act(swiglu, r_up, r_gate, topk) + + c = cute.make_rmem_tensor(r_layout.shape, self.fc1_output_dtype) + if cutlass.const_expr(self.fc1_output_dtype.width == 8): + # Quantized hand-off: fp8 data + E8M0 block scale. + qpvscale = quant_sfd_row(swiglu, c, norm_const, self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype) + if subtile_idx == 0: + rmem_sf[0] = qpvscale + elif subtile_idx == 1: + rmem_sf[1] = qpvscale + elif subtile_idx == 2: + rmem_sf[2] = qpvscale + elif subtile_idx == 3: + rmem_sf[3] = qpvscale + else: + # Plain-data hand-off: direct cast to the fc1 output dtype (the + # fc1_output workspace is reloaded as fc2's A operand). + c.store(swiglu.load().to(self.fc1_output_dtype)) + + thread_in_warp = tidx % WarpThreadCount + if cutlass.const_expr(self._use_stg_fc1): + # Direct STG.256 to GMEM — no SMEM staging or TMA store needed. + token_in_tile = cutlass.Int32(warp_idx * EpilogueTileN) + thread_in_warp + if token_in_tile < work_tile_info.valid_tokens_in_cta_tile: + abs_token = ( + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + cutlass.Int32(warp_idx * EpilogueTileN) + thread_in_warp + ) + # absolute column start (element index in the intermediate axis) + col_elem = ( + (work_tile_info.tile_n_idx * cutlass.Int32(self._subtile_cnt) + subtile_idx) + * cutlass.Int32(Fc1GateUpInterleave) + ) + g_base = cute.local_tile( + real_fc1_output, + (1, 1, 1), + (abs_token, col_elem, cutlass.Int32(0)), + ) + stg_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.fc1_output_dtype, num_bits_per_copy=256, + ) + # col_elem is always a multiple of Fc1GateUpInterleave=32 (FP8 elements) + aligned_iter = cute.make_ptr( + self.fc1_output_dtype, + g_base.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=32, + ) + g_vec = cute.make_tensor(aligned_iter, cute.make_layout(Fc1GateUpInterleave)) + cute.copy(stg_atom, cute.coalesce(c), g_vec) + else: + sC_stage = cute.slice_(smem_fc1_output_buffer, (None, None, warp_idx)) + sC_thread_row = cute.local_tile( + sC_stage, (1, Fc1GateUpInterleave), (thread_in_warp, subtile_idx) + ) + cute.copy(r2s_copy_atom, cute.coalesce(c), cute.coalesce(sC_thread_row)) + + + if cutlass.const_expr(self._generate_c): + if warp_idx == 0: + c_pipeline.producer_acquire() + c_store_bar = pipeline.NamedBarrier( + barrier_id=self._CStoreBarId, + num_threads=EpiWarpCount * WarpThreadCount, + ) + c_store_bar.arrive_and_wait() + + if warp_idx == 0: + iket.range_pop() + + @cute.jit + def _subtile_fc2_tmem_tensor( + self, + tmem_acc_tensor: cute.Tensor, + subtile_idx, + warp_idx, + ) -> cute.Tensor: + """ + Per-warp TMEM view for one fc2 subtile (EpilogueTileN=32 cols). + """ + base = tmem_acc_tensor.iterator + warp_lane_off = warp_idx * WarpThreadCount + subtile_col_off = subtile_idx * EpilogueTileN + total = (warp_lane_off << 16) + subtile_col_off + subtile_ptr = base + cute.assume(total, divby=16) + return cute.make_tensor( + subtile_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ) + + @cute.jit + def _advance_fc2_tmem_tensor( + self, + tmem_tensor: cute.Tensor, + col_offset: int, + ) -> cute.Tensor: + """Advance the fc2 TMEM tensor by col_offset cols (mirrors _subtile_forward_tmem_tensor).""" + new_ptr = tmem_tensor.iterator + cute.assume(col_offset, divby=16) + return cute.make_tensor( + new_ptr, + _TmemTranspose16x32Core._tmem_layout(32, EpilogueTileN), + ) + + @cute.jit + def _acc_pipeline_consumer_release( + self, + acc_pipeline, + acc_consumer_state, + is_release: bool, + ) -> None: + """Release the acc pipeline consumer.""" + if is_release: + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + @cute.jit + def _run_fc2_subtile( + self, + subtile_idx, + tmem_subtile_tensor: cute.Tensor, + real_fc2_output: cute.Tensor, + work_tile_info, + valid_hidden, + warp_idx: int, + tidx, + token_comm_args=None, + rmem_sf_fc2=None, + smem_fc2_tma_buffer=None, + tma_atom_fc2_output=None, + gmem_fc2_tma_output=None, + smem_fc2_reduce_buffer=None, + ) -> None: + """fc2 subtile: LDTM + encode + STG.""" + if warp_idx == 0: + iket.range_push("fc2_epi_subtile") + + fc2_subtile_cnt = self._cta_tile_n // EpilogueTileN # = 8 + hidden_group = ( + work_tile_info.tile_n_idx * cutlass.Int32(fc2_subtile_cnt) + subtile_idx + ) + hidden_col_start = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + + subtile_idx * cutlass.Int32(EpilogueTileN) + ) + r_acc_layout = cute.make_layout((((EpilogueTileN,), 1),), stride=(((1,), 0),)) + atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), self.acc_dtype, + ) + r_acc = cute.make_rmem_tensor(r_acc_layout.shape, self.acc_dtype) + cute.copy(atom_t2r, tmem_subtile_tensor, r_acc) + + thread_in_warp = tidx % WarpThreadCount + token_row_in_cta = cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + if token_row_in_cta < valid_tokens and hidden_col_start < valid_hidden: + if cutlass.const_expr( + token_comm_args is not None + and not self._token_back_by_dispatch + and self._combine_mxfp8 + ): + # MegaMoE Form A, quantized combine: + # 1. Quantize fp32 → fp8 + compute E8M0 block scale. + # 2. STG fp8 data to peer's combine_output. + # 3. Write E8M0 scale to local fc2_output_sf for token-back push. + fp8_dtype = self._combine_format.act_dtype + r_fp8 = cute.make_rmem_tensor(r_acc_layout.shape, fp8_dtype) + qpvscale = quant_sfd_row( + r_acc, r_fp8, 1.0, EpilogueTileN, + cutlass.Float8E8M0FNU, fp8_dtype, + ) + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + metadata_u32 = cute.recast_tensor( + token_comm_args.token_src_metadata, cutlass.Uint32, + ) + fc2_output_dest = Fc2OutputDest( + tensor=token_comm_args.combine_output, + metadata=metadata_u32, + peer_rank_ptr_mapper=token_comm_args.peer_rank_ptr_mapper, + ) + dest_row = fc2_output_dest.resolve_token_row(pool_token_global) + r_fp8_flat = cute.make_tensor( + r_fp8.iterator, cute.make_layout(EpilogueTileN) + ) + dest_fp8_ptr = cute.make_ptr( + fp8_dtype, + dest_row.iterator.toint() + Int64(hidden_col_start), + cute.AddressSpace.gmem, + assumed_align=32, + ) + if cutlass.const_expr(self._fc2_use_ublk): + # UBLK: stage this token's 32 fp8 elems into SMEM + stage_idx = subtile_idx % cutlass.Int32(self._fc2_tma_stages) + smem_row = cute.slice_( + smem_fc2_tma_buffer, (token_row_in_cta, None, stage_idx) + ) + sts_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=256, + ) + cute.copy(sts_fp8_atom, r_fp8_flat, smem_row) + cute.arch.fence_proxy("async.shared", space="cta") + cp_async_bulk_s2g( + dest_fp8_ptr, + smem_row.iterator, + cutlass.Int32(EpilogueTileN * fp8_dtype.width // 8), + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group( + self._fc2_tma_stages - 1, read=True + ) + else: + # STG 32 fp8 elements = 256 bits in one shot. + stg_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, + num_bits_per_copy=256, + ) + cute.copy( + stg_fp8_atom, r_fp8_flat, + cute.make_tensor(dest_fp8_ptr, cute.make_layout(EpilogueTileN)), + ) + # Buffer the E8M0 scale; the whole task tile's 8 scales are + # flushed together by _stg_sf_fc2 (single stg.64 when aligned). + self._write_sf_fc2_buffer(rmem_sf_fc2, subtile_idx, qpvscale) + elif cutlass.const_expr( + self._token_back_by_dispatch and self._combine_mxfp8 + ): + # MegaMoE token-back-by-dispatch + quantized combine: + # Epilogue writes fp8 data to local pool; dispatch warps push + # both data (fc2_output_workspace) and SF (fc2_output_sf) to peers. + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + fp8_dtype = self._combine_format.act_dtype + r_fp8 = cute.make_rmem_tensor(r_acc_layout.shape, fp8_dtype) + qpvscale = quant_sfd_row( + r_acc, r_fp8, 1.0, EpilogueTileN, + cutlass.Float8E8M0FNU, fp8_dtype, + ) + r_fp8_flat = cute.make_tensor(r_fp8.iterator, cute.make_layout(EpilogueTileN)) + if cutlass.const_expr(self._fc2_use_tma): + stage_idx = subtile_idx % cutlass.Int32(self._fc2_tma_stages) + smem_row = cute.slice_( + smem_fc2_tma_buffer, (token_row_in_cta, None, stage_idx) + ) + sts_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=256, + ) + cute.copy(sts_fp8_atom, r_fp8_flat, smem_row) + else: + # Write 32 fp8 elements to local fc2_output_workspace pool. + fp8_byte_addr = ( + token_comm_args.fc2_output_workspace.iterator.toint() + + Int64(pool_token_global) * Int64(self._hidden_fc2) + + Int64(hidden_col_start) + ) + stg_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, + num_bits_per_copy=256, + ) + aligned_fp8_iter = cute.make_ptr( + fp8_dtype, + fp8_byte_addr, + cute.AddressSpace.gmem, + assumed_align=32, + ) + cute.copy( + stg_fp8_atom, r_fp8_flat, + cute.make_tensor(aligned_fp8_iter, cute.make_layout(EpilogueTileN)), + ) + # Buffer the E8M0 scale; flushed together by _stg_sf_fc2 after + # the subtile loop (single stg.64 when hidden-aligned). + self._write_sf_fc2_buffer(rmem_sf_fc2, subtile_idx, qpvscale) + elif cutlass.const_expr( + token_comm_args is None and self._combine_mxfp8 + ): + # Lean path (no token-comm), quantized fc2 output + fp8_dtype = self._combine_format.act_dtype + r_fp8 = cute.make_rmem_tensor(r_acc_layout.shape, fp8_dtype) + qpvscale = quant_sfd_row( + r_acc, r_fp8, 1.0, EpilogueTileN, + cutlass.Float8E8M0FNU, fp8_dtype, + ) + g_fc2_output_tile = cute.local_tile( + real_fc2_output, + (self._cta_tile_m, EpilogueTileN, 1), + (work_tile_info.tile_m_idx, hidden_group, 0), + ) + g_fc2_slice = cute.slice_(g_fc2_output_tile, (None, None, 0)) + g_thread_row = cute.local_tile( + g_fc2_slice, (1, EpilogueTileN), (token_row_in_cta, 0), + ) + g_flat = cute.coalesce(g_thread_row) + aligned_iter = cute.make_ptr( + fp8_dtype, + g_flat.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=32, + ) + stg_fp8_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), fp8_dtype, + num_bits_per_copy=256, + ) + r_fp8_flat = cute.make_tensor( + r_fp8.iterator, cute.make_layout(EpilogueTileN) + ) + cute.copy( + stg_fp8_atom, r_fp8_flat, + cute.make_tensor(aligned_iter, cute.make_layout(EpilogueTileN)), + ) + # Buffer the E8M0 scale; flushed together by _stg_sf_fc2 after + # the subtile loop (single stg.64 when hidden-aligned). + self._write_sf_fc2_buffer(rmem_sf_fc2, subtile_idx, qpvscale) + else: + # BF16 path (default): fp32->bf16, two 256-bit STGs. + r_bf16 = cute.make_rmem_tensor(r_acc_layout.shape, cutlass.BFloat16) + r_bf16.store(r_acc.load().to(cutlass.BFloat16)) + stg_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=256, + ) + if cutlass.const_expr(self._fc2_reduce_coalesce): + # In-kernel reduce: stage the full bf16 token row token-major into + # SMEM (row start is 64 B aligned). + cute.autovec_copy( + cute.make_tensor(r_bf16.iterator, cute.make_layout(EpilogueTileN)), + cute.slice_(smem_fc2_reduce_buffer, (token_row_in_cta, None)), + ) + if cutlass.const_expr( + token_comm_args is not None and not self._token_back_by_dispatch + ): + metadata_u32 = cute.recast_tensor( + token_comm_args.token_src_metadata, cutlass.Uint32, + ) + fc2_output_dest = Fc2OutputDest( + tensor=token_comm_args.combine_output, + metadata=metadata_u32, + peer_rank_ptr_mapper=token_comm_args.peer_rank_ptr_mapper, + reduce_topk_in_kernel=self._fc2_in_kernel_topk_reduce, + ) + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + dest_row = fc2_output_dest.resolve_token_row(pool_token_global) + for stg_half in cutlass.range_constexpr(EpilogueTileN // 16): + reg_view = cute.make_tensor( + r_bf16.iterator + stg_half * 16, + cute.make_layout(16), + ) + if cutlass.const_expr( + token_comm_args is not None and not self._token_back_by_dispatch + ): + hidden_off = hidden_col_start + cutlass.Int32(stg_half * 16) + dest_ptr = cute.make_ptr( + cutlass.BFloat16, + dest_row.iterator.toint() + hidden_off * cutlass.Int64(2), + cute.AddressSpace.gmem, + assumed_align=32, + ) + # Reduce path staged to SMEM above; only the non-reduce Form-A + # combine does the direct peer STG here. + if cutlass.const_expr(not self._fc2_reduce_coalesce): + cute.copy( + stg_atom, reg_view, + cute.make_tensor(dest_ptr, cute.make_layout(16)), + ) + else: + g_fc2_output_tile = cute.local_tile( + real_fc2_output, + (self._cta_tile_m, EpilogueTileN, 1), + (work_tile_info.tile_m_idx, hidden_group, 0), + ) + g_fc2_slice = cute.slice_(g_fc2_output_tile, (None, None, 0)) + g_thread_row = cute.local_tile( + g_fc2_slice, (1, 16), (token_row_in_cta, stg_half), + ) + g_flat = cute.coalesce(g_thread_row) + aligned_iter = cute.make_ptr( + cutlass.BFloat16, + g_flat.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=32, + ) + cute.copy(stg_atom, reg_view, cute.make_tensor(aligned_iter, g_flat.layout)) + + if cutlass.const_expr(self._fc2_use_tma): + self._issue_fc2_tma_store( + subtile_idx=subtile_idx, + smem_fc2_tma_buffer=smem_fc2_tma_buffer, + tma_atom_fc2_output=tma_atom_fc2_output, + gmem_fc2_tma_output=gmem_fc2_tma_output, + work_tile_info=work_tile_info, + warp_idx=warp_idx, + ) + + # In-kernel reduce: the token-major SMEM tile is now populated; all 128 epi + # threads re-partition hidden-major and issue coalesced peer red.add. + if cutlass.const_expr(self._fc2_reduce_coalesce): + self._issue_fc2_reduce_coalesced( + subtile_idx=subtile_idx, + smem_fc2_reduce_buffer=smem_fc2_reduce_buffer, + token_comm_args=token_comm_args, + work_tile_info=work_tile_info, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + ) + + if warp_idx == 0: + iket.range_pop() + + + @cute.jit + def _issue_fc2_tma_store( + self, + subtile_idx, + smem_fc2_tma_buffer: cute.Tensor, + tma_atom_fc2_output: cute.CopyAtom, + gmem_fc2_tma_output: cute.Tensor, + work_tile_info, + warp_idx: int, + ) -> None: + """Issue the FC2 bulk-tensor store for one subtile's staged SMEM tile.""" + stage_idx = subtile_idx % cutlass.Int32(self._fc2_tma_stages) + smem_stage = cute.slice_(smem_fc2_tma_buffer, (None, None, stage_idx)) + token_base = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + ) + hidden_col = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + + subtile_idx * cutlass.Int32(EpilogueTileN) + ) + token_tile_idx = token_base // cutlass.Int32(self._cta_tile_m) + hidden_tile_idx = hidden_col // cutlass.Int32(EpilogueTileN) + tiled_output = cute.flat_divide( + gmem_fc2_tma_output, (self._cta_tile_m, EpilogueTileN) + ) + gmem_subtile = tiled_output[None, None, token_tile_idx, hidden_tile_idx] + bSG_s, bSG_g = cpasync.tma_partition( + tma_atom_fc2_output, + 0, + cute.make_layout(1), + cute.group_modes(smem_stage, 0, 2), + cute.group_modes(gmem_subtile, 0, 2), + ) + cute.arch.fence_proxy("async.shared", space="cta") + bar = pipeline.NamedBarrier( + barrier_id=self._epilog_sync_bar_id, + num_threads=EpiWarpCount * WarpThreadCount, + ) + bar.arrive_and_wait() + if warp_idx == 0: + cute.copy(tma_atom_fc2_output, bSG_s, bSG_g) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(self._fc2_tma_stages - 1, read=True) + bar.arrive_and_wait() + + @cute.jit + def _issue_fc2_reduce_coalesced( + self, + subtile_idx, + smem_fc2_reduce_buffer: cute.Tensor, + token_comm_args, + work_tile_info, + valid_hidden, + warp_idx: int, + tidx, + ) -> None: + """Coalesced peer red.add of one subtile's staged reduce tile.""" + redg_width: cutlass.Constexpr[int] = 4 # bf16 per red.v2.bf16x2 + lanes_per_token: cutlass.Constexpr[int] = EpilogueTileN // redg_width # 8 + tokens_per_warp: cutlass.Constexpr[int] = WarpThreadCount // lanes_per_token # 4 + tokens_per_pass: cutlass.Constexpr[int] = tokens_per_warp * EpiWarpCount # 16 + passes: cutlass.Constexpr[int] = self._cta_tile_m // tokens_per_pass # 8 + + # Rendezvous: make every thread's staged SMEM row visible before the reads. + bar = pipeline.NamedBarrier( + barrier_id=self._epilog_sync_bar_id, + num_threads=EpiWarpCount * WarpThreadCount, + ) + bar.arrive_and_wait() + + thread_in_warp = tidx % WarpThreadCount + token_in_warp = thread_in_warp // cutlass.Int32(lanes_per_token) # 0..tokens_per_warp-1 + chunk = thread_in_warp % cutlass.Int32(lanes_per_token) # 0..lanes_per_token-1 + hidden_off = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + + subtile_idx * cutlass.Int32(EpilogueTileN) + + chunk * cutlass.Int32(redg_width) + ) + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + metadata_u32 = cute.recast_tensor(token_comm_args.token_src_metadata, cutlass.Uint32) + fc2_output_dest = Fc2OutputDest( + tensor=token_comm_args.combine_output, + metadata=metadata_u32, + peer_rank_ptr_mapper=token_comm_args.peer_rank_ptr_mapper, + reduce_topk_in_kernel=True, + ) + for p in cutlass.range_constexpr(passes): + token_in_cta = ( + cutlass.Int32(p * tokens_per_pass) + + cutlass.Int32(warp_idx) * cutlass.Int32(tokens_per_warp) + + token_in_warp + ) + if token_in_cta < valid_tokens and hidden_off < valid_hidden: + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_in_cta + ) + dest_row = fc2_output_dest.resolve_token_row(pool_token_global) + # LDS this lane's 4 staged bf16 (contiguous hidden chunk). + smem_chunk_ptr = cute.slice_( + smem_fc2_reduce_buffer, (token_in_cta, None) + ).iterator + chunk * cutlass.Int32(redg_width) + r4 = cute.make_rmem_tensor((redg_width,), cutlass.BFloat16) + cute.autovec_copy( + cute.make_tensor(smem_chunk_ptr, cute.make_layout(redg_width)), r4 + ) + r4_u32 = cute.recast_tensor(r4, cutlass.Uint32) # 2 u32 = 4 bf16 + dest_ptr = cute.make_ptr( + cutlass.BFloat16, + dest_row.iterator.toint() + Int64(hidden_off) * Int64(2), + cute.AddressSpace.gmem, + assumed_align=8, + ) + _red_add_relaxed_sys_v2_bf16x2( + dest_ptr, + cutlass.Uint32(r4_u32[0]), + cutlass.Uint32(r4_u32[1]), + ) + bar.arrive_and_wait() + + @cute.jit + def _write_sf_fc2_buffer(self, rmem_sf_fc2, subtile_idx, qpvscale) -> None: + """Scatter one subtile's E8M0 scale into the per-tile SF buffer.""" + for j in cutlass.range_constexpr(self._cta_tile_n // EpilogueTileN): + if subtile_idx == cutlass.Int32(j): + rmem_sf_fc2[j] = qpvscale + + @cute.jit + def _stg_sf_fc2( + self, + rmem_sf_fc2: cute.Tensor, + sf_base_addr, + sf_row_stride, + work_tile_info, + valid_hidden, + warp_idx: int, + tidx, + ) -> None: + """Flush a task tile's fc2 E8M0 scales to the local ``fc2_output_sf``.""" + fc2_subtile_cnt = self._cta_tile_n // EpilogueTileN + thread_in_warp = tidx % WarpThreadCount + token_row_in_cta = cutlass.Int32(warp_idx * WarpThreadCount) + thread_in_warp + if token_row_in_cta < work_tile_info.valid_tokens_in_cta_tile: + pool_token_global = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + + token_row_in_cta + ) + hidden_group_base = ( + work_tile_info.tile_n_idx * cutlass.Int32(fc2_subtile_cnt) + ) + sf_byte_addr = ( + sf_base_addr + + Int64(pool_token_global) * Int64(sf_row_stride) + + Int64(hidden_group_base) + ) + if cutlass.const_expr(self._fc2_sf_batch8): + stg_e8m0x8_from_f32( + sf_byte_addr, + rmem_sf_fc2[0], rmem_sf_fc2[1], rmem_sf_fc2[2], rmem_sf_fc2[3], + rmem_sf_fc2[4], rmem_sf_fc2[5], rmem_sf_fc2[6], rmem_sf_fc2[7], + ) + else: + for j in cutlass.range_constexpr(fc2_subtile_cnt): + block_hidden_start = ( + work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) + + cutlass.Int32(j * EpilogueTileN) + ) + if block_hidden_start < valid_hidden: + stg_e8m0_from_f32(sf_byte_addr + Int64(j), rmem_sf_fc2[j]) + + @cute.jit + def _run_fc2_task_tile( + self, + work_tile_info, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + acc_consumer_state, + sched_ext, + gmem_fc2_output: cute.Tensor, + valid_hidden, + warp_idx: int, + tidx, + token_comm_args=None, + gmem_fc2_output_sf=None, + smem_fc2_tma_buffer=None, + tma_atom_fc2_output=None, + gmem_fc2_tma_output=None, + smem_fc2_reduce_buffer=None, + ) -> None: + """FC2 (Linear2) task-tile body using two-stage TMEM accumulation.""" + real_fc2_output, _ = sched_ext.get_gmem_tensor( + "d", gmem_fc2_output, work_tile_info, + ) + acc_pipeline.consumer_wait(acc_consumer_state) + if warp_idx == 0: + iket.range_push("fc2_epi_tile") + + fc2_subtile_cnt = self._cta_tile_n // EpilogueTileN # = 8 + + tmem_t = self._subtile_fc2_tmem_tensor( + tmem_acc_tensor, cutlass.Int32(0), warp_idx, + ) + + tmem_forward_cols = EpilogueTileN + + # Quantized combine: buffer the per-subtile E8M0 scales and flush them in + # one stg.64 after the loop (see _stg_sf_fc2). Indexed by subtile_idx, so + # the reversed odd-turn walk fills the same slots. Both the token-comm + # combine plane and the lean stand-alone fc2_output_sf plane need it. + if cutlass.const_expr(self._combine_mxfp8): + layout_sf_fc2 = cute.make_layout(fc2_subtile_cnt) + rmem_sf_fc2 = cute.make_rmem_tensor(layout_sf_fc2.shape, self.acc_dtype) + else: + rmem_sf_fc2 = None + + for i in cutlass.range(0, fc2_subtile_cnt, 1, unroll=1): + subtile_idx = cutlass.Int32(i) + + self._run_fc2_subtile( + subtile_idx=subtile_idx, + tmem_subtile_tensor=tmem_t, + real_fc2_output=real_fc2_output, + work_tile_info=work_tile_info, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + token_comm_args=token_comm_args, + rmem_sf_fc2=rmem_sf_fc2, + smem_fc2_tma_buffer=smem_fc2_tma_buffer, + tma_atom_fc2_output=tma_atom_fc2_output, + gmem_fc2_tma_output=gmem_fc2_tma_output, + smem_fc2_reduce_buffer=smem_fc2_reduce_buffer, + ) + + tmem_t = self._advance_fc2_tmem_tensor(tmem_t, tmem_forward_cols) + + self._acc_pipeline_consumer_release(acc_pipeline, acc_consumer_state, True) + + # Flush the buffered E8M0 scales + if cutlass.const_expr(self._combine_mxfp8 and token_comm_args is not None): + self._stg_sf_fc2( + rmem_sf_fc2=rmem_sf_fc2, + sf_base_addr=token_comm_args.fc2_output_sf.iterator.toint(), + sf_row_stride=Int64(self._fc2_sf_block_pad), + work_tile_info=work_tile_info, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + ) + elif cutlass.const_expr( + self._combine_mxfp8 and gmem_fc2_output_sf is not None + ): + self._stg_sf_fc2( + rmem_sf_fc2=rmem_sf_fc2, + sf_base_addr=gmem_fc2_output_sf.iterator.toint(), + sf_row_stride=Int64(gmem_fc2_output_sf.stride[0]), + work_tile_info=work_tile_info, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + ) + + if warp_idx == 0: + iket.range_pop() + + + @cute.jit + def _stg_sf_fc1( + self, + rmem_sf_f32: cute.Tensor, + real_fc1_output_sf: cute.Tensor, + work_tile_info, + tidx, + ) -> None: + """Compute gmem SF tile coords and store fc1 scale factors to gmem.""" + bx, _, _ = cute.arch.block_idx() + sf_idx = work_tile_info.tile_n_idx + token_idx = ( + work_tile_info.tile_m_idx * self._cta_tile_m + + tidx + ) + if tidx < work_tile_info.valid_tokens_in_cta_tile: + sf_base = cute.local_tile( + real_fc1_output_sf, + (1, 1, 1), + (token_idx, sf_idx * cutlass.Int32(Fc1EpilogueOutputTileN), cutlass.Int32(0)), + ) + sf_ptr = cute.make_ptr( + self.sf_dtype, + sf_base.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=4, + ) + gmem_sf_f8 = cute.make_tensor(sf_ptr, cute.make_layout(4)) + sf_layout = cute.make_layout(4) + r_sf_f8 = cute.make_rmem_tensor(sf_layout.shape, self.sf_dtype) + r_sf_f8.store(rmem_sf_f32.load().to(self.sf_dtype)) + cute.autovec_copy(r_sf_f8, gmem_sf_f8) + + + @cute.jit + def run( + self, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + sched_consumer, + sched_ext, + smem_fc1_output_buffer: Optional[cute.Tensor], + tma_atom_fc1_output: cute.CopyAtom, + gmem_fc1_output: cute.Tensor, + gmem_topk_scores: cute.Tensor, + gmem_fc2_output: cute.Tensor, + gmem_fc1_done_counter: cute.Tensor, + warp_idx: int, + tidx, + gmem_fc1_output_sf: Optional[cute.Tensor] = None, + alpha=None, + norm_const=None, + token_comm_args=None, + gmem_fc2_output_sf: Optional[cute.Tensor] = None, + smem_c_buffer: cute.Tensor = None, + tma_atom_c: cute.CopyAtom = None, + gmem_c: cute.Tensor = None, + smem_fc2_tma_buffer: Optional[cute.Tensor] = None, + tma_atom_fc2_output: cute.CopyAtom = None, + gmem_fc2_tma_output: cute.Tensor = None, + smem_fc2_reduce_buffer: Optional[cute.Tensor] = None, + ) -> None: + """ + Run the full fc1+fc2-fused epilogue task-tile loop. + """ + assert self._use_stg_fc1 or smem_fc1_output_buffer is not None, ( + "smem_fc1_output_buffer=None requires use_stg_fc1=True (the TMA " + "fc1-output store path consumes the sD staging buffer)" + ) + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._num_acc_pipeline_stages + ) + + if cutlass.const_expr(self._generate_c): + _c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=Fc1CTMAStages, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + EpiWarpCount * WarpThreadCount, + ), + ) + else: + _c_pipeline = None + + task_tile_boundary_bar = pipeline.NamedBarrier( + barrier_id=self._epilog_sync_bar_id, + num_threads=32 * len(self._epilogue_warp_ids), + ) + + valid_hidden = cutlass.Int32(gmem_fc2_output.shape[1]) + + bidx, bidy, bidz = cute.arch.block_idx() + work_tile_info = sched_consumer.consume_work() + + flag_tracker = GpuReleaseFlagBatchTracker( + flag_address=Int64(0), + accumulated_flags=cutlass.Int32(0), + phase=cutlass.Int32(work_tile_info.phase), + thread_idx=tidx % (len(self._epilogue_warp_ids) * WarpThreadCount), + ) + + while work_tile_info.is_valid_tile: + acc_stage_index = acc_consumer_state.index + tmem_acc_stage_tesnor = tmem_acc_tensor[(None, None, None, acc_stage_index)] + + if work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1): + if cutlass.const_expr(self._generate_c): + _smem_c_buf = smem_c_buffer + _tma_atom_c = tma_atom_c + _gmem_c = gmem_c + else: + _smem_c_buf = smem_fc1_output_buffer + _tma_atom_c = tma_atom_fc1_output + _gmem_c = gmem_fc1_output + self._run_fc1_task_tile( + work_tile_info=work_tile_info, + tmem_acc_tensor=tmem_acc_stage_tesnor, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + smem_fc1_output_buffer=smem_fc1_output_buffer, + tma_atom_fc1_output=tma_atom_fc1_output, + sched_ext=sched_ext, + gmem_fc1_output=gmem_fc1_output, + gmem_fc1_output_sf=gmem_fc1_output_sf, + gmem_topk_scores=gmem_topk_scores, + warp_idx=warp_idx, + tidx=tidx, + alpha=alpha, + norm_const=norm_const, + smem_c_buffer=_smem_c_buf, + tma_atom_c=_tma_atom_c, + gmem_c=_gmem_c, + c_pipeline=_c_pipeline, + ) + else: + self._run_fc2_task_tile( + work_tile_info=work_tile_info, + tmem_acc_tensor=tmem_acc_stage_tesnor, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + sched_ext=sched_ext, + gmem_fc2_output=gmem_fc2_output, + valid_hidden=valid_hidden, + warp_idx=warp_idx, + tidx=tidx, + token_comm_args=token_comm_args, + gmem_fc2_output_sf=gmem_fc2_output_sf, + smem_fc2_tma_buffer=smem_fc2_tma_buffer, + tma_atom_fc2_output=tma_atom_fc2_output, + gmem_fc2_tma_output=gmem_fc2_tma_output, + smem_fc2_reduce_buffer=smem_fc2_reduce_buffer, + ) + + acc_consumer_state.advance() + + cur_was_linear1 = work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + cur_fc1_counter_slot = ( + work_tile_info.cumulative_token_block_count + + work_tile_info.tile_m_idx // cutlass.Int32(self._atom_thr_size) + ) + cur_fc2_expert_idx = work_tile_info.expert_idx + + work_tile_info = sched_consumer.consume_work() + + # Drain in-flight bulk stores before publishing the done counter. + if cur_was_linear1 or cutlass.const_expr( + self._fc2_use_tma or self._fc2_use_ublk + ): + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.fence_proxy("async") + cute.arch.fence_acq_rel_gpu() + + task_tile_boundary_bar.arrive_and_wait() + + if cur_was_linear1: + flag_tracker = flag_tracker.accumulate( + work_tile_info.phase, + self._epi_fc1_batch, + (gmem_fc1_done_counter.iterator + cur_fc1_counter_slot).toint(), + ) + else: + _fire_fc2_counter: cutlass.Constexpr = ( + (self._token_back_by_dispatch or self._combine_mxfp8) + and token_comm_args is not None + ) + if cutlass.const_expr(_fire_fc2_counter): + # Fence before (deferred) counter release: make the fc2 + # pool-output STG writes device-visible. + cute.arch.fence_acq_rel_gpu() + fc2_flag_addr = ( + token_comm_args.fc2_done_counter.iterator + cur_fc2_expert_idx + ).toint() + else: + fc2_flag_addr = Int64(0) + no_fire: cutlass.Constexpr = not _fire_fc2_counter + flag_tracker = flag_tracker.accumulate( + work_tile_info.phase, + self._epi_fc2_batch, + fc2_flag_addr, + no_fire, + ) + + flag_tracker.fire() diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_extension.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_extension.py new file mode 100644 index 000000000..4dfd4cd29 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_extension.py @@ -0,0 +1,232 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Scheduling adapter for the MXFP8 GLU FC12 kernel.""" + +import dataclasses +from typing import ClassVar, List, Literal, Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass._mlir import ir +from cutlass.cute.typing import Pointer +from cutlass.cutlass_dsl import Int32, extract_mlir_values, new_from_mlir_values +from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF + +from ......helpers.dsl_helpers import spin_peek, spin_wait +from .....schedulers.fc12_mapping import BlockPhase, NonSwapAbFc12WorkTileInfo, peek_ready_bit + + +# Forward GLU tensor roles: FC1 activation is token-indexed (M), FC1 weight is +# expert-indexed; "d"/"sfd" are the FC1 fp8 output + E8M0 plane, "c" the raw gate/up. +TensorRole = Literal[ + "fc1_activation", + "fc1_weight", + "fc1_activation_sf", + "fc1_weight_sf", + "c", + "d", + "sfd", + "topk", + "fc2_activation", + "fc2_activation_sf", + "fc2_weight", + "fc2_weight_sf", +] + + +@cute.jit +def _rewrite_tensor_shape(tensor: cute.Tensor, new_shape: Tuple) -> cute.Tensor: + return cute.make_tensor(tensor.iterator, cute.make_layout(new_shape, stride=tensor.stride)) + + +@dataclasses.dataclass(frozen=True) +class GluMxFp8Fc12SchedExtension: + """Kernel-owned work-tile preparation and GMEM view adapter (non-swap MXFP8).""" + + work_tile_type: ClassVar[type] = NonSwapAbFc12WorkTileInfo + + sf_vec_size: int + fc1_done_counter_pointer: Pointer + fc2_spin_threshold: Int32 + fc1_ready_counter_pointer: Optional[Pointer] = None + cluster_m: int = 1 + + def __post_init__(self) -> None: + if self.sf_vec_size <= 0: + raise ValueError(f"sf_vec_size must be positive, got {self.sf_vec_size}.") + if self.cluster_m <= 0: + raise ValueError(f"cluster_m must be positive, got {self.cluster_m}.") + object.__setattr__(self, "fc2_spin_threshold", Int32(self.fc2_spin_threshold)) + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + values.extend(extract_mlir_values(self.fc1_done_counter_pointer)) + values.extend(extract_mlir_values(self.fc2_spin_threshold)) + if self.fc1_ready_counter_pointer is not None: + values.extend(extract_mlir_values(self.fc1_ready_counter_pointer)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GluMxFp8Fc12SchedExtension": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + fc1_done_counter_pointer = rebuild(self.fc1_done_counter_pointer) + fc2_spin_threshold = rebuild(self.fc2_spin_threshold) + fc1_ready_counter_pointer = ( + rebuild(self.fc1_ready_counter_pointer) if self.fc1_ready_counter_pointer is not None else None + ) + if value_index != len(values): + raise ValueError( + f"GluMxFp8Fc12SchedExtension MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return type(self)( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter_pointer, + fc2_spin_threshold=fc2_spin_threshold, + fc1_ready_counter_pointer=fc1_ready_counter_pointer, + cluster_m=self.cluster_m, + ) + + @cute.jit + def _counter_slot(self, work_tile: NonSwapAbFc12WorkTileInfo) -> Int32: + # Cluster-granular token-block slot: dispatch_pull increments one counter per + # cluster-level token block, so M-direction tiles fold by cluster_m. + return work_tile.cumulative_token_block_count + work_tile.tile_m_idx // Int32(self.cluster_m) + + @cute.jit + def prepare_work_tile(self, work_tile: NonSwapAbFc12WorkTileInfo) -> NonSwapAbFc12WorkTileInfo: + """Pack kernel readiness observations into the published tile flags.""" + phase_and_flags = work_tile.phase_and_flags + if work_tile.is_valid_tile: + counter_slot = self._counter_slot(work_tile) + is_fc1 = work_tile.phase == Int32(BlockPhase.Linear1) + is_fc2 = work_tile.phase == Int32(BlockPhase.Linear2) + + if cutlass.const_expr(self.fc1_ready_counter_pointer is not None): + if is_fc1: + counter_pointer = self.fc1_ready_counter_pointer + counter_slot + peek_flag = Int32(0) + if spin_peek(counter_pointer, lambda value: value >= work_tile.valid_tokens_in_cluster_tile): + peek_flag = Int32(peek_ready_bit) + phase_and_flags = work_tile.phase_and_flags | peek_flag + + if is_fc2: + counter_pointer = self.fc1_done_counter_pointer + counter_slot + peek_flag = Int32(0) + if spin_peek(counter_pointer, lambda value: value >= self.fc2_spin_threshold): + peek_flag = Int32(peek_ready_bit) + phase_and_flags = work_tile.phase_and_flags | peek_flag + + return NonSwapAbFc12WorkTileInfo( + expert_idx=work_tile.expert_idx, + tile_m_idx=work_tile.tile_m_idx, + tile_n_idx=work_tile.tile_n_idx, + cumulative_data_physical_row=work_tile.cumulative_data_physical_row, + cumulative_sf_physical_row=work_tile.cumulative_sf_physical_row, + cumulative_token_block_count=work_tile.cumulative_token_block_count, + valid_tokens_in_cta_cluster_tile=work_tile.valid_tokens_in_cta_cluster_tile, + phase_and_flags=phase_and_flags, + ) + + @cute.jit + def wait_for_input(self, work_tile: NonSwapAbFc12WorkTileInfo) -> None: + """Wait until this FC1 input tile's cluster-level token count has arrived.""" + if cutlass.const_expr(self.fc1_ready_counter_pointer is not None): + counter_pointer = self.fc1_ready_counter_pointer + self._counter_slot(work_tile) + spin_wait( + counter_pointer, + lambda value: value >= work_tile.valid_tokens_in_cluster_tile, + peek_status=work_tile.peek_ready, + ) + + @cute.jit + def get_gmem_tensor( + self, + tensor_name: TensorRole, + gmem_tensor_in_moe_view: cute.Tensor, + work_tile_info: NonSwapAbFc12WorkTileInfo, + ) -> Tuple[cute.Tensor, Optional[Pointer]]: + """Phase-invariant GMEM slice for the operands.""" + expert_idx = work_tile_info.expert_idx + data_token_offset = work_tile_info.cumulative_data_physical_row + sf_token_offset = work_tile_info.cumulative_sf_physical_row + + shape = gmem_tensor_in_moe_view.shape + stride = gmem_tensor_in_moe_view.stride + c1 = cutlass.Int32(1) + sf_vec_size = self.sf_vec_size + + if cutlass.const_expr(tensor_name == "fc1_activation"): + real = cute.domain_offset((data_token_offset, 0, 0), gmem_tensor_in_moe_view) + return (_rewrite_tensor_shape(real, (shape[0], shape[1], c1)), None) + + elif cutlass.const_expr(tensor_name == "fc1_weight"): + real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view) + return (_rewrite_tensor_shape(real, (shape[0], shape[1], c1)), None) + + elif cutlass.const_expr(tensor_name == "fc1_activation_sf"): + real = cute.domain_offset((sf_token_offset, 0, 0), gmem_tensor_in_moe_view) + per_expert_shape = (shape[0], shape[1], c1) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size) + real = cute.make_tensor(real.iterator, cute.make_layout(sf_layout.shape, stride=stride)) + return (real, None) + + elif cutlass.const_expr(tensor_name == "fc1_weight_sf"): + real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view) + per_expert_shape = (shape[0], shape[1], c1) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size) + real = cute.make_tensor(real.iterator, cute.make_layout(sf_layout.shape, stride=stride)) + return (real, None) + + elif cutlass.const_expr(tensor_name == "c"): + # Raw fc1 accumulator output (gate+up FP32, pre-SwiGLU): token-indexed. + real = cute.domain_offset((data_token_offset, 0, 0), gmem_tensor_in_moe_view) + return (_rewrite_tensor_shape(real, (shape[0], shape[1], c1)), None) + + elif cutlass.const_expr(tensor_name == "d"): + real = cute.domain_offset((data_token_offset, 0, 0), gmem_tensor_in_moe_view) + return (_rewrite_tensor_shape(real, (shape[0], shape[1], c1)), None) + + elif cutlass.const_expr(tensor_name == "sfd"): + real = cute.domain_offset((sf_token_offset, 0, 0), gmem_tensor_in_moe_view) + per_expert_shape = (shape[0], shape[1], c1) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size) + real = cute.make_tensor(real.iterator, cute.make_layout(sf_layout.shape, stride=stride)) + return (real, None) + + elif cutlass.const_expr(tensor_name == "topk"): + real = cute.domain_offset((data_token_offset,), gmem_tensor_in_moe_view) + return (real, None) + + elif cutlass.const_expr(tensor_name == "fc2_activation"): + real = cute.domain_offset((data_token_offset, 0, 0), gmem_tensor_in_moe_view) + return (_rewrite_tensor_shape(real, (shape[0], shape[1], c1)), None) + + elif cutlass.const_expr(tensor_name == "fc2_activation_sf"): + real = cute.domain_offset((sf_token_offset, 0, 0), gmem_tensor_in_moe_view) + per_expert_shape = (shape[0], shape[1], c1) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size) + real = cute.make_tensor(real.iterator, cute.make_layout(sf_layout.shape, stride=stride)) + return (real, None) + + elif cutlass.const_expr(tensor_name == "fc2_weight"): + real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view) + return (_rewrite_tensor_shape(real, (shape[0], shape[1], c1)), None) + + elif cutlass.const_expr(tensor_name == "fc2_weight_sf"): + real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view) + per_expert_shape = (shape[0], shape[1], c1) + sf_layout = tile_atom_to_shape_SF(per_expert_shape, sf_vec_size) + real = cute.make_tensor(real.iterator, cute.make_layout(sf_layout.shape, stride=stride)) + return (real, None) + + raise ValueError(f"Unknown tensor_name: {tensor_name!r}.") + + +__all__ = ["GluMxFp8Fc12SchedExtension", "TensorRole"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_kernel.py new file mode 100644 index 000000000..73f42e30b --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_kernel.py @@ -0,0 +1,2372 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +""" +Fused fc1+fc2 GLU MXFP8 MegaMoE kernel for SM100. +""" + +import dataclasses +from typing import Any, Literal, Optional, Tuple, Type, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, tcgen05, OperandMajorMode +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +import cutlass.utils.rubin_helpers as sm107_utils +from cutlass.cute.nvgpu.tcgen05 import CollectorOp + + +from ..helpers.constants import SupportedMmaTileM, SupportedMmaTileN +from ......helpers.iket_compat import iket +from ......api import ImplDesc, KernelClass, ProblemDesc, StaticOrRuntimeIntegerType +from ......helpers.device_workspace import DeviceWorkspace +from ......helpers.smem_workspace import SmemWorkspace +from ......helpers.dsl_helpers import spin_wait +from ......quant_def import CombineFormat, QuantKind +from ......communication.nvlink_domain.token_comm import TokenCommArgs +from .....schedulers import BlockPhase +from .....schedulers.base import WorkIdAcquisitionMode +from .....schedulers.fc12_scheduler import BlackwellFusedFc12Scheduler +from .glu_mxfp8_fc12_epilogue import GluMxfp8Epilogue +from .glu_mxfp8_fc12_extension import GluMxFp8Fc12SchedExtension + + +@dataclasses.dataclass(frozen=True) +class _EpilogueCommView: + """The exact fields the GLU epilogue reads for cross-rank FC2 routing. + + Built inside the device kernel (so it carries no MLIR-marshaled scalars -- its fields + are already-traced device views) and passed to ``epilogue.run(token_comm_args=...)`` on + the MegaMoE path. Bridges next's Router-push component (``token_src_metadata`` / + ``fc2_done_counter`` come from ``TokenComm`` accessors; ``combine_output`` is the symmetric + ``pre_reduced_activation``) to the epilogue's ``Fc2OutputDest`` peer-store expectations. + """ + + token_src_metadata: Any + combine_output: Any + peer_rank_ptr_mapper: Any + fc2_output_sf: Any = None + fc2_done_counter: Any = None + fc2_output_workspace: Any = None + + +# ============================================================================= +# Sm107Mxfp8GluFc12Kernel +# ============================================================================= + +class Sm107Mxfp8GluFc12Kernel: + + # SMEM budget for all "non-problem-tensor" buffers (mbarriers, sched + # work-tile buffer, TMEM allocator state). Reserved at host side in + # ``_compute_stages``. Bump if ``SharedStorage`` over-allocates SMEM. + _SmemMiscBudget = 1024 + + # Supported (ab_dtype, sf_vec_size) pairings. + # MXFP8 → Float8E4M3FN / Float8E5M2 + sf_vec_size=32 (FP8-E8M0 scales, MmaMXF8Op) + VALID_AB_DTYPE_SF_SIZE: dict = { + 32: (cutlass.Float8E4M3FN, cutlass.Float8E5M2,), + } + + # Interleave granularity for gate and up in SwiGLU / GeGlu + GateUpInterleave: int = 32 + + def __init__( + self, + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mnk: Tuple[int, int, int], + use_2cta_instrs: bool, + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + load_balance_mode: Literal["static", "atomic_counter"] = "static", + static_expert_shape: Optional[Tuple[int, int, int]] = None, + force_static_sched: bool = True, + clc_bundle_size: Optional[int] = None, + num_sched_stages: Optional[int] = None, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + sf_vec_size: int = 32, + ab_dtype: Type[cutlass.Numeric] = cutlass.Float4E2M1FN, + fc2_in_kernel_topk_reduce: bool = False, + token_back_by_dispatch: bool = False, + epi_flag_batch: Tuple[int, int] = (1, 1), + gate_up_clamp: Optional[float] = None, + apply_topk_in_fc1: bool = False, + generate_c: bool = False, + use_stg_fc1: bool = False, + act_func: str = "swiglu", + combine_format: Optional[Any] = None, + fc2_use_bulk: bool = False, + fc2_tma_stages: Optional[int] = None, + ) -> None: + if not force_static_sched: + raise NotImplementedError( + "v1 only implements force_static_sched=True (lean 7-warp). " + "Dynamic CLC (force_static_sched=False) is future work." + ) + + # Validate (ab_dtype, sf_vec_size) pairing. + if sf_vec_size in self.VALID_AB_DTYPE_SF_SIZE: + valid_ab = self.VALID_AB_DTYPE_SF_SIZE[sf_vec_size] + if ab_dtype not in valid_ab: + raise ValueError( + f"ab_dtype={ab_dtype.__name__} is not valid for " + f"sf_vec_size={sf_vec_size}. " + f"Expected one of: {[t.__name__ for t in valid_ab]}." + ) + else: + valid_sf_vec_sizes = tuple(self.VALID_AB_DTYPE_SF_SIZE) + raise NotImplementedError( + f"sf_vec_size must be one of {valid_sf_vec_sizes} (MXFP8); got {sf_vec_size}." + ) + + + if load_balance_mode not in ("static", "atomic_counter"): + raise ValueError( + f"load_balance_mode must be 'static' or 'atomic_counter'; " + f"got {load_balance_mode!r}." + ) + if act_func not in ("swiglu", "geglu"): + raise ValueError( + f"act_func must be 'swiglu' or 'geglu'; got {act_func!r}." + ) + if act_func != "swiglu": + raise NotImplementedError( + f"act_func={act_func!r} is not yet implemented; only " + "'swiglu' is currently supported (geglu support is planned)." + ) + + # Only (M=256, N=256) with 2-CTA instructions is validated now. + m, n, _k = mma_tiler_mnk + if (m, n) != (256, 256) or not use_2cta_instrs: + raise ValueError( + "Sm107Mxfp8GluFc12Kernel only supports mma_tiler (M, N) = " + "(256, 256) with use_2cta_instrs=True; " + f"got mma_tiler_mnk={mma_tiler_mnk}, use_2cta_instrs={use_2cta_instrs}." + ) + + # Store ab_dtype so workspace-size helpers can use it without tensors. + self.ab_dtype = ab_dtype + self.c_dtype = cutlass.BFloat16 + self.act_func = act_func + self.combine_format = combine_format + + self.fc2_in_kernel_topk_reduce = fc2_in_kernel_topk_reduce + self.token_back_by_dispatch = token_back_by_dispatch + self.epi_flag_batch = epi_flag_batch + self.apply_topk_in_fc1 = apply_topk_in_fc1 + self.generate_c = generate_c + self.use_stg_fc1 = use_stg_fc1 + self.fc2_use_bulk = fc2_use_bulk + self._fc2_tma_stages_arg = fc2_tma_stages + self.fc2_tma_stages = fc2_tma_stages if fc2_tma_stages is not None else 0 + self.gate_up_clamp = ( + abs(gate_up_clamp) if gate_up_clamp is not None else None + ) + + self.acc_dtype = acc_dtype + self.mma_tiler_mnk = mma_tiler_mnk + self.cluster_shape_mn = (cluster_shape_mnk[0], cluster_shape_mnk[1]) + self.use_2cta_instrs = use_2cta_instrs + self.force_static_sched = force_static_sched + self.static_expert_shape = static_expert_shape + self.clc_bundle_size = clc_bundle_size + self.num_sched_stages = num_sched_stages + + # Fused fc12 sched-side knobs + self.group_hint = group_hint + self.token_padding_block = token_padding_block + self.sf_padding_block = sf_padding_block + self.load_balance_mode = load_balance_mode + + self.sf_vec_size = sf_vec_size + self.arch = "sm_107" + + self._validate_mma_tiler_and_cluster_shape() + self.mma_tiler = mma_tiler_mnk + + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + # Warp specialization (lean 8-warp / 256 thread) + self.occupancy = 1 + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_a_warp_id = 5 + self.tma_b_warp_id = 6 + self.sched_warp_id = 7 + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_a_warp_id, + self.tma_b_warp_id, + self.sched_warp_id, + *self.epilogue_warp_id, + ) + ) + + # NamedBarrier IDs. + # + # Per-subtile rotated-leader scheme lives inside the epilogue; this + # kernel only owns the four reserved IDs and forwards + # them via ``self.epilog_sync_bar_id`` to the epilogue ctor. + # IDs 8 and 9 are reserved for MXFP8 warp-pair absmax exchange. + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + self.epi_subtile_bar_ids = (4, 5, 6, 7) + + self.enable_token_comm: bool = False + self.dispatch_warp_id: Optional[Tuple[int, int, int, int]] = None + self.token_back_warp_id: Optional[Tuple[int, int, int, int]] = None + self.token_back_standalone: bool = False + + self.smem_capacity = utils.get_smem_capacity_in_bytes() + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols( + self.arch + ) + + def _validate_mma_tiler_and_cluster_shape(self) -> None: + """Validate user-provided geometry against v1 fused-fc12 constraints. + + ``mma_tiler_n`` is restricted to {128, 256}. Short-N is handled by + the swap-AB scheduler via subtile-level early-exit. + """ + m, n, k = self.mma_tiler_mnk + cm, cn = self.cluster_shape_mn + + if m not in SupportedMmaTileM: + raise ValueError( + f"mma_tiler M ({m}) must be one of {SupportedMmaTileM}" + ) + + per_cta_m = m // (2 if self.use_2cta_instrs else 1) + if per_cta_m != 128: + raise ValueError( + f"per-CTA mma_tiler M must be 128, got {per_cta_m} " + f"(mma_tiler_m={m}, use_2cta_instrs={self.use_2cta_instrs})" + ) + + if n not in SupportedMmaTileN: + raise ValueError( + f"mma_tiler N ({n}) must be one of {SupportedMmaTileN} in fused fc12 " + f"(N=64 SFB hack is dropped; swap-AB sched handles short-N " + f"via subtile early-exit)." + ) + + sf_k_granularity = self.sf_vec_size * 4 + if k % sf_k_granularity != 0: + raise ValueError( + f"mma_tiler K ({k}) must be a multiple of " + f"sf_vec_size * 4 = {sf_k_granularity}" + ) + + if cm % (2 if self.use_2cta_instrs else 1) != 0: + raise ValueError( + f"cluster_shape M ({cm}) must be even when use_2cta_instrs=True" + ) + + is_pow2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if cm * cn > 16 or not is_pow2(cm) or not is_pow2(cn) or cm > 4 or cn > 4: + raise ValueError( + f"Invalid cluster_shape ({cm}, {cn}): each dim must be " + f"a power of 2 and <= 4, product must be <= 16" + ) + + if cn > 2: + raise NotImplementedError( + f"cluster_n={cn} is not supported yet (deadlocks in the " + f"mainloop, likely the 4-way scale-factor multicast). " + f"cluster_n in {{1, 2}} is supported and verified." + ) + + if cm > 2: + raise NotImplementedError( + f"cluster_m={cm} is not supported yet (residual fc1->fc2 race " + f"gives ~3-4% mismatch). cluster_m in {{1, 2}} is supported " + f"and verified." + ) + + if cn > 1 and self.static_expert_shape is not None: + cta_tile_n = n # N is not split across the 2-CTA (M) pair + cluster_tile_n = cta_tile_n * cn + _experts, intermediate_gateup, hidden = self.static_expert_shape + if intermediate_gateup % cluster_tile_n != 0: + raise ValueError( + f"cluster_n={cn}: fc1 intermediate_gateup " + f"({intermediate_gateup}) must be a multiple of " + f"cta_tile_n * cluster_n (= {cta_tile_n} * {cn} = " + f"{cluster_tile_n}) to avoid ragged N-peers. Ragged-N " + f"(per-peer N-store predication) is not yet supported." + ) + if hidden % cluster_tile_n != 0: + raise ValueError( + f"cluster_n={cn}: fc2 hidden ({hidden}) must be a multiple " + f"of cta_tile_n * cluster_n (= {cta_tile_n} * {cn} = " + f"{cluster_tile_n}) to avoid ragged N-peers. Ragged-N " + f"(per-peer N-store predication) is not yet supported." + ) + + def _create_tiled_mmas(self) -> Tuple[cute.TiledMma, cute.TiledMma]: + common = ( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + ) + # Rubin: the SM107 blockscaled FP8 MMA op instruction + tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma( + *common, self.cta_group, + (*self.mma_inst_shape_mn, 64), + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma( + *common, tcgen05.CtaGroup.ONE, + (*self.mma_inst_shape_mn_sfb, 64), + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + return tiled_mma, tiled_mma_sfb + + def _build_scheduler( + self, *, expert_cnt, intermediate_gateup, hidden_dim, launch_cluster_count + ) -> None: + """Construct FC12 scheduler and its SMEM/device workspaces.""" + work_id_mode = "grid_stride" if self.load_balance_mode == "static" else "atomic_counter" + # Consumer group = every warp that calls ``consume_work`` (tma_a, tma_b, mma, epilogue). + num_scheduler_consumer_threads = 32 * (len(self.epilogue_warp_id) + 3) + if self.static_expert_shape is not None: + expert_cnt, intermediate_gateup, hidden_dim = self.static_expert_shape + problem_desc = ProblemDesc( + { + "expert_count": expert_cnt, + "intermediate_gateup_size": intermediate_gateup, + "hidden_size": hidden_dim, + } + ) + impl_desc = ImplDesc( + { + "num_scheduler_consumer_threads": num_scheduler_consumer_threads, + "mma_tiler_mnk": self.mma_tiler, + "cluster_shape_mn": self.cluster_shape_mn, + "use_2cta_instrs": self.use_2cta_instrs, + "hint": self.group_hint, + "token_padding_block": self.token_padding_block, + "sf_padding_block": self.sf_padding_block, + "work_id_mode": work_id_mode, + "is_swap_ab": False, + "launch_cluster_count": launch_cluster_count, + } + ) + self.scheduler = BlackwellFusedFc12Scheduler(problem_desc, impl_desc) + + sched_smem_ws = SmemWorkspace() + self.scheduler.register_smem_regions(sched_smem_ws) + sched_smem_ws.finalize(max_bytes=self.smem_capacity) + self.sched_smem_ws = sched_smem_ws + + sched_device_ws = DeviceWorkspace() + self.scheduler.register_device_workspace(sched_device_ws) + sched_device_ws.finalize() + self.sched_device_ws = sched_device_ws + + def _setup_attributes(self) -> None: + """Set up MMA / cluster / tile shapes, SMEM layouts, stage counts. + + The fc12 path shares ``mma_tiler_mnk`` and SMEM layouts across phases. + """ + if self.enable_token_comm: + self.dispatch_warp_id = (8, 9, 10, 11) + num_token_back_warps = ( + len(self.token_back_warp_id) if self.token_back_standalone else 0 + ) + self.threads_per_cta = 32 * ( + len(self.epilogue_warp_id) + + 1 # mma + + 1 # tma_a + + 1 # tma_b + + 1 # sched + + len(self.dispatch_warp_id) + + num_token_back_warps + ) + + self.mma_inst_shape_mn = (self.mma_tiler[0], self.mma_tiler[1]) + self.mma_inst_shape_mn_sfb = ( + self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape_mn[1], 128), + ) + + tiled_mma, tiled_mma_sfb = self._create_tiled_mmas() + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + assert self.mma_tiler[2] % mma_inst_shape_k == 0, ( + f"mma_tiler K ({self.mma_tiler[2]}) must be a multiple of " + f"MMA instruction K ({mma_inst_shape_k})" + ) + + # SFB-specific tiler: rounded-up MN; same K as main tiler. + self.mma_tiler_sfb = ( + self.mma_inst_shape_mn_sfb[0], + self.mma_inst_shape_mn_sfb[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + # Multicast CTA counts + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + # Epilogue owns all epi-side decisions (acc stages, subtile dispatch, + # TMA commit/drain, and piggyback red.add). + _epi_common = dict( + mma_tiler_mnk=self.mma_tiler, + cluster_shape_mn=self.cluster_shape_mn, + use_2cta_instrs=self.use_2cta_instrs, + sf_vec_size=self.sf_vec_size, + fc1_output_dtype=self.fc1_output_dtype, + fc1_output_layout=self.fc1_output_layout, + acc_dtype=self.acc_dtype, + epilog_sync_bar_id=self.epilog_sync_bar_id, + epilogue_warp_ids=self.epilogue_warp_id, + static_expert_shape=self.static_expert_shape, + fc2_in_kernel_topk_reduce=self.fc2_in_kernel_topk_reduce, + token_back_by_dispatch=self.token_back_by_dispatch, + epi_flag_batch=self.epi_flag_batch, + glu_clamp=self.gate_up_clamp, + apply_topk_in_fc1=self.apply_topk_in_fc1, + generate_c=self.generate_c, + use_stg_fc1=self.use_stg_fc1, + combine_format=getattr(self, "combine_format", None), + act_func=self.act_func, + fc2_use_bulk=self.fc2_use_bulk, + fc2_tma_stages=self._fc2_tma_stages_arg, + ) + self.epilogue = GluMxfp8Epilogue(**_epi_common) + + if self.num_sched_stages is None: + self.num_sched_stages = 2 + + self.num_d_stage = self.epilogue.subtile_cnt + # fc1 output (fp8 quantised) SMEM — always present. + d_bytes_total = self.epilogue.bytes_per_stage * self.num_d_stage + # Raw gate+up SMEM (BF16, ping-pong) — only when generate_c=True. + c_bytes_total = 0 + if self.generate_c: + from .glu_mxfp8_fc12_epilogue import Fc1CTMAStages + self.num_c_raw_stage = Fc1CTMAStages + c_bytes_total += self.epilogue.c_bytes_per_stage * self.num_c_raw_stage + else: + self.num_c_raw_stage = 0 + + # FC2 TMASTG staging SMEM + fc2_tma_bytes_total = ( + self.epilogue.fc2_tma_staging_bytes + self.epilogue.fc2_reduce_staging_bytes + ) + + ( + self.num_acc_stage, + self.num_a_stage, + self.num_b_stage, + self.num_sched_stages, + ) = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.sf_dtype, + self.sf_vec_size, + c_bytes_total + d_bytes_total + fc2_tma_bytes_total, + self.smem_capacity, + self.occupancy, + self.num_sched_stages, + ) + + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_a_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_b_stage, + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_a_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_b_stage, + ) + self.d_smem_layout_staged = self.epilogue.staged_smem_layout( + self.num_d_stage, + ) + # Raw gate+up SMEM layout (only meaningful when generate_c=True; pass as + # dummy to kernel when False). + if self.generate_c: + self.c_smem_layout_staged = self.epilogue.staged_c_smem_layout( + self.num_c_raw_stage + ) + else: + self.c_smem_layout_staged = None + + # Read epilogue's accumulator and scale-factor sizing decisions. + self.num_acc_pipeline_stages = self.epilogue.num_acc_pipeline_stages + self.num_acc_stage = self.epilogue.num_acc_stage + self.num_sfa_tmem_cols = self.epilogue.num_sfa_tmem_cols + self.num_sfb_tmem_cols = self.epilogue.num_sfb_tmem_cols + self.num_accumulator_tmem_cols = self.epilogue.num_accumulator_tmem_cols + + # TMA load bytes per stage (A + B + SFA + SFB). + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + self.atom_thr_size = atom_thr_size # store as Python int for use in @cute.kernel + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_a_bytes = (a_copy_size + sfa_copy_size) * atom_thr_size + self.num_tma_load_b_bytes = (b_copy_size + sfb_copy_size) * atom_thr_size + + def _smem_misc_budget_bytes(self) -> int: + """Per-CTA SMEM reserved outside the ABC-stage pipeline.""" + return self._SmemMiscBudget + + def _compute_stages( + self, + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_bytes_total: int, + smem_capacity: int, + occupancy: int, + num_sched_stages: int, + ) -> Tuple[int, int, int]: + """Compute stage counts for ACC, AB+SF, and scheduler.""" + num_acc_stage = 2 + + a_smem_layout_staged_one = sm100_utils.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1, + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 1, + ) + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, mma_tiler_mnk, sf_vec_size, 1, + ) + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, mma_tiler_mnk, sf_vec_size, 1, + ) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_staged_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + b_bytes_per_stage = ( + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + + fixed_overhead = ( + self._smem_misc_budget_bytes() + c_bytes_total + ) + + num_ab_stage = ( + smem_capacity // occupancy - fixed_overhead + ) // ab_bytes_per_stage + num_a_stage = num_ab_stage + num_b_stage = num_ab_stage + + smem_per_cta = smem_capacity // occupancy + unused_smem = smem_per_cta - fixed_overhead - num_ab_stage * ab_bytes_per_stage + if unused_smem > b_bytes_per_stage: + num_b_stage = num_b_stage + 1 + unused_smem = unused_smem - b_bytes_per_stage + print( + f"[fc12 stages] num_ab_stage={num_a_stage, num_b_stage} " + f"ab_bytes_per_stage={ab_bytes_per_stage} " + f"num_acc_stage={num_acc_stage} " + f"misc_budget={self._smem_misc_budget_bytes()} " + f"c_bytes_total={c_bytes_total} " + f"smem_cap={smem_capacity} " + f"unused_smem={unused_smem}" + ) + + return num_acc_stage, num_a_stage, num_b_stage, num_sched_stages + + def get_workspace_size_in_bytes( + self, + fc1_activation_tensor, + fc1_weight_tensor, + ) -> int: + """Compute opaque workspace size for one fused fc1+fc2 launch.""" + sf_padding_block = self.sf_padding_block + sf_vec_size = self.sf_vec_size + + mma_tiler_n = self.mma_tiler_mnk[1] + + data_total_rows, _hidden = fc1_activation_tensor.shape + experts, _hidden_w, intermediate_gateup = fc1_weight_tensor.shape + intermediate_downproj = intermediate_gateup // 2 + + # Conservative upper bound for sf_total_rows. + sf_total_rows_upper = data_total_rows + experts * sf_padding_block + + fc1_output_bytes = ( + data_total_rows * intermediate_downproj * self.ab_dtype.width // 8 + ) + + # fc1_output_sf sf_vec_size matches the kernel's sf_vec_size. + fc1_out_sf_vec_size = self.sf_vec_size + sf_block_cols = ( + (intermediate_downproj // fc1_out_sf_vec_size) + 3 + ) // 4 * 4 + fc1_output_sf_bytes = sf_total_rows_upper * sf_block_cols + + # fc1_done_counter: one Int32 per global token block, plus expert slack. + counter_slots_upper = ( + (data_total_rows + mma_tiler_n - 1) // mma_tiler_n + + experts + ) + fc1_done_counter_bytes = counter_slots_upper * 4 + + # load_balance_counter: Int32 scalar. + if self.load_balance_mode == "atomic_counter": + load_balance_counter_bytes = 4 + else: + load_balance_counter_bytes = 0 + + total = ( + fc1_output_bytes + + fc1_output_sf_bytes + + fc1_done_counter_bytes + + load_balance_counter_bytes + ) + + # 128B align (TMA tensor base address alignment requirement). + alignment = 128 + total = ((total + alignment - 1) // alignment) * alignment + return total + + def mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """SMEM → TMEM tiled copy + partition for SFA / SFB.""" + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor( + tiled_copy_s2t, tCsSF_compact_s2t_ + ) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + def token_comm_extra_smem_storage_class(self) -> type: + """Return a ``@cute.struct`` for dispatch-warp SMEM, or None.""" + return None + + def token_comm_hook_fc1_ready_counter_ptr(self, token_comm_args): + """Return dispatch->fc1 release counter pointer, or None (lean: disabled).""" + return None + + def sched_ext_fc1_peek_threshold(self) -> int: + """Return the fc1 ready-counter peek threshold for GluMxFp8Fc12SchedExtension. + + Must match the spin threshold in ``token_comm_hook_fc1_tma_b_predispatch_spin`` + so that an early peek hit does not skip the spin and expose stale pool rows. + Default 0 → use ``valid_tokens_in_tile`` (no cluster, base class behaviour). + MegaMoE overrides to return ``cluster_tile_tokens`` to match the cluster spin. + """ + return 0 + + def sched_ext_fc1_counter_cumul_scale(self) -> int: + """Return the scale factor for the fc1 ready-counter slot formula. + + Slot = scale * (cumul + fc1_counter_index) + tile_m_idx % scale. + Default 1 = cluster-level granularity (slot = cumul + cluster_token_block_idx). + A MegaMoE subclass can override to cluster_m for per-CTA granularity. + """ + return 1 + + @cute.jit + def token_comm_hook_sched_warp_pre_init_wait(self, token_comm_args): + """Sched warp: wait for dispatch barrier before reading sizes. No-op base.""" + pass + + @cute.jit + def token_comm_hook_fc1_tma_b_predispatch_spin(self, token_comm_args, work_tile_info): + """TMA warp: spin until dispatch-pulled tokens are resident. No-op base.""" + pass + + @cute.jit + def token_comm_hook_dispatch_warp_body( + self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx, + ): + """Body for dispatch warps 8-11 (MegaMoE-only). No-op base.""" + pass + + @cute.jit + def token_comm_hook_token_back_warp_body( + self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx, + ): + """Body for standalone token-back warps 12-15 (MegaMoE-only). No-op base.""" + pass + + @cute.jit + def token_comm_hook_kernel_tail(self, token_comm_args, *, warp_idx, lane_idx, tidx): + """All-warp kernel tail (NVLink release, etc.). No-op base.""" + pass + + @cute.jit + def __call__( + self, + activation: cute.Tensor, # (token_sum_padded, hidden) + fc1_weight: cute.Tensor, # (experts, hidden, intermediate_gateup) + activation_sf: cute.Tensor, # (token_sum_padded_sf, hidden / sf_vec_size) + fc1_weight_sf: cute.Tensor, # (experts, intermediate_gateup_padded * hidden / sf_vec_size) + fc1_output: cute.Tensor, # (token_sum_padded, intermediate_downproj) + fc1_output_sf: cute.Tensor, # (token_sum_padded_sf, intermediate_downproj / sf_vec_size) + fc2_weight: cute.Tensor, # (experts, intermediate_downproj, hidden) + fc2_weight_sf: cute.Tensor, # (experts, hidden_padded * intermediate_downproj / sf_vec_size) + fc2_output: cute.Tensor, # (token_sum_padded, hidden) BFloat16, hidden stride-1 + topk_scores: cute.Tensor, # (token_sum_padded,) Float32 + fc1_done_counter: cute.Tensor, # (max_token_block_per_rank,) Int32 + offs: Optional[cute.Tensor] = None, # (experts,) Int32 cumulative end offsets + max_active_clusters: cutlass.Constexpr = None, + stream: cuda.CUstream = None, + norm_const_tensor: Optional[cute.Tensor] = None, + global_activation_sf: Optional[cute.Tensor] = None, + global_fc1_weight_sf: Optional[cute.Tensor] = None, + load_balance_counter: Optional[cute.Tensor] = None, + expert_token_sizes: Optional[cute.Tensor] = None, + token_comm_args=None, + fc1_c: Optional[cute.Tensor] = None, + # ── Per-rank FC12 overflow output ──────────────────────────────── + overflow_flag: cute.Tensor = None, + fc2_output_sf: Optional[cute.Tensor] = None, + mega_peer_rank_ptr_mapper=None, + mega_local_rank: Optional[cutlass.Int32] = None, + mega_local_workspace: Optional[cute.Pointer] = None, + mega_shared_workspace: Optional[cute.Pointer] = None, + mega_activation: Optional[cute.Tensor] = None, + mega_activation_sf: Optional[cute.Tensor] = None, + mega_pre_reduced_activation: Optional[cute.Tensor] = None, + mega_pre_reduced_activation_sf: Optional[cute.Tensor] = None, + ) -> None: + """Launch the fused fc1+fc2 GLU MXFP8 kernel.""" + + if cutlass.const_expr(self.static_expert_shape is not None): + ( + experts_static, + intermediate_gateup_static, + hidden_static, + ) = self.static_expert_shape + intermediate_downproj_static = intermediate_gateup_static // 2 + + fc1_weight = cute.make_tensor( + fc1_weight.iterator, + cute.make_layout( + (experts_static, hidden_static, intermediate_gateup_static), + stride=fc1_weight.stride, + ), + ) + fc2_weight = cute.make_tensor( + fc2_weight.iterator, + cute.make_layout( + (experts_static, intermediate_downproj_static, hidden_static), + stride=fc2_weight.stride, + ), + ) + activation = cute.make_tensor( + activation.iterator, + cute.make_layout( + (activation.shape[0], hidden_static), + stride=activation.stride, + ), + ) + fc1_output = cute.make_tensor( + fc1_output.iterator, + cute.make_layout( + (fc1_output.shape[0], intermediate_downproj_static), + stride=fc1_output.stride, + ), + ) + # fc2_output is 2D (tokens, hidden) on the lean path and 3D + # (max_tokens, topk, hidden) on the MegaMoE path. + if cutlass.const_expr(len(fc2_output.shape) == 3): + fc2_output = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (fc2_output.shape[0], fc2_output.shape[1], hidden_static), + stride=fc2_output.stride, + ), + ) + else: + fc2_output = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (fc2_output.shape[0], hidden_static), + stride=fc2_output.stride, + ), + ) + + # GEMM-domain transform for fc1 phase + c1 = cutlass.Int32(1) + c0 = cutlass.Int32(0) + + # A_gemm (fc1 activations): (tokens_sum, hidden) -> (M=tokens, K=hidden, L=1). + tokens_sum, hidden = activation.shape + activation_gemm = cute.make_tensor( + activation.iterator, + cute.make_layout( + (tokens_sum, hidden, 1), + stride=(activation.stride[0], activation.stride[1], 0), + ), + ) + + # B_gemm (fc1 weights): (experts, hidden, intermediate_gateup) with hidden stride-1 (K-major) + # -> (N=intermediate_gateup, K=hidden, L=experts). + experts, hidden_b, intermediate_gateup = fc1_weight.shape + fc1_weight_gemm = cute.make_tensor( + fc1_weight.iterator, + cute.make_layout( + (intermediate_gateup, hidden_b, experts), + stride=(fc1_weight.stride[2], fc1_weight.stride[1], fc1_weight.stride[0]), + ), + ) + + # D_gemm is a user-view output tensor; epilogue owns its store path. + intermediate_downproj = fc1_output.shape[1] + fc1_output_gemm = cute.make_tensor( + fc1_output.iterator, + cute.make_layout( + (tokens_sum, intermediate_downproj, 1), + stride=(fc1_output.stride[0], fc1_output.stride[1], 0), + ), + ) + + # SFA / SFB scale tensors (atom-tiled) — fc1 phase. + # SFA (mma M-side) = activation_sf (activation scales, A-side) + # SFB (mma N-side) = fc1_weight_sf (weight scales, B-side) + tokens_sum_padded = activation_sf.shape[0] + hidden_padded = activation_sf.shape[1] * self.sf_vec_size + activation_sf_gemm = cute.make_tensor( + activation_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (tokens_sum_padded, hidden_padded, 1), self.sf_vec_size + ), + ) + intermediate_gateup_padded_mul_hidden_padded = fc1_weight_sf.shape[1] + intermediate_gateup_padded = ( + intermediate_gateup_padded_mul_hidden_padded * self.sf_vec_size + ) // hidden_padded + fc1_weight_sf_gemm = cute.make_tensor( + fc1_weight_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (intermediate_gateup_padded, hidden_padded, experts), + self.sf_vec_size, + ), + ) + + # ── GEMM-domain transform for fc2 phase ── + # + # fc2 roles: M=hidden, N=tokens_sum, K=intermediate_downproj. + + # A_gemm (fc2 weights): (experts, intermediate_downproj, hidden) + # -> (M=hidden, K=intermediate_downproj, L=experts). + experts2, intermediate_downproj_b2, hidden_b2 = fc2_weight.shape + fc2_weight_gemm = cute.make_tensor( + fc2_weight.iterator, + cute.make_layout( + (hidden_b2, intermediate_downproj_b2, experts2), + stride=(fc2_weight.stride[2], fc2_weight.stride[1], fc2_weight.stride[0]), + ), + ) + + if cutlass.const_expr(len(fc2_output.shape) == 3): + fc2_hidden_out = fc2_output.shape[2] + fc2_output_gemm = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (fc2_output.shape[0], fc2_hidden_out, c1), + stride=(fc2_output.stride[0], fc2_output.stride[2], c0), + ), + ) + else: + fc2_hidden_out = fc2_output.shape[1] + fc2_output_gemm = cute.make_tensor( + fc2_output.iterator, + cute.make_layout( + (tokens_sum, fc2_hidden_out, c1), + stride=(fc2_output.stride[0], fc2_output.stride[1], c0), + ), + ) + + # SFA / SFB for fc2: + # SFA (mma M-side) = fc2_weight_sf (fc2 weight scales, sf_vec_size) + # SFB (mma N-side) = fc1_output_sf (fc1 epilogue SFs, uses sf_vec_size) + # Both paths produce SFs with self.sf_vec_size: + fc1_out_sf_vec_size = self.sf_vec_size + tokens_sum_padded_sf = fc1_output_sf.shape[0] + intermediate_downproj_padded = fc1_output_sf.shape[1] * fc1_out_sf_vec_size + fc1_output_sf_gemm_for_fc2_load = cute.make_tensor( + fc1_output_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (tokens_sum_padded_sf, intermediate_downproj_padded, 1), + fc1_out_sf_vec_size, + ), + ) + + hidden_padded_fc2_mul_intermediate_downproj_padded = fc2_weight_sf.shape[1] + hidden_padded_fc2 = ( + hidden_padded_fc2_mul_intermediate_downproj_padded * self.sf_vec_size + ) // intermediate_downproj_padded + fc2_weight_sf_gemm = cute.make_tensor( + fc2_weight_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (hidden_padded_fc2, intermediate_downproj_padded, experts2), + self.sf_vec_size, + ), + ) + + expert_cnt = experts + hidden_dim = hidden + + # ── Infer dtypes and major modes ── + self.a_dtype: Type[cutlass.Numeric] = activation_gemm.element_type + self.b_dtype: Type[cutlass.Numeric] = fc1_weight_gemm.element_type + self.fc1_output_dtype: Type[cutlass.Numeric] = fc1_output_gemm.element_type + self.sf_dtype: Type[cutlass.Numeric] = activation_sf_gemm.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(activation_gemm).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(fc1_weight_gemm).mma_major_mode() + self.fc1_output_layout = utils.LayoutEnum.from_tensor(fc1_output_gemm) + + self._setup_attributes() + tiled_mma, tiled_mma_sfb = self._create_tiled_mmas() + + # ── fc1 TMA atoms ── + + # TMA load A1 + a_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_fc1_activation, tma_tensor_fc1_activation = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + activation_gemm, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load B1 + b_op = sm100_utils.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_fc1_weight, tma_tensor_fc1_weight = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + fc1_weight_gemm, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load SFA1 + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfa_smem_layout = cute.slice_( + self.sfa_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_fc1_activation_sf, tma_tensor_fc1_activation_sf = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + activation_sf_gemm, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64, + ) + + # TMA load SFB1 + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB( + self.cluster_shape_mn, tiled_mma.thr_id + ) + sfb_smem_layout = cute.slice_( + self.sfb_smem_layout_staged, (None, None, None, 0) + ) + tma_atom_fc1_weight_sf, tma_tensor_fc1_weight_sf = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + fc1_weight_sf_gemm, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Uint64, + ) + + # TMA store for fc1 MXFP8 output. + fc1_output_tma_op = cpasync.CopyBulkTensorTileS2GOp() + tma_atom_fc1_output, tma_tensor_fc1_output = cpasync.make_tiled_tma_atom( + fc1_output_tma_op, + fc1_output_gemm, + self.epilogue.smem_layout_one_stage, + self.epilogue.epi_tile, + ) + + # TMA store for raw fc1 accumulator + if cutlass.const_expr(self.generate_c): + c_gemm = cute.make_tensor( + fc1_c.iterator, + cute.make_layout( + (tokens_sum, intermediate_gateup, 1), + stride=(fc1_c.stride[0], fc1_c.stride[1], 0), + ), + ) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c_gemm, + self.epilogue.c_smem_layout_one_stage, + self.epilogue.epi_tile_c, + ) + else: + tma_atom_c = tma_atom_fc1_output + tma_tensor_c = tma_tensor_fc1_output + + # fc1 SFC GMEM tensor (= fc1_output_sf user view). + fc1_output_sf_gemm = cute.make_tensor( + fc1_output_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (tokens_sum_padded, intermediate_downproj, 1), + self.sf_vec_size, + ), + ) + + # ── fc2 TMA atoms: fc1_output → A-side (M=tokens), fc2_weight → B-side (N=hidden) ── + tma_atom_fc2_activation, tma_tensor_fc2_activation = ( + cute.nvgpu.make_tiled_tma_atom_A( + a_op, + fc1_output_gemm, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + ) + tma_atom_fc2_weight, tma_tensor_fc2_weight = ( + cute.nvgpu.make_tiled_tma_atom_B( + b_op, + fc2_weight_gemm, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + ) + tma_atom_fc2_activation_sf, tma_tensor_fc2_activation_sf = ( + cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + fc1_output_sf_gemm_for_fc2_load, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Uint64, + ) + ) + tma_atom_fc2_weight_sf, tma_tensor_fc2_weight_sf = ( + cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + fc2_weight_sf_gemm, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Uint64, + ) + ) + + # ── Scheduler params + grid + launch ── + if cutlass.const_expr(self.load_balance_mode == "atomic_counter"): + if cutlass.const_expr(load_balance_counter is None): + raise ValueError( + "load_balance_counter must be provided when " + "load_balance_mode == 'atomic_counter'" + ) + load_balance_counter_ptr = load_balance_counter.iterator + else: + load_balance_counter_ptr = None + + # On the MegaMoE path the per-expert sizes come from the Router (device-side), so the + # caller supplies neither offs nor expert_token_sizes. + if cutlass.const_expr(not self.enable_token_comm): + if cutlass.const_expr((offs is None) == (expert_token_sizes is None)): + raise ValueError( + "Exactly one of `offs` / `expert_token_sizes` must be " + "provided; got " + f"offs={'set' if offs is not None else 'None'}, " + f"expert_token_sizes=" + f"{'set' if expert_token_sizes is not None else 'None'}." + ) + + self._build_scheduler( + expert_cnt=expert_cnt, + intermediate_gateup=intermediate_gateup, + hidden_dim=hidden_dim, + launch_cluster_count=max_active_clusters, + ) + grid = self.scheduler.get_grid_shape(max_active_clusters=max_active_clusters) + + # FC2 TMASTG (fc2_use_bulk) store atom (host-side) + if cutlass.const_expr(self.enable_token_comm and self.epilogue.fc2_use_tma): + self._mega_device_workspace.assign_device_members( + mega_local_workspace, mega_shared_workspace + ) + _fc2_pool3d = self.token_comm.fc2_activation_tensor(self._mega_device_workspace) + _fc2_pool2d = cute.make_tensor( + _fc2_pool3d.iterator, + cute.make_layout( + (_fc2_pool3d.shape[0], _fc2_pool3d.shape[2]), + stride=(_fc2_pool3d.stride[0], _fc2_pool3d.stride[2]), + ), + ) + tma_atom_fc2_output, fc2_tma_output = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + _fc2_pool2d, + self.epilogue.fc2_tma_smem_layout_one_stage, + self.epilogue.fc2_tma_tile, + ) + else: + tma_atom_fc2_output = None + fc2_tma_output = None + + self.kernel( + tiled_mma, + tiled_mma_sfb, + # fc1 TMA atoms / tensors (A=activations, B=weights) + tma_atom_fc1_activation, + tma_tensor_fc1_activation, + tma_atom_fc1_weight, + tma_tensor_fc1_weight, + tma_atom_fc1_activation_sf, + tma_tensor_fc1_activation_sf, + tma_atom_fc1_weight_sf, + tma_tensor_fc1_weight_sf, + tma_atom_fc1_output, + tma_tensor_fc1_output, + # fc2 TMA atoms / tensors (fc1_output→A, fc2_weight→B) + tma_atom_fc2_activation, + tma_tensor_fc2_activation, + tma_atom_fc2_weight, + tma_tensor_fc2_weight, + tma_atom_fc2_activation_sf, + tma_tensor_fc2_activation_sf, + tma_atom_fc2_weight_sf, + tma_tensor_fc2_weight_sf, + # GEMM-domain tensors (fc1) + activation_gemm, + fc1_weight_gemm, + fc1_output_gemm, + activation_sf_gemm, + fc1_weight_sf_gemm, + fc1_output_sf_gemm, + # GEMM-domain tensors (fc2) + fc2_weight_gemm, + fc2_output_gemm, + fc2_weight_sf_gemm, + fc1_output_sf_gemm_for_fc2_load, + # topk + cross-phase sync workspace + topk_scores, + fc1_done_counter, + # Scheduling + offs, + expert_token_sizes, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + # SMEM layouts + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.d_smem_layout_staged, + self.c_smem_layout_staged, + tma_atom_c, + tma_tensor_c, + overflow_flag, + token_comm_args, + fc2_output_sf, + mega_peer_rank_ptr_mapper, + mega_local_rank, + mega_local_workspace, + mega_shared_workspace, + mega_activation, + mega_activation_sf, + mega_pre_reduced_activation, + mega_pre_reduced_activation_sf, + tma_atom_fc2_output, + fc2_tma_output, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=self.occupancy, + ) + + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + # fc1 TMA atoms / tensors + tma_atom_fc1_activation_1: cute.CopyAtom, + tma_tensor_fc1_activation_1: cute.Tensor, + tma_atom_weight: cute.CopyAtom, + tma_tensor_weight: cute.Tensor, + tma_atom_fc1_activation_1_sf: cute.CopyAtom, + tma_tensor_fc1_activation_1_sf: cute.Tensor, + tma_atom_fc1_weight_sf: cute.CopyAtom, + tma_tensor_fc1_weight_sf: cute.Tensor, + tma_atom_fc1_output: cute.CopyAtom, + tma_tensor_fc1_output: cute.Tensor, + # fc2 TMA atoms / tensors (fc1_output→A, fc2_weight→B) + tma_atom_fc2_activation: cute.CopyAtom, + tma_tensor_fc2_activation: cute.Tensor, + tma_atom_fc2_weight: cute.CopyAtom, + tma_tensor_fc2_weight: cute.Tensor, + tma_atom_fc2_activation_sf: cute.CopyAtom, + tma_tensor_fc2_activation_sf: cute.Tensor, + tma_atom_fc2_weight_sf: cute.CopyAtom, + tma_tensor_fc2_weight_sf: cute.Tensor, + # GEMM-domain tensors (fc1) + activation_gemm: cute.Tensor, + fc1_weight_gemm: cute.Tensor, + fc1_output_gemm: cute.Tensor, + activation_sf_gemm: cute.Tensor, + fc1_weight_sf_gemm: cute.Tensor, + fc1_output_sf_gemm: cute.Tensor, + # GEMM-domain tensors (fc2) + fc2_weight_gemm: cute.Tensor, + fc2_output_gemm: cute.Tensor, + fc2_weight_sf_gemm: cute.Tensor, + fc1_output_sf_gemm_for_fc2_load: cute.Tensor, + # topk + cross-phase sync workspace + topk_scores: cute.Tensor, + fc1_done_counter: cute.Tensor, + # debug: (total_tokens, intermediate_half*3) fp32 for swiglu comparison + # Scheduling + offs: Optional[cute.Tensor], + expert_token_sizes: Optional[cute.Tensor], + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + # SMEM layouts + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + d_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout], + c_smem_layout_staged: Optional[Union[cute.Layout, cute.ComposedLayout]] = None, + tma_atom_c: Optional[cute.CopyAtom] = None, + tma_tensor_c: Optional[cute.Tensor] = None, + overflow_flag: cute.Tensor = None, + token_comm_args=None, + fc2_output_sf: Optional[cute.Tensor] = None, + # MegaMoE push-model token-comm inputs + mega_peer_rank_ptr_mapper=None, + mega_local_rank: Optional[cutlass.Int32] = None, + mega_local_workspace: Optional[cute.Pointer] = None, + mega_shared_workspace: Optional[cute.Pointer] = None, + mega_activation: Optional[cute.Tensor] = None, + mega_activation_sf: Optional[cute.Tensor] = None, + mega_pre_reduced_activation: Optional[cute.Tensor] = None, + mega_pre_reduced_activation_sf: Optional[cute.Tensor] = None, + # FC2 TMASTG (fc2_use_bulk): store atom + token-major pool view (None off-path) + tma_atom_fc2_output: Optional[cute.CopyAtom] = None, + fc2_tma_output: Optional[cute.Tensor] = None, + ): + """Device kernel for fused fc1+fc2 swap-AB GLU MXFP8 grouped GEMM.""" + a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, None, 0)) + b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, None, 0)) + sfa_smem_layout = cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)) + sfb_smem_layout = cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)) + + # MegaMoE (push model) + if cutlass.const_expr(self.enable_token_comm): + self._mega_device_workspace.assign_device_members( + mega_local_workspace, mega_shared_workspace + ) + + # fc2 waits for all fc1 intermediate N-tiles in the same token block. + # Each N-tile is processed by atom_thr_size CTAs (both CTA0 and CTA1 increment + # the counter), so the threshold must account for both CTAs' contributions. + ext_fc2_spin_threshold = ( + fc1_weight_gemm.shape[0] + self.cta_tile_shape_mnk[1] - 1 + ) // self.cta_tile_shape_mnk[1] * self.epilogue._atom_thr_size + + ext = GluMxFp8Fc12SchedExtension( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter.iterator, + fc2_spin_threshold=ext_fc2_spin_threshold, + fc1_ready_counter_pointer=self.token_comm_hook_fc1_ready_counter_ptr( + token_comm_args + ), + cluster_m=self.epilogue._atom_thr_size, + ) + + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + tidx, _, _ = cute.arch.thread_idx() + + # MegaMoE (push model): bind token-comm device members (transfer-warp state + the + # NVLink barrier's peer mapper) before any token_in / token_back / size-wait runs. + if cutlass.const_expr(self.enable_token_comm): + _mega_token_comm_args = TokenCommArgs( + mega_activation, + mega_activation_sf, + mega_pre_reduced_activation, + mega_pre_reduced_activation_sf, + mega_peer_rank_ptr_mapper, + ) + _mega_cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] + _, _, _mega_cluster_idx = cute.arch.block_idx() + _mega_linear_cta_idx = cta_rank_in_cluster + _mega_cluster_idx * _mega_cluster_size + self.token_comm.assign_device_members( + device_workspace=self._mega_device_workspace, + token_comm_args=_mega_token_comm_args, + local_rank=mega_local_rank, + linear_cta_idx=_mega_linear_cta_idx, + ) + + # SharedStorage (mainloop + epilogue SMEM). next's scheduler owns its own + # SMEM workspace, allocated separately below. + @cute.struct + class SharedStorage: + a_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_a_stage * 2] + b_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_b_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_pipeline_stages * 2 + ] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # next scheduler SMEM: a self-contained workspace carved from the same + # allocator; its transport regions resolve against ``sched_smem_base``. + sched_storage = smem.allocate(self.sched_smem_ws.storage_class()) + sched_smem_base = sched_storage.buffer.data_ptr() + + # MegaMoE-only dispatch-warp SMEM (pull_buffer, mbarriers, etc.). + # Kept out of ``SharedStorage`` so the lean path never allocates it. + TokenCommStorageCls = self.token_comm_extra_smem_storage_class() + if cutlass.const_expr(TokenCommStorageCls is not None): + token_comm_storage = smem.allocate(TokenCommStorageCls) + else: + token_comm_storage = None + + # ── Pipelines: separate producer/consumer groups for A and B. ── + + a_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, 1 + ) + a_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mcast_ctas_a + ) + a_producer, a_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.a_full_mbar_ptr.data_ptr(), + num_stages=self.num_a_stage, + producer_group=a_pipeline_producer_group, + consumer_group=a_pipeline_consumer_group, + tx_count=self.num_tma_load_a_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + mcast_mode_mn=(1, 0), + defer_sync=True, + ).make_participants() + b_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, 1 + ) + b_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mcast_ctas_b + ) + b_producer, b_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.b_full_mbar_ptr.data_ptr(), + num_stages=self.num_b_stage, + producer_group=b_pipeline_producer_group, + consumer_group=b_pipeline_consumer_group, + tx_count=self.num_tma_load_b_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + mcast_mode_mn=(0, 1), + defer_sync=True, + ).make_participants() + + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = ( + len(self.epilogue_warp_id) * 32 * (2 if use_2cta_instrs else 1) + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_pipeline_stages, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # TMEM allocator + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, + arch=self.arch, + ) + + scheduler = self.scheduler + if cutlass.const_expr(self.enable_token_comm): + _sched_expert_sizes = self.token_comm.local_expert_sizes( + self._mega_device_workspace, mega_local_rank + ) + _sched_prefix_sum = None + else: + _sched_expert_sizes = expert_token_sizes + _sched_prefix_sum = offs + scheduler.assign_device_members( + expert_token_sizes=_sched_expert_sizes, + expert_token_prefix_sum=_sched_prefix_sum, + actual_expert_shape=None, + block_idx=cute.arch.block_idx(), + smem_workspace=self.sched_smem_ws, + smem_base=sched_smem_base, + device_workspace=self.sched_device_ws, + ) + sched_consumer = scheduler.make_consumer() + + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # ── SMEM tensors A / B / SFA / SFB (shared by fc1 / fc2) ── + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + sSFA = smem.allocate_tensor( + element_type=self.sf_dtype, + layout=sfa_smem_layout_staged, + byte_alignment=128, + ) + sSFB = smem.allocate_tensor( + element_type=self.sf_dtype, + layout=sfb_smem_layout_staged, + byte_alignment=128, + ) + + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + + # acc_fake layout: (MMA, MMA_M, MMA_N, STAGE). + acc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # Cluster wait before TMEM alloc. + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + mma_tiler_k = self.mma_tiler[2] + k_tile_cnt_fc1 = (fc1_weight_gemm.shape[1] + mma_tiler_k - 1) // mma_tiler_k + k_tile_cnt_fc2 = (fc2_weight_gemm.shape[1] + mma_tiler_k - 1) // mma_tiler_k + + cluster_n = self.cluster_shape_mn[1] + _fc1_n_cluster_tile = self.cta_tile_shape_mnk[1] * cluster_n + fc2_spin_threshold = ( + ( + (fc1_weight_gemm.shape[0] + _fc1_n_cluster_tile - 1) + // _fc1_n_cluster_tile + ) + * cluster_n + * self.epilogue._atom_thr_size + ) + + # ════════════════════════════════════════════════════════════════════ + # Scheduler warp (warp 7) — lean path + # ════════════════════════════════════════════════════════════════════ + if warp_idx == self.sched_warp_id: + # MegaMoE: block until the Router has published this rank's per-expert sizes + # (cross-rank), so the lazy size walk in ``gen_next_work`` reads valid counts. + # No-op on the lean path. + self.token_comm_hook_sched_warp_pre_init_wait(token_comm_args) + work_tile = scheduler.gen_next_work() + while work_tile.is_valid_tile: + scheduler.publish_work(ext.prepare_work_tile(work_tile)) + work_tile = scheduler.gen_next_work() + # Sentinel publish (the tile is already invalid here). + scheduler.publish_work(work_tile) + scheduler.produce_tail() + + # ════════════════════════════════════════════════════════════════════ + # TMA load warps (warps 5 / 6) + # ════════════════════════════════════════════════════════════════════ + # + # TMA-A loads activations/SFA into the A pipeline. + # TMA-B loads weights/SFB into the B pipeline and waits for + # fc1 workspace readiness in the fc2 phase. + + # ── TMA-A warp (warp 5) ───────────────────────────────────────────── + if warp_idx == self.tma_a_warp_id: + _iket_active = (tidx == cutlass.Int32(160)) + a_full_mcast_mask = None + sfa_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + + b_full_mcast_mask = None + if cutlass.const_expr(self.is_b_mcast or use_2cta_instrs): + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + sfa_cta_layout = a_cta_layout + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + + work_tile_info = sched_consumer.consume_work() + + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + if is_phase_linear1: + # fc1 phase A-side + if _iket_active: + iket.range_push("tma_token_fc1") + ext.wait_for_input(work_tile_info) + self.token_comm_hook_fc1_tma_b_predispatch_spin( + token_comm_args, work_tile_info, + ) + + k_tile_cnt = k_tile_cnt_fc1 + real_a, desc_ptr_a = ext.get_gmem_tensor( + "fc1_activation", tma_tensor_fc1_activation_1, work_tile_info, + ) + real_sfa, desc_ptr_sfa = ext.get_gmem_tensor( + "fc1_activation_sf", tma_tensor_fc1_activation_1_sf, work_tile_info, + ) + + gA_mkl = cute.local_tile( + real_a, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + gSFA_mkl = cute.local_tile( + real_sfa, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + tCgA = thr_mma.partition_A(gA_mkl) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + + tAsA, tAgA = cpasync.tma_partition( + tma_atom_fc1_activation_1, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_fc1_activation_1_sf, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + mma_tile_m = work_tile_info.tile_m_idx // cute.size( + tiled_mma.thr_id.shape + ) + tAgA_slice = tAgA[(None, mma_tile_m, None, 0)] + tAgSFA_slice = tAgSFA[(None, mma_tile_m, None, 0)] + + a_producer.reset() + peek_a_empty_status = a_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Producer-side backpressure: TMA-A blocked waiting for the + # MMA to free an A SMEM slot. First k-tile only (later tiles + # overlap via peek-ahead). Complements mma_ab_operand_wait. + if _iket_active: + iket.range_push("tma_a_buf_acquire_wait") + handle = a_producer.acquire_and_advance( + peek_a_empty_status + ) + if _iket_active: + iket.range_pop() # tma_a_buf_acquire_wait + peek_a_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_a_empty_status = a_producer.try_acquire() + cute.copy( + tma_atom_fc1_activation_1, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_a, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_fc1_activation_1_sf, + tAgSFA_slice[(None, handle.count)], + tAsSFA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfa, + mcast_mask=sfa_full_mcast_mask, + ) + else: + # fc2 phase A-side: load fc1_output (M=tokens) + wait for fc1 done + if _iket_active: + iket.range_push("tma_token_fc2") + counter_slot = ( + work_tile_info.cumulative_token_block_count + + work_tile_info.tile_m_idx // cutlass.Int32(self.epilogue._atom_thr_size) + ) + counter_ptr = fc1_done_counter.iterator + counter_slot + if _iket_active: + iket.range_push("tma_token_fc2_a_wait") + spin_wait( + counter_ptr, + lambda v: v >= fc2_spin_threshold, + sleep_cycles=20, + ) + if _iket_active: + iket.range_pop() + cute.arch.load(counter_ptr, counter_ptr.dtype, sem="acquire", scope="gpu") + cute.arch.fence_proxy("async") + cute.arch.fence_proxy("async.global") + + if ( + tidx == cutlass.Int32(32 * self.tma_a_warp_id) + and work_tile_info.tile_n_idx == cutlass.Int32(0) + ): + counter_val_post = cute.arch.load( + counter_ptr, counter_ptr.dtype, cop="cg" + ) + fc1_byte_offset = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_m_idx + // cutlass.Int32(self.epilogue._atom_thr_size) + * cutlass.Int32(self.epilogue._cta_tile_m) + ) * fc1_output_gemm.stride[0] + fc1_probe_ptr = cute.make_ptr( + cutlass.Int32, + fc1_output_gemm.iterator.toint() + fc1_byte_offset, + cute.AddressSpace.gmem, + ) + fc1_first_i32 = cute.arch.load(fc1_probe_ptr, cutlass.Int32, cop="cg") + + k_tile_cnt = k_tile_cnt_fc2 + real_a, desc_ptr_a = ext.get_gmem_tensor( + "fc2_activation", tma_tensor_fc2_activation, work_tile_info, + ) + real_sfa, desc_ptr_sfa = ext.get_gmem_tensor( + "fc2_activation_sf", tma_tensor_fc2_activation_sf, work_tile_info, + ) + + gA_mkl = cute.local_tile( + real_a, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + gSFA_mkl = cute.local_tile( + real_sfa, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + tCgA = thr_mma.partition_A(gA_mkl) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + + tAsA, tAgA = cpasync.tma_partition( + tma_atom_fc2_activation, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tAsSFA, tAgSFA = cpasync.tma_partition( + tma_atom_fc2_activation_sf, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + # fc2 A-side = fc1_output (M=tokens). + mma_tile_m = work_tile_info.tile_m_idx // cute.size( + tiled_mma.thr_id.shape + ) + tAgA_slice = tAgA[(None, mma_tile_m, None, 0)] + tAgSFA_slice = tAgSFA[(None, mma_tile_m, None, 0)] + + a_producer.reset() + peek_a_empty_status = a_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Producer-side backpressure: TMA-A blocked waiting for the + # MMA to free an A SMEM slot. First k-tile only (later tiles + # overlap via peek-ahead). Complements mma_ab_operand_wait. + if _iket_active: + iket.range_push("tma_a_buf_acquire_wait") + handle = a_producer.acquire_and_advance( + peek_a_empty_status + ) + if _iket_active: + iket.range_pop() # tma_a_buf_acquire_wait + peek_a_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_a_empty_status = a_producer.try_acquire() + cute.copy( + tma_atom_fc2_activation, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_a, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_fc2_activation_sf, + tAgSFA_slice[(None, handle.count)], + tAsSFA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfa, + mcast_mask=sfa_full_mcast_mask, + ) + + if _iket_active: + iket.range_pop() + work_tile_info = sched_consumer.consume_work() + + a_producer.tail() + + # ── TMA-B warp (warp 6) ───────────────────────────────────────────── + if warp_idx == self.tma_b_warp_id: + _iket_active = (tidx == cutlass.Int32(192)) + b_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr(self.is_b_mcast or use_2cta_instrs): + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, + block_in_cluster_coord_sfb_vmnk, + mcast_mode=1, + ) + + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + sfb_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) + + work_tile_info = sched_consumer.consume_work() + + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + + if is_phase_linear1: + # fc1 phase B-side (fc1_weight, N-side GEMM-B) + if _iket_active: + iket.range_push("tma_weight_fc1") + + k_tile_cnt = k_tile_cnt_fc1 + real_b, desc_ptr_b = ext.get_gmem_tensor( + "fc1_weight", tma_tensor_weight, work_tile_info, + ) + real_sfb, desc_ptr_sfb = ext.get_gmem_tensor( + "fc1_weight_sf", tma_tensor_fc1_weight_sf, work_tile_info, + ) + + # N-K tiling for N-side weight (N=intermediate, K=hidden). + gB_nkl = cute.local_tile( + real_b, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + gSFB_nkl = cute.local_tile( + real_sfb, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + + tCgB = thr_mma.partition_B(gB_nkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + + tBsB, tBgB = cpasync.tma_partition( + tma_atom_weight, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_fc1_weight_sf, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # Use tile_n_idx for N-side weight: invariant across token blocks. + tBgB_slice = tBgB[(None, work_tile_info.tile_n_idx, None, 0)] + tBgSFB_slice = tBgSFB[(None, work_tile_info.tile_n_idx, None, 0)] + + b_producer.reset() + peek_b_empty_status = b_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Producer-side backpressure: TMA-B blocked waiting for the + # MMA to free a B SMEM slot. First k-tile only (later tiles + # overlap via peek-ahead). Complements mma_ab_operand_wait. + if _iket_active: + iket.range_push("tma_b_buf_acquire_wait") + handle = b_producer.acquire_and_advance( + peek_b_empty_status + ) + if _iket_active: + iket.range_pop() # tma_b_buf_acquire_wait + peek_b_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_b_empty_status = b_producer.try_acquire() + cute.copy( + tma_atom_weight, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_b, + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_fc1_weight_sf, + tBgSFB_slice[(None, handle.count)], + tBsSFB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfb, + mcast_mask=sfb_full_mcast_mask, + ) + else: + # fc2 phase B-side: load fc2_weight (N=hidden) + if _iket_active: + iket.range_push("tma_weight_fc2") + k_tile_cnt = k_tile_cnt_fc2 + real_b, desc_ptr_b = ext.get_gmem_tensor( + "fc2_weight", tma_tensor_fc2_weight, work_tile_info, + ) + real_sfb, desc_ptr_sfb = ext.get_gmem_tensor( + "fc2_weight_sf", tma_tensor_fc2_weight_sf, work_tile_info, + ) + + gB_nkl = cute.local_tile( + real_b, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + gSFB_nkl = cute.local_tile( + real_sfb, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + tCgB = thr_mma.partition_B(gB_nkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + + tBsB, tBgB = cpasync.tma_partition( + tma_atom_fc2_weight, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tBsSFB, tBgSFB = cpasync.tma_partition( + tma_atom_fc2_weight_sf, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # fc2 B-side = fc2_weight (N=hidden) + fc2_b_hidden_tile = work_tile_info.tile_n_idx + tBgB_slice = tBgB[(None, fc2_b_hidden_tile, None, 0)] + tBgSFB_slice = tBgSFB[(None, fc2_b_hidden_tile, None, 0)] + + b_producer.reset() + peek_b_empty_status = b_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Producer-side backpressure: TMA-B blocked waiting for the + # MMA to free a B SMEM slot. First k-tile only (later tiles + # overlap via peek-ahead). Complements mma_ab_operand_wait. + if _iket_active: + if k_tile == 0: + iket.range_push("tma_b_buf_acquire_wait") + handle = b_producer.acquire_and_advance( + peek_b_empty_status + ) + if _iket_active: + if k_tile == 0: + iket.range_pop() # tma_b_buf_acquire_wait + peek_b_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_b_empty_status = b_producer.try_acquire() + cute.copy( + tma_atom_fc2_weight, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_b, + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_fc2_weight_sf, + tBgSFB_slice[(None, handle.count)], + tBsSFB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_sfb, + mcast_mask=sfb_full_mcast_mask, + ) + if _iket_active: + iket.range_pop() + work_tile_info = sched_consumer.consume_work() + + b_producer.tail() + + # ════════════════════════════════════════════════════════════════════ + # MMA warp (warp 4) + # ════════════════════════════════════════════════════════════════════ + # + # Both phases share tiled_mma and TMEM; only K-tile count differs. + if warp_idx == self.mma_warp_id: + _iket_active = (tidx == cutlass.Int32(128)) + + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + acc_base = cute.make_tensor(acc_tmem_ptr, acc_fake.layout) + + # SFA TMEM tensor (placed after the acc cols). + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + # SFB TMEM tensor (after acc + SFA cols). + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)), + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + + ( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t, + tCtSFA_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + ( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t, + tCtSFB_compact_s2t, + ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_pipeline_stages + ) + + work_tile_info = sched_consumer.consume_work() + + while work_tile_info.is_valid_tile: + is_phase_linear1 = ( + work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) + ) + # Prebind k_tile_cnt due to DSL AST. + k_tile_cnt = cutlass.Int32(0) + if is_phase_linear1: + k_tile_cnt = k_tile_cnt_fc1 + if _iket_active: + iket.range_push("mma_fc1") + else: + k_tile_cnt = k_tile_cnt_fc2 + if _iket_active: + iket.range_push("mma_fc2") + + acc_stage_index = acc_producer_state.index + + if is_leader_cta: + tCtAcc = acc_base[(None, None, None, acc_stage_index)] + + if _iket_active: + iket.range_push("mma_acc_acquire") + a_consumer.reset() + b_consumer.reset() + peek_a_full_status = cutlass.Boolean(1) + peek_b_full_status = cutlass.Boolean(1) + if k_tile_cnt > 0: + peek_a_full_status = a_consumer.try_wait() + peek_b_full_status = b_consumer.try_wait() + acc_pipeline.producer_acquire(acc_producer_state) + if _iket_active: + iket.range_pop() + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Isolate the real AB-operand stall: the initial k-tile wait + # for A/B to arrive from the TMA warps. Later k-tiles overlap + # with MMA compute via the peek-ahead pipeline, so only the + # first tile's wait reflects true operand-arrival latency. + if _iket_active: + iket.range_push("mma_ab_operand_wait") + handle_a = a_consumer.wait_and_advance(peek_a_full_status) + handle_b = b_consumer.wait_and_advance(peek_b_full_status) + if _iket_active: + iket.range_pop() # mma_ab_operand_wait + peek_a_full_status = cutlass.Boolean(1) + peek_b_full_status = cutlass.Boolean(1) + if handle_a.count + 1 < k_tile_cnt: + peek_a_full_status = a_consumer.try_wait() + peek_b_full_status = b_consumer.try_wait() + + cute.copy( + tiled_copy_s2t_sfa, + tCsSFA_compact_s2t[(None, None, None, None, handle_a.index)], + tCtSFA_compact_s2t, + ) + cute.copy( + tiled_copy_s2t_sfb, + tCsSFB_compact_s2t[(None, None, None, None, handle_b.index)], + tCtSFB_compact_s2t, + ) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + cute.gemm( + tiled_mma, + tCtAcc, + [tCrA[(None, None, None, handle_a.index)], tCtSFA], + [tCrB[(None, None, None, handle_b.index)], tCtSFB], + tCtAcc, + ) + handle_a.release() + handle_b.release() + + if k_tile_cnt > 0: + acc_pipeline.producer_commit(acc_producer_state) + if k_tile_cnt > 0: + acc_producer_state.advance() + + if _iket_active: + iket.range_pop() + + work_tile_info = sched_consumer.consume_work() + + acc_pipeline.producer_tail(acc_producer_state) + + # ── sD SMEM (fc1 output staging; fc2 doesn't use it) ── + sD = smem.allocate_tensor( + element_type=self.fc1_output_dtype, + layout=d_smem_layout_staged.outer, + byte_alignment=128, + swizzle=d_smem_layout_staged.inner, + ) + + # ── sC SMEM (raw gate+up Float32, ping-pong; only when generate_c=True) ── + if cutlass.const_expr(self.generate_c): + sC = smem.allocate_tensor( + element_type=self.epilogue._c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + + # ── sFC2 SMEM (FC2 bulk-store staging; fc2_use_bulk only) ── + if cutlass.const_expr(self.epilogue.fc2_needs_staging): + sFC2 = smem.allocate_tensor( + element_type=self.epilogue._fc2_wire_dtype, + layout=self.epilogue.fc2_tma_staged_smem_layout( + self.epilogue.fc2_tma_stages + ), + byte_alignment=128, + ) + + # ── sRED SMEM (in-kernel reduce coalescing transpose; reduce_topk only) ── + if cutlass.const_expr(self.epilogue.fc2_reduce_coalesce): + sRED = smem.allocate_tensor( + element_type=cutlass.BFloat16, + layout=self.epilogue.fc2_reduce_smem_layout(), + byte_alignment=128, + ) + + # ════════════════════════════════════════════════════════════════════ + # Epilogue warps (warps 0-3) + # ════════════════════════════════════════════════════════════════════ + # + # Fully delegated to ``self.epilogue.run(...)`` -- the epilogue owns + # the entire 2-phase task-tile loop. + if warp_idx < self.mma_warp_id: + epi_warp_idx = warp_idx + + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + acc_tensor = cute.make_tensor(acc_tmem_ptr, acc_fake.layout) + + #acc_tensor = cute.make_tensor( + # acc_tmem_ptr, + # cute.make_layout( + # (((128, self.cta_tile_shape_mnk[1]), 1),), + # stride=(((1 << 16, 1), 0),), + # ), + #) + + # Build common kwargs shared by both epilogue flavours. + if cutlass.const_expr(self.generate_c): + _smem_c_raw_arg = sC + _tma_atom_c_arg = tma_atom_c + _gmem_c_arg = tma_tensor_c + else: + _smem_c_raw_arg = None + _tma_atom_c_arg = None + _gmem_c_arg = None + if cutlass.const_expr(self.use_stg_fc1): + _gmem_fc1_output_arg = fc1_output_gemm + else: + _gmem_fc1_output_arg = tma_tensor_fc1_output + _run_kwargs = dict( + tmem_acc_tensor=acc_tensor, + acc_pipeline=acc_pipeline, + sched_consumer=sched_consumer, + sched_ext=ext, + smem_fc1_output_buffer=sD, + tma_atom_fc1_output=tma_atom_fc1_output, + gmem_fc1_output=_gmem_fc1_output_arg, + gmem_fc1_output_sf=fc1_output_sf_gemm, + gmem_topk_scores=topk_scores, + gmem_fc2_output=fc2_output_gemm, + gmem_fc1_done_counter=fc1_done_counter, + smem_c_buffer=_smem_c_raw_arg, + tma_atom_c=_tma_atom_c_arg, + gmem_c=_gmem_c_arg, + warp_idx=epi_warp_idx, + tidx=tidx, + alpha=cutlass.Float32(1.0), + norm_const=cutlass.Float32(1.0), + ) + + _run_extra = {} + if cutlass.const_expr(fc2_output_sf is not None): + _run_extra["gmem_fc2_output_sf"] = fc2_output_sf + if cutlass.const_expr(self.epilogue.fc2_needs_staging): + # Both bulk paths stage into sFC2; only the dispatch TMASTG path + # also needs the TMA atom + local-pool tensor (UBLK peer-writes). + _run_extra["smem_fc2_tma_buffer"] = sFC2 + if cutlass.const_expr(self.epilogue.fc2_use_tma): + _run_extra["tma_atom_fc2_output"] = tma_atom_fc2_output + _run_extra["gmem_fc2_tma_output"] = fc2_tma_output + if cutlass.const_expr(self.epilogue.fc2_reduce_coalesce): + _run_extra["smem_fc2_reduce_buffer"] = sRED + if cutlass.const_expr(self.enable_token_comm): + # MegaMoE (push model): bridge next's TokenComm accessors + peer mapper into + # the epilogue's Fc2OutputDest peer-store expectations. + _epi_comm = _EpilogueCommView( + token_src_metadata=self.token_comm.token_src_metadata_tensor( + self._mega_device_workspace + ), + combine_output=mega_pre_reduced_activation, + peer_rank_ptr_mapper=mega_peer_rank_ptr_mapper, + fc2_output_sf=self.token_comm.fc2_activation_sf_tensor(self._mega_device_workspace), + fc2_done_counter=self.token_comm.fc2_done_counter_tensor(self._mega_device_workspace), + fc2_output_workspace=self.token_comm.fc2_activation_tensor(self._mega_device_workspace), + ) + self.epilogue.run(**_run_kwargs, **_run_extra, token_comm_args=_epi_comm) + elif cutlass.const_expr(token_comm_args is not None): + self.epilogue.run( + **_run_kwargs, **_run_extra, token_comm_args=token_comm_args + ) + else: + self.epilogue.run(**_run_kwargs, **_run_extra) + + tmem.relinquish_alloc_permit() + tmem.free(acc_tmem_ptr) + if cutlass.const_expr(self.enable_token_comm): + cute.arch.fence_acq_rel_sys() + + # ════════════════════════════════════════════════════════════════════ + # Dispatch warps hook (warps 8-11; MegaMoE-only) + # ════════════════════════════════════════════════════════════════════ + # + # ``enable_token_comm=False`` → warps 8-11 don't exist (threads_per_cta + # = 256), so the guard is const_expr-eliminated in the lean path. + if cutlass.const_expr(self.enable_token_comm): + if warp_idx >= self.dispatch_warp_id[0]: + lane_idx_for_dispatch = cute.arch.lane_idx() + if cutlass.const_expr(self.token_back_standalone): + if warp_idx < self.token_back_warp_id[0]: + self.token_comm_hook_dispatch_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + else: + self.token_comm_hook_token_back_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + else: + self.token_comm_hook_dispatch_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + + # ════════════════════════════════════════════════════════════════════ + # Kernel tail hook (MegaMoE-only; lean base = no-op) + # ════════════════════════════════════════════════════════════════════ + lane_idx = cute.arch.lane_idx() + self.token_comm_hook_kernel_tail( + token_comm_args, + warp_idx=warp_idx, + lane_idx=lane_idx, + tidx=tidx, + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py new file mode 100644 index 000000000..2d39dc12f --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py @@ -0,0 +1,865 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Full MegaMoE (multi-rank) mxfp8 GLU training-forward kernel.""" + +from typing import Any, Literal, Optional, Tuple, Type + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace + +from ......api import ImplDesc, KernelClass, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from ......helpers.device_workspace import DeviceWorkspace +from ......helpers.smem_workspace import SmemWorkspace +from ......helpers.utils import ceil_div, round_up +from ......quant_def import CombineFormat, QuantKind +from ......communication.nvlink_domain.token_comm_deterministic import TokenCommDeterministic +from ..topk_reduce import TopkReduce +from .glu_mxfp8_col_requant import Mxfp8ColRequant +from .glu_mxfp8_fc12_kernel import Sm107Mxfp8GluFc12Kernel + + +_AB_DTYPE_TO_QUANT_KIND = {cutlass.Float8E4M3FN: QuantKind.mxfp8_e4m3, cutlass.Float8E5M2: QuantKind.mxfp8_e5m2} +_QUANT_KIND_TO_AB_DTYPE = {str(k): d for d, k in _AB_DTYPE_TO_QUANT_KIND.items()} + +# TVM-FFI export symbol for the AOT-compiled callable (consumed by ``tester.compiler``). +_aot_symbol_prefix = "rubin_mega_moe_glu_mxfp8_aot" + + +class Sm107MegaMoEMxfp8GluKernel(Sm107Mxfp8GluFc12Kernel, KernelClass): + """Multi-rank MegaMoE wrapper around the lean mxfp8 GLU FC12 kernel.""" + + fc1_output_region = "rubin.glu_mxfp8.mega.fc1_output" + fc1_output_sf_region = "rubin.glu_mxfp8.mega.fc1_output_sf" + fc1_done_counter_region = "rubin.glu_mxfp8.mega.fc1_done_counter" + col_quant_sizes_region = "rubin.glu_mxfp8.mega.col_quant_expert_token_sizes" + + # Reserved on top of the exact token_comm/sched SMEM to cover smem.allocate inter-allocation + # alignment padding that _compute_stages does not model (see _smem_misc_budget_bytes). + _SMEM_ALLOC_MARGIN = 2048 + + @classmethod + def problem_desc_require(cls): + return { + "expert_count": StaticOrRuntimeIntegerType, + "intermediate_gateup_size": StaticOrRuntimeIntegerType, + "hidden_size": StaticOrRuntimeIntegerType, + "quant_kind": str, + "combine_format": CombineFormat, + "world_size": int, + "local_rank": int, + "topk": int, + "max_tokens_per_rank": int, + "max_recv_size_per_rank": int, + "gate_up_clamp": Optional[float], + } + + @classmethod + def impl_desc_require(cls): + return { + "mma_tiler_mnk": tuple, + "cluster_shape_mnk": tuple, + "use_2cta_instrs": bool, + "group_hint": int, + "token_padding_block": int, + "sf_padding_block": int, + "load_balance_mode": str, + "force_static_sched": bool, + "clc_bundle_size": Optional[int], + "num_sched_stages": Optional[int], + "acc_dtype": type, + "sf_vec_size": int, + "launch_cluster_count": int, + "drop_on_overflow": bool, + "fc2_in_kernel_topk_reduce": bool, + "token_back_mode": str, + "epi_flag_batch": tuple, + "flag_batch": int, + "generate_c": bool, + "use_stg_fc1": bool, + "act_func": str, + "fc2_use_bulk": bool, + "fc2_tma_stages": OptionalRequirement(int), + "enable_col_quant": OptionalRequirement(bool), + "col_quant_num_ctas": OptionalRequirement(int), + } + + def name(self) -> str: + return ( + f"sm107_megamoe_glu_{self.quant_kind}_m{self.mma_tiler_mnk[0]}n{self.mma_tiler_mnk[1]}" + f"k{self.mma_tiler_mnk[2]}_e{self.expert_count}_ep{self.world_size}_topk{self.topk}_" + f"h{self.hidden_size}_i{self.intermediate_gateup_size}_combine{self.combine_format}_" + f"clamp{self.gate_up_clamp}_" + f"tokenback{self.token_back_mode}_hint{self.group_hint}_" + f"epi{self.epi_flag_batch[0]}x{self.epi_flag_batch[1]}_tif{self.flag_batch}_" + f"deterministic_mtpr{self.max_tokens_per_rank}_mrpr{self.max_recv_size_per_rank}_" + f"drop{int(self.drop_on_overflow)}_lc{self.launch_cluster_count}_" + f"genc{int(self.generate_c)}_topkfc11_" + f"fc2bulk{int(self.fc2_use_bulk)}x{self.fc2_tma_stages}_" + f"redtopk{int(self.reduce_topk_in_kernel)}" + ) + + def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): + """Compile against fake (metadata-only) inputs; ``out_path=None`` returns the in-memory callable.""" + import math + + from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_ptr + from cutlass.cute.typing import AddressSpace, sym_int64 + from cutlass.cutlass_dsl import Int32, Int64 + + from ......communication.nvlink_domain.symmetric_buffer import SymmetricBufferHost + + def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): + extents = tuple( + sym_int64(divisibility=math.gcd(int(extent), 128)) if axis in dynamic_axes else int(extent) + for axis, extent in enumerate(shape) + ) + return make_fake_compact_tensor(dtype, extents, stride_order=stride_order, assumed_align=alignment) + + tokens = self.max_tokens_per_rank + aux_shapes = self.get_aux_output_shapes() + hidden = self.hidden_size + intermediate_gateup = self.intermediate_gateup_size + intermediate_downproj = intermediate_gateup // 2 + experts = self.num_experts_per_rank + sf_vec_size = self.sf_vec_size + fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up(hidden // sf_vec_size, 4) + fc2_weight_sf_columns = round_up(hidden, 128) * round_up(intermediate_downproj // sf_vec_size, 4) + output_dtype = cutlass.BFloat16 + # Weight SF and activation SF share the E8M0 block-scale dtype for mxfp8. + weight_sf_dtype = self.token_comm.activation_sf_dtype + + fake_arguments = dict( + activation=fake_tensor(self.token_comm.activation_dtype, (tokens, hidden), (1, 0), {0}, 16), + activation_sf=fake_tensor( + self.token_comm.activation_sf_dtype, + (tokens, self.token_comm.activation_sf_hidden_padded), + (1, 0), + {0}, + 16, + ), + topk_indices=fake_tensor(cutlass.Int32, (tokens, self.topk), (1, 0), {0}, 16), + topk_scores=fake_tensor(cutlass.Float32, (tokens, self.topk), (1, 0), {0}, 4), + fc1_weight=fake_tensor(self.ab_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16), + fc1_weight_sf=fake_tensor(weight_sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16), + fc2_weight=fake_tensor(self.ab_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16), + fc2_weight_sf=fake_tensor(weight_sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16), + output_activation=fake_tensor(output_dtype, (tokens, hidden), (1, 0), {0}, 16), + overflow_flag=fake_tensor(cutlass.Int32, (1,), (0,), set(), 4), + local_workspace=make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), + shared_workspace=make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), + peer_rank_ptr_mapper_host=SymmetricBufferHost( + base_address=Int64(0), + offsets=tuple(Int64(0) for _ in range(self.world_size)), + rank=Int32(0), + max_ranks=self.world_size, + ), + stream=make_fake_stream(), + ) + if self.generate_c: + fake_arguments["fc1_c"] = fake_tensor(output_dtype, aux_shapes["fc1_c"], (1, 0), set(), 16) + else: + fake_arguments["fc1_c"] = None + + if self.enable_col_quant: + # col_quant_data shares the dispatch pool's row-major (token, hidden) layout; + # col_quant_sf is flat concat_e [hidden_atom][token_atom] E8M0 bytes. + fake_arguments["col_quant_data"] = fake_tensor( + self.ab_dtype, aux_shapes["col_quant_data"], (1, 0), set(), 16 + ) + fake_arguments["col_quant_sf"] = fake_tensor( + cutlass.Uint8, aux_shapes["col_quant_sf"], (0,), set(), 16 + ) + else: + fake_arguments["col_quant_data"] = None + fake_arguments["col_quant_sf"] = None + + compiled = cute.compile[cute.EnableTVMFFI(True)](self, **fake_arguments) + if out_path is None: + return compiled + compiled.export_to_c(out_path, function_name=_aot_symbol_prefix, export_only_tvm_ffi_symbols=True) + return out_path + + @staticmethod + def load_compiled(path: str): + from cutlass.cute.runtime import load_module + + return load_module(path, enable_tvm_ffi=True)[_aot_symbol_prefix] + + @classmethod + def from_kwargs( + cls, + # Base-class (lean FC12) kwargs. + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mnk: Tuple[int, int, int], + use_2cta_instrs: bool, + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + load_balance_mode: str = "static", + static_expert_shape: Optional[Tuple[int, int, int]] = None, + force_static_sched: bool = True, + clc_bundle_size: Optional[int] = None, + num_sched_stages: Optional[int] = None, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + ab_dtype: Type[cutlass.Numeric] = cutlass.Float8E4M3FN, + sf_vec_size: int = 32, + *, + world_size: int, + local_rank: int, + num_topk: int, + max_tokens_per_rank: int, + max_recv_size_per_rank: int, + hidden: int, + launch_cluster_count: int, + drop_on_overflow: bool, + fc2_in_kernel_topk_reduce: bool = False, + token_back_mode: Literal["epi_warps", "standalone_warps", "reuse_dispatch_warps"] = "epi_warps", + epi_flag_batch: Optional[Tuple[int, int]] = (4, 2), + flag_batch: int = 1, + gate_up_clamp: Optional[float] = None, + generate_c: bool = False, + use_stg_fc1: bool = False, + combine_format: Optional[CombineFormat] = None, + act_func: str = "swiglu", + fc2_use_bulk: bool = False, + fc2_tma_stages: Optional[int] = None, + enable_col_quant: bool = False, + # -1 = let Mxfp8ColRequant derive the grid from the resident-CTA quantum. + # A hardcoded CTA count is not generally a multiple of that quantum, which + # leaves a fractional resident wave. + col_quant_num_ctas: int = -1, + ) -> "Sm107MegaMoEMxfp8GluKernel": + """Build the ``(ProblemDesc, ImplDesc)`` pair from the legacy flat signature.""" + if static_expert_shape is None: + raise NotImplementedError("Sm107MegaMoEMxfp8GluKernel requires a static_expert_shape.") + if hidden != static_expert_shape[2]: + raise ValueError(f"hidden ({hidden}) must equal static_expert_shape[2] ({static_expert_shape[2]}).") + if ab_dtype not in _AB_DTYPE_TO_QUANT_KIND: + raise ValueError(f"ab_dtype {ab_dtype} has no mxfp8 QuantKind.") + num_experts_per_rank, intermediate_gateup, _hidden = static_expert_shape + combine_format = CombineFormat.parse("bf16" if combine_format is None else str(combine_format)) + problem_desc = ProblemDesc( + { + "expert_count": world_size * num_experts_per_rank, + "intermediate_gateup_size": intermediate_gateup, + "hidden_size": hidden, + "quant_kind": str(_AB_DTYPE_TO_QUANT_KIND[ab_dtype]), + "combine_format": combine_format, + "world_size": world_size, + "local_rank": local_rank, + "topk": num_topk, + "max_tokens_per_rank": max_tokens_per_rank, + "max_recv_size_per_rank": max_recv_size_per_rank, + "gate_up_clamp": gate_up_clamp, + } + ) + impl_desc = ImplDesc( + { + "mma_tiler_mnk": tuple(mma_tiler_mnk), + "cluster_shape_mnk": tuple(cluster_shape_mnk), + "use_2cta_instrs": use_2cta_instrs, + "group_hint": group_hint, + "token_padding_block": token_padding_block, + "sf_padding_block": sf_padding_block, + "load_balance_mode": load_balance_mode, + "force_static_sched": force_static_sched, + "clc_bundle_size": clc_bundle_size, + "num_sched_stages": num_sched_stages, + "acc_dtype": acc_dtype, + "sf_vec_size": sf_vec_size, + "launch_cluster_count": launch_cluster_count, + "drop_on_overflow": drop_on_overflow, + "fc2_in_kernel_topk_reduce": fc2_in_kernel_topk_reduce, + "token_back_mode": token_back_mode, + "epi_flag_batch": tuple(epi_flag_batch) if epi_flag_batch is not None else (1, 1), + "flag_batch": flag_batch, + "generate_c": generate_c, + "use_stg_fc1": use_stg_fc1, + "act_func": act_func, + "fc2_use_bulk": fc2_use_bulk, + # OptionalRequirement: present only when set (absent == None). + **({"fc2_tma_stages": fc2_tma_stages} if fc2_tma_stages is not None else {}), + # Col-quant keys present only when enabled + **( + { + "enable_col_quant": True, + "col_quant_num_ctas": col_quant_num_ctas, + } + if enable_col_quant + else {} + ), + } + ) + return cls(problem_desc, impl_desc) + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + # -- Extract descriptors into locals matching the legacy param names so the body below + # is unchanged; derive the base-class flat inputs (static_expert_shape, ab_dtype). -- + world_size = problem_desc["world_size"] + local_rank = problem_desc["local_rank"] + num_topk = problem_desc["topk"] + max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + max_recv_size_per_rank = min( + problem_desc["max_recv_size_per_rank"], world_size * max_tokens_per_rank * num_topk + ) + hidden = problem_desc["hidden_size"] + gate_up_clamp = problem_desc["gate_up_clamp"] + combine_format = problem_desc["combine_format"] + _quant_kind = problem_desc["quant_kind"] + ab_dtype = _QUANT_KIND_TO_AB_DTYPE[_quant_kind] + static_expert_shape = ( + problem_desc["expert_count"] // world_size, + problem_desc["intermediate_gateup_size"], + hidden, + ) + + mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + cluster_shape_mnk = impl_desc["cluster_shape_mnk"] + use_2cta_instrs = impl_desc["use_2cta_instrs"] + group_hint = impl_desc["group_hint"] + token_padding_block = impl_desc["token_padding_block"] + sf_padding_block = impl_desc["sf_padding_block"] + load_balance_mode = impl_desc["load_balance_mode"] + force_static_sched = impl_desc["force_static_sched"] + clc_bundle_size = impl_desc["clc_bundle_size"] + num_sched_stages = impl_desc["num_sched_stages"] + acc_dtype = impl_desc["acc_dtype"] + sf_vec_size = impl_desc["sf_vec_size"] + launch_cluster_count = impl_desc["launch_cluster_count"] + drop_on_overflow = impl_desc["drop_on_overflow"] + fc2_in_kernel_topk_reduce = impl_desc["fc2_in_kernel_topk_reduce"] + token_back_mode = impl_desc["token_back_mode"] + epi_flag_batch = impl_desc["epi_flag_batch"] + flag_batch = impl_desc["flag_batch"] + generate_c = impl_desc["generate_c"] + use_stg_fc1 = impl_desc["use_stg_fc1"] + act_func = impl_desc["act_func"] + fc2_use_bulk = impl_desc["fc2_use_bulk"] + fc2_tma_stages = impl_desc.get("fc2_tma_stages") + self.enable_col_quant = bool(impl_desc.get("enable_col_quant") or False) + self._col_quant_num_ctas = int(impl_desc.get("col_quant_num_ctas") or -1) + + if static_expert_shape is None: + raise NotImplementedError("Sm107MegaMoEMxfp8GluKernel requires a static_expert_shape.") + if hidden != static_expert_shape[2]: + raise ValueError(f"hidden ({hidden}) must equal static_expert_shape[2] ({static_expert_shape[2]}).") + token_back_by_dispatch = token_back_mode != "epi_warps" + + combine_format = CombineFormat.parse("bf16" if combine_format is None else str(combine_format)) + if fc2_in_kernel_topk_reduce and (token_back_by_dispatch or combine_format.is_quantized): + raise ValueError("fc2_in_kernel_topk_reduce requires epi_warps + non-quantized (bf16) combine.") + if token_back_mode not in ("epi_warps", "standalone_warps", "reuse_dispatch_warps"): + raise ValueError(f"unsupported token_back_mode={token_back_mode!r}.") + if ab_dtype not in _AB_DTYPE_TO_QUANT_KIND: + raise ValueError(f"ab_dtype {ab_dtype} has no mxfp8 QuantKind.") + # FC2 bulk store + if fc2_use_bulk and not combine_format.is_quantized: + raise ValueError("fc2_use_bulk currently supports only a quantized (mxfp8) combine format.") + if fc2_tma_stages is not None and not fc2_use_bulk: + raise ValueError("fc2_tma_stages requires fc2_use_bulk=True.") + + super().__init__( + mma_tiler_mnk=mma_tiler_mnk, + cluster_shape_mnk=cluster_shape_mnk, + use_2cta_instrs=use_2cta_instrs, + group_hint=group_hint, + token_padding_block=token_padding_block, + sf_padding_block=sf_padding_block, + load_balance_mode=load_balance_mode, + static_expert_shape=static_expert_shape, + force_static_sched=force_static_sched, + clc_bundle_size=clc_bundle_size, + num_sched_stages=num_sched_stages, + acc_dtype=acc_dtype, + ab_dtype=ab_dtype, + sf_vec_size=sf_vec_size, + fc2_in_kernel_topk_reduce=fc2_in_kernel_topk_reduce, + token_back_by_dispatch=token_back_by_dispatch, + epi_flag_batch=epi_flag_batch, + gate_up_clamp=gate_up_clamp, + apply_topk_in_fc1=True, + generate_c=generate_c, + use_stg_fc1=use_stg_fc1, + act_func=act_func, + fc2_use_bulk=fc2_use_bulk, + fc2_tma_stages=fc2_tma_stages, + ) + + # --- Warp topology: expand to 12 warps (or 16 for standalone token-back). --- + self.enable_token_comm = True + self.dispatch_warp_id = (8, 9, 10, 11) + self.token_back_mode = token_back_mode + self.token_back_standalone = token_back_by_dispatch and token_back_mode == "standalone_warps" + self.token_back_warp_id = (12, 13, 14, 15) if self.token_back_standalone else None + num_token_back_warps = len(self.token_back_warp_id) if self.token_back_standalone else 0 + self.threads_per_cta = 32 * (len(self.epilogue_warp_id) + 4 + len(self.dispatch_warp_id) + num_token_back_warps) + + # --- MegaMoE constants. --- + self.world_size = world_size + self.local_rank = local_rank + self.num_topk = num_topk + self.max_tokens_per_rank = max_tokens_per_rank + self.max_recv_size_per_rank = max_recv_size_per_rank + self.hidden = hidden + self.launch_cluster_count = launch_cluster_count + self.drop_on_overflow = drop_on_overflow + self.combine_format = combine_format + self.num_experts_per_rank = static_expert_shape[0] + self.intermediate_gateup = static_expert_shape[1] + self.intermediate_downproj = self.intermediate_gateup // 2 + self.num_total_experts = world_size * self.num_experts_per_rank + self.reduce_topk_in_kernel = fc2_in_kernel_topk_reduce + self.token_back_schedule_mode = load_balance_mode if load_balance_mode == "atomic_counter" else "static" + + # --- next Router-push token communication component. --- + mma_cta_count = 2 if use_2cta_instrs else 1 + cta_tile_m = mma_tiler_mnk[0] // mma_cta_count + cluster_m, cluster_n = self.cluster_shape_mn + tokens_per_fc1_ready_slot = cta_tile_m * cluster_m + hidden_per_fc2_cluster_tile = cta_tile_m * cluster_m + fc2_done_signals_per_token_tile = ceil_div(hidden, hidden_per_fc2_cluster_tile) * cluster_m * cluster_n + promised_launchable_sm_count = launch_cluster_count * cluster_m * cluster_n + quant_kind = _AB_DTYPE_TO_QUANT_KIND[ab_dtype] + tc_problem_desc = ProblemDesc( + { + "world_size": world_size, + "expert_count": self.num_total_experts, + "topk": num_topk, + "max_tokens_per_rank": max_tokens_per_rank, + "max_recv_size_per_rank": max_recv_size_per_rank, + "hidden_size": hidden, + "quant_kind": str(quant_kind), + "combine_format": combine_format, + "apply_topk_at_fc1": True, + } + ) + tc_impl_desc = ImplDesc( + { + "token_padding_block": token_padding_block, + "sf_padding_block": sf_padding_block, + "tokens_per_fc1_ready_slot": tokens_per_fc1_ready_slot, + "fc2_done_signals_per_token_tile": fc2_done_signals_per_token_tile, + "promised_launchable_sm_count": promised_launchable_sm_count, + "drop_on_overflow": drop_on_overflow, + "token_in_flag_batch": flag_batch, + "token_back_mode": token_back_mode, + "token_back_schedule_mode": self.token_back_schedule_mode, + "reduce_topk_in_kernel": fc2_in_kernel_topk_reduce, + } + ) + self.token_comm = TokenCommDeterministic(tc_problem_desc, tc_impl_desc) + self.pool_token_capacity = self.token_comm.worst_case_token_count + + # --- SMEM sub-buffer for the token_comm transport (allocated in the device kernel). --- + tc_smem_ws = SmemWorkspace() + self.token_comm.register_smem_regions(tc_smem_ws) + tc_smem_ws.finalize(max_bytes=self.smem_capacity) + self.tc_smem_ws = tc_smem_ws + self._token_comm_smem_bytes = tc_smem_ws.total_bytes + + # Build the scheduler + _ec, _ig, _hd = static_expert_shape + self._build_scheduler( + expert_cnt=_ec, intermediate_gateup=_ig, hidden_dim=_hd, launch_cluster_count=launch_cluster_count + ) + self._sched_smem_bytes = self.sched_smem_ws.total_bytes + + # --- Post-kernel top-k reduction (skipped under in-kernel REDG reduce). --- + self._topk_reduce = None if fc2_in_kernel_topk_reduce else TopkReduce(hidden, num_topk, combine_format) + + # --- Device workspace (next model): fc1 pool/output + token_comm regions. --- + self._mega_device_workspace = self._build_megamoe_device_workspace() + + self.expert_count = self.num_total_experts + self.intermediate_gateup_size = self.intermediate_gateup + self.hidden_size = hidden + self.quant_kind = _quant_kind + self.topk = num_topk + self.cluster_shape_mnk = tuple(cluster_shape_mnk) + self.mma_tiler_mnk = tuple(mma_tiler_mnk) + self.group_hint = group_hint + self.token_padding_block = token_padding_block + self.sf_padding_block = sf_padding_block + self.load_balance_mode = load_balance_mode + self.force_static_sched = force_static_sched + self.clc_bundle_size = clc_bundle_size + self.num_sched_stages = num_sched_stages + self.acc_dtype = acc_dtype + self.sf_vec_size = sf_vec_size + self.fc2_in_kernel_topk_reduce = fc2_in_kernel_topk_reduce + self.epi_flag_batch = tuple(epi_flag_batch) + self.flag_batch = flag_batch + self.generate_c = generate_c + self.use_stg_fc1 = use_stg_fc1 + self.act_func = act_func + self.use_2cta_instrs = use_2cta_instrs + self.gate_up_clamp = gate_up_clamp + + # Optional standalone token-axis (column) MXFP8 requantization. + if self.enable_col_quant: + col_quant_type = "mxfp8_e4m3" if ab_dtype is cutlass.Float8E4M3FN else "mxfp8_e5m2" + self.col_quant = Mxfp8ColRequant( + hidden=self.hidden, + num_experts=self.num_experts_per_rank, + max_total_tokens=( + self.world_size + * self.max_tokens_per_rank + * min(self.num_topk, self.num_experts_per_rank) + ), + quant_type=col_quant_type, + num_persistent_ctas=self._col_quant_num_ctas, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + ) + + def get_aux_output_shapes(self) -> dict: + """Return receiver-domain auxiliary output shapes.""" + return { + "fc1_c": (self.pool_token_capacity, self.intermediate_gateup), + "col_quant_data": (self.pool_token_capacity, self.hidden), + "col_quant_sf": ( + self.token_comm.worst_case_sf_token_count + * (self.hidden // self.sf_vec_size), + ), + } + + @cute.jit + def _validate_fixed_matrix(self, tensor: cute.Tensor, dtype, expected_shape) -> None: + if cutlass.const_expr(tensor.element_type is not dtype): + raise TypeError("pool-domain matrix has an unexpected element type.") + if cutlass.const_expr(cute.rank(tensor.layout) != 2): + raise ValueError("pool-domain matrix must be rank 2.") + if cutlass.const_expr( + not isinstance(tensor.shape[0], int) + or not isinstance(tensor.shape[1], int) + or tensor.shape[0] != expected_shape[0] + or tensor.shape[1] != expected_shape[1] + ): + raise ValueError(f"pool-domain matrix must have static shape {expected_shape}.") + if cutlass.const_expr(tensor.stride[0] != expected_shape[1] or tensor.stride[1] != 1): + raise ValueError("pool-domain matrix must be compact row-major.") + + @cute.jit + def _validate_fixed_vector(self, tensor: cute.Tensor, dtype, expected_size: int) -> None: + if cutlass.const_expr(tensor.element_type is not dtype): + raise TypeError("pool-domain vector has an unexpected element type.") + if cutlass.const_expr(cute.rank(tensor.layout) != 1): + raise ValueError("pool-domain vector must be rank 1.") + if cutlass.const_expr( + not isinstance(tensor.shape[0], int) or tensor.shape[0] != expected_size + ): + raise ValueError(f"pool-domain vector must have static size {expected_size}.") + if cutlass.const_expr(tensor.stride[0] != 1): + raise ValueError("pool-domain vector must be contiguous.") + + def _smem_misc_budget_bytes(self) -> int: + """Reserve the token_comm transport SMEM on top of the base misc budget.""" + _sched = getattr(self, "_sched_smem_bytes", 0) + return super()._smem_misc_budget_bytes() + self._token_comm_smem_bytes + _sched + self._SMEM_ALLOC_MARGIN + + def _build_megamoe_device_workspace(self) -> DeviceWorkspace: + """Register the FC1 output/pool + fc1_done_counter + all token_comm regions.""" + sf_dtype = cutlass.Float8E8M0FNU + sf_column_count = round_up(ceil_div(self.intermediate_downproj, self.sf_vec_size), 4) + max_sf_rows = self.token_comm.worst_case_sf_token_count + counter_slot_count = self.token_comm.max_fc1_ready_slot_count + + device_workspace = DeviceWorkspace() + device_workspace.register( + self.fc1_output_region, + self.ab_dtype, + (self.pool_token_capacity, self.intermediate_downproj), + buffer_space="local", + mem_order=(1, 0), + byte_alignment=128, + ) + device_workspace.register( + self.fc1_output_sf_region, + sf_dtype, + (max_sf_rows, sf_column_count), + buffer_space="local", + mem_order=(1, 0), + byte_alignment=128, + ) + device_workspace.register( + self.fc1_done_counter_region, + cutlass.Int32, + (counter_slot_count,), + buffer_space="local", + byte_alignment=16, + reset="tail_reset", + ) + if self.enable_col_quant: + # Persistent per-expert token-count snapshot for the post-kernel col-quant launch. + device_workspace.register( + self.col_quant_sizes_region, + cutlass.Int32, + (self.num_experts_per_rank,), + buffer_space="local", + byte_alignment=16, + ) + self.token_comm.register_device_workspace(device_workspace) + device_workspace.finalize() + return device_workspace + + def get_workspace_sizes(self) -> Tuple[int, int]: + """Return required (local, shared/symmetric) workspace bytes.""" + return self._mega_device_workspace.local_and_shared_bytes + + @property + def require_zero_workspace_leading_bytes(self) -> Tuple[int, int]: + return self._mega_device_workspace.require_zero_workspace_leading_bytes + + # ========================================================================= + # token_comm_hook_* -- the lean device method's integration seams, here + # filled with next's Router-push TokenCommDeterministic calls. + # ========================================================================= + + def token_comm_extra_smem_storage_class(self) -> type: + """SMEM struct for the token_comm transport overlay (allocated in the device kernel).""" + return self.tc_smem_ws.storage_class() + + def token_comm_hook_fc1_ready_counter_ptr(self, token_comm_args): + """Pointer the FC1 scheduler/extension spins on; token_in increments it per ready slot.""" + return self.token_comm.fc1_ready_counter_pointer(self._mega_device_workspace) + + def sched_ext_fc1_peek_threshold(self) -> int: # noqa: D401 - lean hook override point + return super().sched_ext_fc1_peek_threshold() + + @cute.jit + def token_comm_hook_sched_warp_pre_init_wait(self, token_comm_args): + """The scheduler warp must wait for the Router to publish per-expert sizes.""" + self.token_comm.wait_for_sizes_ready(self._mega_device_workspace) + + @cute.jit + def token_comm_hook_fc1_tma_b_predispatch_spin(self, token_comm_args, work_tile_info): + """No-op: FC1 input readiness is enforced by the scheduler extension's fc1_ready spin.""" + pass + + @cute.jit + def token_comm_hook_dispatch_warp_body(self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx): + """Transfer warps (8-11): pull activation from peers into the local FC1 pool.""" + self.token_comm.token_in(self.tc_smem_ws, token_comm_storage.buffer.data_ptr()) + if cutlass.const_expr(self.token_comm.token_back_enabled and not self.token_back_standalone): + self.token_comm.token_back(self.tc_smem_ws, token_comm_storage.buffer.data_ptr()) + + @cute.jit + def token_comm_hook_token_back_warp_body(self, token_comm_args, token_comm_storage, *, warp_idx, lane_idx, tidx): + """Standalone token-back warps (12-15): push FC2 output back to source ranks.""" + self.token_comm.token_back(self.tc_smem_ws, token_comm_storage.buffer.data_ptr()) + + @cute.jit + def token_comm_hook_tail_reset_shared_counters(self, token_comm_args, *, warp_idx, lane_idx, tidx): + """Absorbed into reset_tail (kernel_tail hook).""" + pass + + @cute.jit + def token_comm_hook_kernel_tail(self, token_comm_args, *, warp_idx, lane_idx, tidx): + """Cross-rank drain + workspace tail reset, performed by the transfer warps. + + The whole-CTA barrier is REQUIRED (mirrors inference's mainloop kernel-tail + sync_threads): it forces every compute warp (scheduler / tma / mma / epilogue) to + finish its consume loop -- including any ``ext.wait_for_input`` spin on ``fc1_ready`` + -- before the transfer warps run ``reset_tail``. ``reset_tail`` tail-resets the + (local, GPU-wide) ``fc1_ready`` / ``fc2_done`` counters; without this barrier a fast + CTA's transfer warps can zero a counter that a slower compute warp is still spinning + on, which non-deterministically deadlocks (~1/5 runs). + """ + # Preserve the per-expert dispatch token counts before reset_tail zeros the + # shared sizes_region. + if cutlass.const_expr(self.enable_col_quant): + self._snapshot_col_quant_expert_sizes(tidx) + cute.arch.sync_threads() + # reset_tail must run on EXACTLY the 4 transfer/token_in warps (8-11). For + # standalone_warps (16 warps) the token_back warps (12-15) must be EXCLUDED: their + # thread_idx aliases (thread_idx % transfer_thread_count) back onto transfer warps 0-3, + # so including them double-counts the reset_tail NVLink grid barrier -> deadlock. + if (warp_idx >= self.dispatch_warp_id[0]) & (warp_idx <= self.dispatch_warp_id[-1]): + self.token_comm.reset_tail() + self.token_comm.remove_device_members() + + @cute.jit + def _snapshot_col_quant_expert_sizes(self, tidx) -> None: + from cutlass.cutlass_dsl import Int32, Int64 + """CTA 0 copies this rank's per-expert token counts into the persistent + col-quant region before token_comm's tail_reset zeros ``sizes_region``.""" + dw = self._mega_device_workspace + if self.token_comm._linear_cta_idx == Int32(0): + sizes = self.token_comm.local_expert_sizes(dw, self.token_comm._local_rank) + snapshot = dw.tensor(self.col_quant_sizes_region) + block_dim_x, _, _ = cute.arch.block_dim() + expert_idx = tidx + while expert_idx < Int32(self.num_experts_per_rank): + snapshot[expert_idx] = Int32(sizes[expert_idx]) + expert_idx = expert_idx + block_dim_x + + # ========================================================================= + # Host launch: Router kernel -> fused MegaMoE main kernel -> top-k reduction. + # ========================================================================= + + @cute.jit + def __call__( + self, + activation: cute.Tensor, # (max_tokens_per_rank, hidden) raw per-rank, symmetric heap + activation_sf: cute.Tensor, # (max_tokens_per_rank, hidden // sf_vec_size), symmetric + topk_indices: cute.Tensor, # (max_tokens_per_rank, topk) + topk_scores: cute.Tensor, # (max_tokens_per_rank, topk) Float32 + fc1_weight: cute.Tensor, # (experts_per_rank, hidden, intermediate_gateup) + fc1_weight_sf: cute.Tensor, + fc2_weight: cute.Tensor, # (experts_per_rank, intermediate_downproj, hidden) + fc2_weight_sf: cute.Tensor, + output_activation: cute.Tensor, # (max_tokens_per_rank, topk, hidden) final combined output + fc1_c: Optional[cute.Tensor], # (pool_token_capacity, intermediate_gateup) when generate_c=True + col_quant_data: Optional[cute.Tensor], # dispatch-pool-strided fp8 segments + col_quant_sf: Optional[cute.Tensor], # flat concat_e [hidden_atom][token_atom] E8M0 bytes + overflow_flag: cute.Tensor, # (1,) Int32, per-rank router receive-overflow output + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, # symmetric (NVLink) heap base + peer_rank_ptr_mapper_host, + stream: cuda.CUstream, + ) -> None: + """Launch the Router, then the fused main kernel, then (optionally) the top-k reduce.""" + from cutlass.cutlass_dsl import Int64 + + dw = self._mega_device_workspace + local_rank = peer_rank_ptr_mapper_host.rank + aux_shapes = self.get_aux_output_shapes() + if cutlass.const_expr(self.generate_c): + if cutlass.const_expr(fc1_c is None): + raise ValueError("generate_c=True requires a receiver-domain fc1_c tensor.") + self._validate_fixed_matrix( + fc1_c, cutlass.BFloat16, aux_shapes["fc1_c"] + ) + if cutlass.const_expr(self.enable_col_quant): + if cutlass.const_expr(col_quant_data is None or col_quant_sf is None): + raise ValueError("enable_col_quant=True requires data and scale outputs.") + self._validate_fixed_matrix( + col_quant_data, self.ab_dtype, aux_shapes["col_quant_data"] + ) + self._validate_fixed_vector( + col_quant_sf, cutlass.Uint8, aux_shapes["col_quant_sf"][0] + ) + self.token_comm.launch_router( + topk_indices=topk_indices, + topk_scores=topk_scores, + local_rank=local_rank, + local_workspace=local_workspace, + shared_workspace=shared_workspace, + peer_rank_ptr_mapper_host=peer_rank_ptr_mapper_host, + device_workspace=dw, + overflow_flag=overflow_flag, + stream=stream, + ) + peer_mapper = peer_rank_ptr_mapper_host.make_device_object() + dw.assign_device_members(local_workspace, shared_workspace) + + activation_pool = self.token_comm.fc1_activation_tensor(dw) + _sf_pool_atom = self.token_comm.fc1_activation_sf_tensor(dw) + activation_sf_pool = cute.make_tensor( + _sf_pool_atom.iterator, + cute.make_layout( + (self.token_comm.worst_case_sf_token_count, self.hidden // self.sf_vec_size), + stride=(self.token_comm.activation_sf_hidden_padded, 1), + ), + ) + fc1_output = dw.tensor(self.fc1_output_region) + fc1_output_sf = dw.tensor(self.fc1_output_sf_region) + fc1_done_counter = dw.tensor(self.fc1_done_counter_region) + pool_topk_scores = self.token_comm.fc1_topk_scores_tensor(dw) + + if cutlass.const_expr(self.reduce_topk_in_kernel): + # In-kernel top-k reduce (epi_warps + bf16 combine) + pre_reduced = cute.make_tensor( + output_activation.iterator, + cute.make_layout( + (output_activation.shape[0], 1, output_activation.shape[1]), + stride=(output_activation.stride[0], output_activation.stride[0], output_activation.stride[1]), + ), + ) + pre_reduced_sf = None + else: + pre_reduced = self.token_comm.pre_reduced_activation_tensor(dw) + pre_reduced_sf = self.token_comm.pre_reduced_activation_sf_tensor(dw) + + if cutlass.const_expr(self.token_comm.token_back_push_data): + # token_back-by-dispatch (standalone_warps / reuse_dispatch_warps) + fc2_output = self.token_comm.fc2_activation_tensor(dw) + else: + # epi_warps: the epilogue peer-writes FC2 directly. + _combine_hidden = pre_reduced.shape[2] + fc2_output = cute.make_tensor( + pre_reduced.iterator, + cute.make_layout( + (pre_reduced.shape[0] * pre_reduced.shape[1], _combine_hidden), stride=(_combine_hidden, 1) + ), + ) + + super().__call__( + activation_pool, + fc1_weight, + activation_sf_pool, + fc1_weight_sf, + fc1_output, + fc1_output_sf, + fc2_weight, + fc2_weight_sf, + fc2_output, + pool_topk_scores, + fc1_done_counter, + offs=None, + max_active_clusters=self.launch_cluster_count, + stream=stream, + fc1_c=fc1_c, + overflow_flag=overflow_flag, + mega_peer_rank_ptr_mapper=peer_mapper, + mega_local_rank=local_rank, + mega_local_workspace=local_workspace, + mega_shared_workspace=shared_workspace, + mega_activation=activation, + mega_activation_sf=activation_sf, + mega_pre_reduced_activation=pre_reduced, + mega_pre_reduced_activation_sf=pre_reduced_sf, + ) + + # Top-k weights were already applied before FC1 quantization, so the + # post-kernel reduction is a plain dequantized sum. + if cutlass.const_expr(not self.reduce_topk_in_kernel): + self._topk_reduce(pre_reduced, pre_reduced_sf, output_activation, None, stream) + + # Post-kernel token-axis MXFP8 requantization of the preserved dispatch-local + # FC1 activation pool, stored by the tail snapshot of per-expert token counts. + if cutlass.const_expr(self.enable_col_quant): + lw = local_workspace + data_offset = dw.offset(self.token_comm.fc1_activation_region) + sf_offset = dw.offset(self.token_comm.fc1_activation_sf_region) + sizes_offset = dw.offset(self.col_quant_sizes_region) + sf_pool_bytes = self.token_comm.worst_case_sf_token_count * (self.hidden // self.sf_vec_size) + src_data = cute.make_tensor( + cute.make_ptr( + self.ab_dtype, lw.toint() + Int64(data_offset), AddressSpace.gmem, assumed_align=128 + ), + cute.make_layout( + (self.token_comm.worst_case_token_count, self.hidden), stride=(self.hidden, 1) + ), + ) + src_sf_u8 = cute.make_tensor( + cute.make_ptr(cutlass.Uint8, lw.toint() + Int64(sf_offset), AddressSpace.gmem, assumed_align=16), + cute.make_layout((sf_pool_bytes,)), + ) + expert_sizes = cute.make_tensor( + cute.make_ptr(cutlass.Int32, lw.toint() + Int64(sizes_offset), AddressSpace.gmem, assumed_align=16), + cute.make_layout((self.num_experts_per_rank,)), + ) + self.col_quant( + src_data, + src_sf_u8, + expert_sizes, + col_quant_data, + col_quant_sf, + stream, + ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/__init__.py new file mode 100644 index 000000000..82b548c79 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Training mega helpers (SwiGLU / mxfp8 register-level quant primitives).""" + +from .constants import ( + SupportedMmaTileM, + SupportedMmaTileN, +) +from .utils import dswiglu_act, quant_sfd_col, quant_sfd_row, swiglu_act + +__all__ = [ + "SupportedMmaTileM", + "SupportedMmaTileN", + "dswiglu_act", + "quant_sfd_col", + "quant_sfd_row", + "swiglu_act", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/constants.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/constants.py new file mode 100644 index 000000000..1bd314338 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/constants.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Training-mega-specific numeric constants (MMA tiler extents). + +These are consumed only by the training GLU/dGLU FC12 kernels, so they live in the +training mega ``helpers`` package rather than the cross-inference/training +``next/sources/helpers/constants.py``. +""" + + +# MMA tiler GLU FC12 kernels accept along M and N. +SupportedMmaTileM = (128, 256) +SupportedMmaTileN = (128, 256) + + +__all__ = [ + "SupportedMmaTileM", + "SupportedMmaTileN", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/utils.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/utils.py new file mode 100644 index 000000000..a015fe2a5 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/helpers/utils.py @@ -0,0 +1,322 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +from typing import Optional + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass._mlir.dialects import math as _math +from cutlass.cutlass_dsl import Float32, T, dsl_user_op + +from ......helpers.constants import Fp32Max, Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Log2E +from ......helpers.ptx_helpers import cvt_f32_to_fp8_to_f32, cvt_f32x4_to_f8x4_pack_i32 + + +@dsl_user_op +def zero_unless_equal( + value: Float32, + raw: Float32, + clamped: Float32, + *, + loc=None, + ip=None, +) -> Float32: + """Return ``value`` if ``raw == clamped``, using branch-free PTX.""" + return Float32( + llvm.inline_asm( + T.f32(), + [ + Float32(value).ir_value(loc=loc, ip=ip), + Float32(raw).ir_value(loc=loc, ip=ip), + Float32(clamped).ir_value(loc=loc, ip=ip), + ], + "{\n" + " .reg .pred in_range;\n" + " setp.eq.f32 in_range, $2, $3;\n" + " selp.f32 $0, $1, 0f00000000, in_range;\n" + "}", + "=f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def swiglu_act( + t_swiglu: cute.Tensor, + t_up: cute.Tensor, + t_gate: cute.Tensor, + prob: Optional[Float32] = None, + gate_up_clamp: Optional[Float32] = None, +) -> None: + """SwiGLU with optional gate-upper/up-symmetric clamp.""" + for i in cutlass.range_constexpr(0, cute.size(t_swiglu), 2): + gate = (t_gate[i], t_gate[i + 1]) + up = (t_up[i], t_up[i + 1]) + if cutlass.const_expr(gate_up_clamp is not None): + gate = ( + cute.arch.fmin(gate[0], gate_up_clamp), + cute.arch.fmin(gate[1], gate_up_clamp), + ) + up = ( + cute.arch.fmax(cute.arch.fmin(up[0], gate_up_clamp), -gate_up_clamp), + cute.arch.fmax(cute.arch.fmin(up[1], gate_up_clamp), -gate_up_clamp), + ) + up_gate = cute.arch.mul_packed_f32x2( + up, + gate, + rnd="rn", + ftz=False, + ) + gate_log2e = cute.arch.mul_packed_f32x2( + gate, (-Log2E, -Log2E), rnd="rn", ftz=False, + ) + one_plus_exp = cute.arch.add_packed_f32x2( + ( + cute.math.exp2(gate_log2e[0], fastmath=True), + cute.math.exp2(gate_log2e[1], fastmath=True), + ), + (1.0, 1.0), + ) + sigmoid = (cute.arch.rcp_approx(one_plus_exp[0]), cute.arch.rcp_approx(one_plus_exp[1])) + (t_swiglu[i], t_swiglu[i + 1]) = cute.arch.mul_packed_f32x2( + up_gate, sigmoid, rnd="rn", ftz=False, + ) + if cutlass.const_expr(prob is not None): + (t_swiglu[i], t_swiglu[i + 1]) = cute.arch.mul_packed_f32x2( + (t_swiglu[i], t_swiglu[i + 1]), + (prob, prob), + rnd="rn", + ftz=False, + ) + + +@cute.jit +def quant_sfd_row( + src: cute.Tensor, + dst: cute.Tensor, + norm_const, + sf_vec_size, + sf_dtype, + d_dtype, +): + """Quantize the ``sf_vec_size`` values in ``src`` to ``d_dtype`` with one block scale.""" + rcp_limit = Fp8E4M3RcpLimit if d_dtype == cutlass.Float8E4M3FN else Fp8E5M2RcpLimit + acc_frg = src.load() + abs_acc_frg_ir = _math.absf(acc_frg.ir_value()) + abs_acc_frg = type(acc_frg)(abs_acc_frg_ir, acc_frg.shape, acc_frg.dtype) + # Fuse the two loop-invariant constants into one multiply + rcp_limit_norm = rcp_limit * norm_const + avg_fp32 = abs_acc_frg.reduce(cute.ReductionOp.MAX, Float32(0.0), 0) * rcp_limit_norm + qpvscale_up = cvt_f32_to_fp8_to_f32(avg_fp32, sf_dtype) + acc_scale = norm_const * cute.arch.rcp_approx(qpvscale_up) + acc_scale = cute.arch.fmin(acc_scale, Fp32Max, nan=True) + for ei in cutlass.range_constexpr(0, sf_vec_size, 2): + src[ei], src[ei + 1] = cute.arch.mul_packed_f32x2( + (src[ei], src[ei + 1]), (acc_scale, acc_scale), rnd="rn", ftz=False, + ) + dst_i32 = cute.recast_tensor(dst, cutlass.Int32) + for ei in cutlass.range_constexpr(0, sf_vec_size, 4): + fp32x4 = cute.make_rmem_tensor(4, Float32) + fp32x4[0] = src[ei + 0] + fp32x4[1] = src[ei + 1] + fp32x4[2] = src[ei + 2] + fp32x4[3] = src[ei + 3] + fp8x4_i32 = cvt_f32x4_to_f8x4_pack_i32(fp32x4, d_dtype) + dst_i32[ei // 4] = cutlass.Int32(fp8x4_i32) + return qpvscale_up + + +@cute.jit +def dswiglu_act( + t_dgate: cute.Tensor, + t_dup: cute.Tensor, + t_acc: cute.Tensor, + t_gate: cute.Tensor, + t_up: cute.Tensor, + beta_val: Float32, + prob: Float32, + gate_up_clamp: Optional[Float32] = None, +) -> Float32: + """SwiGLU backward with optional clamp, beta/prob scaling, and dprob. + + Given upstream gradient ``acc``, per-expert scalar ``beta_val``, per-token routing + probability ``prob``, and forward pre-activations ``gate``/``up``:: + + gate_raw = gate * beta_val + up_raw = up * beta_val + gate_b = min(gate_raw, clamp) + up_b = clamp(up_raw, -clamp, clamp) + sig = sigmoid(gate_b) + swish = gate_b * sig + + dprob += acc * up_b * swish (returned to the caller) + d_up = acc * prob * swish * I[-clamp <= up_raw <= clamp] + d_gate = acc * prob * up_b * silu'(gate_b) * I[gate_raw <= clamp] + + The clamp is skipped when ``gate_up_clamp`` is ``None``. Boundary values retain + their gradient, matching ``torch.clamp``. + """ + dprob_acc = Float32(0.0) + for i in cutlass.range_constexpr(0, cute.size(t_acc), 2): + gate_raw = cute.arch.mul_packed_f32x2( + (t_gate[i], t_gate[i + 1]), (beta_val, beta_val), rnd="rn", ftz=False, + ) + up_raw = cute.arch.mul_packed_f32x2( + (t_up[i], t_up[i + 1]), (beta_val, beta_val), rnd="rn", ftz=False, + ) + gate_b = gate_raw + up_b = up_raw + if cutlass.const_expr(gate_up_clamp is not None): + gate_b = ( + cute.arch.fmin(gate_raw[0], gate_up_clamp), + cute.arch.fmin(gate_raw[1], gate_up_clamp), + ) + up_b = ( + cute.arch.fmax(cute.arch.fmin(up_raw[0], gate_up_clamp), -gate_up_clamp), + cute.arch.fmax(cute.arch.fmin(up_raw[1], gate_up_clamp), -gate_up_clamp), + ) + + # sig = 1 / (1 + exp(-gate_b)); exp(-x) = exp2(-Log2E * x) + sig_rcp = cute.arch.mul_packed_f32x2( + gate_b, (-Log2E, -Log2E), rnd="rn", ftz=False, + ) + (sig0, sig1) = cute.arch.add_packed_f32x2( + ( + cute.math.exp2(sig_rcp[0], fastmath=True), + cute.math.exp2(sig_rcp[1], fastmath=True), + ), + (1.0, 1.0), + ) + sig0 = cute.arch.rcp_approx(sig0) + sig1 = cute.arch.rcp_approx(sig1) + + # swish = gate_b * sig + swish = cute.arch.mul_packed_f32x2(gate_b, (sig0, sig1), rnd="rn", ftz=False) + + # dprob += acc * up_b * swish (both lanes into the running scalar) + dp = cute.arch.mul_packed_f32x2( + (t_acc[i], t_acc[i + 1]), (up_b[0], up_b[1]), rnd="rn", ftz=False, + ) + dp = cute.arch.mul_packed_f32x2(dp, swish, rnd="rn", ftz=False) + dprob_acc = dprob_acc + dp[0] + dp[1] + + # acc * prob (shared factor for d_up and d_gate) + acc_prob = cute.arch.mul_packed_f32x2( + (t_acc[i], t_acc[i + 1]), (prob, prob), rnd="rn", ftz=False, + ) + + # d_up = acc * prob * swish + (t_dup[i], t_dup[i + 1]) = cute.arch.mul_packed_f32x2( + acc_prob, swish, rnd="rn", ftz=False, + ) + + # d_gate = acc * prob * up_b * sig * (1 + gate_b * (1 - sig)) + one_minus_sig = cute.arch.add_packed_f32x2( + (1.0, 1.0), (-sig0, -sig1), rnd="rn", ftz=False, + ) + dsig = cute.arch.mul_packed_f32x2(gate_b, one_minus_sig, rnd="rn", ftz=False) + term = cute.arch.add_packed_f32x2( + (dsig[0], dsig[1]), (1.0, 1.0), rnd="rn", ftz=False, + ) + dgate = cute.arch.mul_packed_f32x2( + acc_prob, (up_b[0], up_b[1]), rnd="rn", ftz=False, + ) + dgate = cute.arch.mul_packed_f32x2(dgate, (sig0, sig1), rnd="rn", ftz=False) + (t_dgate[i], t_dgate[i + 1]) = cute.arch.mul_packed_f32x2( + dgate, term, rnd="rn", ftz=False, + ) + if cutlass.const_expr(gate_up_clamp is not None): + t_dgate[i] = zero_unless_equal(t_dgate[i], gate_raw[0], gate_b[0]) + t_dgate[i + 1] = zero_unless_equal(t_dgate[i + 1], gate_raw[1], gate_b[1]) + t_dup[i] = zero_unless_equal(t_dup[i], up_raw[0], up_b[0]) + t_dup[i + 1] = zero_unless_equal(t_dup[i + 1], up_raw[1], up_b[1]) + + return dprob_acc + + +@cute.jit +def quant_sfd_col( + src: cute.Tensor, + dst: cute.Tensor, + norm_const, + sf_vec_size, + sf_dtype, + d_dtype, +): + """Column (cross-thread) block-scale quantize: the amax is a warp reduction.""" + rcp_limit = Fp8E4M3RcpLimit if d_dtype == cutlass.Float8E4M3FN else Fp8E5M2RcpLimit + acc_frg = src.load() + abs_acc_frg_ir = _math.absf(acc_frg.ir_value()) + acc_frg = type(acc_frg)(abs_acc_frg_ir, acc_frg.shape, acc_frg.dtype) + + qpvscale_up = Float32(0.0) + tidx, _, _ = cute.arch.thread_idx() + scale = rcp_limit * norm_const + + for vi in cutlass.range_constexpr(0, sf_vec_size, 4): + # Warp-wide MAX across the 32 rows for each of the 4 lanes. + max_value0 = Float32(cute.arch.warp_redux_sync(acc_frg[vi], "fmax", nan=True)) + max_value1 = Float32(cute.arch.warp_redux_sync(acc_frg[vi + 1], "fmax", nan=True)) + max_value2 = Float32(cute.arch.warp_redux_sync(acc_frg[vi + 2], "fmax", nan=True)) + max_value3 = Float32(cute.arch.warp_redux_sync(acc_frg[vi + 3], "fmax", nan=True)) + + (max_value0, max_value1) = cute.arch.mul_packed_f32x2( + (max_value0, max_value1), (scale, scale), rnd="rn", ftz=False, + ) + (max_value2, max_value3) = cute.arch.mul_packed_f32x2( + (max_value2, max_value3), (scale, scale), rnd="rn", ftz=False, + ) + + max_value0 = cvt_f32_to_fp8_to_f32(max_value0, sf_dtype) + max_value1 = cvt_f32_to_fp8_to_f32(max_value1, sf_dtype) + max_value2 = cvt_f32_to_fp8_to_f32(max_value2, sf_dtype) + max_value3 = cvt_f32_to_fp8_to_f32(max_value3, sf_dtype) + + # Each thread keeps its assigned column's pre-round-trip scale. + if tidx % 32 == vi: + qpvscale_up = max_value0 + if tidx % 32 == vi + 1: + qpvscale_up = max_value1 + if tidx % 32 == vi + 2: + qpvscale_up = max_value2 + if tidx % 32 == vi + 3: + qpvscale_up = max_value3 + + max_value_rcp0 = cute.arch.fmin(cute.arch.rcp_approx(max_value0), Fp32Max, nan=True) + max_value_rcp1 = cute.arch.fmin(cute.arch.rcp_approx(max_value1), Fp32Max, nan=True) + max_value_rcp2 = cute.arch.fmin(cute.arch.rcp_approx(max_value2), Fp32Max, nan=True) + max_value_rcp3 = cute.arch.fmin(cute.arch.rcp_approx(max_value3), Fp32Max, nan=True) + + (acc_scale_col0, acc_scale_col1) = cute.arch.mul_packed_f32x2( + (norm_const, norm_const), (max_value_rcp0, max_value_rcp1), rnd="rn", ftz=False, + ) + (acc_scale_col2, acc_scale_col3) = cute.arch.mul_packed_f32x2( + (norm_const, norm_const), (max_value_rcp2, max_value_rcp3), rnd="rn", ftz=False, + ) + + (src[vi], src[vi + 1]) = cute.arch.mul_packed_f32x2( + (src[vi], src[vi + 1]), (acc_scale_col0, acc_scale_col1), rnd="rn", ftz=False, + ) + (src[vi + 2], src[vi + 3]) = cute.arch.mul_packed_f32x2( + (src[vi + 2], src[vi + 3]), (acc_scale_col2, acc_scale_col3), rnd="rn", ftz=False, + ) + + dst_i32 = cute.recast_tensor(dst, cutlass.Int32) + for ei in cutlass.range_constexpr(0, sf_vec_size, 4): + fp32x4 = cute.make_rmem_tensor(4, Float32) + fp32x4[0] = src[ei + 0] + fp32x4[1] = src[ei + 1] + fp32x4[2] = src[ei + 2] + fp32x4[3] = src[ei + 3] + fp8x4_i32 = cvt_f32x4_to_f8x4_pack_i32(fp32x4, d_dtype) + dst_i32[ei // 4] = cutlass.Int32(fp8x4_i32) + return qpvscale_up + + +__all__ = ["dswiglu_act", "quant_sfd_col", "quant_sfd_row", "swiglu_act"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/tmem_transpose.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/tmem_transpose.py new file mode 100644 index 000000000..12767ec26 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/tmem_transpose.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Rubin-training source-copy shim for the 16x32 TMEM transpose core. + +``_TmemTranspose16x32Core`` is the register-level transpose helper shared with +the Blackwell swap-AB epilogue. It is arch-compatible (identical math), so we +re-export it through a marked import rather than re-porting the transpose, and +rather than reaching into another kernel product's directory at port time -- +the kernel_export script inlines the source here. +""" + +# <<>> +from ....blackwell.inference.mega.block_scaled_swap_ab_fc12_epilogue import ( + _TmemTranspose16x32Core, +) + + +__all__ = ["_TmemTranspose16x32Core"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/topk_reduce.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/topk_reduce.py new file mode 100644 index 000000000..5eb1d6102 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/topk_reduce.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Rubin-training source-copy shim for the compatible Blackwell TopK reduction. + +Identical implementation to the Blackwell / inference ``TopkReduce``; kept as a +marked import so ``rubin.training.mega`` stays a self-contained deliverable and +the kernel_export script can inline the source instead of pulling in another +kernel product's directory. +""" + +# <<>> +from ....blackwell.inference.mega.topk_reduce import TopkReduce + + +__all__ = ["TopkReduce"] From a92b2d5cd340eb24f0db5b1fdb11cbf7c7120e99 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Mon, 24 Aug 2026 16:41:35 -0700 Subject: [PATCH 04/31] feat: add the MoeEp runtime resource layer Manage NVSHMEM lifecycle, symmetric workspaces, capability checks, and execution plans behind a lazy backend seam for multi-rank expert parallelism. --- .../cudnn/moe_ep/_megamoe_backend/README.md | 135 ++++ .../cudnn/moe_ep/_megamoe_backend/__init__.py | 4 + .../moe_ep/_megamoe_backend/_capability.py | 185 ++++++ python/cudnn/moe_ep/_megamoe_backend/_comm.py | 274 ++++++++ python/cudnn/moe_ep/_megamoe_backend/_plan.py | 181 ++++++ .../cudnn/moe_ep/_megamoe_backend/_runtime.py | 611 ++++++++++++++++++ .../moe_ep/_megamoe_backend/_workspace.py | 465 +++++++++++++ 7 files changed, 1855 insertions(+) create mode 100644 python/cudnn/moe_ep/_megamoe_backend/README.md create mode 100644 python/cudnn/moe_ep/_megamoe_backend/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/_capability.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/_comm.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/_plan.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/_runtime.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/_workspace.py diff --git a/python/cudnn/moe_ep/_megamoe_backend/README.md b/python/cudnn/moe_ep/_megamoe_backend/README.md new file mode 100644 index 000000000..8fddc1a52 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/README.md @@ -0,0 +1,135 @@ +# MegaMoE backend capabilities + +This private backend implements the public `cudnn.moe_ep` contract with Rubin +SM107 CuTeDSL products. Public validation and backend capability checks are +separate: a request may be valid for `MoeEp` but unavailable in this backend. + +## Implemented forward paths + +- MXFP8 inputs use the training MegaMoE forward GLU product. +- Plain BF16, FP16, and FP32 operands are quantized into the same logical + MXFP8 representation before launch. +- Block-scaled operands must use the public MXFP8 representation. Native NVFP4 + operands are part of the public contract but are not executable in this + backend. +- Combine format may be BF16 or MXFP8. Final output format is BF16. +- MXFP8 combine quantizes each FP32 route accumulator directly. +- Rubin training execution requires `apply_topk_in_fc1=True`. +- Forward with `generate_c=False` supports CUDA Graph capture after warmup; + `generate_c=True` is eager-only. + +## Explicit Rubin limits + +- CUDA compute capability must be 10.7. +- `max_tokens_per_rank` must be positive and explicit. +- `hidden_size` must be divisible by 128, and `intermediate_size` must be + divisible by 256. +- `top_k` must not exceed 32. +- Forward EP size must not exceed 16 because the validated peer-mapper ABI + carries a fixed 128-byte by-value offset table. + +These are backend limits, not additional public `MoeEp` semantics. They remain +precise, product-specific capability gates rather than hidden padding or a +silent numerical fallback. + +## Backward status + +`MoeEp.backward` has a validated backend seam and requires a forward stash from +`generate_c=True`. In the default `backward_wgrad_mode="none"`, the restricted +Rubin MXFP8 path returns +`(grad_activation, grad_topk_weights)` for EP1/EP2/EP4 with BF16 or MXFP8 +combine, BF16 output, `apply_topk_in_fc1=True`, optional `gate_up_clamp`, and +eager execution. It uses `fc1_c` and `route_metadata` to reconstruct an +external pool-layout `fc1_preact` tensor and converts `grad_output` to FP32 +before re-dispatching it for semantic dprob. The kernel's source-domain dprob +plane is symmetric and reset before every launch; the public router-weight +gradient remains an FP32 semantic recomputation. Default mode does not accept +or retain the forward activation, does not produce FC1/FC2 wgrad operands, and +does not depend on "most recent forward" state. + +### Opt-in grouped-wgrad operands + +Constructing with `backward_wgrad_mode="operands"` requires +`generate_c=True`, `token_padding_size=256`, and `sf_padding_size=128`. +Forward then returns +`(output, fc1_c, route_metadata, wgrad_forward_stash)`. The fourth value is a +`MoeEpWgradForwardStash` for that exact routed call: MXFP8 `x.T` data/scales, +cumulative padded expert offsets, valid route counts, and the same route +metadata. + +Backward takes the stash by keyword: + +```python +grad_activation, grad_topk_weights, operands = op.backward( + grad_output, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + wgrad_forward_stash=wgrad_forward_stash, +) +``` + +`MoeEpWgradOperands` is directly shaped for the grouped-wgrad Tensor2D ABI: + +- `fc1_a=(H,Kp)`, `fc1_b=(Kp,2I)` represent + `dW1 = x.T @ dC`; +- `fc2_a=(I,Kp)`, `fc2_b=(Kp,H)` represent the upstream factorization + `dW2 = (p * h).T @ dY`; +- every local expert's valid rows precede zero padding to 256 routes; + `expert_offsets` contains cumulative padded ends and may repeat for empty + experts, while `valid_route_counts` excludes padding; +- E4M3 data uses unit stride on K. E8M0 scales represent logical 1x32 blocks + assembled into the grouped kernel's physical 128x4 layout. + +The device order is deliberate. Forward first MXFP8-stages `x` along H, then +column-requantizes routed/padded rows along K. Backward follows the upstream +FC2-gradient factorization: the recomputed `h` export carries the route weight, +while the token-axis `grad_y2` export is the unweighted routed `dY`. Their +grouped product is therefore `dW2 = (p*h).T @ dY`. Staged `dY` and `W2.T` +produce `dH`, the route weight is applied before the SwiGLU derivative, and +`dC` is directly column-requantized along K. All three backward auxiliary +scale outputs use the upstream MN-major 128-column by 4-token-block atom +layout. + +The forward stash and backward operand tensors are caller-owned fresh +allocations, not views of reusable execution-plan workspace. Callers must keep +the forward stash alive through its matching backward call and keep returned +operands alive until external grouped-wgrad work completes. Later operator +calls do not overwrite them. Route identity/count validation prevents mixing +stashes from different forwards. + +This mode only produces operands; it does not launch grouped wgrad or return +dense `dW1`/`dW2`. It remains eager-only and inherits the restricted Rubin +MXFP8 backward gates (SM107, EP1/EP2/EP4, BF16 output, +`apply_topk_in_fc1=True`, and BF16/MXFP8 combine). End-to-end operand +production still requires SM107 acceptance. Direct FC1/FC2 consumer execution +has been validated separately on SM100 with reference-generated operands. + +The dGLU product emits BF16 `grad_activation`; the backend converts it to FP32 +for the public return. This is a documented BF16-rounded numerical limitation, +not strict FP32 dgrad parity. `apply_topk_in_fc1=False`, NVFP4 operands or +combine, non-BF16 output, backward CUDA Graph capture, and EP sizes outside +1/2/4 remain capability-gated. + +## Validation boundary + +L0 tests cover public validation, plain-to-MXFP8 staging, compile/cache keys, +workspace sizing, combine semantics, overflow audit behavior, and backward +layout/dispatch capability gates. Current CUDA Graph acceptance covers EP1 +forward. The single-node eager forward suite defines WORLD EP2/EP3/EP4, and +the current SM107 L1 hardware run establishes PASS for all three EP sizes. +The torchrun-native multi-node suite balances NVSHMEM PEs across participating +nodes: EP7 uses a WORLD14 subgroup over seven nodes with two workers per node, +EP12 uses WORLD12 over three nodes with four workers per node, EP15 uses a +WORLD20 subgroup over five nodes with four workers per node, and EP16 uses +WORLD16 over four nodes with four workers per node. Multi-node collection or +skip results do not establish a hardware PASS. Current hardware runs establish +PASS for EP12 and EP16; EP7 and EP15 remain pending. + +End-to-end device forward/backward parity requires SM107 hardware and the +`moe_ep` optional runtime dependencies, including a CuTeDSL installation that +provides `cutlass.utils.rubin_helpers`. Backward acceptance remains limited to +EP1/EP2/EP4 and is not expanded by the multi-node forward suite. diff --git a/python/cudnn/moe_ep/_megamoe_backend/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/__init__.py new file mode 100644 index 000000000..8187472b9 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Private MegaMoE backend implementation for :mod:`cudnn.moe_ep`.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/_capability.py b/python/cudnn/moe_ep/_megamoe_backend/_capability.py new file mode 100644 index 000000000..d8d0124cf --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/_capability.py @@ -0,0 +1,185 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Capability policy for the private MegaMoE execution backend. + +Inputs reaching this module already satisfy the public :mod:`cudnn.moe_ep` +contract. These checks only describe the subset that the current MegaMoE +implementation can execute; they must run before runtime initialization, +allocation, compilation, or collectives. +""" + +from __future__ import annotations + +import torch + +from .._contracts import ( + ForwardConfig, + ValidatedBackwardRequest, + ValidatedForwardRequest, +) +from .._types import BlockScaledTensor, MoeFormat + + +def _validate_operand(name, tensor) -> None: + if isinstance(tensor, BlockScaledTensor): + if tensor.format is MoeFormat.MXFP8: + return + raise NotImplementedError( + "MoeEp training MegaMoE supports only MXFP8 BlockScaledTensor " + f"inputs; {name} has format={tensor.format.value!r}" + ) + if tensor.dtype not in { + torch.bfloat16, + torch.float16, + torch.float32, + }: + raise NotImplementedError( + f"MoeEp MegaMoE {name} staging supports BF16, FP16, " + f"or FP32 plain tensors, got {tensor.dtype}" + ) + + +def _validate_device(device: torch.device) -> None: + if device.type != "cuda": + raise NotImplementedError( + f"MoeEp MegaMoE backend requires a CUDA device, got {device}" + ) + + major, minor = torch.cuda.get_device_capability(device) + if (major, minor) != (10, 7): + raise NotImplementedError( + "MoeEp MegaMoE backend requires Rubin SM107 " + "(compute capability 10.7); " + f"found compute capability {major}.{minor}" + ) + + +def _is_cuda_stream_capturing(device: torch.device) -> bool: + """Return capture state for the request device.""" + + with torch.cuda.device(device): + return torch.cuda.is_current_stream_capturing() + + +def _validate_wgrad_config(config: ForwardConfig) -> None: + if config.backward_wgrad_mode not in ("none", "operands"): + raise ValueError( + "unsupported backward_wgrad_mode " + f"{config.backward_wgrad_mode!r}" + ) + if config.backward_wgrad_mode == "operands": + if not config.generate_c: + raise ValueError( + "backward_wgrad_mode='operands' requires generate_c=True" + ) + if config.token_padding_size != 256: + raise ValueError( + "backward_wgrad_mode='operands' requires " + "token_padding_size=256" + ) + if config.sf_padding_size != 128: + raise ValueError( + "backward_wgrad_mode='operands' requires " + "sf_padding_size=128" + ) + + +def validate_config(config: ForwardConfig) -> None: + """Reject static configurations outside the current MegaMoE milestone.""" + + _validate_wgrad_config(config) + if config.output_format != MoeFormat.BF16.value: + raise NotImplementedError( + "MoeEp training MegaMoE supports output_format='bf16' only" + ) + supported_combine_formats = { + MoeFormat.BF16.value, + MoeFormat.MXFP8.value, + } + if config.combine_format not in supported_combine_formats: + raise NotImplementedError( + "MoeEp training MegaMoE supports combine_format='bf16' " + "or 'mxfp8'" + ) + if config.max_tokens_per_rank is None: + raise NotImplementedError( + "MoeEp MegaMoE backend requires an explicit max_tokens_per_rank" + ) + if config.max_tokens_per_rank == 0: + raise NotImplementedError( + "MoeEp SM107 MXFP8 execution requires " + "max_tokens_per_rank to be positive" + ) + if config.hidden_size % 128: + raise NotImplementedError( + "MoeEp SM107 MXFP8 kernel currently requires hidden_size " + f"to be divisible by 128, got {config.hidden_size}" + ) + if config.intermediate_size % 256: + raise NotImplementedError( + "MoeEp SM107 MXFP8 kernel currently requires intermediate_size " + f"to be divisible by 256, got {config.intermediate_size}" + ) + if config.top_k > 32: + raise NotImplementedError( + "MoeEp SM107 MXFP8 dispatch currently requires top_k <= 32" + ) + if config.ep_size > 16: + raise NotImplementedError( + "MoeEp SM107 MXFP8 execution supports at most EP16 because the " + "validated peer-mapper ABI uses a 128-byte by-value offset table" + ) + if not config.apply_topk_in_fc1: + raise NotImplementedError( + "MoeEp Rubin training MegaMoE requires apply_topk_in_fc1=True" + ) + + +def validate_request(request: ValidatedForwardRequest) -> None: + """Reject valid requests outside the current MegaMoE input/device family.""" + + for name, tensor in ( + ("activation", request.activation), + ("fc1_weight", request.fc1_weight), + ("fc2_weight", request.fc2_weight), + ): + _validate_operand(name, tensor) + + _validate_device(request.device) + + +def validate_backward_request(request: ValidatedBackwardRequest) -> None: + """Reject backward requests outside the Rubin MXFP8 training path.""" + + _validate_wgrad_config(request.config) + for name, tensor in ( + ("fc1_weight", request.fc1_weight), + ("fc2_weight", request.fc2_weight), + ): + _validate_operand(name, tensor) + _validate_device(request.device) + config = request.config + if config.output_format != MoeFormat.BF16.value: + raise NotImplementedError( + "MoeEp MXFP8 backward currently requires output_format='bf16'" + ) + if config.ep_size not in (1, 2, 4): + raise NotImplementedError( + "MoeEp MXFP8 backward currently supports EP1/EP2/EP4" + ) + if not config.apply_topk_in_fc1: + raise NotImplementedError( + "MoeEp MXFP8 backward currently requires apply_topk_in_fc1=True" + ) + if _is_cuda_stream_capturing(request.device): + raise NotImplementedError( + "MoeEp MXFP8 backward does not support CUDA graph capture" + ) + + +__all__ = [ + "validate_backward_request", + "validate_config", + "validate_request", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/_comm.py b/python/cudnn/moe_ep/_megamoe_backend/_comm.py new file mode 100644 index 000000000..b1379e4cd --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/_comm.py @@ -0,0 +1,274 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Symmetric-memory ownership and peer-pointer descriptors for MegaMoE.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Protocol + +import torch + +from ._runtime import RuntimeHandle, RuntimeUnavailableError + + +class SymmetricMemoryProvider(Protocol): + """Injectable allocation boundary for a symmetric root slab.""" + + def allocate(self, nbytes: int, device: torch.device) -> torch.Tensor: ... + + def free(self, tensor: torch.Tensor) -> None: ... + + def peer_address(self, tensor: torch.Tensor, peer: int) -> int: ... + + +class _TorchMemoryProvider: + """CUDA tensor provider for local and single-rank symmetric memory.""" + + def allocate(self, nbytes: int, device: torch.device) -> torch.Tensor: + return torch.empty(nbytes, dtype=torch.uint8, device=device) + + def free(self, tensor: torch.Tensor) -> None: + del tensor + + def peer_address(self, tensor: torch.Tensor, peer: int) -> int: + if peer != 0: + raise ValueError(f"single-rank symmetric memory has no peer {peer}") + return int(tensor.data_ptr()) + + +class _NvshmemMemoryProvider: + """Lazy adapter over NVSHMEM symmetric tensor allocation.""" + + @staticmethod + def _core(): + try: + import nvshmem.core as core + except (ImportError, OSError) as exc: + raise RuntimeUnavailableError( + "symmetric workspace requires nvshmem4py and NVSHMEM libraries" + ) from exc + return core + + def allocate(self, nbytes: int, device: torch.device) -> torch.Tensor: + del device # NVSHMEM allocates on the device bound during runtime init. + try: + return self._core().tensor( + (nbytes,), + dtype=torch.uint8, + release=False, + except_on_del=True, + ) + except Exception as exc: + raise RuntimeUnavailableError( + f"failed to allocate {nbytes} bytes from the NVSHMEM symmetric heap" + ) from exc + + def free(self, tensor: torch.Tensor) -> None: + try: + self._core().free_tensor(tensor) + except Exception as exc: + raise RuntimeUnavailableError( + "failed to free the NVSHMEM symmetric root slab" + ) from exc + + def peer_address(self, tensor: torch.Tensor, peer: int) -> int: + try: + peer_tensor = self._core().get_peer_tensor(tensor, peer) + except Exception as exc: + raise RuntimeUnavailableError( + f"failed to map symmetric root slab for peer {peer}" + ) from exc + return int(peer_tensor.data_ptr()) + + +@dataclass(frozen=True) +class PeerMapping: + """Dense EP-rank peer deltas packed into the vendored kernel mapper ABI.""" + + base_address: int + offsets: tuple[int, ...] + rank: int + + def __post_init__(self) -> None: + if len(self.offsets) == 0: + raise ValueError("peer mapping requires at least one rank") + if self.rank < 0 or self.rank >= len(self.offsets): + raise ValueError( + f"peer mapping rank {self.rank} is outside {len(self.offsets)} ranks" + ) + if self.offsets[self.rank] != 0: + raise ValueError( + f"local peer offset must be zero, got {self.offsets[self.rank]}" + ) + + @property + def world_size(self) -> int: + return len(self.offsets) + + def to_sym_buffer_host(self): + """Build the CuTeDSL host payload lazily at the launch boundary.""" + + from .cutedsl_src.communication.nvlink_domain.symmetric_buffer import ( + SymmetricBufferHost, + ) + + return SymmetricBufferHost( + base_address=self.base_address, + offsets=self.offsets, + rank=self.rank, + max_ranks=self.world_size, + ) + + +class SymmetricSlab: + """One stable root allocation shared by all peer-visible workspace views.""" + + def __init__( + self, + runtime: RuntimeHandle, + nbytes: int, + *, + provider: Optional[SymmetricMemoryProvider] = None, + ) -> None: + if nbytes <= 0: + raise ValueError(f"symmetric slab size must be positive, got {nbytes}") + runtime.ensure_open() + + self._runtime = runtime + self._nbytes = nbytes + self._provider = provider or ( + _NvshmemMemoryProvider() + if runtime.nvshmem_enabled + else _TorchMemoryProvider() + ) + self._root: Optional[torch.Tensor] = None + self._mapping: Optional[PeerMapping] = None + self._cleanup_required = False + + def ensure_allocated(self) -> None: + if self._cleanup_required: + raise RuntimeError( + "symmetric slab requires cleanup before allocation" + ) + if self._root is not None and self._mapping is not None: + return + if self._root is not None: + raise RuntimeError( + "symmetric slab has an allocation pending cleanup" + ) + + root = self._provider.allocate(self._nbytes, self._runtime.device) + if not isinstance(root, torch.Tensor): + raise TypeError( + "symmetric memory provider must return a torch.Tensor" + ) + self._root = root + try: + if root.dtype is not torch.uint8 or root.numel() < self._nbytes: + raise ValueError( + "symmetric root must be a uint8 tensor with at least " + f"{self._nbytes} elements" + ) + if root.device != self._runtime.device: + raise ValueError( + "symmetric root device does not match runtime device: " + f"root={root.device}, runtime={self._runtime.device}" + ) + if not root.is_contiguous(): + raise ValueError("symmetric root tensor must be contiguous") + except Exception: + self._cleanup_required = True + raise + + try: + root.zero_() + + base_address = int(root.data_ptr()) + offsets = [] + for peer in range(self._runtime.world_size): + if peer == self._runtime.rank: + offsets.append(0) + continue + offsets.append( + self._provider.peer_address(root, peer) - base_address + ) + mapping = PeerMapping( + base_address=base_address, + offsets=tuple(offsets), + rank=self._runtime.rank, + ) + if ( + mapping.world_size != self._runtime.world_size + or mapping.rank != self._runtime.rank + ): + raise RuntimeError( + "symmetric peer mapping does not match the EP subgroup" + ) + except Exception: + self._cleanup_required = True + raise + + self._mapping = mapping + + @property + def nbytes(self) -> int: + return self._nbytes + + @property + def closed(self) -> bool: + return self._root is None + + @property + def allocated(self) -> bool: + return ( + not self._cleanup_required + and self._root is not None + and self._mapping is not None + ) + + @property + def mapping(self) -> PeerMapping: + if self._cleanup_required: + raise RuntimeError("symmetric slab requires cleanup") + if self._mapping is None: + raise RuntimeError("symmetric slab is closed") + return self._mapping + + @property + def root(self) -> torch.Tensor: + if self._cleanup_required: + raise RuntimeError("symmetric slab requires cleanup") + if self._root is None: + raise RuntimeError("symmetric slab is closed") + return self._root + + def byte_view(self, offset: int, nbytes: int) -> torch.Tensor: + if offset < 0 or nbytes < 0 or offset + nbytes > self._nbytes: + raise ValueError( + f"byte view [{offset}, {offset + nbytes}) exceeds " + f"symmetric slab size {self._nbytes}" + ) + return self.root.narrow(0, offset, nbytes) + + def close(self) -> None: + root = self._root + if root is None: + self._cleanup_required = False + return + try: + self._provider.free(root) + except Exception: + self._cleanup_required = True + raise + self._root = None + self._mapping = None + self._cleanup_required = False + + +__all__ = [ + "PeerMapping", + "SymmetricMemoryProvider", + "SymmetricSlab", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/_plan.py b/python/cudnn/moe_ep/_megamoe_backend/_plan.py new file mode 100644 index 000000000..a2b72b5fa --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/_plan.py @@ -0,0 +1,181 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy runtime/workspace owner for a compiled MegaMoE execution plan.""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from typing import Optional + +import torch + +from .._contracts import ( + ForwardConfig, + ValidatedBackwardRequest, + ValidatedForwardRequest, +) +from ._comm import SymmetricMemoryProvider +from ._runtime import RuntimeHandle, RuntimeManager, get_runtime_manager +from ._workspace import ( + LocalMemoryProvider, + WorkspaceOwner, + WorkspaceRequirements, + WorkspaceViews, +) + + +@dataclass(frozen=True) +class PreparedResources: + """Resources prepared for staging, compile, and launch integration.""" + + runtime: RuntimeHandle + workspace: WorkspaceViews + + +class ExecutionPlanOwner: + """Own runtime and stable workspace without compiling or launching a kernel.""" + + def __init__( + self, + config: ForwardConfig, + device: torch.device, + requirements: WorkspaceRequirements, + *, + runtime_manager: Optional[RuntimeManager] = None, + symmetric_provider: Optional[SymmetricMemoryProvider] = None, + local_provider: Optional[LocalMemoryProvider] = None, + ) -> None: + if config.max_tokens_per_rank != requirements.max_tokens_per_rank: + raise ValueError( + "workspace capacity must match ForwardConfig.max_tokens_per_rank" + ) + self.config = config + self.device = torch.device(device) + self.requirements = requirements + self._runtime_manager = runtime_manager or get_runtime_manager() + self._symmetric_provider = symmetric_provider + self._local_provider = local_provider + self._runtime: Optional[RuntimeHandle] = None + self._workspace: Optional[WorkspaceOwner] = None + self._closed = False + self._cleanup_required = False + self._lock = threading.RLock() + + @property + def prepared(self) -> bool: + return ( + not self._cleanup_required + and self._runtime is not None + and self._workspace is not None + and self._workspace.allocated + ) + + @property + def cleanup_required(self) -> bool: + return self._cleanup_required + + @property + def closed(self) -> bool: + return self._closed + + def prepare( + self, + request: ValidatedForwardRequest | ValidatedBackwardRequest, + ) -> PreparedResources: + with self._lock: + if self._closed: + raise RuntimeError("MegaMoE execution plan is closed") + if self._cleanup_required: + raise RuntimeError( + "MegaMoE execution plan requires cleanup before prepare" + ) + if request.config is not self.config: + raise ValueError("request does not belong to this static plan") + if torch.device(request.device) != self.device: + raise ValueError( + f"execution plan is bound to {self.device}, got {request.device}" + ) + if request.token_count > self.requirements.max_tokens_per_rank: + raise ValueError( + f"token count {request.token_count} exceeds " + f"max_tokens_per_rank={self.requirements.max_tokens_per_rank}" + ) + if ( + not self.prepared + and torch.cuda.is_current_stream_capturing() + ): + raise RuntimeError( + "MegaMoE runtime/workspace must be warmed up before " + "CUDA graph capture" + ) + + if not self.prepared: + runtime = self._runtime_manager.acquire(self.config, self.device) + self._runtime = runtime + try: + workspace = WorkspaceOwner( + self.requirements, + runtime, + symmetric_provider=self._symmetric_provider, + local_provider=self._local_provider, + ) + self._workspace = workspace + views = workspace.views(request.token_count) + except Exception: + try: + self._cleanup_failed_prepare() + except Exception: + self._cleanup_required = True + raise + raise + else: + assert self._runtime is not None + assert self._workspace is not None + views = self._workspace.views(request.token_count) + + return PreparedResources( + runtime=self._runtime, + workspace=views, + ) + + def _cleanup_failed_prepare(self) -> None: + if self._workspace is not None: + self._workspace.close() + self._workspace = None + if self._runtime is not None: + self._runtime.close() + self._runtime = None + + def close(self) -> None: + with self._lock: + if self._closed: + return + + try: + if self._workspace is not None: + self._workspace.close() + self._workspace = None + if self._runtime is not None: + self._runtime.close() + self._runtime = None + except Exception: + self._cleanup_required = True + raise + self._cleanup_required = False + self._closed = True + + def __enter__(self) -> "ExecutionPlanOwner": + with self._lock: + if self._closed: + raise RuntimeError("MegaMoE execution plan is closed") + return self + + def __exit__(self, exc_type, exc_value, traceback) -> bool: + del exc_type, exc_value, traceback + self.close() + return False + + +__all__ = ["ExecutionPlanOwner", "PreparedResources"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/_runtime.py b/python/cudnn/moe_ep/_megamoe_backend/_runtime.py new file mode 100644 index 000000000..9976acd11 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/_runtime.py @@ -0,0 +1,611 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Process-level runtime ownership for the private MegaMoE backend. + +This module is import-light: importing it does not import CUDA Python or +NVSHMEM and does not initialize CUDA. Optional runtime modules are loaded only +when a distributed runtime is actually acquired. +""" + +from __future__ import annotations + +import logging +import threading +from dataclasses import dataclass +from enum import Enum +from typing import Callable, Optional, Protocol + +import torch +import torch.distributed as dist + +from .._contracts import ForwardConfig + +_logger = logging.getLogger(__name__) + + +class RuntimeUnavailableError(RuntimeError): + """The requested runtime cannot be loaded or initialized.""" + + +class RuntimeInitState(Enum): + """Normalized NVSHMEM initialization state.""" + + NOT_INITIALIZED = "not_initialized" + INITIALIZED = "initialized" + PARTIAL = "partial" + + +@dataclass(frozen=True) +class RuntimeWorld: + """Group-relative geometry and ordered membership used for bootstrap.""" + + rank: int + size: int + group: object + global_ranks: tuple[int, ...] + + @property + def identity(self) -> tuple[int, int, tuple[int, ...]]: + """Stable process-group identity independent of ProcessGroup objects.""" + + return self.rank, self.size, self.global_ranks + + +class NvshmemRuntimeProvider(Protocol): + """Injectable NVSHMEM lifecycle boundary used by :class:`RuntimeManager`.""" + + def initialization_state(self) -> RuntimeInitState: ... + + def initialize(self, device: torch.device, world: RuntimeWorld) -> None: ... + + def rank(self) -> int: ... + + def world_size(self) -> int: ... + + def device(self) -> torch.device: ... + + def finalize(self) -> None: ... + + +def _resolve_world(config: ForwardConfig) -> RuntimeWorld: + if config.ep_group is None: + if ( + config.ep_size != 1 + or config.ep_rank != 0 + or config.ep_global_ranks + ): + raise ValueError( + "ep_group=None requires ep_size=1, ep_rank=0, and no " + "distributed rank membership" + ) + return RuntimeWorld(rank=0, size=1, group=None, global_ranks=()) + + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError( + "distributed MegaMoE runtime requires torch.distributed to be initialized" + ) + + group = config.ep_group + rank = dist.get_rank(group) + size = dist.get_world_size(group) + global_ranks = tuple( + dist.get_global_rank(group, group_rank) + for group_rank in range(size) + ) + if (rank, size) != (config.ep_rank, config.ep_size): + raise RuntimeError( + "ForwardConfig EP geometry does not match its process group: " + f"config=({config.ep_rank}, {config.ep_size}), " + f"runtime=({rank}, {size})" + ) + if global_ranks != config.ep_global_ranks: + raise RuntimeError( + "ForwardConfig EP membership does not match its process group: " + f"config={config.ep_global_ranks}, runtime={global_ranks}" + ) + return RuntimeWorld( + rank=rank, + size=size, + group=group, + global_ranks=global_ranks, + ) + + +def _canonical_cuda_device(device: torch.device) -> torch.device: + device = torch.device(device) + if device.type != "cuda": + raise ValueError(f"MegaMoE runtime requires a CUDA device, got {device}") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + return device + + +def _spans_default_distributed_world(world: RuntimeWorld) -> bool: + """Whether ``world`` has the default world's complete ordered membership.""" + + if not dist.is_available() or not dist.is_initialized(): + return False + return world.global_ranks == tuple(range(dist.get_world_size())) + + +def _load_nvshmem_core(): + try: + import nvshmem.core as core + except (ImportError, OSError) as exc: + raise RuntimeUnavailableError( + "MegaMoE distributed runtime requires nvshmem4py and NVSHMEM libraries" + ) from exc + return core + + +def _normalize_nvshmem_init_state(status) -> RuntimeInitState: + """Normalize enum and integer forms used across nvshmem4py releases.""" + + name = getattr(status, "name", "") + if name.endswith("NOT_INITIALIZED"): + return RuntimeInitState.NOT_INITIALIZED + if name.endswith("IS_INITIALIZED") or name.endswith( + ("LIMITED_MPG", "FULL_MPG") + ): + return RuntimeInitState.INITIALIZED + if name.endswith("IS_BOOTSTRAPPED"): + return RuntimeInitState.PARTIAL + + try: + value = int(getattr(status, "value", status)) + except (TypeError, ValueError): + return RuntimeInitState.PARTIAL + if value == 0: + return RuntimeInitState.NOT_INITIALIZED + if value in {2, 3, 4}: + return RuntimeInitState.INITIALIZED + return RuntimeInitState.PARTIAL + + +class _DefaultNvshmemRuntimeProvider: + """Lazy adapter over the installed ``nvshmem.core`` API.""" + + def initialization_state(self) -> RuntimeInitState: + core = _load_nvshmem_core() + try: + status = core.init_status() + except Exception as exc: + raise RuntimeUnavailableError( + "failed to query NVSHMEM initialization status" + ) from exc + return _normalize_nvshmem_init_state(status) + + def initialize(self, device: torch.device, world: RuntimeWorld) -> None: + if world.size <= 1: + raise ValueError("NVSHMEM initialization requires a distributed subgroup") + if world.group is None: + raise ValueError("NVSHMEM initialization requires a process group") + + core = _load_nvshmem_core() + try: + import numpy as np + + try: + from cuda.core.experimental import Device + except ImportError: + from cuda.core import Device + + torch.cuda.set_device(device) + cuda_device = Device(device.index) + cuda_device.set_current() + + uid = core.get_unique_id(empty=(world.rank != 0)) + uid_bytes = uid._data.view(np.uint8).copy() + uid_tensor = torch.from_numpy(uid_bytes) + group_backend = dist.get_backend(world.group) + if ( + group_backend == dist.Backend.NCCL + or str(group_backend).lower() == "nccl" + ): + uid_tensor = uid_tensor.to(device=device) + root_global_rank = dist.get_global_rank(world.group, 0) + if root_global_rank != world.global_ranks[0]: + raise RuntimeError( + "EP subgroup root changed during NVSHMEM bootstrap" + ) + dist.broadcast( + uid_tensor, + src=root_global_rank, + group=world.group, + ) + dist.barrier(group=world.group) + uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype) + + core.init( + device=cuda_device, + uid=uid, + rank=world.rank, + nranks=world.size, + initializer_method="uid", + ) + except RuntimeUnavailableError: + raise + except Exception as exc: + raise RuntimeUnavailableError( + "failed to initialize the NVSHMEM EP subgroup runtime" + ) from exc + + def rank(self) -> int: + try: + return int(_load_nvshmem_core().my_pe()) + except Exception as exc: + raise RuntimeUnavailableError("failed to query the NVSHMEM PE rank") from exc + + def world_size(self) -> int: + try: + return int(_load_nvshmem_core().n_pes()) + except Exception as exc: + raise RuntimeUnavailableError( + "failed to query the NVSHMEM PE world size" + ) from exc + + def device(self) -> torch.device: + try: + from nvshmem.core.memory import _cached_device + + cached = _cached_device["device"] + if cached is None: + raise RuntimeError("NVSHMEM cached device is empty") + return torch.device("cuda", int(cached.device_id)) + except Exception as exc: + raise RuntimeUnavailableError( + "failed to query the NVSHMEM initialization device" + ) from exc + + def finalize(self) -> None: + try: + _load_nvshmem_core().finalize() + except Exception as exc: + raise RuntimeUnavailableError("failed to finalize NVSHMEM") from exc + + +@dataclass +class _ActiveRuntime: + token: object + device: torch.device + world: RuntimeWorld + provider: Optional[NvshmemRuntimeProvider] + owns_runtime: bool + ref_count: int = 1 + cleanup_required: bool = False + + +class _ProcessRuntimeRegistry: + """State shared by every RuntimeManager instance in this process.""" + + def __init__(self) -> None: + self.lock = threading.RLock() + self.active: Optional[_ActiveRuntime] = None + + +_PROCESS_RUNTIME_REGISTRY = _ProcessRuntimeRegistry() + + +class RuntimeHandle: + """Per-backend lease on process-level runtime state.""" + + def __init__( + self, + manager: "RuntimeManager", + token: object, + device: torch.device, + world: RuntimeWorld, + owns_runtime: bool, + ) -> None: + self._manager = manager + self._token = token + self.device = device + self.rank = world.rank + self.world_size = world.size + self.group = world.group + self.global_ranks = world.global_ranks + self.owns_runtime = owns_runtime + self._closed = False + self._close_lock = threading.Lock() + + @property + def nvshmem_enabled(self) -> bool: + return self.world_size > 1 + + @property + def closed(self) -> bool: + return self._closed + + def ensure_open(self) -> None: + if self._closed: + raise RuntimeError("MegaMoE runtime handle is closed") + + def current_stream(self) -> torch.cuda.Stream: + self.ensure_open() + return torch.cuda.current_stream(self.device) + + def close(self) -> None: + with self._close_lock: + if self._closed: + return + self._manager._release(self._token) + self._closed = True + + +class RuntimeManager: + """Reference-counted owner for one process-global runtime subgroup.""" + + def __init__( + self, + *, + provider_factory: Callable[[], NvshmemRuntimeProvider] = ( + _DefaultNvshmemRuntimeProvider + ), + world_resolver: Callable[[ForwardConfig], RuntimeWorld] = _resolve_world, + ) -> None: + self._provider_factory = provider_factory + self._world_resolver = world_resolver + + @property + def ref_count(self) -> int: + with _PROCESS_RUNTIME_REGISTRY.lock: + active = _PROCESS_RUNTIME_REGISTRY.active + return 0 if active is None else active.ref_count + + @property + def active_device(self) -> Optional[torch.device]: + with _PROCESS_RUNTIME_REGISTRY.lock: + active = _PROCESS_RUNTIME_REGISTRY.active + return None if active is None else active.device + + def acquire( + self, + config: ForwardConfig, + device: torch.device, + ) -> RuntimeHandle: + device = _canonical_cuda_device(device) + world = self._world_resolver(config) + + with _PROCESS_RUNTIME_REGISTRY.lock: + if _PROCESS_RUNTIME_REGISTRY.active is not None: + active = _PROCESS_RUNTIME_REGISTRY.active + if active.cleanup_required: + raise RuntimeError( + "MegaMoE process runtime requires cleanup before reacquire" + ) + if active.device != device: + raise ValueError( + f"MegaMoE process runtime is bound to {active.device}; " + f"cannot acquire it for {device}" + ) + if active.world.identity != world.identity: + raise RuntimeError( + "MegaMoE process runtime is already bound to a different " + "EP subgroup" + ) + active.ref_count += 1 + return RuntimeHandle( + self, + active.token, + active.device, + active.world, + active.owns_runtime, + ) + + provider: Optional[NvshmemRuntimeProvider] = None + owns_runtime = False + if world.size > 1: + provider = self._provider_factory() + status = provider.initialization_state() + if status is RuntimeInitState.PARTIAL: + raise RuntimeError( + "cannot attach to a partially initialized NVSHMEM runtime" + ) + if ( + status is RuntimeInitState.INITIALIZED + and not _spans_default_distributed_world(world) + ): + raise RuntimeError( + "cannot safely attach an externally initialized NVSHMEM " + "runtime to a non-WORLD EP subgroup because its ordered " + "membership cannot be verified" + ) + if status is RuntimeInitState.NOT_INITIALIZED: + try: + provider.initialize(device, world) + except Exception as initialization_error: + self._rollback_failed_initialization( + provider, + device, + world, + initialization_error, + ) + raise + owns_runtime = True + + try: + provider_device = provider.device() + provider_rank = provider.rank() + provider_size = provider.world_size() + except Exception as validation_error: + if owns_runtime: + self._cleanup_owned_runtime_after_error( + provider, + device, + world, + validation_error, + ) + raise + + if provider_device != device: + if owns_runtime: + self._cleanup_owned_runtime_after_error( + provider, + device, + world, + RuntimeError( + "NVSHMEM initialization device does not match " + f"the requested device: nvshmem={provider_device}, " + f"requested={device}" + ), + ) + ownership = ( + "owned" + if owns_runtime + else "externally initialized" + ) + raise RuntimeError( + f"{ownership} NVSHMEM runtime is bound to " + f"{provider_device}, not the requested device {device}" + ) + if (provider_rank, provider_size) != (world.rank, world.size): + if owns_runtime: + self._cleanup_owned_runtime_after_error( + provider, + device, + world, + RuntimeError("NVSHMEM PE geometry mismatch"), + ) + raise RuntimeError( + "NVSHMEM PE geometry does not match the EP subgroup: " + f"nvshmem=({provider_rank}, {provider_size}), " + f"torch=({world.rank}, {world.size})" + ) + + token = object() + _PROCESS_RUNTIME_REGISTRY.active = _ActiveRuntime( + token=token, + device=device, + world=world, + provider=provider, + owns_runtime=owns_runtime, + ) + return RuntimeHandle(self, token, device, world, owns_runtime) + + @staticmethod + def _mark_cleanup_required( + provider: NvshmemRuntimeProvider, + device: torch.device, + world: RuntimeWorld, + ) -> None: + _PROCESS_RUNTIME_REGISTRY.active = _ActiveRuntime( + token=object(), + device=device, + world=world, + provider=provider, + owns_runtime=True, + ref_count=0, + cleanup_required=True, + ) + + @classmethod + def _rollback_failed_initialization( + cls, + provider: NvshmemRuntimeProvider, + device: torch.device, + world: RuntimeWorld, + initialization_error: Exception, + ) -> None: + try: + state = provider.initialization_state() + if state is not RuntimeInitState.NOT_INITIALIZED: + provider.finalize() + except Exception as cleanup_error: + cls._mark_cleanup_required(provider, device, world) + raise RuntimeError( + "NVSHMEM initialization failed and rollback requires retry" + ) from cleanup_error + _logger.debug( + "rolled back failed NVSHMEM initialization: %s", + initialization_error, + ) + + @classmethod + def _cleanup_owned_runtime_after_error( + cls, + provider: NvshmemRuntimeProvider, + device: torch.device, + world: RuntimeWorld, + original_error: Exception, + ) -> None: + try: + provider.finalize() + except Exception as cleanup_error: + cls._mark_cleanup_required(provider, device, world) + raise RuntimeError( + "NVSHMEM validation failed and cleanup requires retry" + ) from cleanup_error + _logger.debug( + "finalized owned NVSHMEM after validation failure: %s", + original_error, + ) + + def retry_cleanup(self) -> None: + """Retry cleanup after an acquire-time rollback failure.""" + + with _PROCESS_RUNTIME_REGISTRY.lock: + active = _PROCESS_RUNTIME_REGISTRY.active + if active is None: + return + if not active.cleanup_required or active.ref_count != 0: + raise RuntimeError( + "MegaMoE process runtime does not have retryable cleanup" + ) + if active.provider is None: + raise RuntimeError( + "retryable MegaMoE runtime cleanup has no provider" + ) + active.provider.finalize() + _PROCESS_RUNTIME_REGISTRY.active = None + + def _release(self, token: object) -> None: + with _PROCESS_RUNTIME_REGISTRY.lock: + active = _PROCESS_RUNTIME_REGISTRY.active + if active is None or active.token is not token: + return + if active.ref_count <= 0: + raise RuntimeError( + "MegaMoE process runtime has invalid release state" + ) + + if active.cleanup_required: + if active.ref_count != 1 or active.provider is None: + raise RuntimeError( + "MegaMoE process runtime has invalid retry state" + ) + active.provider.finalize() + _PROCESS_RUNTIME_REGISTRY.active = None + return + + if active.ref_count > 1: + active.ref_count -= 1 + return + + if active.owns_runtime and active.provider is not None: + try: + active.provider.finalize() + except Exception: + active.cleanup_required = True + raise + _PROCESS_RUNTIME_REGISTRY.active = None + + +_DEFAULT_RUNTIME_MANAGER = RuntimeManager() + + +def get_runtime_manager() -> RuntimeManager: + """Return the process-level manager used by the default MegaMoE backend.""" + + return _DEFAULT_RUNTIME_MANAGER + + +__all__ = [ + "NvshmemRuntimeProvider", + "RuntimeHandle", + "RuntimeInitState", + "RuntimeManager", + "RuntimeUnavailableError", + "RuntimeWorld", + "get_runtime_manager", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/_workspace.py new file mode 100644 index 000000000..97f9ba793 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/_workspace.py @@ -0,0 +1,465 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Stable local and symmetric workspace ownership for MegaMoE.""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from types import MappingProxyType +from typing import Mapping, Optional, Protocol, Sequence + +import torch + +from .._contracts import ForwardConfig +from ._comm import ( + PeerMapping, + SymmetricMemoryProvider, + SymmetricSlab, + _TorchMemoryProvider, +) +from ._runtime import RuntimeHandle + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def padded_mxfp8_scale_columns(hidden: int) -> int: + """Return the E8M0 row width required by Rubin's 16-byte token-in copy.""" + + logical_columns = (hidden + 31) // 32 + return _align_up(logical_columns, 16) + + +@dataclass(frozen=True) +class BufferRegion: + """One named byte region within a stable root allocation.""" + + name: str + nbytes: int + alignment: int = 256 + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("workspace region name must not be empty") + if self.nbytes < 0: + raise ValueError( + f"workspace region {self.name!r} has negative size {self.nbytes}" + ) + if self.alignment <= 0 or self.alignment & (self.alignment - 1): + raise ValueError( + f"workspace region {self.name!r} alignment must be a power of two" + ) + + +@dataclass(frozen=True) +class BufferPlacement: + """Resolved byte offset for one region.""" + + name: str + offset: int + nbytes: int + + +@dataclass(frozen=True) +class BufferLayout: + """Deterministic aligned layout for one root byte allocation.""" + + placements: tuple[BufferPlacement, ...] + total_bytes: int + + @classmethod + def build(cls, regions: Sequence[BufferRegion]) -> "BufferLayout": + names: set[str] = set() + placements = [] + offset = 0 + max_alignment = 1 + for region in regions: + if region.name in names: + raise ValueError(f"duplicate workspace region {region.name!r}") + names.add(region.name) + offset = _align_up(offset, region.alignment) + placements.append( + BufferPlacement( + name=region.name, + offset=offset, + nbytes=region.nbytes, + ) + ) + offset += region.nbytes + max_alignment = max(max_alignment, region.alignment) + return cls( + placements=tuple(placements), + total_bytes=_align_up(offset, max_alignment), + ) + + def placement(self, name: str) -> BufferPlacement: + for placement in self.placements: + if placement.name == name: + return placement + raise KeyError(name) + + +@dataclass(frozen=True) +class WorkspaceRequirements: + """Capacity-driven regions supplied before runtime allocation. + + The executable backend obtains exact Rubin kernel workspace sizes, then + passes them here without making the runtime owner import or instantiate + CuTeDSL kernels. + """ + + max_tokens_per_rank: int + symmetric_regions: tuple[BufferRegion, ...] + local_regions: tuple[BufferRegion, ...] + + def __post_init__(self) -> None: + if self.max_tokens_per_rank < 0: + raise ValueError("max_tokens_per_rank must be non-negative") + symmetric_names = {region.name for region in self.symmetric_regions} + local_names = {region.name for region in self.local_regions} + duplicates = symmetric_names & local_names + if duplicates: + raise ValueError( + "workspace region names must be unique across roots: " + f"{sorted(duplicates)}" + ) + + @classmethod + def for_mxfp8( + cls, + config: ForwardConfig, + *, + kernel_local_workspace_bytes: int, + kernel_shared_workspace_bytes: int, + col_quant_data_bytes: int = 0, + col_quant_sf_bytes: int = 0, + backward_fc1_preact_bytes: int = 0, + backward_dprob_bytes: int = 0, + backward_aux_data_bytes: int = 0, + backward_aux_scale_bytes: int = 0, + ) -> "WorkspaceRequirements": + if config.max_tokens_per_rank is None: + raise ValueError("MXFP8 workspace requires max_tokens_per_rank") + for name, value in ( + ("kernel_local_workspace_bytes", kernel_local_workspace_bytes), + ("kernel_shared_workspace_bytes", kernel_shared_workspace_bytes), + ("col_quant_data_bytes", col_quant_data_bytes), + ("col_quant_sf_bytes", col_quant_sf_bytes), + ("backward_fc1_preact_bytes", backward_fc1_preact_bytes), + ("backward_dprob_bytes", backward_dprob_bytes), + ("backward_aux_data_bytes", backward_aux_data_bytes), + ("backward_aux_scale_bytes", backward_aux_scale_bytes), + ): + if value < 0: + raise ValueError(f"{name} must be non-negative, got {value}") + if bool(col_quant_data_bytes) != bool(col_quant_sf_bytes): + raise ValueError( + "column requant data and scale workspace must be enabled together" + ) + backward_sizes = ( + backward_fc1_preact_bytes, + backward_dprob_bytes, + backward_aux_data_bytes, + backward_aux_scale_bytes, + ) + if any(backward_sizes) and not all(backward_sizes): + raise ValueError( + "backward preactivation, dprob, data, and scale workspace " + "must be enabled together" + ) + + tokens = config.max_tokens_per_rank + hidden = config.hidden_size + top_k = config.top_k + kernel_sf_columns = padded_mxfp8_scale_columns(hidden) + + backward_symmetric_regions = ( + (BufferRegion("backward_dprob", backward_dprob_bytes),) + if backward_dprob_bytes + else () + ) + symmetric_regions = ( + BufferRegion("activation_data", tokens * hidden), + BufferRegion("activation_scale", tokens * kernel_sf_columns), + BufferRegion("topk_weights", tokens * top_k * 4), + BufferRegion("output_data", tokens * hidden * 2), + *backward_symmetric_regions, + BufferRegion( + "kernel_shared_workspace", + kernel_shared_workspace_bytes, + ), + # Test-visible tail canary placed immediately after the opaque + # peer-visible kernel workspace. It does not enter the kernel ABI. + BufferRegion("symmetric_guard", 256, alignment=1), + ) + col_quant_regions = ( + ( + BufferRegion("col_quant_data", col_quant_data_bytes), + BufferRegion("col_quant_sf", col_quant_sf_bytes), + ) + if col_quant_data_bytes + else () + ) + backward_local_regions = ( + ( + BufferRegion( + "backward_fc1_preact", + backward_fc1_preact_bytes, + alignment=128, + ), + BufferRegion( + "backward_aux_data", + backward_aux_data_bytes, + alignment=128, + ), + BufferRegion( + "backward_aux_scale", + backward_aux_scale_bytes, + alignment=128, + ), + ) + if backward_fc1_preact_bytes + else () + ) + local_regions = ( + BufferRegion("topk_idx", tokens * top_k * 4), + BufferRegion("overflow_flag", 4), + *col_quant_regions, + *backward_local_regions, + BufferRegion("kernel_local_workspace", kernel_local_workspace_bytes), + BufferRegion("local_guard", 256, alignment=1), + ) + return cls( + max_tokens_per_rank=tokens, + symmetric_regions=symmetric_regions, + local_regions=local_regions, + ) + +class LocalMemoryProvider(Protocol): + """Injectable local allocation boundary.""" + + def allocate(self, nbytes: int, device: torch.device) -> torch.Tensor: ... + + def free(self, tensor: torch.Tensor) -> None: ... + + +class _LocalSlab: + def __init__( + self, + nbytes: int, + device: torch.device, + provider: LocalMemoryProvider, + ) -> None: + if nbytes <= 0: + raise ValueError(f"local slab size must be positive, got {nbytes}") + self._provider = provider + self._nbytes = nbytes + root = provider.allocate(nbytes, device) + self._root: Optional[torch.Tensor] = None + try: + if not isinstance(root, torch.Tensor): + raise TypeError("local memory provider must return a torch.Tensor") + if root.dtype is not torch.uint8 or root.numel() < nbytes: + raise ValueError( + "local root must be a uint8 tensor with at least " + f"{nbytes} elements" + ) + if root.device != device: + raise ValueError( + "local root device does not match runtime device: " + f"root={root.device}, runtime={device}" + ) + if not root.is_contiguous(): + raise ValueError("local root tensor must be contiguous") + root.zero_() + except Exception: + if isinstance(root, torch.Tensor): + provider.free(root) + raise + self._root = root + + @property + def root(self) -> torch.Tensor: + if self._root is None: + raise RuntimeError("local workspace slab is closed") + return self._root + + def byte_view(self, offset: int, nbytes: int) -> torch.Tensor: + if offset < 0 or nbytes < 0 or offset + nbytes > self._nbytes: + raise ValueError( + f"byte view [{offset}, {offset + nbytes}) exceeds " + f"local slab size {self._nbytes}" + ) + return self.root.narrow(0, offset, nbytes) + + def close(self) -> None: + root = self._root + if root is None: + return + self._provider.free(root) + self._root = None + + +@dataclass(frozen=True) +class WorkspaceViews: + """Stable full-capacity byte views for one prepared request.""" + + token_count: int + symmetric: Mapping[str, torch.Tensor] + local: Mapping[str, torch.Tensor] + peer_mapping: PeerMapping + + +class WorkspaceOwner: + """Own local and symmetric slabs for one static execution plan.""" + + def __init__( + self, + requirements: WorkspaceRequirements, + runtime: RuntimeHandle, + *, + symmetric_provider: Optional[SymmetricMemoryProvider] = None, + local_provider: Optional[LocalMemoryProvider] = None, + ) -> None: + self.requirements = requirements + self.runtime = runtime + self.symmetric_layout = BufferLayout.build( + requirements.symmetric_regions + ) + self.local_layout = BufferLayout.build(requirements.local_regions) + if self.symmetric_layout.total_bytes <= 0: + raise ValueError("workspace requires at least one symmetric byte") + if self.local_layout.total_bytes <= 0: + raise ValueError("workspace requires at least one local byte") + + self._symmetric_provider = symmetric_provider + self._local_provider = local_provider or _TorchMemoryProvider() + self._symmetric: Optional[SymmetricSlab] = None + self._local: Optional[_LocalSlab] = None + self._closed = False + self._cleanup_required = False + self._lock = threading.RLock() + + @property + def allocated(self) -> bool: + return ( + not self._cleanup_required + and self._symmetric is not None + and self._symmetric.allocated + and self._local is not None + ) + + @property + def cleanup_required(self) -> bool: + return self._cleanup_required + + @property + def closed(self) -> bool: + return self._closed + + def ensure_allocated(self) -> None: + with self._lock: + if self._closed: + raise RuntimeError("workspace owner is closed") + if self._cleanup_required: + raise RuntimeError( + "workspace owner requires cleanup before allocation" + ) + if self.allocated: + return + self.runtime.ensure_open() + + local = _LocalSlab( + self.local_layout.total_bytes, + self.runtime.device, + self._local_provider, + ) + self._local = local + try: + symmetric = SymmetricSlab( + self.runtime, + self.symmetric_layout.total_bytes, + provider=self._symmetric_provider, + ) + self._symmetric = symmetric + symmetric.ensure_allocated() + except Exception: + try: + if self._symmetric is not None: + self._symmetric.close() + self._symmetric = None + if self._local is not None: + self._local.close() + self._local = None + except Exception: + self._cleanup_required = True + raise + raise + + def views(self, token_count: int) -> WorkspaceViews: + with self._lock: + if token_count < 0: + raise ValueError( + f"token_count must be non-negative, got {token_count}" + ) + if token_count > self.requirements.max_tokens_per_rank: + raise ValueError( + f"token count {token_count} exceeds " + f"max_tokens_per_rank={self.requirements.max_tokens_per_rank}" + ) + self.ensure_allocated() + assert self._symmetric is not None + assert self._local is not None + + symmetric_views = { + placement.name: self._symmetric.byte_view( + placement.offset, + placement.nbytes, + ) + for placement in self.symmetric_layout.placements + } + local_views = { + placement.name: self._local.byte_view( + placement.offset, + placement.nbytes, + ) + for placement in self.local_layout.placements + } + return WorkspaceViews( + token_count=token_count, + symmetric=MappingProxyType(symmetric_views), + local=MappingProxyType(local_views), + peer_mapping=self._symmetric.mapping, + ) + + def close(self) -> None: + with self._lock: + if self._closed: + return + try: + if self._symmetric is not None: + self._symmetric.close() + self._symmetric = None + if self._local is not None: + self._local.close() + self._local = None + except Exception: + self._cleanup_required = True + raise + self._cleanup_required = False + self._closed = True + + +__all__ = [ + "BufferLayout", + "BufferPlacement", + "BufferRegion", + "LocalMemoryProvider", + "WorkspaceOwner", + "WorkspaceRequirements", + "WorkspaceViews", + "padded_mxfp8_scale_columns", +] From 220f0470df0210f4f271534355b2939f0bae801c Mon Sep 17 00:00:00 2001 From: zhibinz Date: Mon, 24 Aug 2026 16:42:00 -0700 Subject: [PATCH 05/31] feat: implement MXFP8 MoeEp training execution Connect Rubin kernels to validated forward and backward dispatch, including deterministic staging, overflow handling, recomputation stashes, and grouped-wgrad operand export. --- .../moe_ep/_megamoe_backend/mxfp8/__init__.py | 4 + .../moe_ep/_megamoe_backend/mxfp8/_adapter.py | 648 ++++++++++++++++++ .../moe_ep/_megamoe_backend/mxfp8/_backend.py | 334 +++++++++ .../_megamoe_backend/mxfp8/_backward.py | 112 +++ .../mxfp8/_backward_compile.py | 377 ++++++++++ .../mxfp8/_backward_dispatch.py | 204 ++++++ .../_megamoe_backend/mxfp8/_backward_dprob.py | 87 +++ .../mxfp8/_backward_launch.py | 54 ++ .../mxfp8/_backward_layout.py | 82 +++ .../mxfp8/_backward_staging.py | 516 ++++++++++++++ .../mxfp8/_backward_wgrad_export.py | 226 ++++++ .../moe_ep/_megamoe_backend/mxfp8/_compile.py | 373 ++++++++++ .../moe_ep/_megamoe_backend/mxfp8/_config.py | 201 ++++++ .../_megamoe_backend/mxfp8/_fingerprint.py | 150 ++++ .../moe_ep/_megamoe_backend/mxfp8/_formats.py | 23 + .../moe_ep/_megamoe_backend/mxfp8/_launch.py | 180 +++++ .../moe_ep/_megamoe_backend/mxfp8/_stash.py | 302 ++++++++ .../_megamoe_backend/mxfp8/_wgrad_layout.py | 410 +++++++++++ 18 files changed, 4283 insertions(+) create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/__init__.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dispatch.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dprob.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_layout.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_staging.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_wgrad_export.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_fingerprint.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_formats.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_launch.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_stash.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_wgrad_layout.py diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/__init__.py new file mode 100644 index 000000000..a1362f39f --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Private MXFP8 implementation for the MegaMoE execution backend.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py new file mode 100644 index 000000000..ef5d4ee75 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py @@ -0,0 +1,648 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Logical MXFP8 to Rubin SM107 MegaMoE tensor staging.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from ..._contracts import ValidatedForwardRequest +from ..._types import BlockScaledTensor, MoeFormat +from .._plan import PreparedResources +from .._workspace import padded_mxfp8_scale_columns +from ._config import Mxfp8KernelConfig + +_MXFP8_DATA_DTYPE = torch.float8_e4m3fn +_MXFP8_SCALE_DTYPE = torch.float8_e8m0fnu +_GATE_UP_INTERLEAVE = 32 +_WORKSPACE_GUARD_BYTE = 0xA5 + + +def _decode_moe_tensor( + tensor: torch.Tensor | BlockScaledTensor, +) -> torch.Tensor: + """Decode a public MoE tensor to float32 for host-side staging math.""" + + if isinstance(tensor, BlockScaledTensor): + return tensor.dequantize(torch.float32) + return tensor.float() + + +def _quantize_plain_mxfp8( + tensor: torch.Tensor, + *, + axis: int = 1, +) -> BlockScaledTensor: + """Stage a plain floating tensor through the backend's MXFP8 family.""" + + moved = tensor.float().movedim(axis, -1) + logical_extent = moved.shape[-1] + block_count = (logical_extent + 31) // 32 + padded_extent = block_count * 32 + if padded_extent != logical_extent: + moved = torch.nn.functional.pad( + moved, + (0, padded_extent - logical_extent), + ) + blocks = moved.reshape(*moved.shape[:-1], block_count, 32) + raw_scale = blocks.abs().amax(dim=-1) / 448.0 + safe_scale = torch.where(raw_scale > 0, raw_scale, 1.0) + scale_float = torch.where( + raw_scale > 0, + torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), + torch.zeros_like(raw_scale), + ) + scale = scale_float.to(_MXFP8_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(-448.0, 448.0) + data = ( + normalized.to(_MXFP8_DATA_DTYPE) + .reshape(*moved.shape)[..., :logical_extent] + .movedim(-1, axis) + .contiguous() + ) + return BlockScaledTensor( + data=data, + scale=scale.movedim(-1, axis).contiguous(), + format=MoeFormat.MXFP8, + logical_shape=tuple(tensor.shape), + axis=axis, + ) + + +def _as_mxfp8(tensor: torch.Tensor | BlockScaledTensor) -> BlockScaledTensor: + if isinstance(tensor, BlockScaledTensor): + if tensor.format is not MoeFormat.MXFP8: + raise NotImplementedError( + "MXFP8 staging cannot convert " + f"{tensor.format.value!r} block-scaled input" + ) + return tensor + return _quantize_plain_mxfp8(tensor) + + +def _typed_view( + byte_tensor: torch.Tensor, + dtype: torch.dtype, + shape: tuple[int, ...], +) -> torch.Tensor: + expected_bytes = 1 + for extent in shape: + expected_bytes *= extent + expected_bytes *= dtype.itemsize + if byte_tensor.numel() != expected_bytes: + raise ValueError( + f"byte region has {byte_tensor.numel()} bytes, " + f"expected {expected_bytes} for shape={shape}, dtype={dtype}" + ) + return byte_tensor.view(dtype).reshape(shape) + + +def _as_bytes(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(torch.uint8) + + +def _validate_int32_downcast(tensor: torch.Tensor) -> None: + """Validate the only lossy public-to-kernel dtype conversion.""" + + if tensor.dtype is torch.int32: + return + if tensor.dtype is not torch.int64: + raise TypeError( + "topk_idx staging requires torch.int32 or torch.int64, " + f"got {tensor.dtype}" + ) + capturing = ( + tensor.device.type == "cuda" + and torch.cuda.is_current_stream_capturing() + ) + if capturing or tensor.numel() == 0: + # The public validator checked the same tensor before capture. During + # replay callers must preserve its documented expert-id invariant. + return + limits = torch.iinfo(torch.int32) + outside_int32 = (tensor < limits.min) | (tensor > limits.max) + if bool(outside_int32.any().item()): + raise OverflowError("topk_idx contains a value outside the int32 range") + + +def _zero_workspace_prefix( + workspace: torch.Tensor, + nbytes: int, + *, + name: str, +) -> None: + if nbytes < 0 or nbytes > workspace.numel(): + raise ValueError( + f"{name} zero prefix {nbytes} exceeds {workspace.numel()} bytes" + ) + workspace[:nbytes].zero_() + + +def _zero_workspace_range( + workspace: torch.Tensor, + offset: int, + nbytes: int, + *, + name: str, +) -> None: + if offset < 0 or nbytes < 0 or offset + nbytes > workspace.numel(): + raise ValueError( + f"{name} byte range [{offset}, {offset + nbytes}) exceeds " + f"{workspace.numel()} bytes" + ) + workspace.narrow(0, offset, nbytes).zero_() + + +def _interleave_gate_up_rows( + tensor: torch.Tensor, + intermediate: int, +) -> torch.Tensor: + """Convert gate-half/up-half rows to 32-row gate/up pairs.""" + + if intermediate % _GATE_UP_INTERLEAVE: + raise ValueError( + "MXFP8 gate/up interleave requires intermediate_size to be " + f"divisible by {_GATE_UP_INTERLEAVE}, got {intermediate}" + ) + if tensor.ndim != 3 or tensor.shape[1] != 2 * intermediate: + raise ValueError( + f"expected (experts, {2 * intermediate}, K) tensor, " + f"got {tuple(tensor.shape)}" + ) + + experts, _gate_up, reduction = tensor.shape + pairs = intermediate // _GATE_UP_INTERLEAVE + gate = tensor[:, :intermediate].reshape( + experts, + pairs, + _GATE_UP_INTERLEAVE, + reduction, + ) + up = tensor[:, intermediate:].reshape( + experts, + pairs, + _GATE_UP_INTERLEAVE, + reduction, + ) + return ( + torch.stack((gate, up), dim=2) + .reshape(experts, 2 * intermediate, reduction) + .contiguous() + ) + + +def _to_blocked_bytes(scale_2d: torch.Tensor) -> torch.Tensor: + """Apply the kernel's 32x4x4 scale-factor atom swizzle.""" + + if scale_2d.ndim != 2: + raise ValueError(f"expected 2D scale tensor, got {scale_2d.ndim}D") + rows, columns = scale_2d.shape + if rows == 0 or columns == 0: + return scale_2d.new_empty((0,), dtype=torch.uint8) + + row_blocks = (rows + 127) // 128 + column_blocks = (columns + 3) // 4 + padded_rows = row_blocks * 128 + padded_columns = column_blocks * 4 + padded = torch.zeros( + padded_rows, + padded_columns, + dtype=torch.uint8, + device=scale_2d.device, + ) + padded[:rows, :columns].copy_(_as_bytes(scale_2d)) + blocks = padded.view(row_blocks, 128, column_blocks, 4).permute( + 0, + 2, + 1, + 3, + ) + return ( + blocks.reshape(-1, 4, 32, 4) + .transpose(1, 2) + .reshape(-1, 32, 16) + .flatten() + ) + + +def _stack_blocked_scales(raw_scales: torch.Tensor) -> torch.Tensor: + experts = raw_scales.shape[0] + blocked = [_to_blocked_bytes(raw_scales[e]) for e in range(experts)] + if not blocked: + return torch.empty( + (0, 0), + dtype=torch.uint8, + device=raw_scales.device, + ).view(raw_scales.dtype) + flat_size = blocked[0].numel() + output = torch.empty( + experts, + flat_size, + dtype=torch.uint8, + device=raw_scales.device, + ) + for expert, values in enumerate(blocked): + output[expert].copy_(values) + return output.view(raw_scales.dtype) + + +def _prepare_fc1( + tensor: BlockScaledTensor, + intermediate: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Preserve logical bytes while building interleaved K-major FC1 tensors.""" + + payload_nkh = _as_bytes(tensor.data).permute(0, 2, 1).contiguous() + payload_interleaved = _interleave_gate_up_rows( + payload_nkh, + intermediate, + ) + payload = payload_interleaved.view(_MXFP8_DATA_DTYPE).permute(0, 2, 1) + + scales_nk = _as_bytes(tensor.scale).permute(0, 2, 1).contiguous() + scales_interleaved = _interleave_gate_up_rows(scales_nk, intermediate) + scale = _stack_blocked_scales(scales_interleaved) + return payload, scale + + +def _prepare_fc2( + tensor: BlockScaledTensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Preserve logical bytes while building K-major FC2 tensors.""" + + payload_nk = _as_bytes(tensor.data).permute(0, 2, 1).contiguous() + payload = payload_nk.view(_MXFP8_DATA_DTYPE).permute(0, 2, 1) + scales_nk = _as_bytes(tensor.scale).permute(0, 2, 1).contiguous() + scale = _stack_blocked_scales(scales_nk) + return payload, scale + + +def _tensor_fingerprint(tensor: torch.Tensor) -> tuple | None: + try: + version = tensor._version + except RuntimeError: + return None + return ( + tensor.data_ptr(), + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + tensor.device, + version, + ) + + +def _block_scaled_fingerprint(tensor: BlockScaledTensor) -> tuple | None: + data = _tensor_fingerprint(tensor.data) + scale = _tensor_fingerprint(tensor.scale) + if data is None or scale is None: + return None + return data, scale + + +@dataclass(frozen=True) +class Mxfp8Weights: + fc1_weight: torch.Tensor + fc1_weight_sf: torch.Tensor + fc2_weight: torch.Tensor + fc2_weight_sf: torch.Tensor + + +@dataclass(frozen=True) +class Mxfp8LaunchInputs: + activation: torch.Tensor + activation_sf: torch.Tensor + topk_indices: torch.Tensor + topk_scores: torch.Tensor + weights: Mxfp8Weights + fc1_c: torch.Tensor | None + output_data: torch.Tensor + col_quant_data: torch.Tensor | None + col_quant_sf: torch.Tensor | None + overflow_flag: torch.Tensor + local_workspace: torch.Tensor + shared_workspace: torch.Tensor + token_count: int + + +class Mxfp8InputAdapter: + """Stateful staging adapter with mutation-aware weight transforms.""" + + def __init__(self) -> None: + self._weight_key: tuple | None = None + self._weights: Mxfp8Weights | None = None + self._weight_sources: tuple[torch.Tensor, ...] | None = None + self._weight_refresh_count = 0 + self._initialized_workspace_key: tuple[int, int] | None = None + + @property + def weight_refresh_count(self) -> int: + return self._weight_refresh_count + + def has_cached_weights(self, request: ValidatedForwardRequest) -> bool: + key = self._request_weight_key(request) + return key is not None and key == self._weight_key and self._weights is not None + + def weights_have_version_counters( + self, + request: ValidatedForwardRequest, + ) -> bool: + return self._request_weight_key(request) is not None + + @staticmethod + def _request_weight_key( + request: ValidatedForwardRequest, + ) -> tuple | None: + fc1 = ( + _block_scaled_fingerprint(request.fc1_weight) + if isinstance(request.fc1_weight, BlockScaledTensor) + else _tensor_fingerprint(request.fc1_weight) + ) + fc2 = ( + _block_scaled_fingerprint(request.fc2_weight) + if isinstance(request.fc2_weight, BlockScaledTensor) + else _tensor_fingerprint(request.fc2_weight) + ) + if fc1 is None or fc2 is None: + return None + return fc1, fc2 + + def _prepare_weights( + self, + request: ValidatedForwardRequest, + config: Mxfp8KernelConfig, + ) -> Mxfp8Weights: + key = self._request_weight_key(request) + if key is not None and key == self._weight_key and self._weights is not None: + return self._weights + + fc1_source = _as_mxfp8(request.fc1_weight) + fc2_source = _as_mxfp8(request.fc2_weight) + fc1_weight, fc1_weight_sf = _prepare_fc1( + fc1_source, + config.intermediate, + ) + fc2_weight, fc2_weight_sf = _prepare_fc2(fc2_source) + weights = Mxfp8Weights( + fc1_weight=fc1_weight, + fc1_weight_sf=fc1_weight_sf, + fc2_weight=fc2_weight, + fc2_weight_sf=fc2_weight_sf, + ) + self._weight_key = key + self._weights = weights + # Retain the source storages while this entry is cached so allocator + # pointer reuse cannot produce a false cache hit. + self._weight_sources = ( + *( + (request.fc1_weight.data, request.fc1_weight.scale) + if isinstance(request.fc1_weight, BlockScaledTensor) + else (request.fc1_weight,) + ), + *( + (request.fc2_weight.data, request.fc2_weight.scale) + if isinstance(request.fc2_weight, BlockScaledTensor) + else (request.fc2_weight,) + ), + ) + self._weight_refresh_count += 1 + return weights + + def stage( + self, + request: ValidatedForwardRequest, + resources: PreparedResources, + config: Mxfp8KernelConfig, + *, + local_workspace_zero_bytes: int, + shared_workspace_zero_bytes: int, + pre_reduced_activation_offset: int | None, + pre_reduced_activation_bytes_per_token: int, + pre_reduced_activation_sf_offset: int | None, + pre_reduced_activation_sf_bytes_per_token: int, + col_quant_data_rows: int, + col_quant_sf_elements: int, + fc1_c: torch.Tensor | None = None, + ) -> Mxfp8LaunchInputs: + capacity = config.max_tokens_per_rank + token_count = request.token_count + if config.generate_c: + if fc1_c is None: + raise ValueError("generate_c=True requires an fc1_c buffer") + if ( + fc1_c.dtype is not torch.bfloat16 + or fc1_c.device != request.device + or fc1_c.ndim != 2 + or fc1_c.shape[0] <= 0 + or fc1_c.shape[1] != config.fc1_out + or not fc1_c.is_contiguous() + ): + raise ValueError( + "fc1_c buffer must be contiguous BF16 on the request " + f"device with shape (capacity, {config.fc1_out})" + ) + elif fc1_c is not None: + raise ValueError("generate_c=False must not receive an fc1_c buffer") + hidden_sf_columns = (config.hidden + 31) // 32 + padded_sf_columns = padded_mxfp8_scale_columns(config.hidden) + symmetric = resources.workspace.symmetric + local = resources.workspace.local + symmetric_guard = symmetric.get("symmetric_guard") + if symmetric_guard is not None: + symmetric_guard.fill_(_WORKSPACE_GUARD_BYTE) + local_guard = local.get("local_guard") + if local_guard is not None: + local_guard.fill_(_WORKSPACE_GUARD_BYTE) + + activation = _typed_view( + symmetric["activation_data"], + _MXFP8_DATA_DTYPE, + (capacity, config.hidden), + ) + activation_sf = _typed_view( + symmetric["activation_scale"], + _MXFP8_SCALE_DTYPE, + (capacity, padded_sf_columns), + ) + topk_weights = _typed_view( + symmetric["topk_weights"], + torch.float32, + (capacity, config.top_k), + ) + output_data = _typed_view( + symmetric["output_data"], + torch.bfloat16, + (capacity, config.hidden), + ) + topk_indices = _typed_view( + local["topk_idx"], + torch.int32, + (capacity, config.top_k), + ) + overflow_flag = _typed_view( + local["overflow_flag"], + torch.int32, + (1,), + ) + if config.enable_col_quant: + if col_quant_data_rows <= 0 or col_quant_sf_elements <= 0: + raise ValueError( + "enabled column requant requires positive output capacities" + ) + col_quant_data = _typed_view( + local["col_quant_data"], + _MXFP8_DATA_DTYPE, + (col_quant_data_rows, config.hidden), + ) + col_quant_sf = _typed_view( + local["col_quant_sf"], + torch.uint8, + (col_quant_sf_elements,), + ) + else: + if col_quant_data_rows != 0 or col_quant_sf_elements != 0: + raise ValueError( + "disabled column requant must not reserve output capacity" + ) + col_quant_data = None + col_quant_sf = None + local_workspace = local["kernel_local_workspace"] + shared_workspace = symmetric["kernel_shared_workspace"] + + staged_activation = _as_mxfp8(request.activation) + _as_bytes(activation).zero_() + _as_bytes(activation[:token_count]).copy_( + _as_bytes(staged_activation.data) + ) + _as_bytes(activation_sf).zero_() + _as_bytes( + activation_sf[:token_count, :hidden_sf_columns] + ).copy_(_as_bytes(staged_activation.scale)) + _validate_int32_downcast(request.topk_idx) + topk_indices.fill_(-1) + topk_indices[:token_count].copy_(request.topk_idx) + topk_weights.zero_() + topk_weights[:token_count].copy_(request.topk_weights) + _as_bytes(output_data).zero_() + if col_quant_data is not None: + _as_bytes(col_quant_data).zero_() + if col_quant_sf is not None: + col_quant_sf.zero_() + overflow_flag.zero_() + workspace_key = ( + local_workspace.data_ptr(), + shared_workspace.data_ptr(), + ) + if workspace_key != self._initialized_workspace_key: + # This prefix contains both tail-reset regions and persistent + # sense-reversing NVLink barrier counters marked + # zero_on_first_allocate. Re-zeroing it on a later rank-skewed + # launch can erase a peer's signal and deadlock both kernels. + _zero_workspace_prefix( + local_workspace, + local_workspace_zero_bytes, + name="local workspace", + ) + _zero_workspace_prefix( + shared_workspace, + shared_workspace_zero_bytes, + name="shared workspace", + ) + self._initialized_workspace_key = workspace_key + if config.fc2_in_kernel_topk_reduce: + if ( + pre_reduced_activation_offset is not None + or pre_reduced_activation_bytes_per_token != 0 + or pre_reduced_activation_sf_offset is not None + or pre_reduced_activation_sf_bytes_per_token != 0 + ): + raise ValueError( + "in-kernel top-k reduction must not receive a " + "standalone pre-reduced activation workspace" + ) + # output_data is the in-kernel REDG accumulation base and was + # cleared above. + else: + if ( + pre_reduced_activation_offset is None + or pre_reduced_activation_bytes_per_token <= 0 + ): + raise ValueError( + "standalone top-k reduction requires a pre-reduced " + "activation workspace" + ) + # The kernel writes only valid routes into this persistent combine + # plane. Clear the active token rows so dropped routes cannot reuse + # contributions from a previous launch. + _zero_workspace_range( + shared_workspace, + pre_reduced_activation_offset, + token_count * pre_reduced_activation_bytes_per_token, + name="pre-reduced activation workspace", + ) + quantized_combine = config.combine_format != "bf16" + if quantized_combine: + if ( + pre_reduced_activation_sf_offset is None + or pre_reduced_activation_sf_bytes_per_token <= 0 + ): + raise ValueError( + "quantized standalone top-k reduction requires a " + "pre-reduced scale workspace" + ) + _zero_workspace_range( + shared_workspace, + pre_reduced_activation_sf_offset, + token_count + * pre_reduced_activation_sf_bytes_per_token, + name="pre-reduced activation scale workspace", + ) + elif ( + pre_reduced_activation_sf_offset is not None + or pre_reduced_activation_sf_bytes_per_token != 0 + ): + raise ValueError( + "BF16 standalone top-k reduction must not receive a " + "pre-reduced scale workspace" + ) + + weights = self._prepare_weights(request, config) + return Mxfp8LaunchInputs( + activation=activation, + activation_sf=activation_sf, + topk_indices=topk_indices, + topk_scores=topk_weights, + weights=weights, + fc1_c=fc1_c, + output_data=output_data, + col_quant_data=col_quant_data, + col_quant_sf=col_quant_sf, + overflow_flag=overflow_flag, + local_workspace=local_workspace, + shared_workspace=shared_workspace, + token_count=token_count, + ) + + def close(self) -> None: + self._weights = None + self._weight_key = None + self._weight_sources = None + self._initialized_workspace_key = None + + +__all__ = [ + "Mxfp8InputAdapter", + "Mxfp8LaunchInputs", + "Mxfp8Weights", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py new file mode 100644 index 000000000..b95f61415 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py @@ -0,0 +1,334 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Single-rank and EP-subgroup MXFP8 execution backend orchestration.""" + +from __future__ import annotations + +import threading + +import torch +import torch.distributed as dist + +from ..._backend import BackendUnavailableError +from ..._contracts import ( + ForwardConfig, + ValidatedBackwardRequest, + ValidatedForwardRequest, +) +from .._plan import ExecutionPlanOwner +from ._adapter import Mxfp8InputAdapter +from ._backward import Mxfp8BackwardExecutor +from ._compile import ( + CompiledMxfp8Kernel, + PreparedMxfp8Kernel, + compile_or_get, + prepare_kernel, +) +from ._config import Mxfp8KernelConfig +from ._launch import launch_forward +from ._stash import Mxfp8ForwardStash + + +class Mxfp8Backend: + """Own forward/backward executors and per-instance plan resources.""" + + def __init__(self, config: ForwardConfig, device: torch.device) -> None: + self.config = config + self.device = torch.device(device) + self.kernel_config = Mxfp8KernelConfig.from_forward_config(config) + self._adapter = Mxfp8InputAdapter() + self._stash = ( + Mxfp8ForwardStash(config, self.device) + if config.generate_c + else None + ) + self._prepared_kernel: PreparedMxfp8Kernel | None = None + self._compiled: CompiledMxfp8Kernel | None = None + self._plan: ExecutionPlanOwner | None = None + self._warmed_up = False + self._closed = False + self._completion_event: torch.cuda.Event | None = None + self._completion_recorded = False + self._device_work_may_be_pending = False + self._ep_launch_ready = config.ep_size == 1 + self._backward_executor: Mxfp8BackwardExecutor | None = None + self._lock = threading.RLock() + + @property + def warmed_up(self) -> bool: + return self._warmed_up + + @property + def kernel_fingerprint(self) -> dict | None: + """Fingerprint of the callable compiled by the most recent launch.""" + + if self._compiled is None: + return None + return self._compiled.fingerprint + + def _ensure_prepared_kernel(self) -> PreparedMxfp8Kernel: + if self._prepared_kernel is None: + try: + self._prepared_kernel = prepare_kernel( + self.config, + self.kernel_config, + self.device, + ) + except (ImportError, OSError) as exc: + raise BackendUnavailableError( + "MoeEp MXFP8 backend requires the 'moe_ep' optional " + "dependencies and their shared libraries" + ) from exc + return self._prepared_kernel + + def _ensure_ep_launch_ready(self, resources, stream) -> None: + if self._ep_launch_ready: + return + # First subgroup launch only: peer metadata writes begin before the + # kernel's first cross-rank device barrier. Ensure every rank's + # root-zero and staging work has completed before any rank can issue + # those writes. + stream.synchronize() + if resources.runtime.group is None: + raise RuntimeError( + "distributed MXFP8 launch requires a " + "torch.distributed process group" + ) + tuning_signature = self.kernel_config.tuning_signature( + self._ensure_prepared_kernel().launch_cluster_count + ) + rank_tuning_signatures = [None] * resources.runtime.world_size + dist.all_gather_object( + rank_tuning_signatures, + tuning_signature, + group=resources.runtime.group, + ) + if any( + signature != rank_tuning_signatures[0] + for signature in rank_tuning_signatures[1:] + ): + raise RuntimeError( + "MoeEp tuning must match on every expert-parallel rank; " + f"effective signatures by rank: {rank_tuning_signatures}" + ) + dist.barrier(group=resources.runtime.group) + self._ep_launch_ready = True + + def forward(self, request: ValidatedForwardRequest): + with self._lock: + if self._closed: + raise RuntimeError("MoeEp MXFP8 backend is closed") + if request.device != self.device: + raise ValueError( + f"MoeEp MXFP8 backend is bound to {self.device}, " + f"got {request.device}" + ) + + with torch.cuda.device(self.device): + capturing = torch.cuda.is_current_stream_capturing() + if capturing and self._stash is not None: + raise NotImplementedError( + "MoeEp generate_c=True is eager-only and does not " + "support CUDA graph capture" + ) + if ( + capturing + and not self._adapter.weights_have_version_counters( + request + ) + ): + raise NotImplementedError( + "CUDA graph capture does not support inference tensor " + "weights without version counters; eager calls remain " + "supported and repack those weights on every call" + ) + if capturing and ( + not self._warmed_up + or not self._adapter.has_cached_weights(request) + ): + raise RuntimeError( + "MoeEp MXFP8 backend and weights must be warmed up " + "before CUDA graph capture" + ) + + stream = torch.cuda.current_stream(self.device) + if self._device_work_may_be_pending: + torch.cuda.synchronize(self.device) + self._device_work_may_be_pending = False + if self._completion_event is None: + self._completion_event = torch.cuda.Event() + elif self._completion_recorded and not capturing: + stream.wait_event(self._completion_event) + + prepared = self._ensure_prepared_kernel() + if self._plan is None: + self._plan = ExecutionPlanOwner( + self.config, + self.device, + prepared.workspace_requirements, + ) + device_work_attempted = False + try: + # Allocation zeroing, input staging, weight transforms, + # compilation, and launch can all enqueue device work. + # Record one completion event even if a later step fails so + # a retry on another stream cannot race those writes. + device_work_attempted = True + resources = self._plan.prepare(request) + stash_plan = ( + None + if self._stash is None + else self._stash.prepare( + request, + pool_token_capacity=prepared.pool_token_capacity, + ) + ) + inputs = self._adapter.stage( + request, + resources, + self.kernel_config, + local_workspace_zero_bytes=( + prepared.local_workspace_zero_bytes + ), + shared_workspace_zero_bytes=( + prepared.shared_workspace_zero_bytes + ), + pre_reduced_activation_offset=( + prepared.pre_reduced_activation_offset + ), + pre_reduced_activation_bytes_per_token=( + prepared.pre_reduced_activation_bytes_per_token + ), + pre_reduced_activation_sf_offset=( + prepared.pre_reduced_activation_sf_offset + ), + pre_reduced_activation_sf_bytes_per_token=( + prepared.pre_reduced_activation_sf_bytes_per_token + ), + col_quant_data_rows=prepared.col_quant_data_rows, + col_quant_sf_elements=prepared.col_quant_sf_elements, + fc1_c=( + None + if stash_plan is None + else stash_plan.buffer + ), + ) + self._compiled = compile_or_get( + prepared, + inputs, + resources, + ) + self._ensure_ep_launch_ready(resources, stream) + output = launch_forward( + self._compiled, + inputs, + resources, + ) + if self._stash is not None: + assert stash_plan is not None + ( + fc1_c, + route_metadata, + wgrad_stash, + ) = self._stash.materialize( + stash_plan, + inputs, + prepared, + ) + if wgrad_stash is None: + output = (output, fc1_c, route_metadata) + else: + output = ( + output, + fc1_c, + route_metadata, + wgrad_stash, + ) + except (ImportError, OSError) as exc: + raise BackendUnavailableError( + "MoeEp MXFP8 backend requires the 'moe_ep' optional " + "dependencies and their shared libraries" + ) from exc + finally: + if device_work_attempted and not capturing: + try: + self._completion_event.record(stream) + self._completion_recorded = True + self._device_work_may_be_pending = False + except Exception: + self._completion_recorded = False + self._device_work_may_be_pending = True + raise + + self._warmed_up = True + return output + + def backward(self, request: ValidatedBackwardRequest): + """Run the restricted explicit dgrad/dprob Rubin MXFP8 path.""" + + with self._lock: + if self._closed: + raise RuntimeError("MoeEp MXFP8 backend is closed") + if request.device != self.device: + raise ValueError( + f"MoeEp MXFP8 backend is bound to {self.device}, " + f"got {request.device}" + ) + stream = torch.cuda.current_stream(self.device) + if self._device_work_may_be_pending: + torch.cuda.synchronize(self.device) + self._device_work_may_be_pending = False + if self._completion_event is None: + self._completion_event = torch.cuda.Event() + elif self._completion_recorded: + stream.wait_event(self._completion_event) + if self._backward_executor is None: + self._backward_executor = Mxfp8BackwardExecutor( + self.config, + self.device, + ) + try: + result = self._backward_executor.run(request) + finally: + try: + self._completion_event.record(stream) + self._completion_recorded = True + self._device_work_may_be_pending = False + except Exception: + self._completion_recorded = False + self._device_work_may_be_pending = True + raise + return result + + def close(self) -> None: + with self._lock: + if self._closed: + return + with torch.cuda.device(self.device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "MoeEp MXFP8 backend cannot be closed during " + "CUDA graph capture" + ) + if self._plan is not None or self._backward_executor is not None: + torch.cuda.synchronize(self.device) + self._adapter.close() + if self._backward_executor is not None: + self._backward_executor.close() + self._backward_executor = None + if self._stash is not None: + self._stash.close() + if self._plan is not None: + self._plan.close() + self._plan = None + self._prepared_kernel = None + self._compiled = None + self._completion_event = None + self._completion_recorded = False + self._device_work_may_be_pending = False + self._ep_launch_ready = self.config.ep_size == 1 + self._closed = True + + +__all__ = ["Mxfp8Backend"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward.py new file mode 100644 index 000000000..2d642f373 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward.py @@ -0,0 +1,112 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Explicit dgrad/dprob Rubin MXFP8 backward orchestration.""" + +from __future__ import annotations + +import torch +import torch.distributed as dist + +from ..._backend import BackendUnavailableError +from ..._contracts import ForwardConfig, ValidatedBackwardRequest +from .._plan import ExecutionPlanOwner +from ._backward_compile import ( + CompiledMxfp8BackwardKernel, + PreparedMxfp8BackwardKernel, + compile_backward_or_get, + prepare_backward_kernel, +) +from ._backward_dispatch import Mxfp8BackwardRedispatch +from ._backward_dprob import return_grad_topk_weights +from ._backward_launch import launch_backward_dglu +from ._backward_layout import Mxfp8BackwardLayout +from ._backward_staging import stage_backward +from ._backward_wgrad_export import export_wgrad_operands +from ._config import Mxfp8KernelConfig + + +class Mxfp8BackwardExecutor: + """Own only compiled products and reusable capacity workspaces.""" + + def __init__(self, config: ForwardConfig, device: torch.device) -> None: + self.config = config + self.device = torch.device(device) + self.kernel_config = Mxfp8KernelConfig.from_forward_config(config) + self._prepared: PreparedMxfp8BackwardKernel | None = None + self._compiled: CompiledMxfp8BackwardKernel | None = None + self._plan: ExecutionPlanOwner | None = None + self._ep_launch_ready = config.ep_size == 1 + + def _ensure_prepared(self) -> PreparedMxfp8BackwardKernel: + if self._prepared is None: + try: + self._prepared = prepare_backward_kernel( + self.config, + self.kernel_config, + self.device, + ) + except (ImportError, OSError) as exc: + raise BackendUnavailableError( + "MoeEp MXFP8 backward requires the 'moe_ep' optional " + "dependencies and their shared libraries" + ) from exc + return self._prepared + + def _ensure_ep_launch_ready(self) -> None: + if self._ep_launch_ready: + return + if self.config.ep_group is None: + raise RuntimeError( + "distributed MXFP8 backward requires an EP process group" + ) + torch.cuda.current_stream(self.device).synchronize() + dist.barrier(group=self.config.ep_group) + self._ep_launch_ready = True + + def run( + self, + request: ValidatedBackwardRequest, + ): + prepared = self._ensure_prepared() + if self._plan is None: + self._plan = ExecutionPlanOwner( + self.config, + self.device, + prepared.workspace_requirements, + ) + + layout = Mxfp8BackwardLayout.from_request(request) + redispatched = Mxfp8BackwardRedispatch(request).run() + resources = self._plan.prepare(request) + inputs = stage_backward(request, layout, prepared, resources) + self._compiled = compile_backward_or_get( + prepared, + inputs, + resources, + ) + self._ensure_ep_launch_ready() + dglu = launch_backward_dglu( + self._compiled, + inputs, + resources, + ) + grad_topk_weights = return_grad_topk_weights( + request, + redispatched.grad_output, + ) + if request.config.backward_wgrad_mode == "operands": + operands = export_wgrad_operands(request, dglu) + return dglu.grad_activation, grad_topk_weights, operands + return dglu.grad_activation, grad_topk_weights + + def close(self) -> None: + if self._plan is not None: + self._plan.close() + self._plan = None + self._prepared = None + self._compiled = None + self._ep_launch_ready = self.config.ep_size == 1 + + +__all__ = ["Mxfp8BackwardExecutor"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py new file mode 100644 index 000000000..77530ae39 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py @@ -0,0 +1,377 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Compilation and ABI metadata for Rubin MXFP8 dGLU backward.""" + +from __future__ import annotations + +import math +import os +import threading +from dataclasses import dataclass +from typing import Any + +import torch + +from ..._contracts import ForwardConfig +from .._plan import PreparedResources +from .._workspace import WorkspaceRequirements +from ._compile import ( + _pre_reduced_sf_workspace_metadata, + _pre_reduced_workspace_metadata, +) +from ._config import Mxfp8KernelConfig +from ._formats import combine_wire_format +from ._launch import _to_cute, _to_cute_ptr + + +@dataclass(frozen=True) +class PreparedMxfp8BackwardKernel: + config: Mxfp8KernelConfig + device: torch.device + architecture: tuple[int, int] + kernel: Any + launch_cluster_count: int + workspace_requirements: WorkspaceRequirements + pool_token_capacity: int + pre_reduced_activation_offset: int + pre_reduced_activation_bytes_per_token: int + pre_reduced_activation_sf_offset: int | None + pre_reduced_activation_sf_bytes_per_token: int + local_workspace_zero_bytes: int + shared_workspace_zero_bytes: int + dfc2_recompute: bool + dfc2_col_output: bool + enable_grad_y2_col_quant: bool + + +@dataclass(frozen=True) +class Mxfp8BackwardLaunchInputs: + grad_out: torch.Tensor + grad_out_sf: torch.Tensor + topk_idx: torch.Tensor + topk_weights: torch.Tensor + fc1_weight: torch.Tensor + fc1_weight_sf: torch.Tensor + fc2_weight: torch.Tensor + fc2_weight_sf: torch.Tensor + beta: torch.Tensor + fc1_preact: torch.Tensor + output_activation: torch.Tensor + overflow_flag: torch.Tensor + dprob: torch.Tensor + fc1_recompute: torch.Tensor + fc1_recompute_sf: torch.Tensor + fc1_col_output: torch.Tensor + fc1_col_output_sf: torch.Tensor + grad_y2: torch.Tensor + grad_y2_sf: torch.Tensor + local_workspace: torch.Tensor + shared_workspace: torch.Tensor + token_count: int + + +@dataclass(frozen=True) +class CompiledMxfp8BackwardKernel: + key: tuple + callable: Any + + +_COMPILE_LOCK = threading.RLock() +_COMPILE_CACHE: dict[tuple, CompiledMxfp8BackwardKernel] = {} + + +def prepare_backward_kernel( + forward_config: ForwardConfig, + config: Mxfp8KernelConfig, + device: torch.device, +) -> PreparedMxfp8BackwardKernel: + """Instantiate the fixed Rubin dGLU specialization.""" + + torch.cuda.set_device(device) + architecture = torch.cuda.get_device_capability(device) + if architecture != (10, 7): + raise RuntimeError( + "Rubin MXFP8 backward requires compute capability (10, 7), " + f"got {architecture}" + ) + configured_architecture = os.environ.get("CUTE_DSL_ARCH") + if configured_architecture is None: + os.environ["CUTE_DSL_ARCH"] = "sm_107a" + elif configured_architecture not in ("sm_107", "sm_107a"): + raise RuntimeError( + "CUTE_DSL_ARCH must target SM107 for the Rubin MXFP8 backward" + ) + import cutlass + import cutlass.utils as utils + + from ..cutedsl_src.kernel_src.rubin.training.mega.bwd_dglu import ( + Sm107MegaMoEMxfp8DgluKernel, + ) + from ..cutedsl_src.quant_def import CombineFormat + + launch_cluster_count = int( + utils.HardwareInfo().get_max_active_clusters(config.cluster_size) + ) + if launch_cluster_count <= 0: + raise RuntimeError( + "hardware occupancy query returned no launchable Rubin clusters" + ) + group_hint = ( + launch_cluster_count + if config.group_hint is None + else config.group_hint + ) + operands_mode = forward_config.backward_wgrad_mode == "operands" + dfc2_recompute = operands_mode + dfc2_col_output = operands_mode + enable_grad_y2_col_quant = operands_mode + kernel = Sm107MegaMoEMxfp8DgluKernel.from_kwargs( + mma_tiler_mnk=config.mma_tiler_mnk, + cluster_shape_mnk=config.cluster_shape_mnk, + use_2cta_instrs=config.use_2cta_instrs, + group_hint=group_hint, + token_padding_block=config.token_padding_block, + sf_padding_block=config.sf_padding_block, + load_balance_mode=config.load_balance_mode, + static_expert_shape=( + config.num_experts, + config.intermediate, + config.hidden, + ), + force_static_sched=config.force_static_sched, + clc_bundle_size=config.clc_bundle_size, + num_sched_stages=config.num_sched_stages, + ab_dtype=cutlass.Float8E4M3FN, + sf_vec_size=config.sf_vec_size, + world_size=config.world_size, + local_rank=0, + num_topk=config.top_k, + max_tokens_per_rank=config.max_tokens_per_rank, + max_recv_size_per_rank=( + config.world_size + * config.max_tokens_per_rank + * config.top_k + ), + hidden=config.hidden, + launch_cluster_count=launch_cluster_count, + drop_on_overflow=True, + fc2_in_kernel_topk_reduce=False, + token_back_mode="epi_warps", + epi_flag_batch=config.epi_flag_batch, + flag_batch=config.flag_batch, + combine_format=CombineFormat.parse( + combine_wire_format(forward_config.combine_format) + ), + act_func=config.act_func, + gate_up_clamp=config.gate_up_clamp, + dfc2_recompute=dfc2_recompute, + dfc2_col_output=dfc2_col_output, + enable_grad_y2_col_quant=enable_grad_y2_col_quant, + num_ctas_grad_y2_col_quant=config.col_quant_num_ctas, + ) + local_bytes, shared_bytes = kernel.get_workspace_sizes() + local_zero, shared_zero = kernel.require_zero_workspace_leading_bytes + device_workspace = kernel._mega_device_workspace + pool_capacity = int(kernel.pool_token_capacity) + fc1_preact_shape = tuple( + int(extent) for extent in kernel.get_fc1_preact_shape() + ) + expected_preact_shape = ( + pool_capacity, + 2 * config.intermediate, + ) + if fc1_preact_shape != expected_preact_shape: + raise RuntimeError( + "Rubin dGLU fc1_preact shape mismatch: " + f"{fc1_preact_shape} != {expected_preact_shape}" + ) + aux_shapes = { + name: tuple(int(extent) for extent in shape) + for name, shape in kernel.get_aux_output_shapes().items() + } + fc1_preact_bytes = ( + math.prod(fc1_preact_shape) * torch.bfloat16.itemsize + ) + dprob_bytes = math.prod(aux_shapes["dprob"]) * torch.float32.itemsize + aux_data_bytes = max( + math.prod(aux_shapes["fc1_recompute"]), + math.prod(aux_shapes["fc1_col_output"]), + math.prod(aux_shapes["grad_y2"]), + ) * torch.float8_e4m3fn.itemsize + aux_scale_bytes = max( + math.prod(aux_shapes["fc1_recompute_sf"]), + math.prod(aux_shapes["fc1_col_output_sf"]), + math.prod(aux_shapes["grad_y2_sf"]), + ) * torch.float8_e8m0fnu.itemsize + requirements = WorkspaceRequirements.for_mxfp8( + forward_config, + kernel_local_workspace_bytes=local_bytes, + kernel_shared_workspace_bytes=shared_bytes, + backward_fc1_preact_bytes=fc1_preact_bytes, + backward_dprob_bytes=dprob_bytes, + backward_aux_data_bytes=aux_data_bytes, + backward_aux_scale_bytes=aux_scale_bytes, + ) + pre_reduced_offset, pre_reduced_bytes_per_token = ( + _pre_reduced_workspace_metadata( + device_workspace, + config, + shared_bytes, + ) + ) + if pre_reduced_offset is None or pre_reduced_bytes_per_token <= 0: + raise RuntimeError( + "Rubin MXFP8 backward requires standalone pre-reduced activation" + ) + pre_reduced_sf_offset, pre_reduced_sf_bytes_per_token = ( + _pre_reduced_sf_workspace_metadata( + device_workspace, + config, + shared_bytes, + ) + ) + return PreparedMxfp8BackwardKernel( + config=config, + device=torch.device(device), + architecture=architecture, + kernel=kernel, + launch_cluster_count=launch_cluster_count, + workspace_requirements=requirements, + pool_token_capacity=pool_capacity, + pre_reduced_activation_offset=pre_reduced_offset, + pre_reduced_activation_bytes_per_token=pre_reduced_bytes_per_token, + pre_reduced_activation_sf_offset=pre_reduced_sf_offset, + pre_reduced_activation_sf_bytes_per_token=( + pre_reduced_sf_bytes_per_token + ), + local_workspace_zero_bytes=int(local_zero), + shared_workspace_zero_bytes=int(shared_zero), + dfc2_recompute=dfc2_recompute, + dfc2_col_output=dfc2_col_output, + enable_grad_y2_col_quant=enable_grad_y2_col_quant, + ) + + +def _layout_signature(inputs: Mxfp8BackwardLaunchInputs) -> tuple: + tensors = tuple( + value + for value in inputs.__dict__.values() + if isinstance(value, torch.Tensor) + ) + return tuple( + (tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype) + for tensor in tensors + ) + + +def build_backward_runtime_kwargs( + inputs: Mxfp8BackwardLaunchInputs, + resources: PreparedResources, +) -> dict[str, Any]: + import cuda.bindings.driver as cuda + + stream = resources.runtime.current_stream() + return { + "grad_out": _to_cute(inputs.grad_out), + "grad_out_sf": _to_cute(inputs.grad_out_sf), + "topk_idx": _to_cute(inputs.topk_idx), + "topk_weights": _to_cute(inputs.topk_weights, assumed_align=4), + "fc1_weight": _to_cute(inputs.fc1_weight), + "fc1_weight_sf": _to_cute(inputs.fc1_weight_sf), + "fc2_weight": _to_cute(inputs.fc2_weight), + "fc2_weight_sf": _to_cute(inputs.fc2_weight_sf), + "beta": _to_cute(inputs.beta, assumed_align=4), + "fc1_preact": _to_cute( + inputs.fc1_preact, + assumed_align=128, + dynamic_layout=False, + ), + "output_activation": _to_cute(inputs.output_activation), + "overflow_flag": _to_cute( + inputs.overflow_flag, + assumed_align=4, + dynamic_layout=False, + ), + "dprob": _to_cute(inputs.dprob, dynamic_layout=False), + "fc1_recompute": _to_cute( + inputs.fc1_recompute, + assumed_align=128, + dynamic_layout=False, + ), + "fc1_recompute_sf": _to_cute( + inputs.fc1_recompute_sf, + assumed_align=128, + dynamic_layout=False, + ), + "fc1_col_output": _to_cute( + inputs.fc1_col_output, + assumed_align=128, + dynamic_layout=False, + ), + "fc1_col_output_sf": _to_cute( + inputs.fc1_col_output_sf, + assumed_align=128, + dynamic_layout=False, + ), + "grad_y2": _to_cute( + inputs.grad_y2, + assumed_align=128, + dynamic_layout=False, + ), + "grad_y2_sf": _to_cute( + inputs.grad_y2_sf, + dynamic_layout=False, + ), + "local_workspace": _to_cute_ptr(inputs.local_workspace), + "shared_workspace": _to_cute_ptr(inputs.shared_workspace), + "peer_rank_ptr_mapper_host": ( + resources.workspace.peer_mapping.to_sym_buffer_host() + ), + "stream": cuda.CUstream(stream.cuda_stream), + } + + +def compile_backward_or_get( + prepared: PreparedMxfp8BackwardKernel, + inputs: Mxfp8BackwardLaunchInputs, + resources: PreparedResources, +) -> CompiledMxfp8BackwardKernel: + signature = _layout_signature(inputs) + key = ( + prepared.config, + prepared.device.index, + prepared.architecture, + prepared.launch_cluster_count, + prepared.dfc2_recompute, + prepared.dfc2_col_output, + prepared.enable_grad_y2_col_quant, + signature, + ) + with _COMPILE_LOCK: + cached = _COMPILE_CACHE.get(key) + if cached is not None: + return cached + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "MXFP8 backward kernel must be compiled before capture" + ) + import cutlass.cute as cute + + runtime_kwargs = build_backward_runtime_kwargs(inputs, resources) + compiled = CompiledMxfp8BackwardKernel( + key=key, + callable=cute.compile(prepared.kernel, **runtime_kwargs), + ) + _COMPILE_CACHE[key] = compiled + return compiled + + +__all__ = [ + "CompiledMxfp8BackwardKernel", + "Mxfp8BackwardLaunchInputs", + "PreparedMxfp8BackwardKernel", + "build_backward_runtime_kwargs", + "compile_backward_or_get", + "prepare_backward_kernel", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dispatch.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dispatch.py new file mode 100644 index 000000000..f073e17c1 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dispatch.py @@ -0,0 +1,204 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Explicit grad-output re-dispatch for semantic router gradients.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Sequence + +import torch +import torch.distributed as dist + +from ..._contracts import ValidatedBackwardRequest + + +@dataclass(frozen=True) +class _DispatchPlan: + send_token: torch.Tensor + send_slot: torch.Tensor + send_local_expert: torch.Tensor + send_counts: tuple[int, ...] + recv_counts: tuple[int, ...] + + +@dataclass(frozen=True) +class RedispatchedGradOutput: + """Route rows in the compact public ``route_metadata`` order.""" + + grad_output: torch.Tensor + + +class Mxfp8BackwardRedispatch: + """Recreate the identical-route grad-output exchange for dprob.""" + + def __init__(self, request: ValidatedBackwardRequest) -> None: + self.request = request + self.config = request.config + + def _collective_device(self, device: torch.device) -> torch.device: + if ( + device.type != "cpu" + and self.config.ep_size > 1 + and dist.get_backend(self.config.ep_group) == "gloo" + ): + return torch.device("cpu") + return device + + def _exchange_counts(self, send_counts: torch.Tensor) -> torch.Tensor: + if self.config.ep_size == 1: + return send_counts.clone() + staged = send_counts.to(self._collective_device(send_counts.device)) + recv_counts = torch.empty_like(staged) + dist.all_to_all_single( + recv_counts, + staged, + group=self.config.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.config.ep_size == 1: + return send.clone() + comm_device = self._collective_device(send.device) + staged = send.contiguous().to(comm_device) + recv = torch.empty( + (sum(recv_counts), *send.shape[1:]), + 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.config.ep_group, + ) + return recv.to(send.device) + + def _plan(self) -> _DispatchPlan: + config = self.config + flat_expert = self.request.topk_idx.reshape(-1).to(torch.int64) + valid = flat_expert != -1 + token = torch.arange( + self.request.token_count, + dtype=torch.int64, + device=self.request.device, + ).repeat_interleave(config.top_k) + slot = torch.arange( + config.top_k, + dtype=torch.int64, + device=self.request.device, + ).repeat(self.request.token_count) + expert = flat_expert[valid] + destination = torch.div( + expert, + config.experts_per_rank, + rounding_mode="floor", + ) + order = torch.argsort(destination, stable=True) + destination = destination.index_select(0, order) + send_counts_tensor = torch.bincount( + destination, + minlength=config.ep_size, + ).to(torch.int64) + recv_counts_tensor = self._exchange_counts(send_counts_tensor) + return _DispatchPlan( + send_token=token[valid].index_select(0, order), + send_slot=slot[valid].index_select(0, order), + send_local_expert=expert.index_select(0, order).remainder( + config.experts_per_rank + ), + send_counts=tuple( + int(value) for value in send_counts_tensor.cpu().tolist() + ), + recv_counts=tuple( + int(value) for value in recv_counts_tensor.cpu().tolist() + ), + ) + + def run(self) -> RedispatchedGradOutput: + config = self.config + plan = self._plan() + + recv_token = self._all_to_all( + plan.send_token, + plan.send_counts, + plan.recv_counts, + ) + recv_slot = self._all_to_all( + plan.send_slot, + plan.send_counts, + plan.recv_counts, + ) + recv_expert = self._all_to_all( + plan.send_local_expert, + plan.send_counts, + plan.recv_counts, + ) + recv_grad_output = self._all_to_all( + self.request.grad_output.index_select(0, plan.send_token).float(), + plan.send_counts, + plan.recv_counts, + ) + recv_rank = torch.repeat_interleave( + torch.arange( + config.ep_size, + dtype=torch.int64, + device=self.request.device, + ), + torch.tensor( + plan.recv_counts, + dtype=torch.int64, + device=self.request.device, + ), + output_size=self.request.local_routes, + ) + if recv_grad_output.shape[0] != self.request.local_routes: + raise ValueError( + "route_metadata row count does not match the routes received " + "from the re-supplied topk_idx" + ) + + key = ( + ( + ( + recv_expert * config.ep_size + + recv_rank + ) + * int(config.max_tokens_per_rank) + + recv_token + ) + * config.top_k + + recv_slot + ) + compact_order = torch.argsort(key, stable=True) + actual_metadata = torch.stack( + ( + recv_expert.index_select(0, compact_order), + recv_rank.index_select(0, compact_order), + recv_token.index_select(0, compact_order), + recv_slot.index_select(0, compact_order), + ), + dim=1, + ).to(torch.int32) + if not torch.equal(actual_metadata, self.request.route_metadata): + raise ValueError( + "route_metadata does not match the re-supplied forward routes" + ) + + return RedispatchedGradOutput( + grad_output=recv_grad_output.index_select(0, compact_order), + ) + + +__all__ = [ + "Mxfp8BackwardRedispatch", + "RedispatchedGradOutput", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dprob.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dprob.py new file mode 100644 index 000000000..1cf9bf44a --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dprob.py @@ -0,0 +1,87 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Return pool-ordered dprob to the source ``(token, top-k)`` plane.""" + +from __future__ import annotations + +import torch +import torch.distributed as dist + +from ..._contracts import ValidatedBackwardRequest +from ._adapter import _decode_moe_tensor + + +def return_grad_topk_weights( + request: ValidatedBackwardRequest, + redispatched_grad_output: torch.Tensor, +) -> torch.Tensor: + """Compute semantic route gradients and return them to every source. + + The fused dGLU kernel consumes an MXFP8 materialization of ``grad_output``. + Using its in-kernel dprob would therefore expose quantization error through + an operation whose public contract specifies straight-through semantics. + Recompute only this scalar gradient from the original FP32 dY, the BF16 + forward stash, and the decoded FC2 weight. + """ + + config = request.config + metadata = request.route_metadata.to(torch.int64) + local_dprob = torch.zeros( + (request.local_routes,), + dtype=torch.float32, + device=request.device, + ) + fc2_weight = _decode_moe_tensor(request.fc2_weight) + gate, up = request.fc1_c.float().split( + config.intermediate_size, + dim=-1, + ) + if config.gate_up_clamp is not None: + gate = gate.clamp(max=config.gate_up_clamp) + up = up.clamp( + min=-config.gate_up_clamp, + max=config.gate_up_clamp, + ) + hidden = (gate * torch.sigmoid(gate)) * up + local_expert = metadata[:, 0] + for expert in range(config.experts_per_rank): + positions = torch.nonzero( + local_expert == expert, + as_tuple=False, + ).flatten() + if positions.numel() == 0: + continue + grad_output = redispatched_grad_output.index_select(0, positions) + expert_hidden = hidden.index_select(0, positions) + grad_hidden = grad_output @ fc2_weight[expert].transpose(0, 1) + local_dprob.index_copy_( + 0, + positions, + (grad_hidden * expert_hidden).sum(dim=-1), + ) + + global_dprob = torch.zeros( + ( + config.ep_size, + int(config.max_tokens_per_rank), + config.top_k, + ), + dtype=torch.float32, + device=request.device, + ) + if request.local_routes: + global_dprob[ + metadata[:, 1], + metadata[:, 2], + metadata[:, 3], + ] = local_dprob + if config.ep_size > 1: + dist.all_reduce(global_dprob, group=config.ep_group) + return global_dprob[ + config.ep_rank, + : request.token_count, + ] + + +__all__ = ["return_grad_topk_weights"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py new file mode 100644 index 000000000..9cf069be5 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Current-stream launch of the Rubin MXFP8 dGLU backward product.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from .._plan import PreparedResources +from ._backward_compile import ( + CompiledMxfp8BackwardKernel, + Mxfp8BackwardLaunchInputs, + build_backward_runtime_kwargs, +) +from ._launch import _check_overflow + + +@dataclass(frozen=True) +class Mxfp8DgluResult: + grad_activation: torch.Tensor + fc1_recompute: torch.Tensor + fc1_recompute_sf: torch.Tensor + fc1_col_output: torch.Tensor + fc1_col_output_sf: torch.Tensor + grad_y2: torch.Tensor + grad_y2_sf: torch.Tensor + + +def launch_backward_dglu( + compiled: CompiledMxfp8BackwardKernel, + inputs: Mxfp8BackwardLaunchInputs, + resources: PreparedResources, +) -> Mxfp8DgluResult: + runtime_kwargs = build_backward_runtime_kwargs(inputs, resources) + compiled.callable(**runtime_kwargs) + _check_overflow(inputs.overflow_flag) + + return Mxfp8DgluResult( + grad_activation=inputs.output_activation[ + : inputs.token_count + ].float(), + fc1_recompute=inputs.fc1_recompute, + fc1_recompute_sf=inputs.fc1_recompute_sf, + fc1_col_output=inputs.fc1_col_output, + fc1_col_output_sf=inputs.fc1_col_output_sf, + grad_y2=inputs.grad_y2, + grad_y2_sf=inputs.grad_y2_sf, + ) + + +__all__ = ["Mxfp8DgluResult", "launch_backward_dglu"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_layout.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_layout.py new file mode 100644 index 000000000..e1ae287d5 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_layout.py @@ -0,0 +1,82 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Stateless lowering from the public backward stash to Rubin pool rows.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from ..._contracts import ForwardConfig, ValidatedBackwardRequest + +@dataclass(frozen=True) +class Mxfp8BackwardLayout: + """Stateless lowering from compact route metadata to the dGLU pool LUT.""" + + preact_row_lut: torch.Tensor + + @classmethod + def from_request( + cls, + request: ValidatedBackwardRequest, + ) -> "Mxfp8BackwardLayout": + config = request.config + metadata = request.route_metadata + if config.max_tokens_per_rank is None: + raise ValueError("MXFP8 backward requires max_tokens_per_rank") + + bounds = ( + (metadata[:, 0], 0, config.experts_per_rank, "local expert"), + (metadata[:, 1], 0, config.ep_size, "source rank"), + ( + metadata[:, 2], + 0, + config.max_tokens_per_rank, + "source token", + ), + (metadata[:, 3], 0, config.top_k, "source top-k slot"), + ) + for values, lower, upper, name in bounds: + if values.numel() and bool( + ((values < lower) | (values >= upper)).any().item() + ): + raise ValueError( + f"route_metadata contains an out-of-range {name}" + ) + + preact_row_lut = cls._build_preact_row_lut(config, metadata) + return cls( + preact_row_lut=preact_row_lut, + ) + + @staticmethod + def _build_preact_row_lut( + config: ForwardConfig, + metadata: torch.Tensor, + ) -> torch.Tensor: + lut = torch.full( + ( + config.ep_size, + int(config.max_tokens_per_rank), + config.top_k, + ), + -1, + dtype=torch.int32, + device=metadata.device, + ) + if metadata.shape[0]: + compact_rows = torch.arange( + metadata.shape[0], + dtype=torch.int32, + device=metadata.device, + ) + lut[ + metadata[:, 1].to(torch.int64), + metadata[:, 2].to(torch.int64), + metadata[:, 3].to(torch.int64), + ] = compact_rows + return lut + +__all__ = ["Mxfp8BackwardLayout"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_staging.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_staging.py new file mode 100644 index 000000000..4225aeda7 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_staging.py @@ -0,0 +1,516 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Tensor staging for the explicit Rubin MXFP8 backward invocation.""" + +from __future__ import annotations + +import math + +import torch + +from ..._contracts import ValidatedBackwardRequest +from .._plan import PreparedResources +from .._workspace import padded_mxfp8_scale_columns +from ._adapter import ( + _GATE_UP_INTERLEAVE, + _decode_moe_tensor, + _interleave_gate_up_rows, + _quantize_plain_mxfp8, + _stack_blocked_scales, + _typed_view, + _zero_workspace_prefix, + _zero_workspace_range, +) +from ._backward_compile import ( + Mxfp8BackwardLaunchInputs, + PreparedMxfp8BackwardKernel, +) +from ._backward_layout import Mxfp8BackwardLayout + +_DATA_DTYPE = torch.float8_e4m3fn +_SCALE_DTYPE = torch.float8_e8m0fnu + + +def _typed_prefix_view( + byte_tensor: torch.Tensor, + dtype: torch.dtype, + shape: tuple[int, ...], +) -> torch.Tensor: + """Return a compact typed view of a prefix of a reusable byte region.""" + + nbytes = math.prod(shape) * dtype.itemsize + if nbytes > byte_tensor.numel(): + raise ValueError( + f"byte region has {byte_tensor.numel()} bytes, " + f"cannot provide {nbytes} bytes for shape={shape}, dtype={dtype}" + ) + return _typed_view(byte_tensor.narrow(0, 0, nbytes), dtype, shape) + + +def _k_major(tensor: torch.Tensor) -> torch.Tensor: + """Return the same logical ``(E,K,N)`` tensor with K stride one.""" + + return tensor.permute(0, 2, 1).contiguous().permute(0, 2, 1) + + +def _prepare_backward_weights( + request: ValidatedBackwardRequest, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize W2^T and W1^T along their backward reduction axes.""" + + config = request.config + intermediate = config.intermediate_size + + w2_t = _decode_moe_tensor(request.fc2_weight).transpose(1, 2) + q_w2_t = _quantize_plain_mxfp8(w2_t, axis=1) + fc1_weight = _k_major(q_w2_t.data) + fc1_weight_sf = _stack_blocked_scales( + q_w2_t.scale.permute(0, 2, 1).contiguous() + ) + + w1_t = _decode_moe_tensor(request.fc1_weight).transpose(1, 2) + q_w1_t = _quantize_plain_mxfp8(w1_t, axis=1) + interleaved_w1 = _interleave_gate_up_rows( + q_w1_t.data, + intermediate, + ) + fc2_weight = _k_major(interleaved_w1) + + scale_blocks = intermediate // 32 + gate_sf = q_w1_t.scale[:, :scale_blocks, :] + up_sf = q_w1_t.scale[:, scale_blocks:, :] + interleaved_sf = ( + torch.stack( + ( + gate_sf.view(torch.uint8), + up_sf.view(torch.uint8), + ), + dim=2, + ) + .reshape( + q_w1_t.scale.shape[0], + 2 * scale_blocks, + q_w1_t.scale.shape[2], + ) + .view(_SCALE_DTYPE) + ) + fc2_weight_sf = _stack_blocked_scales( + interleaved_sf.permute(0, 2, 1).contiguous() + ) + return ( + fc1_weight, + fc1_weight_sf, + fc2_weight, + fc2_weight_sf, + ) + + +def _router_order_key( + source_token: torch.Tensor, + source_slot: torch.Tensor, + request: ValidatedBackwardRequest, + prepared: PreparedMxfp8BackwardKernel, +) -> tuple[torch.Tensor, int]: + """Return each route's position in the deterministic router's source run.""" + + elements_per_vector = 4 # The staged top-k index tensor is Int32. + token_comm = prepared.kernel.token_comm + router_ctas = int(token_comm.router_data_cta_count) + router_warps = int(token_comm.router_warps_per_cta) + threads_per_cta = router_warps * 32 + grid_threads = router_ctas * threads_per_cta + tile_span = elements_per_vector * grid_threads + maximum_elements = ( + int(request.config.max_tokens_per_rank) * request.config.top_k + ) + load_rounds = ( + maximum_elements + tile_span - 1 + ) // tile_span + elements_per_thread = load_rounds * elements_per_vector + + flat_index = source_token * request.config.top_k + source_slot + load_round = torch.div( + flat_index, + tile_span, + rounding_mode="floor", + ) + in_tile = flat_index.remainder(tile_span) + grid_thread = torch.div( + in_tile, + elements_per_vector, + rounding_mode="floor", + ) + register_index = ( + load_round * elements_per_vector + + in_tile.remainder(elements_per_vector) + ) + cta = torch.div( + grid_thread, + threads_per_cta, + rounding_mode="floor", + ) + thread_in_cta = grid_thread.remainder(threads_per_cta) + warp = torch.div(thread_in_cta, 32, rounding_mode="floor") + lane = thread_in_cta.remainder(32) + order_key = ( + ( + (cta * router_warps + warp) * elements_per_thread + + register_index + ) + * 32 + + lane + ) + order_span = ( + router_ctas * router_warps * elements_per_thread * 32 + ) + return order_key, order_span + + +def _stage_fc1_preact( + request: ValidatedBackwardRequest, + layout: Mxfp8BackwardLayout, + prepared: PreparedMxfp8BackwardKernel, + fc1_preact: torch.Tensor, +) -> None: + """Lower compact public stash rows into the upstream dGLU pool layout.""" + + config = prepared.config + expected_shape = ( + prepared.pool_token_capacity, + 2 * config.intermediate, + ) + if ( + fc1_preact.dtype is not torch.bfloat16 + or tuple(fc1_preact.shape) != expected_shape + or not fc1_preact.is_contiguous() + ): + raise ValueError( + "backward fc1_preact must be contiguous BF16 with shape " + f"{expected_shape}, got shape={tuple(fc1_preact.shape)}, " + f"dtype={fc1_preact.dtype}" + ) + fc1_preact.zero_() + + metadata = request.route_metadata.to(torch.int64) + compact_rows = layout.preact_row_lut[ + metadata[:, 1], + metadata[:, 2], + metadata[:, 3], + ].to(torch.int64) + if compact_rows.numel() and bool((compact_rows < 0).any().item()): + raise RuntimeError("backward preactivation LUT is incomplete") + + if config.intermediate % _GATE_UP_INTERLEAVE: + raise RuntimeError( + "backward preactivation requires intermediate_size divisible by " + f"{_GATE_UP_INTERLEAVE}" + ) + gate, up = request.fc1_c.split(config.intermediate, dim=1) + pairs = config.intermediate // _GATE_UP_INTERLEAVE + interleaved_preact = torch.stack( + ( + gate.reshape(-1, pairs, _GATE_UP_INTERLEAVE), + up.reshape(-1, pairs, _GATE_UP_INTERLEAVE), + ), + dim=2, + ).reshape(-1, 2 * config.intermediate) + + physical_offset = 0 + for expert in range(config.num_experts): + positions = torch.nonzero( + metadata[:, 0] == expert, + as_tuple=False, + ).flatten() + count = int(positions.numel()) + if count: + # Receiver pools concatenate source ranks in a destination-relative + # ring: local rank first, then increasing ranks with wraparound. + # The public stash is source-rank sorted, so restore pool order. + source_rank = metadata.index_select(0, positions)[:, 1] + source_token = metadata.index_select(0, positions)[:, 2] + source_slot = metadata.index_select(0, positions)[:, 3] + source_order, source_order_span = _router_order_key( + source_token, + source_slot, + request, + prepared, + ) + ring_position = ( + source_rank - request.config.ep_rank + ) % request.config.ep_size + route_key = ring_position * source_order_span + source_order + positions = positions.index_select( + 0, + torch.argsort(route_key, stable=True), + ) + if physical_offset + count > prepared.pool_token_capacity: + raise RuntimeError( + "backward preactivation rows exceed Rubin pool capacity" + ) + destination = fc1_preact.narrow(0, physical_offset, count) + destination.copy_( + interleaved_preact.index_select( + 0, + compact_rows.index_select(0, positions), + ) + ) + physical_offset += ( + count + config.token_padding_block - 1 + ) // config.token_padding_block * config.token_padding_block + + if physical_offset > prepared.pool_token_capacity: + raise RuntimeError( + "backward preactivation rows exceed Rubin pool capacity" + ) + + +def stage_backward( + request: ValidatedBackwardRequest, + layout: Mxfp8BackwardLayout, + prepared: PreparedMxfp8BackwardKernel, + resources: PreparedResources, +) -> Mxfp8BackwardLaunchInputs: + config = prepared.config + capacity = config.max_tokens_per_rank + token_count = request.token_count + hidden = config.hidden + top_k = config.top_k + + symmetric = resources.workspace.symmetric + local = resources.workspace.local + grad_out = _typed_view( + symmetric["activation_data"], + _DATA_DTYPE, + (capacity, hidden), + ) + grad_out_sf = _typed_view( + symmetric["activation_scale"], + _SCALE_DTYPE, + (capacity, padded_mxfp8_scale_columns(hidden)), + ) + topk_weights = _typed_view( + symmetric["topk_weights"], + torch.float32, + (capacity, top_k), + ) + topk_idx = _typed_view( + local["topk_idx"], + torch.int32, + (capacity, top_k), + ) + output_activation = _typed_view( + symmetric["output_data"], + torch.bfloat16, + (capacity, hidden), + ) + overflow_flag = _typed_view( + local["overflow_flag"], + torch.int32, + (1,), + ) + fc1_preact_shape = tuple( + int(extent) for extent in prepared.kernel.get_fc1_preact_shape() + ) + fc1_preact = _typed_view( + local["backward_fc1_preact"], + torch.bfloat16, + fc1_preact_shape, + ) + local_workspace = local["kernel_local_workspace"] + shared_workspace = symmetric["kernel_shared_workspace"] + + _zero_workspace_prefix( + local_workspace, + prepared.local_workspace_zero_bytes, + name="MXFP8 backward local workspace", + ) + _zero_workspace_prefix( + shared_workspace, + prepared.shared_workspace_zero_bytes, + name="MXFP8 backward shared workspace", + ) + _stage_fc1_preact( + request, + layout, + prepared, + fc1_preact, + ) + _zero_workspace_range( + shared_workspace, + prepared.pre_reduced_activation_offset, + token_count * prepared.pre_reduced_activation_bytes_per_token, + name="MXFP8 backward pre-reduced activation workspace", + ) + quantized_combine = config.combine_format != "bf16" + if quantized_combine: + if ( + prepared.pre_reduced_activation_sf_offset is None + or prepared.pre_reduced_activation_sf_bytes_per_token <= 0 + ): + raise RuntimeError( + "MXFP8 backward quantized combine requires scale workspace" + ) + _zero_workspace_range( + shared_workspace, + prepared.pre_reduced_activation_sf_offset, + token_count + * prepared.pre_reduced_activation_sf_bytes_per_token, + name="MXFP8 backward pre-reduced activation scale workspace", + ) + elif ( + prepared.pre_reduced_activation_sf_offset is not None + or prepared.pre_reduced_activation_sf_bytes_per_token != 0 + ): + raise RuntimeError( + "MXFP8 backward BF16 combine must not expose scale workspace" + ) + quantized_grad = _quantize_plain_mxfp8( + request.grad_output, + axis=1, + ) + grad_out.zero_() + grad_out[:token_count].copy_(quantized_grad.data) + grad_out_sf.zero_() + grad_out_sf[ + :token_count, + : quantized_grad.scale.shape[1], + ].copy_(quantized_grad.scale) + topk_idx.fill_(-1) + topk_idx[:token_count].copy_( + request.topk_idx.to(torch.int32) + ) + topk_weights.zero_() + topk_weights[:token_count].copy_( + request.topk_weights.float() + ) + output_activation.zero_() + overflow_flag.zero_() + + aux_shapes = { + name: tuple(int(extent) for extent in shape) + for name, shape in prepared.kernel.get_aux_output_shapes().items() + } + dprob = _typed_view( + symmetric["backward_dprob"], + torch.float32, + aux_shapes["dprob"], + ) + dprob.zero_() + operands_mode = request.config.backward_wgrad_mode == "operands" + expected_flags = ( + prepared.dfc2_recompute, + prepared.dfc2_col_output, + prepared.enable_grad_y2_col_quant, + ) + if operands_mode: + if expected_flags != (True, True, True): + raise RuntimeError( + "wgrad operands require every backward auxiliary output" + ) + # These allocations are intentionally not execution-plan workspace: + # the returned operand bundle must remain valid after later calls. + fc1_recompute = torch.zeros( + aux_shapes["fc1_recompute"], + dtype=_DATA_DTYPE, + device=request.device, + ) + fc1_recompute_sf = torch.full( + aux_shapes["fc1_recompute_sf"], + 127, + dtype=torch.uint8, + device=request.device, + ).view(_SCALE_DTYPE) + fc1_col_output = torch.zeros( + aux_shapes["fc1_col_output"], + dtype=_DATA_DTYPE, + device=request.device, + ) + fc1_col_output_sf = torch.full( + aux_shapes["fc1_col_output_sf"], + 127, + dtype=torch.uint8, + device=request.device, + ).view(_SCALE_DTYPE) + grad_y2 = torch.zeros( + aux_shapes["grad_y2"], + dtype=_DATA_DTYPE, + device=request.device, + ) + grad_y2_sf = torch.full( + aux_shapes["grad_y2_sf"], + 127, + dtype=torch.uint8, + device=request.device, + ) + else: + if expected_flags != (False, False, False): + raise RuntimeError( + "default backward must disable wgrad auxiliary outputs" + ) + # Preserve the existing mode-none allocation/performance behavior: + # disabled fixed-ABI arguments alias reusable plan scratch. + fc1_recompute = _typed_prefix_view( + local["backward_aux_data"], + _DATA_DTYPE, + aux_shapes["fc1_recompute"], + ) + fc1_col_output = _typed_prefix_view( + local["backward_aux_data"], + _DATA_DTYPE, + aux_shapes["fc1_col_output"], + ) + fc1_recompute_sf = _typed_prefix_view( + local["backward_aux_scale"], + _SCALE_DTYPE, + aux_shapes["fc1_recompute_sf"], + ) + fc1_col_output_sf = _typed_prefix_view( + local["backward_aux_scale"], + _SCALE_DTYPE, + aux_shapes["fc1_col_output_sf"], + ) + grad_y2 = _typed_prefix_view( + local["backward_aux_data"], + _DATA_DTYPE, + aux_shapes["grad_y2"], + ) + grad_y2_sf = _typed_prefix_view( + local["backward_aux_scale"], + torch.uint8, + aux_shapes["grad_y2_sf"], + ) + fc1_weight, fc1_weight_sf, fc2_weight, fc2_weight_sf = ( + _prepare_backward_weights(request) + ) + return Mxfp8BackwardLaunchInputs( + grad_out=grad_out, + grad_out_sf=grad_out_sf, + topk_idx=topk_idx, + topk_weights=topk_weights, + fc1_weight=fc1_weight, + fc1_weight_sf=fc1_weight_sf, + fc2_weight=fc2_weight, + fc2_weight_sf=fc2_weight_sf, + beta=torch.ones( + (config.num_experts,), + dtype=torch.float32, + device=request.device, + ), + fc1_preact=fc1_preact, + output_activation=output_activation, + overflow_flag=overflow_flag, + dprob=dprob, + fc1_recompute=fc1_recompute, + fc1_recompute_sf=fc1_recompute_sf, + fc1_col_output=fc1_col_output, + fc1_col_output_sf=fc1_col_output_sf, + grad_y2=grad_y2, + grad_y2_sf=grad_y2_sf, + local_workspace=local_workspace, + shared_workspace=shared_workspace, + token_count=token_count, + ) + + +__all__ = ["stage_backward"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_wgrad_export.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_wgrad_export.py new file mode 100644 index 000000000..6146cf296 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_wgrad_export.py @@ -0,0 +1,226 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Materialize caller-owned grouped-wgrad operands from backward auxiliaries.""" + +from __future__ import annotations + +import torch + +from ..._contracts import ValidatedBackwardRequest +from ..._types import MoeEpWgradOperands +from ._backward_launch import Mxfp8DgluResult +from ._wgrad_layout import ( + assemble_dfc2_atom_scales, + assemble_discrete_col_requant_scales, + deinterleave_gate_up_columns, + pool_data_as_wgrad_a, + pool_data_as_wgrad_b, +) + + +def export_wgrad_operands( + request: ValidatedBackwardRequest, + dglu: Mxfp8DgluResult, +) -> MoeEpWgradOperands: + """Convert physical Rubin pools to the grouped-wgrad Tensor2D ABI.""" + + if request.config.backward_wgrad_mode != "operands": + raise ValueError("wgrad operands can only be exported in operands mode") + stash = request.wgrad_forward_stash + if stash is None: + raise ValueError("wgrad operand export requires a forward stash") + + padded_ends = tuple( + int(value) for value in stash.expert_offsets.detach().cpu().tolist() + ) + valid_counts = tuple( + int(value) + for value in stash.valid_route_counts.detach().cpu().tolist() + ) + padded_routes = padded_ends[-1] if padded_ends else 0 + config = request.config + + _validate_aux_shapes(dglu, padded_routes, config) + + # dC is emitted in the kernel's 32-column gate/up strip order. Upstream + # exports route-weighted h and unweighted dY, whose product is the same + # dW2 contract previously represented as unweighted h and weighted dY. + # Every conversion allocates fresh storage, so no returned tensor aliases + # the reusable execution plan. + dc_pool = deinterleave_gate_up_columns( + dglu.fc1_col_output[:padded_routes], + config.intermediate_size, + ) + fc1_b = pool_data_as_wgrad_b(dc_pool, padded_routes) + fc1_sfb = assemble_dfc2_atom_scales( + dglu.fc1_col_output_sf, + valid_counts, + padded_ends, + 2 * config.intermediate_size, + config.sf_padding_size, + deinterleave_gate_up=config.intermediate_size, + ) + + fc2_a = pool_data_as_wgrad_a( + dglu.fc1_recompute, + padded_routes, + ) + fc2_sfa = assemble_dfc2_atom_scales( + dglu.fc1_recompute_sf, + valid_counts, + padded_ends, + config.intermediate_size, + config.sf_padding_size, + ) + fc2_b = pool_data_as_wgrad_b( + dglu.grad_y2, + padded_routes, + ) + fc2_sfb = assemble_discrete_col_requant_scales( + dglu.grad_y2_sf, + valid_counts, + padded_ends, + config.hidden_size, + config.sf_padding_size, + ) + + operands = MoeEpWgradOperands( + fc1_a=stash.fc1_a, + fc1_sfa=stash.fc1_sfa, + fc1_b=fc1_b, + fc1_sfb=fc1_sfb, + fc2_a=fc2_a, + fc2_sfa=fc2_sfa, + fc2_b=fc2_b, + fc2_sfb=fc2_sfb, + expert_offsets=stash.expert_offsets, + valid_route_counts=stash.valid_route_counts, + route_metadata=stash.route_metadata, + ) + _validate_grouped_wgrad_abi(operands, config, padded_routes) + return operands + + +def _validate_aux_shapes( + dglu: Mxfp8DgluResult, + padded_routes: int, + config, +) -> None: + expected_columns = ( + ("fc1_recompute", dglu.fc1_recompute, config.intermediate_size), + ( + "fc1_col_output", + dglu.fc1_col_output, + 2 * config.intermediate_size, + ), + ( + "grad_y2", + dglu.grad_y2, + config.hidden_size, + ), + ) + for name, tensor, columns in expected_columns: + if tensor.ndim != 2 or tensor.shape[1] != columns: + raise RuntimeError( + f"{name} must have shape (pool_capacity, {columns}), " + f"got {tuple(tensor.shape)}" + ) + if tensor.shape[0] < padded_routes or not tensor.is_contiguous(): + raise RuntimeError( + f"{name} does not contain a contiguous {padded_routes}-row " + "pool prefix" + ) + if tensor.dtype is not torch.float8_e4m3fn: + raise TypeError(f"{name} must have dtype torch.float8_e4m3fn") + _require_alignment(name, tensor, 16) + + for name, tensor in ( + ("fc1_recompute_sf", dglu.fc1_recompute_sf), + ("fc1_col_output_sf", dglu.fc1_col_output_sf), + ): + if tensor.ndim != 2 or tensor.dtype is not torch.float8_e8m0fnu: + raise TypeError(f"{name} must be a rank-2 E8M0 tensor") + _require_alignment(name, tensor, 16) + if dglu.grad_y2_sf.ndim != 1 or dglu.grad_y2_sf.dtype is not torch.uint8: + raise TypeError("grad_y2_sf must be a rank-1 uint8 tensor") + _require_alignment("grad_y2_sf", dglu.grad_y2_sf, 16) + + +def _validate_grouped_wgrad_abi( + operands: MoeEpWgradOperands, + config, + padded_routes: int, +) -> None: + rounded_hidden = _round_up(config.hidden_size, 128) + rounded_intermediate = _round_up(config.intermediate_size, 128) + rounded_gate_up = _round_up(2 * config.intermediate_size, 128) + sf_columns = _round_up(padded_routes // 32, 4) + expected = ( + ("fc1_a", operands.fc1_a, (config.hidden_size, padded_routes)), + ("fc1_sfa", operands.fc1_sfa, (rounded_hidden, sf_columns)), + ( + "fc1_b", + operands.fc1_b, + (padded_routes, 2 * config.intermediate_size), + ), + ("fc1_sfb", operands.fc1_sfb, (rounded_gate_up, sf_columns)), + ( + "fc2_a", + operands.fc2_a, + (config.intermediate_size, padded_routes), + ), + ( + "fc2_sfa", + operands.fc2_sfa, + (rounded_intermediate, sf_columns), + ), + ("fc2_b", operands.fc2_b, (padded_routes, config.hidden_size)), + ("fc2_sfb", operands.fc2_sfb, (rounded_hidden, sf_columns)), + ) + for name, tensor, shape in expected: + if tuple(tensor.shape) != shape: + raise RuntimeError( + f"grouped-wgrad {name} shape must be {shape}, " + f"got {tuple(tensor.shape)}" + ) + _require_alignment(name, tensor, 16) + + for name in ("fc1_a", "fc2_a"): + tensor = getattr(operands, name) + expected_stride = (padded_routes, 1) + if padded_routes and tensor.stride() != expected_stride: + raise RuntimeError( + f"grouped-wgrad {name} strides must be " + f"{expected_stride}, got {tensor.stride()}" + ) + for name in ("fc1_b", "fc2_b"): + tensor = getattr(operands, name) + expected_stride = (1, padded_routes) + if padded_routes and tensor.stride() != expected_stride: + raise RuntimeError( + f"grouped-wgrad {name} strides must be " + f"{expected_stride}, got {tensor.stride()}" + ) + for name in ("fc1_sfa", "fc1_sfb", "fc2_sfa", "fc2_sfb"): + if not getattr(operands, name).is_contiguous(): + raise RuntimeError(f"grouped-wgrad {name} must be contiguous") + _require_alignment("expert_offsets", operands.expert_offsets, 4) + + +def _require_alignment( + name: str, + tensor: torch.Tensor, + alignment: int, +) -> None: + if tensor.data_ptr() % alignment: + raise RuntimeError( + f"{name} address is not {alignment}-byte aligned" + ) + + +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +__all__ = ["export_wgrad_operands"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py new file mode 100644 index 000000000..9cbd18fc3 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py @@ -0,0 +1,373 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""In-process JIT cache for the vendored Rubin SM107 MXFP8 kernel.""" + +from __future__ import annotations + +import os +import threading +from dataclasses import dataclass +from typing import Any + +import torch + +from ..._contracts import ForwardConfig +from .._plan import PreparedResources +from .._workspace import WorkspaceRequirements +from ._adapter import Mxfp8LaunchInputs +from ._config import Mxfp8KernelConfig +from ._fingerprint import build_kernel_fingerprint +from ._launch import build_runtime_kwargs, layout_signature + + +@dataclass(frozen=True) +class PreparedMxfp8Kernel: + config: Mxfp8KernelConfig + device: torch.device + architecture: tuple[int, int] + kernel: Any + launch_cluster_count: int + workspace_requirements: WorkspaceRequirements + pool_token_capacity: int + col_quant_data_rows: int + col_quant_sf_elements: int + token_src_metadata_offset: int + token_src_metadata_bytes: int + pre_reduced_activation_offset: int | None + pre_reduced_activation_bytes_per_token: int + pre_reduced_activation_sf_offset: int | None + pre_reduced_activation_sf_bytes_per_token: int + local_workspace_zero_bytes: int + shared_workspace_zero_bytes: int + + +@dataclass(frozen=True) +class CompiledMxfp8Kernel: + key: tuple + callable: Any + fingerprint: dict[str, Any] + + +_COMPILE_LOCK = threading.RLock() +_COMPILE_CACHE: dict[tuple, CompiledMxfp8Kernel] = {} +_TOKEN_SRC_METADATA_REGION = "nvlink.token_comm.token_src_metadata" +_PRE_REDUCED_ACTIVATION_REGION = ( + "nvlink.token_comm.pre_reduced_activation" +) +_PRE_REDUCED_ACTIVATION_SF_REGION = ( + "nvlink.token_comm.pre_reduced_activation_sf" +) + + +def _compile_kernel(kernel: Any, compile_kwargs: dict[str, Any]) -> Any: + """Import CuTeDSL only on a cache miss and compile one callable.""" + + import cutlass.cute as cute + + return cute.compile(kernel, **compile_kwargs) + + +def _pre_reduced_workspace_metadata( + device_workspace: Any, + config: Mxfp8KernelConfig, + shared_bytes: int, +) -> tuple[int | None, int]: + """Describe the standalone combine plane, if this kernel has one.""" + if config.fc2_in_kernel_topk_reduce: + return None, 0 + + region = device_workspace.region(_PRE_REDUCED_ACTIVATION_REGION) + if region.buffer_space != "shared": + raise RuntimeError( + "Rubin pre_reduced_activation must reside in shared workspace" + ) + offset = int( + device_workspace.offset(_PRE_REDUCED_ACTIVATION_REGION) + ) + nbytes = int( + device_workspace.nbytes(_PRE_REDUCED_ACTIVATION_REGION) + ) + wire_bits_per_element = { + "bf16": 16, + "32e4m3xe8m0": 8, + } + try: + element_bits = wire_bits_per_element[config.combine_format] + except KeyError as exc: + raise ValueError( + f"unsupported combine wire format {config.combine_format!r}" + ) from exc + wire_bits_per_token = config.top_k * config.hidden * element_bits + if wire_bits_per_token % 8: + raise RuntimeError("combine wire row is not byte aligned") + bytes_per_token = wire_bits_per_token // 8 + expected_bytes = config.max_tokens_per_rank * bytes_per_token + if nbytes != expected_bytes: + raise RuntimeError( + "Rubin pre_reduced_activation size does not match " + f"combine_format={config.combine_format!r}: {nbytes} bytes, " + f"expected {expected_bytes}" + ) + if offset + nbytes > shared_bytes: + raise RuntimeError( + "Rubin pre_reduced_activation region exceeds shared workspace" + ) + return offset, bytes_per_token + + +def _pre_reduced_sf_workspace_metadata( + device_workspace: Any, + config: Mxfp8KernelConfig, + shared_bytes: int, +) -> tuple[int | None, int]: + """Describe the standalone quantized-combine scale plane, if present.""" + if config.fc2_in_kernel_topk_reduce or config.combine_format == "bf16": + return None, 0 + + region = device_workspace.region(_PRE_REDUCED_ACTIVATION_SF_REGION) + if region.buffer_space != "shared": + raise RuntimeError( + "Rubin pre_reduced_activation_sf must reside in shared workspace" + ) + offset = int(device_workspace.offset(_PRE_REDUCED_ACTIVATION_SF_REGION)) + nbytes = int(device_workspace.nbytes(_PRE_REDUCED_ACTIVATION_SF_REGION)) + if nbytes % config.max_tokens_per_rank: + raise RuntimeError( + "Rubin pre_reduced_activation_sf size is not token aligned" + ) + if offset + nbytes > shared_bytes: + raise RuntimeError( + "Rubin pre_reduced_activation_sf region exceeds shared workspace" + ) + return offset, nbytes // config.max_tokens_per_rank + + +def prepare_kernel( + forward_config: ForwardConfig, + config: Mxfp8KernelConfig, + device: torch.device, +) -> PreparedMxfp8Kernel: + """Instantiate the kernel and derive exact allocation requirements.""" + + torch.cuda.set_device(device) + architecture = torch.cuda.get_device_capability(device) + if architecture != (10, 7): + raise RuntimeError( + "Rubin MXFP8 kernel preparation requires compute capability " + f"(10, 7), got {architecture}" + ) + configured_architecture = os.environ.get("CUTE_DSL_ARCH") + if configured_architecture is None: + os.environ["CUTE_DSL_ARCH"] = "sm_107a" + elif configured_architecture not in ("sm_107", "sm_107a"): + raise RuntimeError( + "CUTE_DSL_ARCH must target SM107 for the Rubin MXFP8 backend, " + f"got {configured_architecture!r}" + ) + + import cutlass + import cutlass.utils as utils + + from ..cutedsl_src.kernel_src.rubin.training.mega.fwd_glu import ( + Sm107MegaMoEMxfp8GluKernel, + ) + from ..cutedsl_src.quant_def import CombineFormat + + launch_cluster_count = int( + utils.HardwareInfo().get_max_active_clusters(config.cluster_size) + ) + if launch_cluster_count <= 0: + raise RuntimeError( + "hardware occupancy query returned no launchable Rubin clusters" + ) + group_hint = ( + launch_cluster_count + if config.group_hint is None + else config.group_hint + ) + kernel_kwargs = dict( + mma_tiler_mnk=config.mma_tiler_mnk, + cluster_shape_mnk=config.cluster_shape_mnk, + use_2cta_instrs=config.use_2cta_instrs, + group_hint=group_hint, + token_padding_block=config.token_padding_block, + sf_padding_block=config.sf_padding_block, + load_balance_mode=config.load_balance_mode, + static_expert_shape=( + config.num_experts, + config.fc1_out, + config.hidden, + ), + force_static_sched=config.force_static_sched, + clc_bundle_size=config.clc_bundle_size, + num_sched_stages=config.num_sched_stages, + ab_dtype=cutlass.Float8E4M3FN, + sf_vec_size=config.sf_vec_size, + world_size=config.world_size, + # Runtime rank is carried by SymmetricBufferHost. Keeping this + # descriptor rank-independent allows every EP rank to compile the + # same Rubin kernel. + local_rank=0, + num_topk=config.top_k, + max_tokens_per_rank=config.max_tokens_per_rank, + max_recv_size_per_rank=config.max_recv_size_per_rank, + hidden=config.hidden, + launch_cluster_count=launch_cluster_count, + drop_on_overflow=config.drop_on_overflow, + fc2_in_kernel_topk_reduce=config.fc2_in_kernel_topk_reduce, + token_back_mode=config.token_back_mode, + epi_flag_batch=config.epi_flag_batch, + flag_batch=config.flag_batch, + gate_up_clamp=config.gate_up_clamp, + generate_c=config.generate_c, + combine_format=CombineFormat.parse(config.combine_format), + act_func=config.act_func, + fc2_use_bulk=config.fc2_use_bulk, + fc2_tma_stages=config.fc2_tma_stages, + enable_col_quant=config.enable_col_quant, + col_quant_num_ctas=config.col_quant_num_ctas, + ) + kernel = Sm107MegaMoEMxfp8GluKernel.from_kwargs(**kernel_kwargs) + local_bytes, shared_bytes = kernel.get_workspace_sizes() + local_zero_bytes, shared_zero_bytes = ( + kernel.require_zero_workspace_leading_bytes + ) + for name, zero_bytes, total_bytes in ( + ("local", local_zero_bytes, local_bytes), + ("shared", shared_zero_bytes, shared_bytes), + ): + if zero_bytes < 0 or zero_bytes > total_bytes: + raise RuntimeError( + f"Rubin kernel {name} zero prefix {zero_bytes} exceeds " + f"workspace size {total_bytes}" + ) + device_workspace = kernel._mega_device_workspace + metadata_region = device_workspace.region(_TOKEN_SRC_METADATA_REGION) + if metadata_region.buffer_space != "shared": + raise RuntimeError( + "Rubin token_src_metadata must reside in shared workspace" + ) + token_src_metadata_offset = int( + device_workspace.offset(_TOKEN_SRC_METADATA_REGION) + ) + token_src_metadata_bytes = int( + device_workspace.nbytes(_TOKEN_SRC_METADATA_REGION) + ) + if token_src_metadata_offset + token_src_metadata_bytes > shared_bytes: + raise RuntimeError( + "Rubin token_src_metadata region exceeds shared workspace" + ) + pool_token_capacity = int(kernel.pool_token_capacity) + if token_src_metadata_bytes != pool_token_capacity * 8: + raise RuntimeError( + "Rubin token_src_metadata must contain one Int64 per pool token" + ) + col_quant_data_rows = ( + pool_token_capacity if config.enable_col_quant else 0 + ) + col_quant_sf_elements = ( + int(kernel.token_comm.worst_case_sf_token_count) + * (config.hidden // config.sf_vec_size) + if config.enable_col_quant + else 0 + ) + requirements = WorkspaceRequirements.for_mxfp8( + forward_config, + kernel_local_workspace_bytes=local_bytes, + kernel_shared_workspace_bytes=shared_bytes, + col_quant_data_bytes=col_quant_data_rows * config.hidden, + col_quant_sf_bytes=col_quant_sf_elements, + ) + ( + pre_reduced_activation_offset, + pre_reduced_activation_bytes_per_token, + ) = _pre_reduced_workspace_metadata( + device_workspace, + config, + shared_bytes, + ) + ( + pre_reduced_activation_sf_offset, + pre_reduced_activation_sf_bytes_per_token, + ) = _pre_reduced_sf_workspace_metadata( + device_workspace, + config, + shared_bytes, + ) + return PreparedMxfp8Kernel( + config=config, + device=torch.device(device), + architecture=architecture, + kernel=kernel, + launch_cluster_count=launch_cluster_count, + workspace_requirements=requirements, + pool_token_capacity=pool_token_capacity, + col_quant_data_rows=col_quant_data_rows, + col_quant_sf_elements=col_quant_sf_elements, + token_src_metadata_offset=token_src_metadata_offset, + token_src_metadata_bytes=token_src_metadata_bytes, + pre_reduced_activation_offset=pre_reduced_activation_offset, + pre_reduced_activation_bytes_per_token=( + pre_reduced_activation_bytes_per_token + ), + pre_reduced_activation_sf_offset=pre_reduced_activation_sf_offset, + pre_reduced_activation_sf_bytes_per_token=( + pre_reduced_activation_sf_bytes_per_token + ), + local_workspace_zero_bytes=int(local_zero_bytes), + shared_workspace_zero_bytes=int(shared_zero_bytes), + ) + + +def compile_or_get( + prepared: PreparedMxfp8Kernel, + inputs: Mxfp8LaunchInputs, + resources: PreparedResources, +) -> CompiledMxfp8Kernel: + signature = layout_signature(inputs) + key = ( + *prepared.config.compile_key( + prepared.device, + prepared.architecture, + prepared.launch_cluster_count, + signature, + ), + ) + with _COMPILE_LOCK: + cached = _COMPILE_CACHE.get(key) + if cached is not None: + return cached + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("MXFP8 kernel must be compiled before CUDA graph capture") + + compile_kwargs = build_runtime_kwargs( + inputs, + resources, + ) + compiled = CompiledMxfp8Kernel( + key=key, + callable=_compile_kernel(prepared.kernel, compile_kwargs), + fingerprint=build_kernel_fingerprint( + prepared, + signature, + ), + ) + _COMPILE_CACHE[key] = compiled + return compiled + + +def clear_compile_cache() -> None: + """Drop process-local compiled callable references.""" + + with _COMPILE_LOCK: + _COMPILE_CACHE.clear() + + +__all__ = [ + "CompiledMxfp8Kernel", + "PreparedMxfp8Kernel", + "clear_compile_cache", + "compile_or_get", + "prepare_kernel", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py new file mode 100644 index 000000000..647ef07f0 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py @@ -0,0 +1,201 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Import-light static configuration for the Rubin SM107 MXFP8 kernel.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from ..._contracts import ForwardConfig +from ._formats import combine_wire_format + + +@dataclass(frozen=True) +class Mxfp8KernelConfig: + """Code-generation constants for one dense EP subgroup kernel.""" + + num_experts: int + world_size: int + local_rank: int + hidden: int + intermediate: int + top_k: int + max_tokens_per_rank: int + apply_topk_in_fc1: bool + gate_up_clamp: float | None + generate_c: bool + max_recv_size_per_rank: int | None = None + drop_on_overflow: bool = True + enable_col_quant: bool = False + col_quant_num_ctas: int = 2368 + mma_tiler_mnk: tuple[int, int, int] = (256, 256, 128) + cluster_shape_mnk: tuple[int, int, int] = (2, 1, 1) + use_2cta_instrs: bool = True + load_balance_mode: str = "static" + force_static_sched: bool = True + clc_bundle_size: int | None = None + num_sched_stages: int | None = None + token_padding_block: int = 128 + sf_padding_block: int = 128 + sf_vec_size: int = 32 + group_hint: int | None = None + token_back_mode: str = "epi_warps" + epi_flag_batch: tuple[int, int] = (1, 1) + flag_batch: int = 1 + fc2_in_kernel_topk_reduce: bool = False + act_func: str = "swiglu" + combine_format: str = "bf16" + fc2_use_bulk: bool = False + fc2_tma_stages: int | None = None + + def __post_init__(self) -> None: + if ( + self.max_recv_size_per_rank is not None + and self.max_recv_size_per_rank <= 0 + ): + raise ValueError("max_recv_size_per_rank must be positive") + if self.col_quant_num_ctas <= 0: + raise ValueError("col_quant_num_ctas must be positive") + + @classmethod + def from_forward_config(cls, config: ForwardConfig) -> "Mxfp8KernelConfig": + if config.ep_size < 1 or config.ep_size > 16: + raise NotImplementedError( + "MXFP8 execution supports EP subgroup sizes from 1 through 16" + ) + if config.ep_rank < 0 or config.ep_rank >= config.ep_size: + raise ValueError( + f"ep_rank {config.ep_rank} is outside EP size {config.ep_size}" + ) + if config.max_tokens_per_rank is None: + raise ValueError("MXFP8 execution requires max_tokens_per_rank") + max_recv_size_per_rank = ( + config.ep_size * config.max_tokens_per_rank * config.top_k + ) + if max_recv_size_per_rank <= 0: + raise ValueError("max_recv_size_per_rank must be positive") + return cls( + num_experts=config.experts_per_rank, + world_size=config.ep_size, + local_rank=config.ep_rank, + hidden=config.hidden_size, + intermediate=config.intermediate_size, + top_k=config.top_k, + max_tokens_per_rank=config.max_tokens_per_rank, + apply_topk_in_fc1=config.apply_topk_in_fc1, + gate_up_clamp=config.gate_up_clamp, + generate_c=config.generate_c, + max_recv_size_per_rank=max_recv_size_per_rank, + combine_format=combine_wire_format(config.combine_format), + enable_col_quant=( + config.backward_wgrad_mode == "operands" + ), + token_padding_block=( + config.token_padding_size + if config.backward_wgrad_mode == "operands" + else 128 if config.generate_c else config.token_padding_size + ), + sf_padding_block=config.sf_padding_size, + group_hint=config.tuning.group_hint, + token_back_mode=config.tuning.token_back_mode, + epi_flag_batch=config.tuning.epi_flag_batch, + flag_batch=config.tuning.token_in_flag_batch, + fc2_in_kernel_topk_reduce=( + config.tuning.reduce_topk_in_kernel + ), + ) + + @property + def fc1_out(self) -> int: + return 2 * self.intermediate + + @property + def cluster_size(self) -> int: + return self.cluster_shape_mnk[0] * self.cluster_shape_mnk[1] + + def tuning_signature( + self, + launch_cluster_count: int, + ) -> tuple[str, tuple[int, int], int, int, bool]: + """Return the effective rank-independent transport/scheduler knobs.""" + + group_hint = ( + launch_cluster_count + if self.group_hint is None + else self.group_hint + ) + return ( + self.token_back_mode, + self.epi_flag_batch, + self.flag_batch, + group_hint, + self.fc2_in_kernel_topk_reduce, + ) + + def effective_config(self, launch_cluster_count: int) -> dict[str, object]: + """Return the complete JSON-safe compile-time configuration.""" + + effective_group_hint = ( + launch_cluster_count + if self.group_hint is None + else self.group_hint + ) + return { + "num_experts_per_rank": self.num_experts, + "world_size": self.world_size, + "hidden": self.hidden, + "intermediate": self.intermediate, + "top_k": self.top_k, + "max_tokens_per_rank": self.max_tokens_per_rank, + "max_recv_size_per_rank": self.max_recv_size_per_rank, + "drop_on_overflow": self.drop_on_overflow, + "apply_topk_in_fc1": self.apply_topk_in_fc1, + "gate_up_clamp": self.gate_up_clamp, + "generate_c": self.generate_c, + "enable_col_quant": self.enable_col_quant, + "col_quant_num_ctas": self.col_quant_num_ctas, + "combine_format": self.combine_format, + "mma_tiler_mnk": list(self.mma_tiler_mnk), + "cluster_shape_mnk": list(self.cluster_shape_mnk), + "use_2cta_instrs": self.use_2cta_instrs, + "load_balance_mode": self.load_balance_mode, + "force_static_sched": self.force_static_sched, + "clc_bundle_size": self.clc_bundle_size, + "num_sched_stages": self.num_sched_stages, + "token_padding_block": self.token_padding_block, + "sf_padding_block": self.sf_padding_block, + "sf_vec_size": self.sf_vec_size, + "effective_group_hint": effective_group_hint, + "token_back_mode": self.token_back_mode, + "epi_flag_batch": list(self.epi_flag_batch), + "token_in_flag_batch": self.flag_batch, + "fc2_in_kernel_topk_reduce": self.fc2_in_kernel_topk_reduce, + "act_func": self.act_func, + "fc2_use_bulk": self.fc2_use_bulk, + "fc2_tma_stages": self.fc2_tma_stages, + "launch_cluster_count": launch_cluster_count, + } + + def compile_key( + self, + device: torch.device, + architecture: tuple[int, int], + launch_cluster_count: int, + layout_signature: tuple, + ) -> tuple: + """Return a pointer/stream-independent in-process JIT cache key.""" + + canonical_device = torch.device(device) + return ( + self, + canonical_device.index, + architecture, + launch_cluster_count, + layout_signature, + ) + + +__all__ = ["Mxfp8KernelConfig"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_fingerprint.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_fingerprint.py new file mode 100644 index 000000000..d7600d7b4 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_fingerprint.py @@ -0,0 +1,150 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Stable, machine-readable fingerprints for compiled MXFP8 kernels.""" + +from __future__ import annotations + +import hashlib +import importlib.metadata +import json +import os +from pathlib import Path +from typing import Any + + +FINGERPRINT_SCHEMA_VERSION = 1 +_KERNEL_IDENTITY_FIELDS = ( + "kernel_name", + "kernel_source", + "effective_config", + "cutlass_version", + "source_tree_sha256", + "source_git_revision", + "launch_geometry", +) + + +def canonical_json_sha256(value: object) -> str: + payload = json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def kernel_identity_sha256(fingerprint: dict[str, Any]) -> str: + """Hash fields that must match across AOT and in-process JIT paths.""" + + return canonical_json_sha256( + { + field: fingerprint.get(field) + for field in _KERNEL_IDENTITY_FIELDS + } + ) + + +def source_tree_sha256(root: Path) -> str: + """Hash Python source content and relative paths in deterministic order.""" + + resolved = root.expanduser().resolve() + if not resolved.is_dir(): + raise RuntimeError(f"kernel source tree does not exist: {resolved}") + digest = hashlib.sha256() + sources = sorted( + path for path in resolved.rglob("*.py") if path.is_file() + ) + if not sources: + raise RuntimeError(f"kernel source tree contains no Python files: {resolved}") + for path in sources: + relative = path.relative_to(resolved).as_posix().encode("utf-8") + digest.update(len(relative).to_bytes(8, "little")) + digest.update(relative) + payload = path.read_bytes() + digest.update(len(payload).to_bytes(8, "little")) + digest.update(payload) + return digest.hexdigest() + + +def _cutlass_version() -> str: + for distribution in ( + "nvidia-cutlass-dsl", + "nvidia-cutlass-dsl-internal", + ): + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + continue + import cutlass + + return str(getattr(cutlass, "__version__", "unknown")) + + +def _json_layout_signature(signature: tuple) -> list[object]: + return [ + None + if entry is None + else { + "shape": list(entry[0]), + "stride": list(entry[1]), + "dtype": str(entry[2]), + } + for entry in signature + ] + + +def build_kernel_fingerprint( + prepared: Any, + layout_signature: tuple, + *, + compiled_binary_sha256: str | None = None, +) -> dict[str, Any]: + """Describe exactly what was compiled and how the main kernel launches.""" + + kernel = prepared.kernel + source_root = Path(__file__).resolve().parents[1] / "cutedsl_src" + effective_config = prepared.config.effective_config( + prepared.launch_cluster_count + ) + launch_geometry = { + "grid": [ + prepared.config.cluster_shape_mnk[0], + prepared.config.cluster_shape_mnk[1], + prepared.launch_cluster_count, + ], + "block": [int(kernel.threads_per_cta), 1, 1], + "cluster": list(prepared.config.cluster_shape_mnk), + "min_blocks_per_mp": int(getattr(kernel, "occupancy", 1)), + "dynamic_shared_memory_bytes": int( + getattr(kernel, "smem_capacity", 0) + ), + } + layout = _json_layout_signature(layout_signature) + fingerprint = { + "schema_version": FINGERPRINT_SCHEMA_VERSION, + "kernel_name": str(kernel.name()), + "kernel_source": "vendored-training-mega", + "effective_config": effective_config, + "cutlass_version": _cutlass_version(), + "source_tree_sha256": source_tree_sha256(source_root), + "source_git_revision": os.environ.get("MOE_EP_SOURCE_GIT_REVISION"), + "launch_geometry": launch_geometry, + "layout_signature_sha256": canonical_json_sha256(layout), + "compiled_binary_sha256": compiled_binary_sha256, + } + fingerprint["kernel_identity_sha256"] = kernel_identity_sha256( + fingerprint + ) + fingerprint["fingerprint_sha256"] = canonical_json_sha256(fingerprint) + return fingerprint + + +__all__ = [ + "FINGERPRINT_SCHEMA_VERSION", + "build_kernel_fingerprint", + "canonical_json_sha256", + "kernel_identity_sha256", + "source_tree_sha256", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_formats.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_formats.py new file mode 100644 index 000000000..c3d4df4f4 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_formats.py @@ -0,0 +1,23 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Public MoE format names to Rubin combine-wire encodings.""" + +from __future__ import annotations + +from ..._types import MoeFormat, parse_format + + +_COMBINE_WIRE_FORMATS = { + MoeFormat.BF16: "bf16", + MoeFormat.MXFP8: "32e4m3xe8m0", +} + + +def combine_wire_format(value: MoeFormat | str) -> str: + """Return the kernel encoding for one public combine format.""" + + return _COMBINE_WIRE_FORMATS[parse_format(value)] + + +__all__ = ["combine_wire_format"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_launch.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_launch.py new file mode 100644 index 000000000..98d9c24e1 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_launch.py @@ -0,0 +1,180 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CuTe tensor conversion and current-stream MXFP8 launch.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from .._plan import PreparedResources +from ._adapter import Mxfp8LaunchInputs + + +def _to_cute( + tensor: torch.Tensor, + assumed_align: int = 16, + *, + dynamic_layout: bool = True, +): + import cutlass.torch as cutlass_torch + + cute_tensor = cutlass_torch.from_dlpack( + tensor, + assumed_align=assumed_align, + enable_tvm_ffi=True, + ) + if not dynamic_layout: + return cute_tensor + return cute_tensor.mark_layout_dynamic( + leading_dim=cutlass_torch.get_leading_dim(tensor) + ) + + +def _to_cute_ptr(tensor: torch.Tensor, assumed_align: int = 128): + """Build the opaque byte-pointer ABI used by Rubin workspaces.""" + + import cutlass + from cutlass.cute.runtime import make_ptr + from cutlass.cute.typing import AddressSpace + + address = int(tensor.data_ptr()) + if address % assumed_align: + raise ValueError( + f"Rubin workspace address {address:#x} is not " + f"{assumed_align}-byte aligned" + ) + return make_ptr( + cutlass.Uint8, + address, + AddressSpace.gmem, + assumed_align=assumed_align, + ) + + +def build_runtime_kwargs( + inputs: Mxfp8LaunchInputs, + resources: PreparedResources, +) -> dict[str, Any]: + import cuda.bindings.driver as cuda + + stream = resources.runtime.current_stream() + weights = inputs.weights + kwargs = { + "activation": _to_cute(inputs.activation), + "activation_sf": _to_cute(inputs.activation_sf), + "topk_indices": _to_cute(inputs.topk_indices), + "topk_scores": _to_cute(inputs.topk_scores, assumed_align=4), + "fc1_weight": _to_cute(weights.fc1_weight), + "fc1_weight_sf": _to_cute(weights.fc1_weight_sf), + "fc2_weight": _to_cute(weights.fc2_weight), + "fc2_weight_sf": _to_cute(weights.fc2_weight_sf), + "fc1_c": ( + None + if inputs.fc1_c is None + else _to_cute(inputs.fc1_c, dynamic_layout=False) + ), + "output_activation": _to_cute(inputs.output_data), + "col_quant_data": ( + None + if inputs.col_quant_data is None + else _to_cute( + inputs.col_quant_data, + assumed_align=128, + dynamic_layout=False, + ) + ), + "col_quant_sf": ( + None + if inputs.col_quant_sf is None + else _to_cute( + inputs.col_quant_sf, + dynamic_layout=False, + ) + ), + "overflow_flag": _to_cute( + inputs.overflow_flag, + assumed_align=4, + dynamic_layout=False, + ), + "local_workspace": _to_cute_ptr(inputs.local_workspace), + "shared_workspace": _to_cute_ptr(inputs.shared_workspace), + "peer_rank_ptr_mapper_host": ( + resources.workspace.peer_mapping.to_sym_buffer_host() + ), + "stream": cuda.CUstream(stream.cuda_stream), + } + return kwargs + + +def layout_signature(inputs: Mxfp8LaunchInputs) -> tuple: + tensors = ( + inputs.activation, + inputs.activation_sf, + inputs.topk_indices, + inputs.topk_scores, + inputs.weights.fc1_weight, + inputs.weights.fc1_weight_sf, + inputs.weights.fc2_weight, + inputs.weights.fc2_weight_sf, + inputs.fc1_c, + inputs.col_quant_data, + inputs.col_quant_sf, + inputs.output_data, + inputs.overflow_flag, + inputs.local_workspace, + inputs.shared_workspace, + ) + return tuple( + None + if tensor is None + else (tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype) + for tensor in tensors + ) + + +def _check_overflow(overflow_flag: torch.Tensor) -> None: + message = ( + "Rubin MegaMoE receive route-pool overflow; the output is invalid for " + "this routing distribution" + ) + assert_async = getattr(torch, "_assert_async", None) + if assert_async is not None: + assert_async(overflow_flag == 0, message) + return + if torch.cuda.is_current_stream_capturing(): + raise NotImplementedError( + "CUDA graph capture requires torch._assert_async to surface " + "Rubin MegaMoE overflow" + ) + # Compatibility fallback for PyTorch builds without a device-side assert. + value = int(overflow_flag.item()) + if value != 0: + raise RuntimeError(f"{message} (overflow_flag={value})") + + +def launch_forward( + compiled, + inputs: Mxfp8LaunchInputs, + resources: PreparedResources, +) -> torch.Tensor: + runtime_kwargs = build_runtime_kwargs(inputs, resources) + compiled.callable(**runtime_kwargs) + _check_overflow(inputs.overflow_flag) + + output_data = torch.empty( + (inputs.token_count, inputs.output_data.shape[1]), + dtype=inputs.output_data.dtype, + device=inputs.output_data.device, + ) + output_data.copy_(inputs.output_data[: inputs.token_count]) + return output_data + + +__all__ = [ + "build_runtime_kwargs", + "launch_forward", + "layout_signature", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_stash.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_stash.py new file mode 100644 index 000000000..27409e837 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_stash.py @@ -0,0 +1,302 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Eager-only ownership and materialization for the MXFP8 FC1 stash.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.distributed as dist + +from ..._contracts import ForwardConfig, ValidatedForwardRequest +from ..._types import MoeEpWgradForwardStash +from .._workspace import _align_up +from ._adapter import _GATE_UP_INTERLEAVE, Mxfp8LaunchInputs +from ._compile import PreparedMxfp8Kernel +from ._wgrad_layout import ( + assemble_discrete_col_requant_scales, + cumulative_padded_offsets, + pool_data_as_wgrad_a, +) + +_TOKEN_SRC_METADATA_BYTES = 8 + + +@dataclass(frozen=True) +class Mxfp8StashPlan: + """One launch's logical counts over an instance-owned raw C buffer.""" + + buffer: torch.Tensor + expert_counts: tuple[int, ...] + expert_offsets: tuple[int, ...] + local_routes: int + + +class Mxfp8ForwardStash: + """Own a padded high-watermark C buffer and compact public results.""" + + def __init__(self, config: ForwardConfig, device: torch.device) -> None: + if not config.generate_c: + raise ValueError("Mxfp8ForwardStash requires generate_c=True") + self.config = config + self.device = torch.device(device) + self._buffer: torch.Tensor | None = None + self._allocation_count = 0 + self._closed = False + + @property + def capacity(self) -> int: + return 0 if self._buffer is None else int(self._buffer.shape[0]) + + @property + def allocation_count(self) -> int: + return self._allocation_count + + def _local_expert_counts( + self, + request: ValidatedForwardRequest, + ) -> tuple[int, ...]: + flat_experts = request.topk_idx.reshape(-1).to(torch.int64) + valid_experts = flat_experts[flat_experts >= 0] + counts = torch.bincount( + valid_experts, + minlength=self.config.num_experts, + ).to(torch.int64) + if self.config.ep_size > 1: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError( + "distributed generate_c route counting requires an " + "initialized torch.distributed process group" + ) + dist.all_reduce( + counts, + op=dist.ReduceOp.SUM, + group=self.config.ep_group, + ) + begin = self.config.ep_rank * self.config.experts_per_rank + end = begin + self.config.experts_per_rank + return tuple(int(value) for value in counts[begin:end].cpu().tolist()) + + def prepare( + self, + request: ValidatedForwardRequest, + *, + pool_token_capacity: int, + ) -> Mxfp8StashPlan: + if self._closed: + raise RuntimeError("MXFP8 forward stash is closed") + if request.device != self.device: + raise ValueError( + f"MXFP8 forward stash is bound to {self.device}, " + f"got {request.device}" + ) + if torch.cuda.is_current_stream_capturing(): + raise NotImplementedError( + "MoeEp generate_c=True is eager-only and does not support " + "CUDA graph capture" + ) + if pool_token_capacity <= 0: + raise ValueError("pool_token_capacity must be positive") + + expert_counts = self._local_expert_counts(request) + expert_offsets = [] + padded_routes = 0 + token_padding = ( + self.config.token_padding_size + if self.config.backward_wgrad_mode == "operands" + else 128 + ) + for count in expert_counts: + expert_offsets.append(padded_routes) + padded_routes += _align_up( + count, + token_padding, + ) + + if padded_routes > pool_token_capacity: + raise RuntimeError( + "forward stash route layout exceeds Rubin pool capacity: " + f"{padded_routes} > {pool_token_capacity}" + ) + # The upstream training kernel validates the receiver-domain C tensor + # against its complete pool shape, even though only active expert rows + # are materialized for the public stash. + required_capacity = pool_token_capacity + if self._buffer is None or self.capacity < required_capacity: + self._buffer = torch.empty( + required_capacity, + 2 * self.config.intermediate_size, + dtype=torch.bfloat16, + device=self.device, + ) + self._allocation_count += 1 + + return Mxfp8StashPlan( + buffer=self._buffer, + expert_counts=expert_counts, + expert_offsets=tuple(expert_offsets), + local_routes=sum(expert_counts), + ) + + def materialize( + self, + plan: Mxfp8StashPlan, + inputs: Mxfp8LaunchInputs, + prepared: PreparedMxfp8Kernel, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + MoeEpWgradForwardStash | None, + ]: + """Compact padded kernel rows into the documented public stash.""" + + if self._closed: + raise RuntimeError("MXFP8 forward stash is closed") + if inputs.fc1_c is not plan.buffer: + raise ValueError("launch inputs do not use this stash plan's buffer") + if prepared.token_src_metadata_bytes != ( + prepared.pool_token_capacity * _TOKEN_SRC_METADATA_BYTES + ): + raise RuntimeError("unexpected token_src_metadata byte size") + + local_routes = plan.local_routes + if local_routes == 0: + fc1_c = torch.empty( + (0, 2 * self.config.intermediate_size), + dtype=torch.bfloat16, + device=self.device, + ) + route_metadata = torch.empty( + (0, 4), + dtype=torch.int32, + device=self.device, + ) + else: + physical_rows = torch.cat( + tuple( + torch.arange( + count, + dtype=torch.int64, + device=self.device, + ) + + offset + for count, offset in zip( + plan.expert_counts, + plan.expert_offsets, + ) + if count + ) + ) + local_experts = torch.repeat_interleave( + torch.arange( + self.config.experts_per_rank, + dtype=torch.int64, + device=self.device, + ), + torch.tensor( + plan.expert_counts, + dtype=torch.int64, + device=self.device, + ), + output_size=local_routes, + ) + + metadata_region = inputs.shared_workspace.narrow( + 0, + prepared.token_src_metadata_offset, + prepared.token_src_metadata_bytes, + ) + packed_metadata = metadata_region.view(torch.int64).index_select( + 0, + physical_rows, + ) + src_tokens = packed_metadata & 0xFFFFFFFF + high = packed_metadata >> 32 + src_ranks = (high >> 16) & 0xFFFF + src_slots = high & 0xFFFF + + order_key = ( + ( + local_experts * self.config.ep_size + + src_ranks + ) + * int(self.config.max_tokens_per_rank) + + src_tokens + ) * self.config.top_k + src_slots + order = torch.argsort(order_key, stable=True) + physical_rows = physical_rows.index_select(0, order) + local_experts = local_experts.index_select(0, order) + src_ranks = src_ranks.index_select(0, order) + src_tokens = src_tokens.index_select(0, order) + src_slots = src_slots.index_select(0, order) + + raw_fc1_c = plan.buffer.index_select(0, physical_rows) + pairs = self.config.intermediate_size // _GATE_UP_INTERLEAVE + gate_up_blocks = raw_fc1_c.reshape( + local_routes, + pairs, + 2, + _GATE_UP_INTERLEAVE, + ) + gate = gate_up_blocks[:, :, 0, :].reshape( + local_routes, + self.config.intermediate_size, + ) + up = gate_up_blocks[:, :, 1, :].reshape( + local_routes, + self.config.intermediate_size, + ) + fc1_c = torch.cat((gate, up), dim=1) + route_metadata = torch.stack( + (local_experts, src_ranks, src_tokens, src_slots), + dim=1, + ).to(torch.int32) + + wgrad_stash = None + if self.config.backward_wgrad_mode == "operands": + if inputs.col_quant_data is None or inputs.col_quant_sf is None: + raise RuntimeError( + "wgrad operand mode requires forward column-requant outputs" + ) + padded_ends, expert_offsets = cumulative_padded_offsets( + plan.expert_counts, + self.config.token_padding_size, + self.device, + ) + padded_routes = padded_ends[-1] if padded_ends else 0 + fc1_a = pool_data_as_wgrad_a( + inputs.col_quant_data, + padded_routes, + ) + fc1_sfa = assemble_discrete_col_requant_scales( + inputs.col_quant_sf, + plan.expert_counts, + padded_ends, + self.config.hidden_size, + self.config.sf_padding_size, + ) + valid_route_counts = torch.tensor( + plan.expert_counts, + dtype=torch.int32, + device=self.device, + ) + wgrad_stash = MoeEpWgradForwardStash( + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + expert_offsets=expert_offsets, + valid_route_counts=valid_route_counts, + route_metadata=route_metadata, + ) + return fc1_c, route_metadata, wgrad_stash + + def close(self) -> None: + self._buffer = None + self._closed = True + + +__all__ = [ + "Mxfp8ForwardStash", + "Mxfp8StashPlan", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_wgrad_layout.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_wgrad_layout.py new file mode 100644 index 000000000..cb2a2fa03 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_wgrad_layout.py @@ -0,0 +1,410 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""MXFP8 pool-to-grouped-wgrad data and scale layout conversion.""" + +from __future__ import annotations + +import torch + +from .._workspace import _align_up + +_SF_VEC_SIZE = 32 +_SF_ATOM_ROWS = 128 +_SF_ATOM_COLUMNS = 4 +_SF_ATOM_BYTES = _SF_ATOM_ROWS * _SF_ATOM_COLUMNS +_GATE_UP_INTERLEAVE = 32 +_NEUTRAL_E8M0 = 127 + + +def cumulative_padded_offsets( + valid_counts: tuple[int, ...], + padding: int, + device: torch.device, +) -> tuple[tuple[int, ...], torch.Tensor]: + """Return cumulative padded expert ends as host values and device Int32.""" + + ends = [] + total = 0 + for count in valid_counts: + if count < 0: + raise ValueError("expert route counts must be non-negative") + total += _align_up(count, padding) + ends.append(total) + return tuple(ends), torch.tensor( + ends, + dtype=torch.int32, + device=device, + ) + + +def pool_data_as_wgrad_a( + pool_data: torch.Tensor, + padded_routes: int, +) -> torch.Tensor: + """Copy pool ``(K,M)`` bytes into contiguous logical wgrad A ``(M,K)``.""" + + _validate_pool_prefix(pool_data, padded_routes) + return pool_data[:padded_routes].transpose(0, 1).contiguous() + + +def pool_data_as_wgrad_b( + pool_data: torch.Tensor, + padded_routes: int, +) -> torch.Tensor: + """Copy pool ``(K,N)`` bytes into a K-major logical wgrad B ``(K,N)``.""" + + _validate_pool_prefix(pool_data, padded_routes) + return ( + pool_data[:padded_routes] + .transpose(0, 1) + .contiguous() + .transpose(0, 1) + ) + + +def deinterleave_gate_up_columns( + tensor: torch.Tensor, + intermediate: int, +) -> torch.Tensor: + """Convert 32-column ``gate,up`` strips to logical ``gate || up``.""" + + expected = 2 * intermediate + if tensor.ndim != 2 or tensor.shape[1] != expected: + raise ValueError( + f"gate/up tensor must have shape (rows, {expected}), " + f"got {tuple(tensor.shape)}" + ) + if intermediate % _GATE_UP_INTERLEAVE: + raise ValueError( + "intermediate size must be divisible by " + f"{_GATE_UP_INTERLEAVE}" + ) + pairs = intermediate // _GATE_UP_INTERLEAVE + blocks = tensor.reshape( + tensor.shape[0], + pairs, + 2, + _GATE_UP_INTERLEAVE, + ) + gate = blocks[:, :, 0, :].reshape(tensor.shape[0], intermediate) + up = blocks[:, :, 1, :].reshape(tensor.shape[0], intermediate) + return torch.cat((gate, up), dim=1) + + +def assemble_discrete_col_requant_scales( + packed_scales: torch.Tensor, + valid_counts: tuple[int, ...], + padded_ends: tuple[int, ...], + non_k_size: int, + sf_padding: int, +) -> torch.Tensor: + """Assemble upstream col-requant SF atoms for grouped wgrad. + + Since upstream revision 71d5fc1, each expert already emits + ``(non-K/128, K/128, 512)`` atoms, which is grouped-wgrad order. + """ + + return _assemble_atom_scales( + packed_scales, + valid_counts, + padded_ends, + non_k_size, + sf_padding, + source_hidden_major=True, + source_name="col-requant", + ) + + +def assemble_dfc2_atom_scales( + packed_scales: torch.Tensor, + valid_counts: tuple[int, ...], + padded_ends: tuple[int, ...], + non_k_size: int, + sf_padding: int, + *, + deinterleave_gate_up: int | None = None, +) -> torch.Tensor: + """Reorder dFC2 epilogue atoms from token-major to grouped-wgrad order.""" + + return _assemble_atom_scales( + packed_scales, + valid_counts, + padded_ends, + non_k_size, + sf_padding, + source_hidden_major=False, + source_name="dFC2", + deinterleave_gate_up=deinterleave_gate_up, + ) + + +def _assemble_atom_scales( + packed_scales: torch.Tensor, + valid_counts: tuple[int, ...], + padded_ends: tuple[int, ...], + non_k_size: int, + sf_padding: int, + *, + source_hidden_major: bool, + source_name: str, + deinterleave_gate_up: int | None = None, +) -> torch.Tensor: + """Expand compact per-expert SF atoms to the data-padded K extent.""" + + if len(valid_counts) != len(padded_ends): + raise ValueError("expert count and padded offset lengths must match") + _validate_padded_ends(padded_ends) + if sf_padding % _SF_ATOM_ROWS: + raise ValueError("scale padding must be divisible by 128") + padded_non_k = _align_up(non_k_size, _SF_ATOM_ROWS) + non_k_atoms = padded_non_k // _SF_ATOM_ROWS + flat_u8 = packed_scales.view(torch.uint8).reshape(-1) + expert_parts = [] + previous_end = 0 + source_byte_offset = 0 + for count, end in zip(valid_counts, padded_ends): + target_extent = end - previous_end + if count < 0 or count > target_extent: + raise ValueError("valid expert routes exceed their padded extent") + if target_extent % _SF_ATOM_ROWS: + raise ValueError( + "data-padded expert extents must be multiples of 128" + ) + source_extent = _align_up(count, sf_padding) + source_token_atoms = source_extent // _SF_ATOM_ROWS + target_token_atoms = target_extent // _SF_ATOM_ROWS + if source_token_atoms > target_token_atoms: + raise ValueError("scale-padded extent exceeds data-padded extent") + source_byte_count = ( + source_token_atoms * non_k_atoms * _SF_ATOM_BYTES + ) + if source_byte_offset + source_byte_count > flat_u8.numel(): + raise ValueError( + f"{source_name} scale output is smaller than its layout" + ) + target_raw = torch.full( + (padded_non_k, target_token_atoms * _SF_ATOM_COLUMNS), + _NEUTRAL_E8M0, + dtype=torch.uint8, + device=flat_u8.device, + ) + if source_byte_count: + source = flat_u8.narrow( + 0, + source_byte_offset, + source_byte_count, + ) + if source_hidden_major: + source = source.reshape( + non_k_atoms, + source_token_atoms, + _SF_ATOM_BYTES, + ) + else: + source = ( + source.reshape( + source_token_atoms, + non_k_atoms, + _SF_ATOM_BYTES, + ) + .permute(1, 0, 2) + .contiguous() + ) + source_raw = _from_blocked_bytes( + source.reshape(-1), + padded_non_k, + source_token_atoms * _SF_ATOM_COLUMNS, + ) + target_raw[:, : source_raw.shape[1]].copy_(source_raw) + if deinterleave_gate_up is not None: + target_raw = ( + deinterleave_gate_up_columns( + target_raw.transpose(0, 1), + deinterleave_gate_up, + ) + .transpose(0, 1) + .contiguous() + ) + expert_parts.append(_to_blocked_bytes(target_raw)) + source_byte_offset += source_byte_count + previous_end = end + + total_routes = padded_ends[-1] if padded_ends else 0 + scale_columns = _align_up(total_routes // _SF_VEC_SIZE, 4) + if expert_parts: + assembled_u8 = torch.cat(expert_parts) + else: + assembled_u8 = flat_u8.new_empty((0,)) + expected = padded_non_k * scale_columns + if assembled_u8.numel() != expected: + raise RuntimeError( + f"assembled {source_name} scale size mismatch: " + f"{assembled_u8.numel()} != {expected}" + ) + return assembled_u8.reshape(padded_non_k, scale_columns).view( + torch.float8_e8m0fnu + ) + + +def assemble_plain_col_scales( + col_scales: torch.Tensor, + valid_counts: tuple[int, ...], + padded_ends: tuple[int, ...], + non_k_size: int, + sf_padding: int, + *, + deinterleave_gate_up: int | None = None, +) -> torch.Tensor: + """Assemble plain ``(K/32,N)`` col scales into grouped-wgrad SF atoms.""" + + if len(valid_counts) != len(padded_ends): + raise ValueError("expert count and padded offset lengths must match") + _validate_padded_ends(padded_ends) + if col_scales.ndim != 2 or col_scales.shape[1] != non_k_size: + raise ValueError( + "plain column scales must have shape " + f"(rows, {non_k_size}), got {tuple(col_scales.shape)}" + ) + if sf_padding % _SF_VEC_SIZE: + raise ValueError("scale padding must be divisible by 32") + + source_u8 = col_scales.view(torch.uint8) + expert_parts = [] + previous_end = 0 + sf_row = 0 + for count, end in zip(valid_counts, padded_ends): + padded_extent = end - previous_end + if padded_extent % _SF_VEC_SIZE: + raise ValueError("padded expert extents must be divisible by 32") + valid_sf_rows = (count + _SF_VEC_SIZE - 1) // _SF_VEC_SIZE + padded_sf_rows = padded_extent // _SF_VEC_SIZE + if count < 0 or count > padded_extent: + raise ValueError("valid expert routes exceed their padded extent") + if sf_row + valid_sf_rows > source_u8.shape[0]: + raise ValueError("plain column scale output is too short") + + raw = torch.full( + (non_k_size, padded_sf_rows), + _NEUTRAL_E8M0, + dtype=torch.uint8, + device=col_scales.device, + ) + if valid_sf_rows: + source = source_u8[ + sf_row : sf_row + valid_sf_rows, + :, + ] + if deinterleave_gate_up is not None: + source = deinterleave_gate_up_columns( + source, + deinterleave_gate_up, + ) + raw[:, :valid_sf_rows].copy_(source.transpose(0, 1)) + expert_parts.append(_to_blocked_bytes(raw)) + sf_row += _align_up(count, sf_padding) // _SF_VEC_SIZE + previous_end = end + + padded_non_k = _align_up(non_k_size, _SF_ATOM_ROWS) + total_routes = padded_ends[-1] if padded_ends else 0 + scale_columns = _align_up(total_routes // _SF_VEC_SIZE, 4) + if expert_parts: + assembled_u8 = torch.cat(expert_parts) + else: + assembled_u8 = source_u8.new_empty((0,)) + expected = padded_non_k * scale_columns + if assembled_u8.numel() != expected: + raise RuntimeError( + "assembled plain column scale size mismatch: " + f"{assembled_u8.numel()} != {expected}" + ) + return assembled_u8.reshape(padded_non_k, scale_columns).view( + torch.float8_e8m0fnu + ) + + +def _to_blocked_bytes(raw_scale: torch.Tensor) -> torch.Tensor: + rows, columns = raw_scale.shape + if rows == 0 or columns == 0: + return raw_scale.new_empty((0,), dtype=torch.uint8) + padded_rows = _align_up(rows, _SF_ATOM_ROWS) + padded_columns = _align_up(columns, _SF_ATOM_COLUMNS) + padded = torch.full( + (padded_rows, padded_columns), + _NEUTRAL_E8M0, + dtype=torch.uint8, + device=raw_scale.device, + ) + padded[:rows, :columns].copy_(raw_scale) + blocks = padded.view( + padded_rows // _SF_ATOM_ROWS, + _SF_ATOM_ROWS, + padded_columns // _SF_ATOM_COLUMNS, + _SF_ATOM_COLUMNS, + ).permute(0, 2, 1, 3) + return ( + blocks.reshape(-1, 4, 32, 4) + .transpose(1, 2) + .reshape(-1) + ) + + +def _from_blocked_bytes( + packed_scale: torch.Tensor, + rows: int, + columns: int, +) -> torch.Tensor: + """Invert the grouped-wgrad 128x4 scale-atom swizzle.""" + + padded_rows = _align_up(rows, _SF_ATOM_ROWS) + padded_columns = _align_up(columns, _SF_ATOM_COLUMNS) + expected = padded_rows * padded_columns + flat = packed_scale.view(torch.uint8).reshape(-1) + if flat.numel() != expected: + raise ValueError( + f"blocked scale has {flat.numel()} bytes, expected {expected}" + ) + if expected == 0: + return flat.new_empty((rows, columns)) + row_atoms = padded_rows // _SF_ATOM_ROWS + column_atoms = padded_columns // _SF_ATOM_COLUMNS + raw = ( + flat.reshape(row_atoms * column_atoms, 32, 4, 4) + .transpose(1, 2) + .reshape(row_atoms, column_atoms, _SF_ATOM_ROWS, _SF_ATOM_COLUMNS) + .permute(0, 2, 1, 3) + .reshape(padded_rows, padded_columns) + ) + return raw[:rows, :columns] + + +def _validate_pool_prefix( + pool_data: torch.Tensor, + padded_routes: int, +) -> None: + if pool_data.ndim != 2 or not pool_data.is_contiguous(): + raise ValueError("pool data must be a contiguous rank-2 tensor") + if padded_routes < 0 or padded_routes > pool_data.shape[0]: + raise ValueError( + f"padded route count {padded_routes} exceeds pool capacity " + f"{pool_data.shape[0]}" + ) + + +def _validate_padded_ends(padded_ends: tuple[int, ...]) -> None: + previous = 0 + for end in padded_ends: + if end < previous: + raise ValueError("expert offsets must be non-decreasing") + previous = end + + +__all__ = [ + "assemble_dfc2_atom_scales", + "assemble_discrete_col_requant_scales", + "assemble_plain_col_scales", + "cumulative_padded_offsets", + "deinterleave_gate_up_columns", + "pool_data_as_wgrad_a", + "pool_data_as_wgrad_b", +] From 1846033836efa212226808568069ff1944e1ed72 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Mon, 24 Aug 2026 16:42:28 -0700 Subject: [PATCH 06/31] test: add semantic MoeEp test support Provide reusable references, quantized input builders, and distributed workers so forward and backward behavior can be validated consistently across execution modes. --- test/python/moe_ep/moe_ep_backward_support.py | 212 +++ .../moe_ep/moe_ep_distributed_workers.py | 506 +++++++ test/python/moe_ep/moe_ep_forward_support.py | 319 +++++ test/python/moe_ep/moe_ep_reference.py | 1192 +++++++++++++++++ test/python/moe_ep/moe_ep_test_data.py | 168 +++ test/python/pytest.ini | 1 + 6 files changed, 2398 insertions(+) create mode 100644 test/python/moe_ep/moe_ep_backward_support.py create mode 100644 test/python/moe_ep/moe_ep_distributed_workers.py create mode 100644 test/python/moe_ep/moe_ep_forward_support.py create mode 100644 test/python/moe_ep/moe_ep_reference.py create mode 100644 test/python/moe_ep/moe_ep_test_data.py diff --git a/test/python/moe_ep/moe_ep_backward_support.py b/test/python/moe_ep/moe_ep_backward_support.py new file mode 100644 index 000000000..390daa42d --- /dev/null +++ b/test/python/moe_ep/moe_ep_backward_support.py @@ -0,0 +1,212 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared reference helpers for MoE EP backward tests.""" + +from __future__ import annotations + +import torch + +from moe_ep.moe_ep_forward_support import _reference_args +from moe_ep.moe_ep_reference import MoeEpReference + +__all__ = [ + "_assert_backward_matches", + "_dense_wgrads_from_operands", + "_expected_backward", + "_grad_output", + "_reference_backward", +] + + +_BACKWARD_CLOSE_KWARGS = ( + {"rtol": 0.15, "atol": 0.125}, # grad_activation is BF16-rounded. + {"rtol": 0.15, "atol": 0.125}, # router-weight gradient. +) + + +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +def _unpack_wgrad_scale_part( + packed: torch.Tensor, + rows: int, + columns: int, +) -> torch.Tensor: + """Invert grouped-wgrad's 128x4 scale-atom swizzle.""" + + padded_rows = _round_up(rows, 128) + padded_columns = _round_up(columns, 4) + row_atoms = padded_rows // 128 + column_atoms = padded_columns // 4 + atom_count = row_atoms * column_atoms + expected = padded_rows * padded_columns + if packed.numel() != expected: + raise ValueError( + f"packed scale part has {packed.numel()} bytes, expected {expected}" + ) + blocked = ( + packed.reshape(atom_count, 32, 4, 4) + .transpose(1, 2) + .reshape(row_atoms, column_atoms, 128, 4) + .permute(0, 2, 1, 3) + .reshape(padded_rows, padded_columns) + ) + return blocked[:rows, :columns].view(torch.float8_e8m0fnu).float() + + +def _dequantize_wgrad_operand( + data: torch.Tensor, + scales: torch.Tensor, + expert_offsets: torch.Tensor, + *, + k_dim: int, +) -> torch.Tensor: + """Decode one public grouped-wgrad operand without launching a GEMM.""" + + if data.ndim != 2 or k_dim not in (0, 1): + raise ValueError("wgrad operand must be rank 2 with k_dim 0 or 1") + non_k = int(data.shape[1 - k_dim]) + padded_non_k = _round_up(non_k, 128) + flat_scales = scales.view(torch.uint8).reshape(-1) + output = torch.empty(data.shape, dtype=torch.float32, device=data.device) + ends = [int(value) for value in expert_offsets.detach().cpu().tolist()] + previous = 0 + scale_byte_offset = 0 + for end in ends: + extent = end - previous + scale_columns = _round_up(extent // 32, 4) + scale_byte_count = padded_non_k * scale_columns + part = flat_scales.narrow( + 0, + scale_byte_offset, + scale_byte_count, + ) + logical_scale = _unpack_wgrad_scale_part( + part, + non_k, + extent // 32, + ) + if k_dim == 1: + expanded_scale = logical_scale.repeat_interleave(32, dim=1) + output[:, previous:end] = ( + data[:, previous:end].float() * expanded_scale + ) + else: + expanded_scale = logical_scale.repeat_interleave( + 32, + dim=1, + ).transpose(0, 1) + output[previous:end, :] = ( + data[previous:end, :].float() * expanded_scale + ) + previous = end + scale_byte_offset += scale_byte_count + if previous != data.shape[k_dim]: + raise ValueError("expert offsets do not cover the operand K dimension") + if scale_byte_offset != flat_scales.numel(): + raise ValueError("expert offsets do not cover the scale tensor") + return output + + +def _dense_wgrads_from_operands(operands): + """Reference grouped matmuls over the exported operand ABI.""" + + fc1_a = _dequantize_wgrad_operand( + operands.fc1_a, + operands.fc1_sfa, + operands.expert_offsets, + k_dim=1, + ) + fc1_b = _dequantize_wgrad_operand( + operands.fc1_b, + operands.fc1_sfb, + operands.expert_offsets, + k_dim=0, + ) + fc2_a = _dequantize_wgrad_operand( + operands.fc2_a, + operands.fc2_sfa, + operands.expert_offsets, + k_dim=1, + ) + fc2_b = _dequantize_wgrad_operand( + operands.fc2_b, + operands.fc2_sfb, + operands.expert_offsets, + k_dim=0, + ) + fc1_parts = [] + fc2_parts = [] + previous = 0 + for end_value in operands.expert_offsets.detach().cpu().tolist(): + end = int(end_value) + fc1_parts.append( + fc1_a[:, previous:end] @ fc1_b[previous:end, :] + ) + fc2_parts.append( + fc2_a[:, previous:end] @ fc2_b[previous:end, :] + ) + previous = end + return torch.stack(fc1_parts), torch.stack(fc2_parts) + + +def _reference_backward(config) -> MoeEpReference: + options = dict(config) + options.pop("tuning", None) + options.pop("sf_padding_size", None) + options["intermediate_format"] = "mxfp8" + options["backward_operand_format"] = "mxfp8" + return MoeEpReference(**options) + + +def _grad_output( + device: torch.device, + token_count: int, + *, + seed: int, +) -> torch.Tensor: + generator = torch.Generator(device=device).manual_seed(seed) + return ( + torch.randn( + token_count, + 128, + generator=generator, + dtype=torch.float32, + device=device, + ) + / 8 + ) + + +def _expected_backward(reference, grad_output, args, stash): + return reference.backward( + grad_output, + *_reference_args(args)[1:], + *stash, + ) + + +def _assert_backward_matches(actual, expected, topk_idx) -> None: + assert len(actual) == len(expected) == 2 + for name, gradient, reference, close_kwargs in zip( + ("grad_activation", "grad_topk_weights"), + actual, + expected, + _BACKWARD_CLOSE_KWARGS, + ): + assert gradient.shape == reference.shape + assert gradient.dtype == torch.float32 + assert torch.isfinite(gradient).all() + torch.testing.assert_close( + gradient, + reference, + msg=lambda default, name=name: ( + f"{name} does not match the backward reference\n{default}" + ), + **close_kwargs, + ) + + dropped = topk_idx == -1 + assert actual[1][dropped].eq(0).all() diff --git a/test/python/moe_ep/moe_ep_distributed_workers.py b/test/python/moe_ep/moe_ep_distributed_workers.py new file mode 100644 index 000000000..9c8e7b5a0 --- /dev/null +++ b/test/python/moe_ep/moe_ep_distributed_workers.py @@ -0,0 +1,506 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Picklable multiprocessing workers for distributed MoE EP tests.""" + +from __future__ import annotations + +from datetime import timedelta + +import torch +import torch.distributed as dist + +from moe_ep.moe_ep_backward_support import ( + _assert_backward_matches, + _dense_wgrads_from_operands, + _expected_backward, + _grad_output, + _reference_backward, +) +from moe_ep.moe_ep_forward_support import ( + _assert_matches_reference, + _forward_config, + _output_as_float, + _reference_args, + _reference_forward, +) +from moe_ep.moe_ep_test_data import make_distributed_forward_inputs + +__all__ = [ + "_distributed_backward_worker", + "_distributed_output_worker", + "_distributed_subgroup_output_worker", + "_distributed_wgrad_worker", + "_run_forward_output_case", + "_run_wgrad_operand_case", +] + + +def _run_forward_output_case( + *, + device: torch.device, + ep_group, + ep_rank: int, + ep_size: int, + combine_format: str = "bf16", + expected_global_ranks: tuple[int, ...] | None = None, +) -> None: + """Run forward parity and dropped-route checks on one initialized EP group.""" + + from cudnn import MoeEp + + args = make_distributed_forward_inputs(ep_rank, ep_size, device) + config = _forward_config( + num_experts=2 * ep_size, + ep_group=ep_group, + max_tokens_per_rank=8, + combine_format=combine_format, + ) + expected = _reference_forward(args, **config) + op = MoeEp(**config) + try: + assert op.ep_rank == ep_rank + if expected_global_ranks is not None: + assert op.ep_global_ranks == expected_global_ranks + + actual = op(*args) + torch.cuda.synchronize(device) + _assert_matches_reference(actual, expected) + + args[3].fill_(-1) + dropped = op(*args) + torch.cuda.synchronize(device) + assert _output_as_float(dropped).eq(0).all() + + dist.barrier(group=ep_group) + op.close() + op = None + dist.barrier(group=ep_group) + finally: + if op is not None: + op.close() + + +def _distributed_output_worker( + rank: int, + world_size: int, + init_file: str, + combine_format: str = "bf16", +) -> None: + device = torch.device("cuda", rank) + torch.cuda.set_device(device) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + device_id=device, + timeout=timedelta(seconds=180), + ) + try: + _run_forward_output_case( + device=device, + ep_group=dist.group.WORLD, + ep_rank=rank, + ep_size=world_size, + combine_format=combine_format, + expected_global_ranks=tuple(range(world_size)), + ) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _distributed_subgroup_output_worker( + global_rank: int, + global_world_size: int, + init_file: str, +) -> None: + """Run one of two disjoint, non-contiguous EP2 groups inside WORLD4.""" + + device = torch.device("cuda", global_rank) + torch.cuda.set_device(device) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=global_rank, + world_size=global_world_size, + device_id=device, + timeout=timedelta(seconds=180), + ) + try: + subgroup_memberships = ((0, 2), (1, 3)) + subgroups = [ + dist.new_group(list(members), backend="nccl") + for members in subgroup_memberships + ] + subgroup_index = global_rank % 2 + ep_group = subgroups[subgroup_index] + ep_rank = dist.get_rank(ep_group) + ep_size = dist.get_world_size(ep_group) + assert tuple( + dist.get_global_rank(ep_group, group_rank) + for group_rank in range(ep_size) + ) == subgroup_memberships[subgroup_index] + + _run_forward_output_case( + device=device, + ep_group=ep_group, + ep_rank=ep_rank, + ep_size=ep_size, + expected_global_ranks=subgroup_memberships[subgroup_index], + ) + dist.barrier() + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _distributed_backward_worker( + rank: int, + world_size: int, + init_file: str, + combine_format: str, + gate_up_clamp: float | None = None, +) -> None: + """Run distributed forward stashing and backward reference parity.""" + + from cudnn import MoeEp + + device = torch.device("cuda", rank) + torch.cuda.set_device(device) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + device_id=device, + timeout=timedelta(seconds=300), + ) + op = None + try: + args = make_distributed_forward_inputs(rank, world_size, device) + config = _forward_config( + num_experts=2 * world_size, + ep_group=dist.group.WORLD, + max_tokens_per_rank=8, + generate_c=True, + combine_format=combine_format, + gate_up_clamp=gate_up_clamp, + ) + reference = _reference_backward(config) + grad_output = _grad_output( + device, + args[3].shape[0], + seed=20260820 + rank, + ) + op = MoeEp(**config) + _, fc1_c, route_metadata = op(*args) + stash = (fc1_c, route_metadata) + expected = _expected_backward(reference, grad_output, args, stash) + + first = op.backward(grad_output, *args[1:], *stash) + second = op.backward(grad_output, *args[1:], *stash) + torch.cuda.synchronize(device) + + # Complete collective work before local assertions. A failure before + # this barrier would leave peer ranks waiting for process-group timeout. + dist.barrier() + _assert_backward_matches(first, expected, args[3]) + _assert_backward_matches(second, expected, args[3]) + + op.close() + op = None + finally: + if op is not None: + op.close() + if dist.is_initialized(): + dist.destroy_process_group() + + +def _make_wgrad_inputs( + rank: int, + world_size: int, + device: torch.device, +): + """Build routes with negative weights, drops, and one empty local expert.""" + + args = list(make_distributed_forward_inputs(rank, world_size, device)) + token_count = args[3].shape[0] + local_expert = 2 * rank + remote_expert = 2 * ((rank + 1) % world_size) + topk_idx = torch.full( + (token_count, 2), + -1, + dtype=torch.int32, + device=device, + ) + topk_weights = torch.empty( + (token_count, 2), + dtype=torch.bfloat16, + device=device, + ) + weight_rows = ( + (0.5, -0.25), + (7.0, -1.25), + (1.5, -9.0), + ) + for token in range(token_count): + pattern = token % 3 + if pattern == 0: + topk_idx[token] = torch.tensor( + (local_expert, remote_expert), + dtype=torch.int32, + device=device, + ) + elif pattern == 1: + topk_idx[token, 1] = local_expert + else: + topk_idx[token, 0] = remote_expert + topk_weights[token] = torch.tensor( + weight_rows[pattern], + dtype=torch.bfloat16, + device=device, + ) + args[3] = topk_idx + args[4] = topk_weights + return tuple(args) + + +def _source_route_expert( + source_rank: int, + token: int, + slot: int, + world_size: int, +) -> int: + pattern = token % 3 + if pattern == 0: + return ( + 2 * source_rank + if slot == 0 + else 2 * ((source_rank + 1) % world_size) + ) + if pattern == 1: + return 2 * source_rank if slot == 1 else -1 + return 2 * ((source_rank + 1) % world_size) if slot == 0 else -1 + + +def _assert_local_operand_metadata( + operands, + reference_operands, + *, + rank: int, + world_size: int, +) -> None: + local_experts = 2 + assert operands.expert_offsets.shape == (local_experts,) + assert operands.valid_route_counts.shape == (local_experts,) + assert torch.equal( + operands.route_metadata, + reference_operands.route_metadata, + ) + assert torch.equal( + operands.valid_route_counts, + reference_operands.valid_route_counts, + ) + assert torch.equal( + operands.expert_offsets, + reference_operands.expert_offsets, + ) + + metadata = operands.route_metadata + if metadata.numel(): + assert metadata[:, 0].ge(0).all() + assert metadata[:, 0].lt(local_experts).all() + counts = torch.bincount( + metadata[:, 0].to(torch.int64), + minlength=local_experts, + ).to(torch.int32) + assert torch.equal(operands.valid_route_counts, counts) + assert counts[0] > 0 + assert counts[1] == 0 + + expected_offsets = [] + padded_end = 0 + for count in counts.tolist(): + padded_end += ((count + 255) // 256) * 256 + expected_offsets.append(padded_end) + assert operands.expert_offsets.tolist() == expected_offsets + + for local_expert, source_rank, token, slot in metadata.tolist(): + global_expert = _source_route_expert( + source_rank, + token, + slot, + world_size, + ) + assert global_expert != -1 + assert global_expert // local_experts == rank + assert global_expert % local_experts == local_expert + + +def _run_grouped_wgrad( + operands, + prefix: str, + *, + wgrad_tensor=None, + accumulate_on_output: bool = False, +): + import cudnn + + return cudnn.grouped_gemm_wgrad_wrapper_sm100( + a_tensor=getattr(operands, f"{prefix}_a"), + b_tensor=getattr(operands, f"{prefix}_b"), + sfa_tensor=getattr(operands, f"{prefix}_sfa"), + sfb_tensor=getattr(operands, f"{prefix}_sfb"), + offsets_tensor=operands.expert_offsets, + output_mode="dense", + wgrad_tensor=wgrad_tensor, + wgrad_dtype=torch.bfloat16, + acc_dtype=torch.float32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=32, + accumulate_on_output=accumulate_on_output, + )["wgrad_tensor"] + + +def _run_wgrad_operand_case( + *, + device: torch.device, + ep_group, + rank: int, + world_size: int, +) -> None: + """Exercise production FC1/FC2 operands through grouped wgrad.""" + + from cudnn import MoeEp, MoeEpWgradOperands + + args = _make_wgrad_inputs(rank, world_size, device) + config = _forward_config( + num_experts=2 * world_size, + ep_group=ep_group, + max_tokens_per_rank=8, + generate_c=True, + backward_wgrad_mode="operands", + token_padding_size=256, + sf_padding_size=128, + ) + reference = _reference_backward(config) + grad_output = _grad_output( + device, + args[3].shape[0], + seed=20260821 + rank, + ) + + with MoeEp(**config) as op: + output, fc1_c, route_metadata, forward_stash = op(*args) + ( + reference_output, + reference_fc1_c, + reference_metadata, + reference_stash, + ) = reference(*_reference_args(args)) + _assert_matches_reference(output, reference_output) + assert torch.equal(route_metadata, reference_metadata) + assert forward_stash.route_metadata is route_metadata + + backward = op.backward( + grad_output, + *args[1:], + fc1_c, + route_metadata, + wgrad_forward_stash=forward_stash, + ) + reference_backward = reference.backward( + grad_output, + *_reference_args(args)[1:], + reference_fc1_c, + reference_metadata, + wgrad_forward_stash=reference_stash, + ) + + _assert_backward_matches(backward[:2], reference_backward[:2], args[3]) + operands = backward[2] + reference_operands = reference_backward[2] + assert isinstance(operands, MoeEpWgradOperands) + assert operands.route_metadata is forward_stash.route_metadata + assert operands.expert_offsets is forward_stash.expert_offsets + assert operands.valid_route_counts is forward_stash.valid_route_counts + _assert_local_operand_metadata( + operands, + reference_operands, + rank=rank, + world_size=world_size, + ) + + expected_wgrads = reference_operands.dense_wgrads() + decoded_wgrads = _dense_wgrads_from_operands(operands) + for actual, expected in zip(decoded_wgrads, expected_wgrads): + torch.testing.assert_close( + actual, + expected, + rtol=0.15, + atol=0.125, + ) + assert actual[1].eq(0).all() + + # The in-tree grouped-wgrad implementation is an SM100 kernel. Rubin + # validates the producer numerics above; the direct SM100 ABI test covers + # execution and accumulate_on_output with the same public bundle layout. + if torch.cuda.get_device_capability(device) != (10, 0): + return + + for prefix, expected in zip(("fc1", "fc2"), expected_wgrads): + actual = _run_grouped_wgrad(operands, prefix) + initial = torch.full_like(actual, 0.25) + accumulated = _run_grouped_wgrad( + operands, + prefix, + wgrad_tensor=initial, + accumulate_on_output=True, + ) + torch.cuda.synchronize(device) + assert accumulated is initial + torch.testing.assert_close( + actual.float(), + expected, + rtol=0.15, + atol=0.125, + ) + torch.testing.assert_close( + accumulated.float(), + expected + 0.25, + rtol=0.15, + atol=0.125, + ) + assert actual[1].eq(0).all() + assert accumulated[1].eq(0.25).all() + + +def _distributed_wgrad_worker( + rank: int, + world_size: int, + init_file: str, +) -> None: + device = torch.device("cuda", rank) + torch.cuda.set_device(device) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + device_id=device, + timeout=timedelta(seconds=600), + ) + try: + _run_wgrad_operand_case( + device=device, + ep_group=dist.group.WORLD, + rank=rank, + world_size=world_size, + ) + dist.barrier() + finally: + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/test/python/moe_ep/moe_ep_forward_support.py b/test/python/moe_ep/moe_ep_forward_support.py new file mode 100644 index 000000000..cd70583bf --- /dev/null +++ b/test/python/moe_ep/moe_ep_forward_support.py @@ -0,0 +1,319 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared support for MoE EP forward tests.""" + +from __future__ import annotations + +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from moe_ep.moe_ep_reference import ( + BlockScaledTensor as ReferenceBlockScaledTensor, + MoeEpReference, + MoeFormat, + forward_combine_round_trip, + quantize_blockwise, +) +from moe_ep.moe_ep_test_data import quantize_mxfp8 + +_DEFAULT_FORWARD_CONFIG = { + "num_experts": 2, + "hidden_size": 128, + "intermediate_size": 256, + "top_k": 2, + "max_tokens_per_rank": 5, + "apply_topk_in_fc1": True, + "combine_format": "bf16", + "output_format": "bf16", +} +_REFERENCE_CLOSE_KWARGS = {"rtol": 0.05, "atol": 0.0625} + +__all__ = [ + "_assert_matches_reference", + "_forward_config", + "_make_forward_case", + "_naive_reference", + "_output_as_float", + "_reference_forward", + "_replay_cuda_graph", + "_require_distributed_sm107", + "_sm107_device", + "_stress_backend_reuse", +] + + +def _forward_config(**overrides): + return {**_DEFAULT_FORWARD_CONFIG, **overrides} + + +def _output_as_float(output): + if isinstance(output, torch.Tensor): + return output.float() + return output.dequantize() + + +def _assert_matches_reference(actual, expected): + torch.testing.assert_close( + _output_as_float(actual), + _output_as_float(expected), + **_REFERENCE_CLOSE_KWARGS, + ) + + +def _naive_reference( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + *, + apply_topk_in_fc1, + clamp=None, + combine_format=MoeFormat.BF16, + intermediate_format=None, + apply_topk_after_combine=False, +): + 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, + device=activation.device, + ) + 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 + if intermediate_format is not None: + intermediate = quantize_blockwise( + intermediate, + intermediate_format, + ).dequantize() + result = intermediate @ fc2_weight[expert].float() + if not apply_topk_in_fc1 and not apply_topk_after_combine: + result = result * route_weight + result = forward_combine_round_trip(result, combine_format) + if not apply_topk_in_fc1 and apply_topk_after_combine: + result = result * route_weight + combine[token, slot] = result + return combine.sum(dim=1).to(torch.bfloat16) + + +def _as_reference_tensor(tensor): + if isinstance(tensor, torch.Tensor): + return tensor + return ReferenceBlockScaledTensor( + data=tensor.data, + scale=tensor.scale, + format=tensor.format.value, + logical_shape=tensor.logical_shape, + axis=tensor.axis, + ) + + +def _reference_args(args): + return ( + _as_reference_tensor(args[0]), + _as_reference_tensor(args[1]), + _as_reference_tensor(args[2]), + args[3], + args[4], + ) + + +def _reference_forward(args, **overrides): + # Rubin's fused FC1 epilogue stores the post-SwiGLU intermediate as MXFP8 + # before FC2 consumes it. Keep MoeEpReference's default raw semantics for + # its standalone tests, but model the device precision for API comparisons. + config = _forward_config(**overrides) + config.pop("tuning", None) + config.setdefault("intermediate_format", "mxfp8") + return MoeEpReference(**config)(*_reference_args(args)) + + +def _sm107_device() -> torch.device: + if not torch.cuda.is_available(): + pytest.skip("Rubin MXFP8 forward requires CUDA") + device = torch.device("cuda", 0) + if torch.cuda.get_device_capability(device) != (10, 7): + pytest.skip("Rubin MXFP8 forward requires exactly SM107 (compute capability 10.7)") + return device + + +def _require_distributed_sm107(world_size: int) -> None: + if not dist.is_available() or not dist.is_nccl_available(): + pytest.skip("multi-GPU Rubin MXFP8 forward requires NCCL") + if torch.cuda.device_count() < world_size: + pytest.skip(f"multi-GPU Rubin MXFP8 forward requires {world_size} GPUs") + if any( + torch.cuda.get_device_capability(index) != (10, 7) + for index in range(world_size) + ): + pytest.skip( + "multi-GPU Rubin MXFP8 forward requires exactly SM107 " + "(compute capability 10.7) on every rank" + ) + try: + import nvshmem.core # noqa: F401 + except (ImportError, OSError): + pytest.skip("multi-GPU Rubin MXFP8 forward requires NVSHMEM") + + +def _make_forward_case( + device: torch.device, + *, + experts: int, + tokens: int, + hidden: int, + intermediate: int, + top_k: int, + index_dtype: torch.dtype, + weight_dtype: torch.dtype, +): + """Build a deterministic supported case for the shape/format matrix.""" + + seed = ( + 20260811 + + experts * 1009 + + tokens * 101 + + hidden * 11 + + intermediate + + top_k + ) + generator = torch.Generator(device=device).manual_seed(seed) + activation = quantize_mxfp8( + torch.randn(tokens, hidden, generator=generator, device=device), + axis=1, + ) + fc1_weight = quantize_mxfp8( + torch.randn( + experts, + hidden, + 2 * intermediate, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + fc2_weight = quantize_mxfp8( + torch.randn( + experts, + intermediate, + hidden, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + topk_idx = ( + torch.arange(tokens * top_k, device=device) + .reshape(tokens, top_k) + .remainder(experts) + .to(index_dtype) + ) + topk_weights = torch.arange( + 1, + tokens * top_k + 1, + dtype=torch.float32, + device=device, + ).reshape(tokens, top_k) + topk_weights /= topk_weights.sum(dim=1, keepdim=True) + return ( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights.to(weight_dtype), + ) + + +def _stress_backend_reuse( + op, + args, + original_topk_idx, + original_topk_weights, + device, + *, + check_weight_refresh, +): + backend = op._forward_backend + assert backend is not None + compiled = backend._compiled + plan_workspace = backend._plan._workspace + weight_refresh_count = ( + backend._adapter.weight_refresh_count if check_weight_refresh else None + ) + alternate_stream = torch.cuda.Stream(device=device) + + for iteration in range(100): + args[3].copy_(original_topk_idx) + args[4].copy_( + original_topk_weights * float((iteration % 7) + 1) / 7.0 + ) + if iteration % 10 == 0: + args[3].fill_(-1) + stream = ( + torch.cuda.current_stream(device) + if iteration % 2 == 0 + else alternate_stream + ) + with torch.cuda.stream(stream): + stressed = op(*args) + stream.synchronize() + if iteration % 10 == 0: + assert _output_as_float(stressed).eq(0).all() + else: + assert torch.isfinite(_output_as_float(stressed)).all() + assert backend._compiled is compiled + assert backend._plan._workspace is plan_workspace + if weight_refresh_count is not None: + assert backend._adapter.weight_refresh_count == weight_refresh_count + + +def _replay_cuda_graph( + op, + args, + original_topk_idx, + expected, + device, + *, + synchronize_ranks=None, +): + synchronize_ranks = synchronize_ranks or (lambda: None) + op.warmup(*args) + synchronize_ranks() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = op(*args) + synchronize_ranks() + + for replay in range(20): + if replay % 2: + args[3].fill_(-1) + else: + args[3].copy_(original_topk_idx) + synchronize_ranks() + graph.replay() + torch.cuda.synchronize(device) + if replay % 2: + assert _output_as_float(graph_output).eq(0).all() + else: + _assert_matches_reference(graph_output, expected) diff --git a/test/python/moe_ep/moe_ep_reference.py b/test/python/moe_ep/moe_ep_reference.py new file mode 100644 index 000000000..3c56db8b9 --- /dev/null +++ b/test/python/moe_ep/moe_ep_reference.py @@ -0,0 +1,1192 @@ +# 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 WgradForwardStashReference: + """Logical reference for the caller-owned forward wgrad stash. + + Unlike the production object, ``fc1_a`` bundles its logical E8M0 scales + with the E4M3 payload. It represents the padded, expert-concatenated + ``x.T`` operand after input MXFP8 staging and token-axis requantization. + """ + + fc1_a: BlockScaledTensor + expert_offsets: torch.Tensor + valid_route_counts: torch.Tensor + route_metadata: torch.Tensor + + +@dataclass(frozen=True) +class WgradOperandsReference: + """Logical MXFP8 operands and dense expert-weight-gradient oracle. + + The K dimension is a concatenation of local experts. Each expert's valid + routes come first, followed by zero rows up to its 256-route boundary. + Production scale tensors use a blocked physical layout; these reference + tensors keep ordinary logical scales so their represented values are easy + to inspect. + """ + + fc1_a: BlockScaledTensor + fc1_b: BlockScaledTensor + fc2_a: BlockScaledTensor + fc2_b: BlockScaledTensor + expert_offsets: torch.Tensor + valid_route_counts: torch.Tensor + route_metadata: torch.Tensor + + def dense_wgrads(self) -> Tuple[torch.Tensor, torch.Tensor]: + """Return dense ``dW1=x.T@dC`` and ``dW2=(p*h).T@dY`` per expert.""" + + a1 = self.fc1_a.dequantize() + b1 = self.fc1_b.dequantize() + a2 = self.fc2_a.dequantize() + b2 = self.fc2_b.dequantize() + expert_count = int(self.expert_offsets.numel()) + dw1 = torch.zeros( + (expert_count, a1.shape[0], b1.shape[1]), + dtype=torch.float32, + device=a1.device, + ) + dw2 = torch.zeros( + (expert_count, a2.shape[0], b2.shape[1]), + dtype=torch.float32, + device=a2.device, + ) + begin = 0 + for expert, end_tensor in enumerate(self.expert_offsets): + end = int(end_tensor.item()) + if end > begin: + dw1[expert] = a1[:, begin:end] @ b1[begin:end] + dw2[expert] = a2[:, begin:end] @ b2[begin:end] + begin = end + return dw1, dw2 + + +@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_axis( + tensor: torch.Tensor, + format: MoeFormat, + *, + axis: int, +) -> torch.Tensor: + if format is MoeFormat.BF16: + return tensor.to(torch.bfloat16).float() + return quantize_blockwise(tensor, format, axis=axis).dequantize() + + +def _format_round_trip(tensor: torch.Tensor, format: MoeFormat) -> torch.Tensor: + return _format_round_trip_axis(tensor, format, axis=-1) + + +def forward_combine_round_trip( + tensor: torch.Tensor, + format: MoeFormat, +) -> torch.Tensor: + """Model GLU combine conversion directly from its FP32 accumulator.""" + + return _format_round_trip(tensor, format) + + +def backward_combine_round_trip( + tensor: torch.Tensor, + format: MoeFormat, +) -> torch.Tensor: + """Model dGLU combine conversion directly from its FP32 accumulator.""" + + return _format_round_trip(tensor, format) + + +def _padded_expert_rows( + rows: torch.Tensor, + expert_rows: torch.Tensor, + valid_counts: Sequence[int], + padded_ends: Sequence[int], +) -> torch.Tensor: + """Place compact expert-grouped rows at the start of padded ranges.""" + + padded_extent = int(padded_ends[-1]) if padded_ends else 0 + padded = torch.zeros( + (padded_extent, *rows.shape[1:]), + dtype=rows.dtype, + device=rows.device, + ) + begin = 0 + for expert, (count, end) in enumerate(zip(valid_counts, padded_ends)): + positions = torch.nonzero( + expert_rows == expert, + as_tuple=False, + ).flatten() + if int(positions.numel()) != int(count): + raise ValueError( + f"expert {expert} has {positions.numel()} rows, expected {count}" + ) + if count: + padded[begin : begin + count].copy_( + rows.index_select(0, positions) + ) + begin = int(end) + return padded + + +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. + + ``intermediate_format`` optionally applies a post-SwiGLU, pre-FC2 format + round trip to model fused kernels that materialize their FC2 input in low + precision. ``None`` preserves the raw mathematical reference semantics. + ``backward_operand_format`` additionally models dGLU staging of grad-output + and transposed weights along their backward reduction dimensions. + """ + + 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, + intermediate_format: Optional[Union[MoeFormat, str]] = None, + backward_operand_format: Optional[Union[MoeFormat, str]] = None, + apply_topk_in_fc1: bool = True, + gate_up_clamp: Optional[float] = None, + generate_c: bool = False, + backward_wgrad_mode: str = "none", + token_padding_size: int = 128, + ) -> 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 backward_wgrad_mode not in ("none", "operands"): + raise ValueError( + "backward_wgrad_mode must be 'none' or 'operands'" + ) + if backward_wgrad_mode == "operands" and not generate_c: + raise ValueError( + "backward_wgrad_mode='operands' requires generate_c=True" + ) + if not isinstance(token_padding_size, int) or token_padding_size <= 0: + raise ValueError("token_padding_size must be a positive integer") + if ( + backward_wgrad_mode == "operands" + and token_padding_size != 256 + ): + raise ValueError( + "backward_wgrad_mode='operands' requires " + "token_padding_size=256" + ) + + 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.intermediate_format = ( + None if intermediate_format is None else _parse_format(intermediate_format) + ) + self.backward_operand_format = ( + None + if backward_operand_format is None + else _parse_format(backward_operand_format) + ) + self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) + self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) + self.generate_c = bool(generate_c) + self.backward_wgrad_mode = backward_wgrad_mode + self.token_padding_size = token_padding_size + + 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 + if self.intermediate_format is not None: + intermediate = _format_round_trip( + intermediate, + self.intermediate_format, + ) + expert_output = intermediate @ fc2_weight[expert] + expert_output = forward_combine_round_trip( + expert_output, + self.combine_format, + ) + if not self.apply_topk_in_fc1: + # The upstream training kernel leaves scores out of dispatch + # and applies them in standalone TopkReduce after the combine + # wire-format round trip. + expert_output = expert_output * weights + 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], + Tuple[ + MoeTensor, + torch.Tensor, + torch.Tensor, + WgradForwardStashReference, + ], + ]: + """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``. In wgrad operand mode, a + fourth :class:`WgradForwardStashReference` item is returned. ``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, + ) + # The Rubin path first stages plain activation along H, then its + # forward column requantization forms x.T scales along routed K. + wgrad_activation_float = None + if self.backward_wgrad_mode == "operands": + wgrad_activation_float = _format_round_trip( + activation_float, + MoeFormat.MXFP8, + ) + 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, + ) + if self.backward_wgrad_mode == "operands": + # Plain forward operands are staged to the same public MXFP8 + # reduction-axis representation before the Rubin GEMMs. + fc1_float = _format_round_trip_axis( + fc1_float, + MoeFormat.MXFP8, + axis=1, + ) + fc2_float = _format_round_trip_axis( + fc2_float, + MoeFormat.MXFP8, + 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 + forward_activation_float = ( + wgrad_activation_float + if wgrad_activation_float is not None + else activation_float + ) + send_tokens = forward_activation_float.index_select( + 0, + send_token_idx, + ) + + recv_tokens = self._all_to_all(send_tokens, send_counts, recv_counts) + recv_wgrad_tokens = None + if wgrad_activation_float is not None: + recv_wgrad_tokens = self._all_to_all( + wgrad_activation_float.index_select(0, send_token_idx), + 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 + fc1_c_order = 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: + if self.backward_wgrad_mode == "operands": + if recv_wgrad_tokens is None or fc1_c_order is None: + raise RuntimeError("wgrad forward staging was not built") + valid_counts = tuple( + int(value) + for value in torch.bincount( + recv_expert, + minlength=self.experts_per_rank, + ).cpu().tolist() + ) + padded_ends = [] + total = 0 + for count in valid_counts: + total += _ceil_div( + count, + self.token_padding_size, + ) * self.token_padding_size + padded_ends.append(total) + ordered_tokens = recv_wgrad_tokens.index_select( + 0, + fc1_c_order, + ) + metadata_experts = route_metadata[:, 0].to(torch.int64) + padded_x = _padded_expert_rows( + ordered_tokens, + metadata_experts, + valid_counts, + padded_ends, + ) + wgrad_stash = WgradForwardStashReference( + fc1_a=quantize_blockwise( + padded_x.transpose(0, 1), + MoeFormat.MXFP8, + axis=1, + ), + expert_offsets=torch.tensor( + padded_ends, + dtype=torch.int32, + device=device, + ), + valid_route_counts=torch.tensor( + valid_counts, + dtype=torch.int32, + device=device, + ), + route_metadata=route_metadata, + ) + return output, fc1_c, route_metadata, wgrad_stash + return output, fc1_c, route_metadata + return output + + def backward( + self, + grad_output: torch.Tensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + fc1_c: torch.Tensor, + route_metadata: torch.Tensor, + *, + wgrad_forward_stash: Optional[ + WgradForwardStashReference + ] = None, + ) -> Union[ + Tuple[torch.Tensor, torch.Tensor], + Tuple[ + torch.Tensor, + torch.Tensor, + WgradOperandsReference, + ], + ]: + """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_topk_weights)`` in float32. In wgrad + operand mode, a third :class:`WgradOperandsReference` item models the + caller-owned grouped-GEMM operands. + """ + + if not self.generate_c: + raise RuntimeError("backward requires the operator to be constructed with generate_c=True") + if self.backward_wgrad_mode == "operands": + if not isinstance( + wgrad_forward_stash, + WgradForwardStashReference, + ): + raise TypeError( + "wgrad_forward_stash must be a " + "WgradForwardStashReference" + ) + if not torch.equal( + wgrad_forward_stash.route_metadata, + route_metadata, + ): + raise ValueError( + "wgrad_forward_stash route identity does not match " + "route_metadata" + ) + elif wgrad_forward_stash is not None: + raise ValueError( + "wgrad_forward_stash is only accepted in operands mode" + ) + 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(fc1_weight) + two_i = 2 * self.intermediate_size + 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, + ) + semantic_fc2_float = fc2_float + effective_backward_format = self.backward_operand_format + if ( + effective_backward_format is None + and self.backward_wgrad_mode == "operands" + ): + effective_backward_format = MoeFormat.MXFP8 + if effective_backward_format is not None: + # The dGLU adapter requantizes both transposed weights along the + # backward GEMM reduction dimension. + fc1_float = _format_round_trip_axis( + fc1_float.transpose(1, 2), + effective_backward_format, + axis=1, + ).transpose(1, 2) + fc2_float = _format_round_trip_axis( + fc2_float.transpose(1, 2), + effective_backward_format, + axis=1, + ).transpose(1, 2) + 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 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 + semantic_grad_output = grad_output.float() + grad_output_float = semantic_grad_output + if effective_backward_format is not None: + grad_output_float = _format_round_trip( + grad_output_float, + effective_backward_format, + ) + 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) + recv_semantic_grad = self._all_to_all( + semantic_grad_output.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) + 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) + semantic_dy_rows = torch.empty_like(recv_semantic_grad) + semantic_dy_rows.index_copy_(0, perm, recv_semantic_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) + weighted_h_rows = torch.zeros( + (local_routes, self.intermediate_size), + dtype=torch.float32, + device=device, + ) + wgrad_dy_rows = torch.zeros( + (local_routes, self.hidden_size), + dtype=torch.float32, + device=device, + ) + dc_rows = torch.zeros( + (local_routes, two_i), + dtype=torch.float32, + device=device, + ) + 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) + w = w_rows.index_select(0, positions).unsqueeze(-1) + d_y = dy_rows.index_select(0, positions) + semantic_d_y = semantic_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 + weighted_h_rows.index_copy_(0, positions, h * w) + wgrad_dy_rows.index_copy_(0, positions, d_y) + + if self.apply_topk_in_fc1: + d_y_pre = d_y + else: + d_y_pre = d_y * w + d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) + if self.apply_topk_in_fc1: + d_h = d_h_fc2 * w + semantic_d_h = ( + semantic_d_y + @ semantic_fc2_float[expert].transpose(0, 1) + ) + d_w_rows[positions] = (semantic_d_h * h).sum(dim=-1) + else: + d_h = d_h_fc2 + d_w_rows[positions] = ( + semantic_d_y * (h @ semantic_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) + dc_rows.index_copy_(0, positions, d_c) + if self.intermediate_format is not None: + d_c = _format_round_trip(d_c, self.intermediate_format) + d_x = d_c @ fc1_float[expert].transpose(0, 1) + d_x_rows.index_copy_( + 0, + positions, + backward_combine_round_trip(d_x, self.combine_format), + ) + + # 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) + grad_topk_weights = grad_topk_weights.view( + token_count, + self.top_k, + ) + if self.backward_wgrad_mode == "operands": + stash = wgrad_forward_stash + padded_ends = tuple( + int(value) + for value in stash.expert_offsets.cpu().tolist() + ) + valid_counts = tuple( + int(value) + for value in stash.valid_route_counts.cpu().tolist() + ) + padded_dc = _padded_expert_rows( + dc_rows, + expert_rows, + valid_counts, + padded_ends, + ) + padded_weighted_h = _padded_expert_rows( + weighted_h_rows, + expert_rows, + valid_counts, + padded_ends, + ) + padded_wgrad_dy = _padded_expert_rows( + wgrad_dy_rows, + expert_rows, + valid_counts, + padded_ends, + ) + operands = WgradOperandsReference( + fc1_a=stash.fc1_a, + fc1_b=quantize_blockwise( + padded_dc, + MoeFormat.MXFP8, + axis=0, + ), + fc2_a=quantize_blockwise( + padded_weighted_h.transpose(0, 1), + MoeFormat.MXFP8, + axis=1, + ), + fc2_b=quantize_blockwise( + padded_wgrad_dy, + MoeFormat.MXFP8, + axis=0, + ), + expert_offsets=stash.expert_offsets, + valid_route_counts=stash.valid_route_counts, + route_metadata=stash.route_metadata, + ) + return grad_activation, grad_topk_weights, operands + return grad_activation, grad_topk_weights + + +__all__ = [ + "BlockScaledTensor", + "MoeEpReference", + "MoeFormat", + "MoeTensor", + "WgradForwardStashReference", + "WgradOperandsReference", + "backward_combine_round_trip", + "forward_combine_round_trip", + "quantize_blockwise", +] \ No newline at end of file diff --git a/test/python/moe_ep/moe_ep_test_data.py b/test/python/moe_ep/moe_ep_test_data.py new file mode 100644 index 000000000..bcd6c6000 --- /dev/null +++ b/test/python/moe_ep/moe_ep_test_data.py @@ -0,0 +1,168 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Deterministic input data and quantization helpers for MoE EP tests.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def make_forward_inputs(device: torch.device): + """Build one deterministic MXFP8 forward case.""" + + generator = torch.Generator(device=device).manual_seed(20260811) + experts, tokens, hidden, intermediate = 2, 5, 128, 256 + activation = quantize_mxfp8( + torch.randn(tokens, hidden, generator=generator, device=device), + axis=1, + ) + fc1_weight = quantize_mxfp8( + torch.randn( + experts, + hidden, + 2 * intermediate, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + fc2_weight = quantize_mxfp8( + torch.randn( + experts, + intermediate, + hidden, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + topk_idx = torch.tensor( + [[0, 1], [1, 0], [0, -1], [1, 0], [0, 1]], + dtype=torch.int32, + device=device, + ) + topk_weights = torch.tensor( + [ + [0.75, 0.25], + [0.625, 0.375], + [1.0, 0.0], + [0.5, 0.5], + [0.875, 0.125], + ], + dtype=torch.bfloat16, + device=device, + ) + return activation, fc1_weight, fc2_weight, topk_idx, topk_weights + + +def make_distributed_forward_inputs( + rank: int, + world_size: int, + device: torch.device, +): + """Build rank-local inputs with one local and one remote route per token.""" + + generator = torch.Generator(device=device).manual_seed(20260811 + rank) + # Vary local shapes without exceeding the distributed tests' + # max_tokens_per_rank=8 contract at EP sizes above seven. + local_experts, tokens, hidden, intermediate = ( + 2, + rank % 7 + 2, + 128, + 256, + ) + activation = quantize_mxfp8( + torch.randn(tokens, hidden, generator=generator, device=device), + axis=1, + ) + fc1_weight = quantize_mxfp8( + torch.randn( + local_experts, + hidden, + 2 * intermediate, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + fc2_weight = quantize_mxfp8( + torch.randn( + local_experts, + intermediate, + hidden, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + remote_rank = (rank + 1) % world_size + topk_idx = torch.tensor( + [ + [ + rank * local_experts + token % local_experts, + remote_rank * local_experts + (token + 1) % local_experts, + ] + for token in range(tokens) + ], + dtype=torch.int32, + device=device, + ) + topk_weights = torch.tensor( + [[0.625, 0.375]], + dtype=torch.bfloat16, + device=device, + ).expand(tokens, -1).contiguous() + return activation, fc1_weight, fc2_weight, topk_idx, topk_weights + + +def quantize_mxfp8(tensor: torch.Tensor, *, axis: int = -1): + """Return a public logical MXFP8 tensor (E4M3 payload + E8M0 scales).""" + + from cudnn import BlockScaledTensor + + axis = axis % tensor.ndim + logical_shape = tuple(tensor.shape) + logical_extent = logical_shape[axis] + moved = tensor.float().movedim(axis, -1) + block_count = (logical_extent + 31) // 32 + padded_extent = block_count * 32 + if padded_extent != logical_extent: + moved = F.pad(moved, (0, padded_extent - logical_extent)) + + blocks = moved.reshape(*moved.shape[:-1], block_count, 32) + raw_scale = blocks.abs().amax(dim=-1) / 448.0 + safe_scale = torch.where(raw_scale > 0, raw_scale, 1.0) + power_of_two_scale = torch.where( + raw_scale > 0, + torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), + torch.zeros_like(raw_scale), + ) + scale = power_of_two_scale.to(torch.float8_e8m0fnu) + reciprocal = torch.where(scale.float() > 0, scale.float().reciprocal(), 0.0) + payload = ( + (blocks * reciprocal.unsqueeze(-1)) + .clamp(-448.0, 448.0) + .to(torch.float8_e4m3fn) + .reshape(*moved.shape)[..., :logical_extent] + ) + + return BlockScaledTensor( + data=payload.movedim(-1, axis).contiguous(), + scale=scale.movedim(-1, axis).contiguous(), + format="mxfp8", + logical_shape=logical_shape, + axis=axis, + ) + + +__all__ = [ + "make_distributed_forward_inputs", + "make_forward_inputs", + "quantize_mxfp8", +] diff --git a/test/python/pytest.ini b/test/python/pytest.ini index 412f6d38a..7ab68cea5 100644 --- a/test/python/pytest.ini +++ b/test/python/pytest.ini @@ -6,5 +6,6 @@ markers = L3: specifies L3 level (use -m L3) L4: specifies L4 level (use -m L4) gpu_exclusive: tests that require exclusive GPU access (no concurrent kernels from other processes) + moe_ep_multinode: torchrun-native multi-node MoE EP tests addopts = -m L0 --tb=short --no-header From e49f2cc307ca9c85653928a4bf561707f78576e7 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Mon, 24 Aug 2026 16:42:38 -0700 Subject: [PATCH 07/31] test: cover MoeEp forward execution Exercise API validation, semantic numerics, arbitrary subgroups, quantized outputs, and single- and multi-node distributed forward paths. --- test/python/moe_ep/test_moe_ep_forward.py | 1641 +++++++++++++++++ .../moe_ep/test_moe_ep_forward_multinode.py | 209 +++ 2 files changed, 1850 insertions(+) create mode 100644 test/python/moe_ep/test_moe_ep_forward.py create mode 100644 test/python/moe_ep/test_moe_ep_forward_multinode.py diff --git a/test/python/moe_ep/test_moe_ep_forward.py b/test/python/moe_ep/test_moe_ep_forward.py new file mode 100644 index 000000000..8ca620596 --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_forward.py @@ -0,0 +1,1641 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Core MoE EP forward contract, parity, runtime, and distributed tests.""" + +from __future__ import annotations + +import os +import sys +from dataclasses import replace +from types import ModuleType, SimpleNamespace + +import numpy as np +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from moe_ep.moe_ep_distributed_workers import ( + _distributed_output_worker, + _distributed_subgroup_output_worker, +) +from moe_ep.moe_ep_forward_support import ( + _assert_matches_reference, + _forward_config, + _make_forward_case, + _naive_reference, + _output_as_float, + _reference_forward, + _replay_cuda_graph, + _require_distributed_sm107, + _sm107_device, + _stress_backend_reuse, +) +from moe_ep.moe_ep_reference import ( + BlockScaledTensor as ReferenceBlockScaledTensor, + MoeEpReference, + MoeFormat, + forward_combine_round_trip, + quantize_blockwise, +) +from moe_ep.moe_ep_test_data import ( + make_forward_inputs, + quantize_mxfp8, +) + + + +def _public_nvfp4(data, scale, logical_shape): + from cudnn import BlockScaledTensor + + return BlockScaledTensor( + data=data, + scale=scale, + format="nvfp4", + logical_shape=logical_shape, + axis=1, + ) + + +def _request(activation, fc1, fc2): + return SimpleNamespace( + activation=activation, + fc1_weight=fc1, + fc2_weight=fc2, + ) + + +# Public API, capability, layout, and workspace contracts. + + +@pytest.mark.L0 +def test_moe_ep_finalizer_warns_without_retaining_failed_backend(): + import cudnn.moe_ep.api as api_module + from cudnn import MoeEp + + class Backend: + close_calls = 0 + + def close(self): + self.close_calls += 1 + raise RuntimeError("cleanup failed") + + operator = MoeEp(**_forward_config()) + backend = Backend() + operator._forward_backend = backend + + with pytest.warns(ResourceWarning, match="cleanup failed"): + operator.__del__() + + assert backend.close_calls == 1 + assert not hasattr(api_module, "_FAILED_FINALIZER_BACKENDS") + operator._forward_backend = None + operator._closed = True + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("field", "value"), + [ + *( + ("token_back_mode", value) + for value in ( + "epi_warps", + "standalone_warps", + "reuse_dispatch_warps", + ) + ), + *( + ("epi_flag_batch", value) + for value in ( + (4, 2), + (1, 1), + (1, 2), + (1, 4), + (2, 1), + (2, 2), + (2, 4), + (4, 4), + ) + ), + *( + ("token_in_flag_batch", value) + for value in (1, 2, 4, 8, 16) + ), + *( + ("group_hint", value) + for value in (None, 64, 128, 256, 512, 768, 1024) + ), + ("reduce_topk_in_kernel", False), + ("reduce_topk_in_kernel", True), + ], +) +def test_moe_ep_tuning_accepts_candidate_values(field, value): + from cudnn import MoeEpTuningConfig + + tuning = MoeEpTuningConfig(**{field: value}) + assert getattr(tuning, field) == value + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "kwargs", + [ + {"token_back_mode": "unknown"}, + {"token_back_mode": []}, + {"epi_flag_batch": [1, 1]}, + {"epi_flag_batch": (3, 3)}, + {"token_in_flag_batch": 3}, + {"token_in_flag_batch": True}, + {"group_hint": 0}, + {"group_hint": True}, + {"reduce_topk_in_kernel": 1}, + { + "token_back_mode": "standalone_warps", + "reduce_topk_in_kernel": True, + }, + ], +) +def test_moe_ep_tuning_rejects_unvalidated_values(kwargs): + from cudnn import MoeEpTuningConfig + + with pytest.raises(ValueError): + MoeEpTuningConfig(**kwargs) + + +@pytest.mark.L0 +def test_moe_ep_tuning_public_contract_mapping_and_cache_key(): + from cudnn import MoeEp, MoeEpTuningConfig + from cudnn.moe_ep import ( + MoeEpTuningConfig as PackageMoeEpTuningConfig, + ) + from cudnn.moe_ep._megamoe_backend.mxfp8._config import ( + Mxfp8KernelConfig, + ) + + assert PackageMoeEpTuningConfig is MoeEpTuningConfig + tuning = MoeEpTuningConfig( + token_back_mode="epi_warps", + epi_flag_batch=(4, 2), + token_in_flag_batch=4, + group_hint=768, + reduce_topk_in_kernel=True, + ) + with MoeEp( + **_forward_config(), + token_padding_size=64, + sf_padding_size=256, + tuning=tuning, + ) as op: + assert op.tuning is tuning + kernel_config = Mxfp8KernelConfig.from_forward_config( + op._forward_config + ) + + assert kernel_config.token_back_mode == "epi_warps" + assert kernel_config.epi_flag_batch == (4, 2) + assert kernel_config.flag_batch == 4 + assert kernel_config.group_hint == 768 + assert kernel_config.token_padding_block == 64 + assert kernel_config.sf_padding_block == 256 + assert kernel_config.tuning_signature(123) == ( + "epi_warps", + (4, 2), + 4, + 768, + True, + ) + effective = kernel_config.effective_config(123) + assert effective["token_padding_block"] == 64 + assert effective["sf_padding_block"] == 256 + assert effective["effective_group_hint"] == 768 + assert effective["fc2_in_kernel_topk_reduce"] is True + assert effective["launch_cluster_count"] == 123 + assert effective["drop_on_overflow"] is True + assert effective["enable_col_quant"] is False + assert "output_format" not in effective + + with MoeEp(**_forward_config()) as default_op: + default_config = Mxfp8KernelConfig.from_forward_config( + default_op._forward_config + ) + assert default_config.tuning_signature(123) == ( + "epi_warps", + (1, 1), + 1, + 123, + False, + ) + key_args = ( + torch.device("cuda", 0), + (10, 7), + 123, + (), + ) + assert kernel_config.compile_key(*key_args) != default_config.compile_key( + *key_args + ) + + +@pytest.mark.L0 +def test_internal_column_requant_config_is_disabled_by_default_and_cache_distinct(): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend.mxfp8._config import ( + Mxfp8KernelConfig, + ) + + with MoeEp(**_forward_config()) as op: + default_forward = op._forward_config + default_config = Mxfp8KernelConfig.from_forward_config(default_forward) + enabled_config = replace( + default_config, + enable_col_quant=True, + col_quant_num_ctas=512, + ) + + assert default_config.enable_col_quant is False + assert default_config.max_recv_size_per_rank == ( + default_forward.ep_size + * default_forward.max_tokens_per_rank + * default_forward.top_k + ) + assert enabled_config.enable_col_quant is True + assert enabled_config.col_quant_num_ctas == 512 + with pytest.raises(ValueError, match="max_recv_size_per_rank"): + replace(default_config, max_recv_size_per_rank=0) + with pytest.raises(ValueError, match="col_quant_num_ctas"): + replace(default_config, col_quant_num_ctas=0) + key_args = (torch.device("cuda", 0), (10, 7), 123, ()) + assert default_config.compile_key(*key_args) != enabled_config.compile_key( + *key_args + ) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("public_format", "wire_format"), + [ + ("bf16", "bf16"), + ("mxfp8", "32e4m3xe8m0"), + ], +) +def test_combine_format_maps_to_contract_wire(public_format, wire_format): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend.mxfp8._config import ( + Mxfp8KernelConfig, + ) + + with MoeEp( + **_forward_config(combine_format=public_format) + ) as op: + kernel_config = Mxfp8KernelConfig.from_forward_config( + op._forward_config + ) + + assert kernel_config.combine_format == wire_format + + +@pytest.mark.L0 +@pytest.mark.parametrize("combine_format", ["bf16", "mxfp8"]) +def test_megamoe_capability_enables_combine_formats(combine_format): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend._capability import validate_config + + with MoeEp(**_forward_config(combine_format=combine_format)) as op: + validate_config(op._forward_config) + + +@pytest.mark.L0 +def test_distributed_topk_can_exceed_local_expert_count(): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend._capability import validate_config + + with MoeEp( + **_forward_config( + num_experts=4, + top_k=3, + ) + ) as op: + distributed = replace( + op._forward_config, + experts_per_rank=2, + ep_size=2, + ep_global_ranks=(0, 1), + ) + + validate_config(distributed) + + +@pytest.mark.L0 +def test_overflow_check_prefers_device_assert(monkeypatch): + from cudnn.moe_ep._megamoe_backend.mxfp8._launch import ( + _check_overflow, + ) + + calls = [] + + def record_assert(condition, message): + calls.append((condition.clone(), message)) + + monkeypatch.setattr(torch, "_assert_async", record_assert) + flag = torch.zeros(1, dtype=torch.int32) + _check_overflow(flag) + + assert len(calls) == 1 + assert bool(calls[0][0]) + assert "route-pool overflow" in calls[0][1] + + +@pytest.mark.L0 +def test_overflow_check_fallback_rejects_nonzero_flag(monkeypatch): + from cudnn.moe_ep._megamoe_backend.mxfp8._launch import ( + _check_overflow, + ) + + monkeypatch.setattr(torch, "_assert_async", None) + monkeypatch.setattr( + torch.cuda, + "is_current_stream_capturing", + lambda: False, + ) + with pytest.raises(RuntimeError, match="receive route-pool overflow"): + _check_overflow(torch.ones(1, dtype=torch.int32)) + + +@pytest.mark.L0 +def test_in_kernel_topk_reduce_omits_standalone_combine_workspace(): + from cudnn.moe_ep._megamoe_backend.mxfp8._compile import ( + _pre_reduced_workspace_metadata, + ) + + class NoStandaloneWorkspace: + def region(self, _name): + raise AssertionError("in-kernel reduction must not query region") + + def offset(self, _name): + raise AssertionError("in-kernel reduction must not query offset") + + def nbytes(self, _name): + raise AssertionError("in-kernel reduction must not query size") + + config = SimpleNamespace( + fc2_in_kernel_topk_reduce=True, + top_k=6, + hidden=7168, + max_tokens_per_rank=4096, + combine_format="bf16", + ) + + assert _pre_reduced_workspace_metadata( + NoStandaloneWorkspace(), + config, + shared_bytes=0, + ) == (None, 0) + + bytes_per_token = config.top_k * config.hidden * 2 + total_bytes = config.max_tokens_per_rank * bytes_per_token + + class StandaloneWorkspace: + def region(self, _name): + return SimpleNamespace(buffer_space="shared") + + def offset(self, _name): + return 256 + + def nbytes(self, _name): + return total_bytes + + config.fc2_in_kernel_topk_reduce = False + assert _pre_reduced_workspace_metadata( + StandaloneWorkspace(), + config, + shared_bytes=256 + total_bytes, + ) == (256, bytes_per_token) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("combine_format", "bits_per_element"), + [ + ("bf16", 16), + ("32e4m3xe8m0", 8), + ], +) +def test_standalone_combine_workspace_tracks_wire_width( + combine_format, + bits_per_element, +): + from cudnn.moe_ep._megamoe_backend.mxfp8._compile import ( + _pre_reduced_workspace_metadata, + ) + + config = SimpleNamespace( + fc2_in_kernel_topk_reduce=False, + top_k=2, + hidden=128, + max_tokens_per_rank=5, + combine_format=combine_format, + ) + bytes_per_token = config.top_k * config.hidden * bits_per_element // 8 + total_bytes = config.max_tokens_per_rank * bytes_per_token + + class StandaloneWorkspace: + def region(self, _name): + return SimpleNamespace(buffer_space="shared") + + def offset(self, _name): + return 128 + + def nbytes(self, _name): + return total_bytes + + assert _pre_reduced_workspace_metadata( + StandaloneWorkspace(), + config, + shared_bytes=128 + total_bytes, + ) == (128, bytes_per_token) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("combine_format", "expected"), + [ + ("bf16", (None, 0)), + ("32e4m3xe8m0", (128, 64)), + ], +) +def test_standalone_combine_scale_workspace_metadata( + combine_format, + expected, +): + from cudnn.moe_ep._megamoe_backend.mxfp8._compile import ( + _pre_reduced_sf_workspace_metadata, + ) + + config = SimpleNamespace( + fc2_in_kernel_topk_reduce=False, + max_tokens_per_rank=5, + combine_format=combine_format, + ) + + class StandaloneWorkspace: + def region(self, _name): + return SimpleNamespace(buffer_space="shared") + + def offset(self, _name): + return 128 + + def nbytes(self, _name): + return config.max_tokens_per_rank * 64 + + workspace = StandaloneWorkspace() + if combine_format == "bf16": + class NoScaleWorkspace: + def region(self, _name): + raise AssertionError("BF16 combine must not query scale region") + + workspace = NoScaleWorkspace() + + assert _pre_reduced_sf_workspace_metadata( + workspace, + config, + shared_bytes=128 + config.max_tokens_per_rank * 64, + ) == expected + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "kwargs", + [ + {"token_padding_size": True}, + {"token_padding_size": 0}, + {"token_padding_size": -1}, + {"token_padding_size": 64.0}, + {"sf_padding_size": True}, + {"sf_padding_size": 0}, + {"sf_padding_size": 64}, + {"sf_padding_size": 128.0}, + ], +) +def test_moe_ep_rejects_invalid_padding(kwargs): + from cudnn import MoeEp + + with pytest.raises(ValueError): + MoeEp(**_forward_config(), **kwargs) + + +@pytest.mark.L0 +def test_moe_ep_rejects_untyped_tuning(): + from cudnn import MoeEp + + with pytest.raises(TypeError, match="MoeEpTuningConfig"): + MoeEp(**_forward_config(), tuning={"group_hint": 768}) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "kwargs", + [ + {"combine_format": "mxfp8"}, + {"output_format": "mxfp8"}, + {"apply_topk_in_fc1": False}, + ], +) +def test_moe_ep_rejects_incompatible_in_kernel_topk_reduce(kwargs): + from cudnn import MoeEp, MoeEpTuningConfig + + config = _forward_config() + config.update(kwargs) + with pytest.raises(ValueError, match="reduce_topk_in_kernel requires"): + MoeEp( + **config, + tuning=MoeEpTuningConfig(reduce_topk_in_kernel=True), + ) + + +@pytest.mark.L0 +def test_distributed_launch_rejects_mismatched_tuning_before_barrier( + monkeypatch, +): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend.mxfp8._backend import ( + Mxfp8Backend, + ) + + with MoeEp(**_forward_config()) as op: + backend = Mxfp8Backend( + op._forward_config, + torch.device("cuda", 0), + ) + backend._ep_launch_ready = False + + stream = SimpleNamespace(synchronize=lambda: None) + resources = SimpleNamespace( + runtime=SimpleNamespace(group=object(), world_size=2) + ) + prepared = SimpleNamespace(launch_cluster_count=123) + monkeypatch.setattr( + backend, + "_ensure_prepared_kernel", + lambda: prepared, + ) + + def mismatched_all_gather(output, signature, *, group): + assert group is resources.runtime.group + output[:] = [ + signature, + ("standalone_warps", (1, 1), 1, 123), + ] + + barrier_called = False + + def unexpected_barrier(*, group): + nonlocal barrier_called + barrier_called = True + + monkeypatch.setattr(dist, "all_gather_object", mismatched_all_gather) + monkeypatch.setattr(dist, "barrier", unexpected_barrier) + + with pytest.raises(RuntimeError, match="MoeEp tuning must match"): + backend._ensure_ep_launch_ready(resources, stream) + assert not barrier_called + assert not backend._ep_launch_ready + + +@pytest.mark.L0 +def test_api_allocates_fresh_bf16_outputs_with_logical_shape(): + from cudnn import MoeEp + + device = _sm107_device() + args = make_forward_inputs(device) + activation, fc1_weight, fc2_weight = args[:3] + + assert activation.logical_shape == (5, 128) + assert fc1_weight.logical_shape == (2, 128, 512) + assert fc2_weight.logical_shape == (2, 256, 128) + + with MoeEp(**_forward_config()) as op: + first = op(*args) + snapshot = first.clone() + second = op(*args) + torch.cuda.synchronize(device) + + assert isinstance(first, torch.Tensor) + assert isinstance(second, torch.Tensor) + assert first.shape == second.shape == (5, 128) + assert first.dtype == second.dtype == torch.bfloat16 + assert first.device == second.device == device + assert first is not second + assert first.data_ptr() != second.data_ptr() + torch.testing.assert_close(first, snapshot, rtol=0, atol=0) + torch.testing.assert_close(first, second, rtol=0, atol=0) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "kwargs", + [ + {"combine_format": "nvfp4"}, + {"output_format": "mxfp8"}, + {"output_format": "nvfp4"}, + {"apply_topk_in_fc1": False}, + ], +) +def test_training_megamoe_rejects_unsupported_config_before_backend(kwargs): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend._capability import validate_config + + with MoeEp(**_forward_config(**kwargs)) as op: + with pytest.raises(NotImplementedError, match="training MegaMoE"): + validate_config(op._forward_config) + + +@pytest.mark.L0 +def test_training_megamoe_rejects_nvfp4_operand_before_cuda_query(monkeypatch): + from cudnn.moe_ep._megamoe_backend._capability import validate_request + + operand = _public_nvfp4( + torch.zeros(2, 64, dtype=torch.uint8), + torch.ones(2, 8).to(torch.float8_e4m3fn), + (2, 128), + ) + request = _request(operand, operand, operand) + request.device = torch.device("cuda", 0) + + monkeypatch.setattr( + torch.cuda, + "get_device_capability", + lambda _device: pytest.fail("CUDA capability queried too early"), + ) + with pytest.raises(NotImplementedError, match="only MXFP8"): + validate_request(request) + + +# Single-rank and distributed forward numerical parity. + + +@pytest.mark.L0 +def test_fp8_activation_bf16_combine_forward_single_gpu(): + from cudnn import MoeEp + + device = _sm107_device() + args = make_forward_inputs(device) + expected = _reference_forward(args) + + with MoeEp(**_forward_config()) as op: + actual = op(*args) + torch.cuda.synchronize(device) + + args[3].fill_(-1) + dropped = op(*args) + torch.cuda.synchronize(device) + + assert actual.shape == (5, 128) + assert actual.dtype == torch.bfloat16 + _assert_matches_reference(actual, expected) + assert dropped.eq(0).all() + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_mxfp8_combine_matches_direct_fp32_training_reference(): + from cudnn import MoeEp + + device = _sm107_device() + args = make_forward_inputs(device) + config = _forward_config( + combine_format="mxfp8", + ) + expected = _reference_forward(args, **config) + + with MoeEp(**config) as op: + actual = op(*args) + torch.cuda.synchronize(device) + + _assert_matches_reference(actual, expected) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +@pytest.mark.parametrize( + "plain_mask", + [ + (True, False, False), + (False, True, False), + (False, False, True), + (True, True, False), + (True, False, True), + (False, True, True), + (True, True, True), + ], +) +@pytest.mark.parametrize( + "plain_dtype", + [torch.bfloat16, torch.float16, torch.float32], +) +def test_plain_and_mixed_inputs_match_staged_reference( + plain_mask, + plain_dtype, +): + from cudnn import MoeEp + + device = _sm107_device() + args = list(make_forward_inputs(device)) + for index, make_plain in enumerate(plain_mask): + if make_plain: + args[index] = args[index].dequantize(dtype=plain_dtype) + args = tuple(args) + expected = _reference_forward(args) + + with MoeEp(**_forward_config()) as op: + actual = op(*args) + torch.cuda.synchronize(device) + + _assert_matches_reference(actual, expected) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_one_operator_switches_mxfp8_and_plain_weight_families(): + from cudnn import MoeEp + + device = _sm107_device() + quantized_args = make_forward_inputs(device) + plain_args = ( + quantized_args[0].dequantize(dtype=torch.bfloat16), + quantized_args[1].dequantize(dtype=torch.bfloat16), + quantized_args[2].dequantize(dtype=torch.bfloat16), + *quantized_args[3:], + ) + expected_quantized = _reference_forward(quantized_args) + expected_plain = _reference_forward(plain_args) + + with MoeEp(**_forward_config()) as op: + quantized = op(*quantized_args) + backend = op._forward_backend + refresh_before = backend._adapter.weight_refresh_count + plain = op(*plain_args) + refresh_after = backend._adapter.weight_refresh_count + torch.cuda.synchronize(device) + + assert op._forward_backend is None + assert refresh_after == refresh_before + 1 + _assert_matches_reference(quantized, expected_quantized) + _assert_matches_reference(plain, expected_plain) + + +@pytest.mark.L0 +def test_nondefault_moe_ep_tuning_matches_reference_and_reuses_plan(): + from cudnn import MoeEp, MoeEpTuningConfig + + device = _sm107_device() + args = make_forward_inputs(device) + expected = _reference_forward(args) + tuning = MoeEpTuningConfig( + token_back_mode="standalone_warps", + epi_flag_batch=(4, 2), + token_in_flag_batch=4, + group_hint=64, + ) + + with MoeEp(**_forward_config(), tuning=tuning) as op: + first = op(*args) + backend = op._forward_backend + assert backend is not None + compiled = backend._compiled + workspace = backend._plan._workspace + second = op(*args) + torch.cuda.synchronize(device) + + assert backend._compiled is compiled + assert backend._plan._workspace is workspace + assert backend.kernel_config.tuning_signature( + backend._prepared_kernel.launch_cluster_count + ) == ("standalone_warps", (4, 2), 4, 64, False) + + _assert_matches_reference(first, expected) + _assert_matches_reference(second, expected) + + +@pytest.mark.L0 +def test_gate_up_clamp_matches_moe_ep_reference(): + from cudnn import MoeEp + + device = _sm107_device() + args = make_forward_inputs(device) + clamp = 0.5 + expected = _reference_forward(args, gate_up_clamp=clamp) + unclamped = _reference_forward(args) + + with MoeEp(**_forward_config(gate_up_clamp=clamp)) as op: + actual = op(*args) + torch.cuda.synchronize(device) + + assert not torch.equal(expected, unclamped) + _assert_matches_reference(actual, expected) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_generate_c_outputs_fc1_c_and_route_metadata(): + from cudnn import MoeEp + + device = _sm107_device() + args = make_forward_inputs(device) + config = _forward_config( + gate_up_clamp=1.25, + generate_c=True, + ) + expected_output, expected_fc1_c, expected_metadata = _reference_forward( + args, + **config, + ) + + with MoeEp(**config) as op: + first = op(*args) + output, fc1_c, route_metadata = first + fc1_c_snapshot = fc1_c.clone() + metadata_snapshot = route_metadata.clone() + + scaled_args = (*args[:4], args[4] * 0.25) + _, scaled_fc1_c, scaled_metadata = op(*scaled_args) + torch.cuda.synchronize(device) + + assert isinstance(first, tuple) + assert len(first) == 3 + assert output.shape == (5, 128) + assert output.dtype == torch.bfloat16 + assert fc1_c.shape == (9, 512) + assert fc1_c.dtype == torch.bfloat16 + assert route_metadata.shape == (9, 4) + assert route_metadata.dtype == torch.int32 + _assert_matches_reference(output, expected_output) + torch.testing.assert_close( + _output_as_float(fc1_c), + _output_as_float(expected_fc1_c), + rtol=0.01, + atol=0.01, + ) + torch.testing.assert_close( + route_metadata, + expected_metadata, + rtol=0, + atol=0, + ) + + # FC1 C is captured before clamp/SwiGLU and does not include router weights. + torch.testing.assert_close(scaled_fc1_c, fc1_c_snapshot, rtol=0, atol=0) + torch.testing.assert_close(scaled_metadata, metadata_snapshot, rtol=0, atol=0) + torch.testing.assert_close(fc1_c, fc1_c_snapshot, rtol=0, atol=0) + torch.testing.assert_close(route_metadata, metadata_snapshot, rtol=0, atol=0) + assert scaled_fc1_c is not fc1_c + assert scaled_metadata is not route_metadata + assert scaled_fc1_c.data_ptr() != fc1_c.data_ptr() + assert scaled_metadata.data_ptr() != route_metadata.data_ptr() + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +@pytest.mark.parametrize( + ( + "experts", + "tokens", + "hidden", + "intermediate", + "top_k", + "index_dtype", + "weight_dtype", + ), + [ + pytest.param( + 2, + 3, + 128, + 256, + 1, + torch.int32, + torch.bfloat16, + id="topk1-h128-i256-int32-bf16", + ), + pytest.param( + 2, + 5, + 128, + 256, + 2, + torch.int64, + torch.float32, + id="topk2-h128-i256-int64-fp32", + ), + pytest.param( + 4, + 3, + 256, + 256, + 4, + torch.int32, + torch.float16, + id="topk4-h256-i256-int32-fp16", + ), + pytest.param( + 32, + 1, + 128, + 256, + 32, + torch.int64, + torch.float32, + id="topk32-boundary-int64-fp32", + ), + ], +) +def test_supported_topk_shape_and_routing_format_matrix( + experts, + tokens, + hidden, + intermediate, + top_k, + index_dtype, + weight_dtype, +): + from cudnn import MoeEp + + device = _sm107_device() + args = _make_forward_case( + device, + experts=experts, + tokens=tokens, + hidden=hidden, + intermediate=intermediate, + top_k=top_k, + index_dtype=index_dtype, + weight_dtype=weight_dtype, + ) + config = _forward_config( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, + max_tokens_per_rank=tokens, + ) + expected = _reference_forward(args, **config) + + with MoeEp(**config) as op: + actual = op(*args) + torch.cuda.synchronize(device) + + assert actual.shape == (tokens, hidden) + assert actual.dtype == torch.bfloat16 + _assert_matches_reference(actual, expected) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_single_gpu_stress_and_cuda_graph_replay(): + from cudnn import MoeEp + + device = _sm107_device() + args = make_forward_inputs(device) + original_topk_idx = args[3].clone() + original_topk_weights = args[4].clone() + config = _forward_config() + expected = _reference_forward(args, **config) + + with MoeEp(**config) as op: + op.warmup(*args) + _stress_backend_reuse( + op, + args, + original_topk_idx, + original_topk_weights, + device, + check_weight_refresh=True, + ) + + args[3].copy_(original_topk_idx) + args[4].copy_(original_topk_weights) + eager = op(*args) + torch.cuda.synchronize(device) + _assert_matches_reference(eager, expected) + _replay_cuda_graph(op, args, original_topk_idx, expected, device) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_nondefault_tuning_warmup_and_cuda_graph_replay(): + from cudnn import MoeEp, MoeEpTuningConfig + + device = _sm107_device() + args = make_forward_inputs(device) + original_topk_idx = args[3].clone() + expected = _reference_forward(args) + tuning = MoeEpTuningConfig( + token_back_mode="reuse_dispatch_warps", + epi_flag_batch=(2, 2), + token_in_flag_batch=2, + group_hint=128, + ) + + with MoeEp(**_forward_config(), tuning=tuning) as op: + _replay_cuda_graph( + op, + args, + original_topk_idx, + expected, + device, + ) + + +@pytest.mark.L0 +def test_reference_apply_topk_after_fc2_weights_after_combine_rounding(): + """Keep post-combine router weighting in reference-only semantics.""" + + device = torch.device("cpu") + args = make_forward_inputs(device) + # Duplicate routes make pre/post-combine weighting observably different. + args[3].copy_( + torch.tensor( + [[0, 0], [1, 1], [0, 0], [1, 1], [0, 0]], + dtype=torch.int32, + device=device, + ) + ) + args[4].copy_( + torch.tensor( + [[256.0, -255.0]] * 5, + dtype=torch.bfloat16, + device=device, + ) + ) + decoded_args = ( + args[0].dequantize(), + args[1].dequantize(), + args[2].dequantize(), + args[3], + args[4], + ) + expected = _naive_reference( + *decoded_args, + apply_topk_in_fc1=False, + intermediate_format=MoeFormat.MXFP8, + apply_topk_after_combine=True, + ) + pre_combine_weighting = _naive_reference( + *decoded_args, + apply_topk_in_fc1=False, + intermediate_format=MoeFormat.MXFP8, + ) + assert not torch.equal(expected, pre_combine_weighting) + +@pytest.mark.L0 +def test_forward_mxfp8_combine_is_direct_fp32(): + generator = torch.Generator().manual_seed(20260819) + accumulator = torch.randn(4, 128, generator=generator) * 3.25 + + forward = forward_combine_round_trip(accumulator, MoeFormat.MXFP8) + direct_fp32 = quantize_blockwise( + accumulator, + MoeFormat.MXFP8, + ).dequantize() + bf16_preround = quantize_blockwise( + accumulator.to(torch.bfloat16).float(), + MoeFormat.MXFP8, + ).dequantize() + + torch.testing.assert_close( + forward, + direct_fp32, + rtol=0, + atol=0, + ) + assert not torch.equal(forward, bf16_preround) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +@pytest.mark.parametrize( + "world_size", + [2, 3, 4], + ids=["ep2", "ep3", "ep4"], +) +@pytest.mark.parametrize("combine_format", ["bf16", "mxfp8"]) +def test_mxfp8_forward_multi_gpu_matches_reference( + world_size, + combine_format, + tmp_path, +): + _require_distributed_sm107(world_size) + os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") + init_file = tmp_path / f"{combine_format}_combine_ep{world_size}.init" + mp.spawn( + _distributed_output_worker, + args=(world_size, str(init_file), combine_format), + nprocs=world_size, + join=True, + ) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_mxfp8_forward_noncontiguous_ep_subgroups(tmp_path): + global_world_size = 4 + _require_distributed_sm107(global_world_size) + os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") + init_file = tmp_path / "two_noncontiguous_ep2.init" + mp.spawn( + _distributed_subgroup_output_worker, + args=(global_world_size, str(init_file)), + nprocs=global_world_size, + join=True, + ) + + +# Input staging and workspace layout. + + +@pytest.mark.L0 +def test_plain_tensor_staging_matches_logical_mxfp8_quantization(): + if not torch.cuda.is_available(): + pytest.skip("MXFP8 staging test requires CUDA") + from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( + _quantize_plain_mxfp8, + ) + + device = torch.device("cuda", 0) + plain = torch.randn(2, 128, 3, device=device).to(torch.bfloat16) + actual = _quantize_plain_mxfp8(plain, axis=1) + expected = quantize_mxfp8(plain, axis=1) + + torch.testing.assert_close( + actual.data.view(torch.uint8), + expected.data.view(torch.uint8), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + actual.scale.view(torch.uint8), + expected.scale.view(torch.uint8), + rtol=0, + atol=0, + ) + + +@pytest.mark.L0 +def test_intermediate_requires_full_mma_n_tile(): + from cudnn import MoeEp + + args = _make_forward_case( + torch.device("cpu"), + experts=2, + tokens=3, + hidden=128, + intermediate=128, + top_k=2, + index_dtype=torch.int32, + weight_dtype=torch.bfloat16, + ) + with MoeEp( + **_forward_config(intermediate_size=128, max_tokens_per_rank=3) + ) as op: + with pytest.raises( + NotImplementedError, + match=r"intermediate_size .*divisible by 256", + ): + op(*args) + assert op._forward_backend is None + + +@pytest.mark.L0 +def test_activation_scale_rows_are_padded_to_16_bytes(): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend._workspace import ( + WorkspaceRequirements, + padded_mxfp8_scale_columns, + ) + + assert padded_mxfp8_scale_columns(128) == 16 + assert padded_mxfp8_scale_columns(512) == 16 + assert padded_mxfp8_scale_columns(640) == 32 + + with MoeEp(**_forward_config()) as op: + requirements = WorkspaceRequirements.for_mxfp8( + op._forward_config, + kernel_local_workspace_bytes=128, + kernel_shared_workspace_bytes=128, + ) + activation_scale = next( + region + for region in requirements.symmetric_regions + if region.name == "activation_scale" + ) + assert activation_scale.nbytes == 5 * 16 + + +@pytest.mark.L0 +def test_column_requant_workspace_is_allocated_only_when_enabled(): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend._workspace import WorkspaceRequirements + + with MoeEp(**_forward_config()) as op: + disabled = WorkspaceRequirements.for_mxfp8( + op._forward_config, + kernel_local_workspace_bytes=128, + kernel_shared_workspace_bytes=128, + ) + enabled = WorkspaceRequirements.for_mxfp8( + op._forward_config, + kernel_local_workspace_bytes=128, + kernel_shared_workspace_bytes=128, + col_quant_data_bytes=640, + col_quant_sf_bytes=80, + ) + + disabled_names = {region.name for region in disabled.local_regions} + assert "col_quant_data" not in disabled_names + assert "col_quant_sf" not in disabled_names + enabled_sizes = { + region.name: region.nbytes for region in enabled.local_regions + } + assert enabled_sizes["col_quant_data"] == 640 + assert enabled_sizes["col_quant_sf"] == 80 + + with pytest.raises(ValueError, match="must be enabled together"): + WorkspaceRequirements.for_mxfp8( + op._forward_config, + kernel_local_workspace_bytes=128, + kernel_shared_workspace_bytes=128, + col_quant_data_bytes=640, + ) + + +# Reference and quantization self-checks. + + +@pytest.mark.L0 +def test_mxfp8_activation_representation(): + activation = quantize_mxfp8(torch.randn(3, 128), axis=1) + + assert activation.format.value == "mxfp8" + assert activation.logical_shape == (3, 128) + assert activation.axis == 1 + assert activation.data.shape == (3, 128) + assert activation.data.dtype == torch.float8_e4m3fn + assert activation.scale.shape == (3, 4) + assert activation.scale.dtype == torch.float8_e8m0fnu + assert torch.isfinite(activation.dequantize()).all() + + +@pytest.mark.L0 +def test_reference_mxfp8_block_scaled_round_trip(): + values = torch.linspace(-4.0, 4.0, 3 * 64).reshape(3, 64) + quantized = quantize_blockwise(values, MoeFormat.MXFP8) + + assert isinstance(quantized, ReferenceBlockScaledTensor) + assert quantized.format is MoeFormat.MXFP8 + assert quantized.logical_shape == (3, 64) + assert tuple(quantized.data.shape) == (3, 64) + assert tuple(quantized.scale.shape) == (3, 2) + assert quantized.scale.dtype == torch.float8_e8m0fnu + assert quantized.dequantize().shape == values.shape + assert torch.isfinite(quantized.dequantize()).all() + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "intermediate_format", + [None, MoeFormat.MXFP8], + ids=["fp32-intermediate", "mxfp8-intermediate"], +) +def test_reference_mxfp8_inputs_bf16_combine_matches_naive( + intermediate_format, +): + torch.manual_seed(19) + experts, tokens, hidden, intermediate = 2, 3, 128, 128 + 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, MoeFormat.MXFP8, axis=1) + q_fc1 = quantize_blockwise(fc1_weight, MoeFormat.MXFP8, axis=1) + q_fc2 = quantize_blockwise(fc2_weight, MoeFormat.MXFP8, 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, + combine_format="bf16", + output_format="bf16", + intermediate_format=intermediate_format, + ) + + 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, + combine_format=MoeFormat.BF16, + intermediate_format=intermediate_format, + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +# Host-side EP topology and runtime bootstrap. + + +@pytest.mark.L0 +def test_resolve_ep_topology_preserves_group_rank_order(monkeypatch): + from cudnn.moe_ep.api import _resolve_ep_topology + + group = object() + monkeypatch.setattr(dist, "is_available", lambda: True) + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda selected=None: 2) + monkeypatch.setattr( + dist, + "get_rank", + lambda selected=None: 1, + ) + monkeypatch.setattr( + dist, + "get_global_rank", + lambda selected, group_rank: (3, 1)[group_rank], + ) + + assert _resolve_ep_topology(group) == (2, 1, (3, 1)) + + +@pytest.mark.L0 +def test_resolve_ep_topology_rejects_nonmember(monkeypatch): + from cudnn.moe_ep.api import _resolve_ep_topology + + monkeypatch.setattr(dist, "is_available", lambda: True) + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda group: 2) + monkeypatch.setattr(dist, "get_rank", lambda group: -1) + + with pytest.raises(ValueError, match="must be a member"): + _resolve_ep_topology(object()) + + +@pytest.mark.L0 +def test_resolve_runtime_world_revalidates_ordered_membership(monkeypatch): + from cudnn.moe_ep._megamoe_backend._runtime import _resolve_world + + group = object() + monkeypatch.setattr(dist, "is_available", lambda: True) + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda selected: 2) + monkeypatch.setattr(dist, "get_rank", lambda selected: 1) + monkeypatch.setattr( + dist, + "get_global_rank", + lambda selected, group_rank: (3, 1)[group_rank], + ) + config = SimpleNamespace( + ep_group=group, + ep_size=2, + ep_rank=1, + ep_global_ranks=(3, 1), + ) + + world = _resolve_world(config) + assert world.identity == (1, 2, (3, 1)) + + config.ep_global_ranks = (1, 3) + with pytest.raises(RuntimeError, match="membership does not match"): + _resolve_world(config) + + +@pytest.mark.L0 +def test_megamoe_capability_accepts_nonworld_subgroup_config(): + from cudnn.moe_ep._contracts import ForwardConfig + from cudnn.moe_ep._megamoe_backend._capability import validate_config + from cudnn.moe_ep._tuning import MoeEpTuningConfig + + config = ForwardConfig( + num_experts=4, + hidden_size=128, + intermediate_size=256, + top_k=2, + experts_per_rank=2, + ep_size=2, + ep_rank=0, + ep_group=object(), + ep_global_ranks=(1, 3), + max_tokens_per_rank=8, + output_format="bf16", + combine_format="bf16", + apply_topk_in_fc1=True, + gate_up_clamp=None, + generate_c=False, + token_padding_size=128, + sf_padding_size=128, + tuning=MoeEpTuningConfig(), + ) + + validate_config(config) + + +@pytest.fixture +def runtime_module(): + from cudnn.moe_ep._megamoe_backend import _runtime + + with _runtime._PROCESS_RUNTIME_REGISTRY.lock: + _runtime._PROCESS_RUNTIME_REGISTRY.active = None + yield _runtime + with _runtime._PROCESS_RUNTIME_REGISTRY.lock: + _runtime._PROCESS_RUNTIME_REGISTRY.active = None + + +class _FakeRuntimeProvider: + def __init__(self, runtime_module, state=None): + self._runtime_module = runtime_module + self._state = state or runtime_module.RuntimeInitState.NOT_INITIALIZED + self._world = None + self.finalize_count = 0 + + def initialization_state(self): + return self._state + + def initialize(self, device, world): + del device + self._world = world + self._state = self._runtime_module.RuntimeInitState.INITIALIZED + + def rank(self): + return self._world.rank + + def world_size(self): + return self._world.size + + def device(self): + return torch.device("cuda", 0) + + def finalize(self): + self.finalize_count += 1 + self._state = self._runtime_module.RuntimeInitState.NOT_INITIALIZED + + +@pytest.mark.L0 +def test_runtime_manager_shares_only_identical_subgroup(runtime_module): + world = runtime_module.RuntimeWorld( + rank=1, + size=2, + group=object(), + global_ranks=(1, 3), + ) + provider = _FakeRuntimeProvider(runtime_module) + manager = runtime_module.RuntimeManager( + provider_factory=lambda: provider, + world_resolver=lambda config: world, + ) + + first = manager.acquire(object(), torch.device("cuda", 0)) + second = manager.acquire(object(), torch.device("cuda", 0)) + assert manager.ref_count == 2 + assert second.global_ranks == (1, 3) + + second.close() + assert manager.ref_count == 1 + first.close() + assert manager.ref_count == 0 + assert provider.finalize_count == 1 + + +@pytest.mark.L0 +def test_runtime_manager_rejects_different_same_geometry_subgroup( + runtime_module, +): + first_world = runtime_module.RuntimeWorld( + rank=0, + size=2, + group=object(), + global_ranks=(0, 2), + ) + second_world = runtime_module.RuntimeWorld( + rank=0, + size=2, + group=object(), + global_ranks=(0, 3), + ) + provider = _FakeRuntimeProvider(runtime_module) + first_manager = runtime_module.RuntimeManager( + provider_factory=lambda: provider, + world_resolver=lambda config: first_world, + ) + second_manager = runtime_module.RuntimeManager( + provider_factory=lambda: provider, + world_resolver=lambda config: second_world, + ) + + handle = first_manager.acquire(object(), torch.device("cuda", 0)) + with pytest.raises(RuntimeError, match="different EP subgroup"): + second_manager.acquire(object(), torch.device("cuda", 0)) + handle.close() + + +@pytest.mark.L0 +def test_runtime_manager_rejects_unverifiable_external_subgroup( + runtime_module, + monkeypatch, +): + world = runtime_module.RuntimeWorld( + rank=0, + size=2, + group=object(), + global_ranks=(1, 3), + ) + provider = _FakeRuntimeProvider( + runtime_module, + runtime_module.RuntimeInitState.INITIALIZED, + ) + provider._world = world + manager = runtime_module.RuntimeManager( + provider_factory=lambda: provider, + world_resolver=lambda config: world, + ) + monkeypatch.setattr( + runtime_module, + "_spans_default_distributed_world", + lambda selected: False, + ) + + with pytest.raises(RuntimeError, match="cannot safely attach"): + manager.acquire(object(), torch.device("cuda", 0)) + + +@pytest.mark.L0 +def test_nvshmem_uid_broadcast_uses_subgroup_root_global_rank( + runtime_module, + monkeypatch, +): + class _FakeDevice: + def __init__(self, index): + self.index = index + + def set_current(self): + return None + + cuda_module = ModuleType("cuda") + cuda_core_module = ModuleType("cuda.core") + cuda_experimental_module = ModuleType("cuda.core.experimental") + cuda_experimental_module.Device = _FakeDevice + cuda_core_module.experimental = cuda_experimental_module + cuda_module.core = cuda_core_module + monkeypatch.setitem(sys.modules, "cuda", cuda_module) + monkeypatch.setitem(sys.modules, "cuda.core", cuda_core_module) + monkeypatch.setitem( + sys.modules, + "cuda.core.experimental", + cuda_experimental_module, + ) + + init_args = {} + + class _FakeUid: + def __init__(self): + self._data = np.arange(16, dtype=np.uint8) + + core = SimpleNamespace( + get_unique_id=lambda empty: _FakeUid(), + init=lambda **kwargs: init_args.update(kwargs), + ) + monkeypatch.setattr(runtime_module, "_load_nvshmem_core", lambda: core) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + + group = object() + broadcast_args = {} + monkeypatch.setattr(dist, "get_backend", lambda selected: "gloo") + monkeypatch.setattr( + dist, + "get_global_rank", + lambda selected, group_rank: (1, 3)[group_rank], + ) + + def _broadcast(tensor, *, src, group): + broadcast_args.update(tensor=tensor, src=src, group=group) + + monkeypatch.setattr(dist, "broadcast", _broadcast) + monkeypatch.setattr(dist, "barrier", lambda *, group: None) + + world = runtime_module.RuntimeWorld( + rank=0, + size=2, + group=group, + global_ranks=(1, 3), + ) + runtime_module._DefaultNvshmemRuntimeProvider().initialize( + torch.device("cuda", 0), + world, + ) + + assert broadcast_args["src"] == 1 + assert broadcast_args["group"] is group + assert broadcast_args["tensor"].device.type == "cpu" + assert init_args["rank"] == 0 + assert init_args["nranks"] == 2 diff --git a/test/python/moe_ep/test_moe_ep_forward_multinode.py b/test/python/moe_ep/test_moe_ep_forward_multinode.py new file mode 100644 index 000000000..46efb85a4 --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_forward_multinode.py @@ -0,0 +1,209 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Torchrun-native multi-node MoE EP forward acceptance tests.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist + +from moe_ep.moe_ep_distributed_workers import ( + _run_forward_output_case, +) + + +pytestmark = [ + pytest.mark.L1, + pytest.mark.gpu_exclusive, + pytest.mark.moe_ep_multinode, +] + +_TORCHRUN_ENV = ("LOCAL_RANK", "LOCAL_WORLD_SIZE", "RANK", "WORLD_SIZE") +_PROCESS_GROUP_TIMEOUT = timedelta(minutes=10) + + +def _bind_torchrun_device_before_pytest_fixtures() -> None: + """Bind before the root conftest creates its session CUDA handle.""" + + value = os.environ.get("LOCAL_RANK") + if value is None or not torch.cuda.is_available(): + return + local_rank = int(value) + if 0 <= local_rank < torch.cuda.device_count(): + torch.cuda.set_device(local_rank) + + +_bind_torchrun_device_before_pytest_fixtures() + + +@dataclass(frozen=True) +class _TorchrunWorld: + rank: int + world_size: int + local_rank: int + local_world_size: int + device: torch.device + + +def _require_torchrun_environment() -> tuple[int, int, int, int]: + missing = [name for name in _TORCHRUN_ENV if name not in os.environ] + if missing: + pytest.skip( + "multi-node MoE EP forward requires torchrun environment variables: " + + ", ".join(missing) + ) + return ( + int(os.environ["RANK"]), + int(os.environ["WORLD_SIZE"]), + int(os.environ["LOCAL_RANK"]), + int(os.environ["LOCAL_WORLD_SIZE"]), + ) + + +@pytest.fixture(scope="session") +def torchrun_world(): + if not dist.is_available() or not dist.is_nccl_available(): + pytest.skip("multi-node Rubin MXFP8 forward requires NCCL") + + rank, world_size, local_rank, local_world_size = ( + _require_torchrun_environment() + ) + if local_rank < 0 or local_rank >= torch.cuda.device_count(): + pytest.skip( + f"torchrun LOCAL_RANK={local_rank} is not backed by a visible GPU" + ) + + device = torch.device("cuda", local_rank) + if torch.cuda.get_device_capability(device) != (10, 7): + pytest.skip( + "multi-node Rubin MXFP8 forward requires exactly SM107 " + "(compute capability 10.7) on every rank" + ) + try: + import nvshmem.core # noqa: F401 + except (ImportError, OSError): + pytest.skip("multi-node Rubin MXFP8 forward requires NVSHMEM") + + os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") + torch.cuda.set_device(device) + if dist.is_initialized(): + if dist.get_rank() != rank or dist.get_world_size() != world_size: + raise RuntimeError( + "existing process group does not match torchrun RANK/WORLD_SIZE" + ) + else: + dist.init_process_group( + backend="nccl", + init_method="env://", + device_id=device, + timeout=_PROCESS_GROUP_TIMEOUT, + ) + + context = _TorchrunWorld( + rank=rank, + world_size=world_size, + local_rank=local_rank, + local_world_size=local_world_size, + device=device, + ) + try: + yield context + finally: + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() + + +@pytest.mark.parametrize( + ( + "ep_size", + "required_world_size", + "required_local_world_size", + "ep_global_ranks", + ), + [ + pytest.param( + 7, + 14, + 2, + tuple(range(0, 14, 2)), + id="ep7-world14", + ), + pytest.param( + 12, + 12, + 4, + tuple(range(12)), + id="ep12-world12", + ), + pytest.param( + 15, + 20, + 4, + tuple(rank for rank in range(20) if rank % 4 < 3), + id="ep15-world20", + ), + pytest.param( + 16, + 16, + 4, + tuple(range(16)), + id="ep16-world16", + ), + ], +) +@pytest.mark.parametrize("combine_format", ["bf16", "mxfp8"]) +def test_mxfp8_forward_multinode_matches_reference( + torchrun_world, + ep_size, + required_world_size, + required_local_world_size, + ep_global_ranks, + combine_format, +): + world = torchrun_world + if ( + world.world_size != required_world_size + or world.local_world_size != required_local_world_size + ): + pytest.skip( + f"EP{ep_size} requires torchrun WORLD_SIZE={required_world_size}, " + f"LOCAL_WORLD_SIZE={required_local_world_size}; got " + f"WORLD_SIZE={world.world_size}, " + f"LOCAL_WORLD_SIZE={world.local_world_size}" + ) + + if ep_size == world.world_size: + ep_group = dist.group.WORLD + else: + # All WORLD ranks must create subgroups in the same order, including + # idle ranks that are not members of this balanced EP group. + ep_group = dist.new_group( + list(ep_global_ranks), + backend="nccl", + timeout=_PROCESS_GROUP_TIMEOUT, + ) + + is_ep_member = world.rank in ep_global_ranks + try: + if is_ep_member: + ep_rank = dist.get_rank(ep_group) + _run_forward_output_case( + device=world.device, + ep_group=ep_group, + ep_rank=ep_rank, + ep_size=ep_size, + combine_format=combine_format, + expected_global_ranks=ep_global_ranks, + ) + dist.barrier() + finally: + if ep_group is not dist.group.WORLD and is_ep_member: + dist.destroy_process_group(ep_group) + dist.barrier() From 0e7402add32ed7a8f257a1f42b0492a89650d886 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Mon, 24 Aug 2026 16:42:47 -0700 Subject: [PATCH 08/31] test: cover MoeEp backward and wgrad contracts Validate dGLU execution, routing-weight gradients, source invariants, stash layouts, and exported grouped-wgrad operands across supported distributed configurations. --- test/python/moe_ep/test_moe_ep_backward.py | 1166 ++++++++++++++++ .../test_moe_ep_cutedsl_grad_y2_source.py | 63 + .../moe_ep/test_moe_ep_wgrad_contract.py | 1170 +++++++++++++++++ 3 files changed, 2399 insertions(+) create mode 100644 test/python/moe_ep/test_moe_ep_backward.py create mode 100644 test/python/moe_ep/test_moe_ep_cutedsl_grad_y2_source.py create mode 100644 test/python/moe_ep/test_moe_ep_wgrad_contract.py diff --git a/test/python/moe_ep/test_moe_ep_backward.py b/test/python/moe_ep/test_moe_ep_backward.py new file mode 100644 index 000000000..804c87cee --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_backward.py @@ -0,0 +1,1166 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Core MoE EP backward contract, parity, and distributed tests.""" + +from __future__ import annotations + +import inspect +import os +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch +import torch.multiprocessing as mp + +from cudnn.moe_ep import MoeEp +from cudnn.moe_ep._contracts import ForwardConfig +from cudnn.moe_ep._megamoe_backend import _capability +from cudnn.moe_ep._megamoe_backend.mxfp8._backend import Mxfp8Backend +from cudnn.moe_ep._megamoe_backend.mxfp8._backward_layout import ( + Mxfp8BackwardLayout, +) +from cudnn.moe_ep._megamoe_backend.mxfp8._backward_staging import ( + _stage_fc1_preact, + stage_backward, +) +from cudnn.moe_ep._megamoe_backend._workspace import WorkspaceRequirements +from cudnn.moe_ep._megamoe_backend.mxfp8._backward_dispatch import ( + Mxfp8BackwardRedispatch, +) +from cudnn.moe_ep._megamoe_backend.mxfp8._backward_dprob import ( + return_grad_topk_weights, +) +from cudnn.moe_ep._tuning import MoeEpTuningConfig +from cudnn.moe_ep._validation import validate_backward +from moe_ep.moe_ep_backward_support import ( + _assert_backward_matches, + _expected_backward, + _grad_output, + _reference_backward, +) +from moe_ep.moe_ep_distributed_workers import ( + _distributed_backward_worker, +) +from moe_ep.moe_ep_forward_support import ( + _forward_config, + _require_distributed_sm107, + _sm107_device, +) +from moe_ep.moe_ep_reference import ( + MoeFormat, + backward_combine_round_trip, + forward_combine_round_trip, +) +from moe_ep.moe_ep_test_data import ( + make_forward_inputs, + quantize_mxfp8, +) + + +def _config(**overrides) -> ForwardConfig: + values = { + "num_experts": 2, + "hidden_size": 128, + "intermediate_size": 256, + "top_k": 2, + "experts_per_rank": 2, + "ep_size": 1, + "ep_rank": 0, + "ep_group": None, + "ep_global_ranks": (), + "max_tokens_per_rank": 4, + "output_format": "bf16", + "combine_format": "bf16", + "apply_topk_in_fc1": True, + "gate_up_clamp": None, + "generate_c": True, + "token_padding_size": 128, + "sf_padding_size": 128, + "tuning": MoeEpTuningConfig(), + } + values.update(overrides) + return ForwardConfig(**values) + + +def _inputs(): + activation = torch.randn(2, 128, dtype=torch.bfloat16) + fc1_weight = torch.randn(2, 128, 512, dtype=torch.bfloat16) + fc2_weight = torch.randn(2, 256, 128, dtype=torch.bfloat16) + topk_idx = torch.tensor([[0, -1], [1, 0]], dtype=torch.int32) + topk_weights = torch.randn(2, 2, dtype=torch.float32) + return activation, fc1_weight, fc2_weight, topk_idx, topk_weights + + +def _validate_backward( + config, + grad_output, + args, + fc1_c, + route_metadata, +): + return validate_backward( + config, + grad_output, + *args[1:], + fc1_c, + route_metadata, + ) + + +# Validation, layout, staging, and backend contracts. + + +@pytest.mark.L0 +def test_validate_backward_builds_typed_request_and_checks_stash(): + config = _config() + args = _inputs() + grad_output = torch.randn(2, 128) + fc1_c = torch.randn(3, 512, dtype=torch.bfloat16) + route_metadata = torch.tensor( + [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], + dtype=torch.int32, + ) + + request = _validate_backward( + config, + grad_output, + args, + fc1_c, + route_metadata, + ) + + assert request.config is config + assert request.fc1_weight is args[1] + assert request.topk_idx is args[3] + assert request.local_routes == 3 + with pytest.raises(ValueError, match="fc1_c shape must be"): + _validate_backward( + config, + grad_output, + args, + fc1_c[:2], + route_metadata, + ) + with pytest.raises(TypeError, match="route_metadata must have dtype"): + _validate_backward( + config, + grad_output, + args, + fc1_c, + route_metadata.to(torch.int64), + ) + + +@pytest.mark.L0 +def test_mxfp8_backward_layout_builds_public_preactivation_lut(): + config = _config() + args = _inputs() + route_metadata = torch.tensor( + [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], + dtype=torch.int32, + ) + request = _validate_backward( + config, + torch.randn(2, 128), + args, + torch.randn(3, 512, dtype=torch.bfloat16), + route_metadata, + ) + + layout = Mxfp8BackwardLayout.from_request(request) + + assert layout.preact_row_lut[0, 0, 0].item() == 0 + assert layout.preact_row_lut[0, 1, 1].item() == 1 + assert layout.preact_row_lut[0, 1, 0].item() == 2 + + +@pytest.mark.L0 +def test_mxfp8_backward_stages_compact_preactivation_into_pool_rows(): + config = _config() + args = _inputs() + route_metadata = torch.tensor( + [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], + dtype=torch.int32, + ) + fc1_c = torch.arange( + 3 * 512, + dtype=torch.float32, + ).reshape(3, 512).to(torch.bfloat16) + request = _validate_backward( + config, + torch.randn(2, 128), + args, + fc1_c, + route_metadata, + ) + layout = Mxfp8BackwardLayout.from_request(request) + pool_capacity = 256 + fc1_preact = torch.empty( + pool_capacity, + 512, + dtype=torch.bfloat16, + ) + prepared = SimpleNamespace( + config=SimpleNamespace( + intermediate=256, + num_experts=2, + token_padding_block=128, + ), + kernel=SimpleNamespace( + token_comm=SimpleNamespace( + router_data_cta_count=1, + router_warps_per_cta=4, + ) + ), + pool_token_capacity=pool_capacity, + ) + + _stage_fc1_preact(request, layout, prepared, fc1_preact) + + staged = fc1_preact + gate, up = fc1_c.split(256, dim=1) + expected = torch.stack( + (gate.reshape(3, 8, 32), up.reshape(3, 8, 32)), + dim=2, + ).reshape(3, 512) + torch.testing.assert_close(staged[0], expected[0], rtol=0, atol=0) + torch.testing.assert_close(staged[1], expected[1], rtol=0, atol=0) + torch.testing.assert_close(staged[128], expected[2], rtol=0, atol=0) + assert staged[2:128].eq(0).all() + assert staged[129:].eq(0).all() + + +@pytest.mark.L0 +def test_mxfp8_backward_stages_sources_in_destination_ring_order(): + metadata = torch.tensor( + [[0, 0, 0, 0], [0, 1, 0, 0]], + dtype=torch.int64, + ) + fc1_c = torch.stack( + ( + torch.cat( + ( + torch.full((32,), 10, dtype=torch.bfloat16), + torch.full((32,), 11, dtype=torch.bfloat16), + ) + ), + torch.cat( + ( + torch.full((32,), 20, dtype=torch.bfloat16), + torch.full((32,), 21, dtype=torch.bfloat16), + ) + ), + ) + ) + request = SimpleNamespace( + config=SimpleNamespace( + ep_rank=1, + ep_size=2, + max_tokens_per_rank=1, + top_k=1, + ), + route_metadata=metadata, + fc1_c=fc1_c, + ) + layout = SimpleNamespace( + preact_row_lut=torch.tensor([[[0]], [[1]]], dtype=torch.int32) + ) + pool_capacity = 128 + fc1_preact = torch.empty( + pool_capacity, + 64, + dtype=torch.bfloat16, + ) + prepared = SimpleNamespace( + config=SimpleNamespace( + intermediate=32, + num_experts=1, + token_padding_block=128, + ), + kernel=SimpleNamespace( + token_comm=SimpleNamespace( + router_data_cta_count=1, + router_warps_per_cta=4, + ) + ), + pool_token_capacity=pool_capacity, + ) + + _stage_fc1_preact(request, layout, prepared, fc1_preact) + + staged = fc1_preact + # Destination rank 1 receives source rank 1 before wrapped source rank 0. + expected = torch.stack( + (fc1_c[:, :32], fc1_c[:, 32:]), + dim=1, + ).reshape(2, 64) + torch.testing.assert_close(staged[0], expected[1], rtol=0, atol=0) + torch.testing.assert_close(staged[1], expected[0], rtol=0, atol=0) + + +@pytest.mark.L0 +def test_mxfp8_backward_stages_source_routes_in_router_vector_order(): + metadata = torch.tensor( + [ + [0, 0, 0, 0], + [0, 0, 1, 1], + [0, 0, 2, 0], + [0, 0, 3, 1], + [0, 0, 4, 0], + ], + dtype=torch.int64, + ) + fc1_c = torch.arange(5, dtype=torch.bfloat16).view(5, 1).expand( + 5, + 64, + ).contiguous() + request = SimpleNamespace( + config=SimpleNamespace( + ep_rank=0, + ep_size=1, + max_tokens_per_rank=5, + top_k=2, + ), + route_metadata=metadata, + fc1_c=fc1_c, + ) + preact_row_lut = torch.full((1, 5, 2), -1, dtype=torch.int32) + preact_row_lut[ + metadata[:, 1], + metadata[:, 2], + metadata[:, 3], + ] = torch.arange(5, dtype=torch.int32) + layout = SimpleNamespace(preact_row_lut=preact_row_lut) + pool_capacity = 128 + fc1_preact = torch.empty( + pool_capacity, + 64, + dtype=torch.bfloat16, + ) + prepared = SimpleNamespace( + config=SimpleNamespace( + intermediate=32, + num_experts=1, + token_padding_block=128, + ), + kernel=SimpleNamespace( + token_comm=SimpleNamespace( + router_data_cta_count=1, + router_warps_per_cta=4, + ) + ), + pool_token_capacity=pool_capacity, + ) + + _stage_fc1_preact(request, layout, prepared, fc1_preact) + + staged = fc1_preact + # Int32 router loads four adjacent routes per thread, then stable-sorts by + # register round and lane: flat routes 0,4,8 precede 3,7. + assert staged[:5, 0].tolist() == [0, 2, 4, 1, 3] + + +@pytest.mark.L0 +def test_mxfp8_backward_workspace_regions_are_explicit_and_symmetric(): + requirements = WorkspaceRequirements.for_mxfp8( + _config(), + kernel_local_workspace_bytes=64, + kernel_shared_workspace_bytes=128, + backward_fc1_preact_bytes=1024, + backward_dprob_bytes=32, + backward_aux_data_bytes=512, + backward_aux_scale_bytes=256, + ) + symmetric = { + region.name: region for region in requirements.symmetric_regions + } + local = {region.name: region for region in requirements.local_regions} + + assert symmetric["backward_dprob"].nbytes == 32 + assert local["backward_fc1_preact"].nbytes == 1024 + assert local["backward_fc1_preact"].alignment == 128 + assert local["backward_aux_data"].nbytes == 512 + assert local["backward_aux_scale"].nbytes == 256 + + with pytest.raises(ValueError, match="must be enabled together"): + WorkspaceRequirements.for_mxfp8( + _config(), + kernel_local_workspace_bytes=64, + kernel_shared_workspace_bytes=128, + backward_dprob_bytes=32, + ) + + +@pytest.mark.L0 +def test_rubin_adapter_source_tracks_current_kernel_signatures(): + from cudnn.moe_ep._megamoe_backend.mxfp8 import ( + _backward_compile, + _compile, + ) + + forward_source = inspect.getsource(_compile.prepare_kernel) + backward_source = inspect.getsource(_backward_compile.prepare_backward_kernel) + runtime_source = inspect.getsource( + _backward_compile.build_backward_runtime_kwargs + ) + + assert "apply_topk_in_fc1=config.apply_topk_in_fc1" not in forward_source + assert "gate_up_clamp=config.gate_up_clamp" in backward_source + assert "dfc2_recompute=dfc2_recompute" in backward_source + assert "dfc2_col_output=dfc2_col_output" in backward_source + assert "enable_grad_y2_col_quant=enable_grad_y2_col_quant" in backward_source + assert '"fc1_preact":' in runtime_source + overflow_runtime_source = runtime_source.split( + '"overflow_flag":', + 1, + )[1].split('"dprob":', 1)[0] + assert "dynamic_layout=False" in overflow_runtime_source + for output_name in ( + "dprob", + "fc1_recompute", + "fc1_recompute_sf", + "fc1_col_output", + "fc1_col_output_sf", + "grad_y2", + "grad_y2_sf", + ): + assert f'"{output_name}":' in runtime_source + + +@pytest.mark.L0 +def test_stage_backward_exposes_fixed_aux_shapes_and_resets_symmetric_dprob(): + config = _config() + args = _inputs() + route_metadata = torch.tensor( + [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], + dtype=torch.int32, + ) + request = _validate_backward( + config, + torch.randn(2, 128), + args, + torch.randn(3, 512, dtype=torch.bfloat16), + route_metadata, + ) + layout = Mxfp8BackwardLayout.from_request(request) + aux_shapes = { + "dprob": (4, 2), + "fc1_recompute": (8, 256), + "fc1_recompute_sf": (1, 256), + "fc1_col_output": (8, 512), + "fc1_col_output_sf": (1, 512), + "grad_y2": (8, 128), + "grad_y2_sf": (32,), + } + kernel = SimpleNamespace( + token_comm=SimpleNamespace( + router_data_cta_count=1, + router_warps_per_cta=4, + ), + get_fc1_preact_shape=lambda: (256, 512), + get_aux_output_shapes=lambda: aux_shapes, + ) + prepared = SimpleNamespace( + config=SimpleNamespace( + max_tokens_per_rank=4, + hidden=128, + top_k=2, + intermediate=256, + num_experts=2, + combine_format="bf16", + token_padding_block=128, + ), + kernel=kernel, + pool_token_capacity=256, + pre_reduced_activation_offset=0, + pre_reduced_activation_bytes_per_token=4, + pre_reduced_activation_sf_offset=None, + pre_reduced_activation_sf_bytes_per_token=0, + local_workspace_zero_bytes=0, + shared_workspace_zero_bytes=0, + dfc2_recompute=False, + dfc2_col_output=False, + enable_grad_y2_col_quant=False, + ) + symmetric_dprob = torch.full((32,), 0x7F, dtype=torch.uint8) + resources = SimpleNamespace( + workspace=SimpleNamespace( + symmetric={ + "activation_data": torch.empty(4 * 128, dtype=torch.uint8), + "activation_scale": torch.empty(4 * 16, dtype=torch.uint8), + "topk_weights": torch.empty(4 * 2 * 4, dtype=torch.uint8), + "output_data": torch.empty(4 * 128 * 2, dtype=torch.uint8), + "backward_dprob": symmetric_dprob, + "kernel_shared_workspace": torch.empty( + 64, + dtype=torch.uint8, + ), + }, + local={ + "topk_idx": torch.empty(4 * 2 * 4, dtype=torch.uint8), + "overflow_flag": torch.empty(4, dtype=torch.uint8), + "backward_fc1_preact": torch.empty( + 256 * 512 * 2, + dtype=torch.uint8, + ), + "backward_aux_data": torch.empty( + 8 * 512, + dtype=torch.uint8, + ), + "backward_aux_scale": torch.empty( + 512, + dtype=torch.uint8, + ), + "kernel_local_workspace": torch.empty( + 64, + dtype=torch.uint8, + ), + }, + ) + ) + + inputs = stage_backward(request, layout, prepared, resources) + + assert inputs.fc1_preact.shape == (256, 512) + assert inputs.fc1_preact.dtype is torch.bfloat16 + assert inputs.dprob.shape == (4, 2) + assert inputs.dprob.dtype is torch.float32 + assert inputs.dprob.eq(0).all() + assert inputs.fc1_recompute.shape == aux_shapes["fc1_recompute"] + assert inputs.fc1_recompute.dtype is torch.float8_e4m3fn + assert inputs.fc1_recompute_sf.shape == aux_shapes["fc1_recompute_sf"] + assert inputs.fc1_recompute_sf.dtype is torch.float8_e8m0fnu + assert inputs.fc1_col_output.shape == aux_shapes["fc1_col_output"] + assert inputs.fc1_col_output_sf.shape == aux_shapes["fc1_col_output_sf"] + assert ( + inputs.fc1_recompute.data_ptr() + == inputs.fc1_col_output.data_ptr() + ) + assert ( + inputs.fc1_recompute_sf.data_ptr() + == inputs.fc1_col_output_sf.data_ptr() + ) + assert inputs.grad_y2.shape == aux_shapes["grad_y2"] + assert inputs.grad_y2_sf.shape == aux_shapes["grad_y2_sf"] + assert inputs.grad_y2.data_ptr() == inputs.fc1_recompute.data_ptr() + assert inputs.grad_y2_sf.data_ptr() == inputs.fc1_recompute_sf.data_ptr() + + operands_config = replace( + config, + backward_wgrad_mode="operands", + token_padding_size=256, + ) + operands_request = replace(request, config=operands_config) + operands_aux_shapes = { + "dprob": (4, 2), + "fc1_recompute": (512, 256), + "fc1_recompute_sf": (8, 256), + "fc1_col_output": (512, 512), + "fc1_col_output_sf": (8, 512), + "grad_y2": (512, 128), + "grad_y2_sf": (256 // 32 * 128,), + } + operands_kernel = SimpleNamespace( + token_comm=kernel.token_comm, + get_fc1_preact_shape=lambda: (512, 512), + get_aux_output_shapes=lambda: operands_aux_shapes, + ) + operands_prepared = SimpleNamespace( + **{ + **vars(prepared), + "config": SimpleNamespace( + **{ + **vars(prepared.config), + "token_padding_block": 256, + } + ), + "kernel": operands_kernel, + "pool_token_capacity": 512, + "dfc2_recompute": True, + "dfc2_col_output": True, + "enable_grad_y2_col_quant": True, + } + ) + operands_local = dict(resources.workspace.local) + operands_local["backward_fc1_preact"] = torch.empty( + 512 * 512 * 2, + dtype=torch.uint8, + ) + operands_resources = SimpleNamespace( + workspace=SimpleNamespace( + symmetric=resources.workspace.symmetric, + local=operands_local, + ) + ) + + operand_inputs = stage_backward( + operands_request, + Mxfp8BackwardLayout.from_request(operands_request), + operands_prepared, + operands_resources, + ) + + assert operand_inputs.fc1_recompute.shape == (512, 256) + assert operand_inputs.fc1_col_output.shape == (512, 512) + assert operand_inputs.grad_y2.shape == (512, 128) + assert operand_inputs.grad_y2_sf.shape == ( + 256 // 32 * 128, + ) + assert ( + operand_inputs.fc1_recompute.data_ptr() + != operands_local["backward_aux_data"].data_ptr() + ) + assert ( + operand_inputs.fc1_col_output.data_ptr() + != operands_local["backward_aux_data"].data_ptr() + ) + assert operand_inputs.fc1_recompute.eq(0).all() + assert operand_inputs.fc1_col_output.eq(0).all() + assert operand_inputs.grad_y2.eq(0).all() + assert operand_inputs.fc1_preact[128:256].eq(0).all() + assert torch.equal( + operand_inputs.fc1_preact[256], + torch.stack( + ( + request.fc1_c[2, :256].reshape(8, 32), + request.fc1_c[2, 256:].reshape(8, 32), + ), + dim=1, + ).reshape(512), + ) + + +@pytest.mark.L0 +@pytest.mark.parametrize("apply_topk_in_fc1", [True, False]) +def test_mxfp8_backward_recomputes_semantic_grad_topk_weights( + apply_topk_in_fc1, +): + args = list(_inputs()) + fc2_weight = torch.zeros_like(args[2]) + fc2_weight[0, 0, 0] = 2 + fc2_weight[1, 0, 0] = 4 + args[2] = fc2_weight + config = _config(apply_topk_in_fc1=apply_topk_in_fc1) + route_metadata = torch.tensor( + [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], + dtype=torch.int32, + ) + fc1_c = torch.zeros(3, 512, dtype=torch.bfloat16) + fc1_c[:, 0] = 1 + fc1_c[:, 256] = 1 + request = _validate_backward( + config, + torch.randn(2, 128), + args, + fc1_c, + route_metadata, + ) + redispatched_grad_output = torch.zeros(3, 128) + redispatched_grad_output[:, 0] = torch.tensor([1.0, 2.0, 3.0]) + + grad_topk = return_grad_topk_weights( + request, + redispatched_grad_output, + ) + silu_one = torch.sigmoid(torch.tensor(1.0)) + torch.testing.assert_close( + grad_topk, + silu_one * torch.tensor([[2.0, 0.0], [12.0, 4.0]]), + rtol=1e-6, + atol=1e-6, + ) + + +@pytest.mark.L0 +def test_mxfp8_backward_dprob_recompute_applies_gate_up_clamp(): + args = list(_inputs()) + fc2_weight = torch.zeros_like(args[2]) + fc2_weight[0, 0, 0] = 2 + fc2_weight[1, 0, 0] = 4 + args[2] = fc2_weight + config = _config(gate_up_clamp=0.5) + route_metadata = torch.tensor( + [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], + dtype=torch.int32, + ) + fc1_c = torch.zeros(3, 512, dtype=torch.bfloat16) + fc1_c[:, 0] = 2 + fc1_c[:, 256] = 2 + request = _validate_backward( + config, + torch.randn(2, 128), + args, + fc1_c, + route_metadata, + ) + redispatched_grad_output = torch.zeros(3, 128) + redispatched_grad_output[:, 0] = torch.tensor([1.0, 2.0, 3.0]) + + grad_topk = return_grad_topk_weights( + request, + redispatched_grad_output, + ) + + clamped_hidden = 0.25 * torch.sigmoid(torch.tensor(0.5)) + torch.testing.assert_close( + grad_topk, + clamped_hidden * torch.tensor([[2.0, 0.0], [12.0, 4.0]]), + rtol=1e-6, + atol=1e-6, + ) + + +@pytest.mark.L0 +def test_mxfp8_grad_output_redispatch_uses_public_route_order(): + config = _config() + args = _inputs() + route_metadata = torch.tensor( + [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], + dtype=torch.int32, + ) + request = _validate_backward( + config, + torch.randn(2, 128), + args, + torch.randn(3, 512, dtype=torch.bfloat16), + route_metadata, + ) + + actual = Mxfp8BackwardRedispatch(request).run() + expected_rows = torch.tensor([0, 1, 1], dtype=torch.int64) + + torch.testing.assert_close( + actual.grad_output, + request.grad_output.index_select(0, expected_rows).float(), + rtol=0, + atol=0, + ) + + +@pytest.mark.L0 +def test_moe_ep_backward_delegates_validated_request(monkeypatch): + import cudnn.moe_ep._backend as backend_seam + + args = _inputs() + grad_output = torch.randn(2, 128) + fc1_c = torch.randn(3, 512, dtype=torch.bfloat16) + route_metadata = torch.tensor( + [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], + dtype=torch.int32, + ) + expected = ( + torch.empty(2, 128, dtype=torch.float32), + torch.empty(2, 2, dtype=torch.float32), + ) + + class Backend: + request = None + + def backward(self, request): + self.request = request + return expected + + def close(self): + pass + + instance = Backend() + monkeypatch.setattr(backend_seam, "validate_config", lambda config: None) + monkeypatch.setattr( + backend_seam, + "validate_backward_request", + lambda request: None, + ) + monkeypatch.setattr( + backend_seam, + "create_backend", + lambda config, device: instance, + ) + + operator = MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=4, + generate_c=True, + ) + actual = operator.backward( + grad_output, + *args[1:], + fc1_c, + route_metadata, + ) + + assert actual is expected + assert len(actual) == 2 + assert instance.request is not None + assert instance.request.local_routes == 3 + assert instance.request.fc1_c is fc1_c + assert instance.request.fc1_weight is args[1] + + +@pytest.mark.L0 +def test_moe_ep_backward_accepts_explicit_stashes_in_reordered_calls( + monkeypatch, +): + import cudnn.moe_ep._backend as backend_seam + + calls = [] + + class Backend: + def backward(self, request): + calls.append(request) + marker = request.fc1_c[0, 0].float().reshape(1) + return marker, marker + + def close(self): + pass + + backend = Backend() + monkeypatch.setattr(backend_seam, "validate_config", lambda config: None) + monkeypatch.setattr( + backend_seam, + "validate_backward_request", + lambda request: None, + ) + monkeypatch.setattr( + backend_seam, + "create_backend", + lambda config, device: backend, + ) + operator = MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=4, + generate_c=True, + ) + args = _inputs() + route_metadata = torch.tensor( + [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], + dtype=torch.int32, + ) + stash_a = torch.full((3, 512), 1.0, dtype=torch.bfloat16) + stash_b = torch.full((3, 512), 2.0, dtype=torch.bfloat16) + + markers = [] + for stash in (stash_a, stash_b, stash_b, stash_a): + result = operator.backward( + torch.randn(2, 128), + *args[1:], + stash, + route_metadata, + ) + markers.append(result[0].item()) + + assert markers == [1.0, 2.0, 2.0, 1.0] + assert calls[0].fc1_c is stash_a + assert calls[1].fc1_c is stash_b + assert calls[2].fc1_c is stash_b + assert calls[3].fc1_c is stash_a + + +@pytest.mark.L0 +def test_moe_ep_backward_requires_generate_c(): + args = _inputs() + operator = MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=4, + generate_c=False, + ) + + with pytest.raises(RuntimeError, match="generate_c=True"): + operator.backward( + torch.randn(2, 128), + *args[1:], + torch.randn(3, 512, dtype=torch.bfloat16), + torch.zeros(3, 4, dtype=torch.int32), + ) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"output_format": "mxfp8"}, "output_format='bf16'"), + ({"ep_size": 8}, "supports EP1/EP2/EP4"), + ({"apply_topk_in_fc1": False}, "apply_topk_in_fc1=True"), + ], +) +def test_mxfp8_backward_capability_rejects_unsupported_config( + monkeypatch, + overrides, + message, +): + config = _config(**overrides) + args = _inputs() + request = _validate_backward( + config, + torch.randn(2, 128), + args, + torch.randn(3, 512, dtype=torch.bfloat16), + torch.zeros(3, 4, dtype=torch.int32), + ) + monkeypatch.setattr(_capability, "_validate_device", lambda device: None) + monkeypatch.setattr( + _capability, + "_is_cuda_stream_capturing", + lambda device: False, + ) + + with pytest.raises( + NotImplementedError, + match=message, + ): + _capability.validate_backward_request(request) + + +@pytest.mark.L0 +def test_mxfp8_backward_capability_accepts_gate_up_clamp(monkeypatch): + config = _config(gate_up_clamp=1.0) + args = _inputs() + request = _validate_backward( + config, + torch.randn(2, 128), + args, + torch.randn(3, 512, dtype=torch.bfloat16), + torch.zeros(3, 4, dtype=torch.int32), + ) + monkeypatch.setattr(_capability, "_validate_device", lambda device: None) + monkeypatch.setattr( + _capability, + "_is_cuda_stream_capturing", + lambda device: False, + ) + + _capability.validate_backward_request(request) + + +@pytest.mark.L0 +def test_mxfp8_backward_capability_rejects_cuda_graph_capture(monkeypatch): + config = _config() + args = _inputs() + request = _validate_backward( + config, + torch.randn(2, 128), + args, + torch.randn(3, 512, dtype=torch.bfloat16), + torch.zeros(3, 4, dtype=torch.int32), + ) + monkeypatch.setattr(_capability, "_validate_device", lambda device: None) + monkeypatch.setattr( + _capability, + "_is_cuda_stream_capturing", + lambda device: True, + ) + + with pytest.raises(NotImplementedError, match="CUDA graph capture"): + _capability.validate_backward_request(request) + + +@pytest.mark.L0 +def test_mxfp8_backward_delegates_to_explicit_executor(monkeypatch): + import cudnn.moe_ep._megamoe_backend.mxfp8._backend as backend_module + + config = _config() + args = _inputs() + request = _validate_backward( + config, + torch.randn(2, 128), + args, + torch.randn(3, 512, dtype=torch.bfloat16), + torch.tensor( + [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], + dtype=torch.int32, + ), + ) + expected = tuple(torch.empty(0) for _ in range(2)) + + class Executor: + def __init__(self, actual_config, actual_device): + assert actual_config is config + assert actual_device == torch.device("cpu") + + def run(self, actual_request): + assert actual_request is request + return expected + + def close(self): + pass + + monkeypatch.setattr( + backend_module, + "Mxfp8BackwardExecutor", + Executor, + ) + monkeypatch.setattr( + torch.cuda, + "is_current_stream_capturing", + lambda: False, + ) + class Stream: + def wait_event(self, event): + del event + + class Event: + def record(self, stream): + del stream + + monkeypatch.setattr(torch.cuda, "current_stream", lambda device: Stream()) + monkeypatch.setattr(torch.cuda, "Event", Event) + backend = Mxfp8Backend(config, torch.device("cpu")) + assert backend.backward(request) is expected + + +# Single-rank and distributed backward numerical parity. + + +def _make_reentrant_case_b(args, device): + activation = quantize_mxfp8( + args[0].dequantize(dtype=torch.float32) + 0.25, + axis=1, + ) + topk_idx = torch.tensor( + [[0, 0], [-1, -1], [0, -1], [0, 0], [0, -1]], + dtype=torch.int32, + device=device, + ) + topk_weights = torch.tensor( + [ + [0.625, 0.375], + [0.0, 0.0], + [1.0, 0.0], + [0.75, 0.25], + [1.0, 0.0], + ], + dtype=torch.bfloat16, + device=device, + ) + return activation, args[1], args[2], topk_idx, topk_weights + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +@pytest.mark.parametrize("combine_format", ["bf16", "mxfp8"]) +@pytest.mark.parametrize( + "gate_up_clamp", + [None, 0.5], + ids=["unclamped", "clamped"], +) +def test_mxfp8_backward_ep1_matches_reference_and_resets_workspace( + combine_format, + gate_up_clamp, +): + device = _sm107_device() + args = make_forward_inputs(device) + config = _forward_config( + generate_c=True, + combine_format=combine_format, + gate_up_clamp=gate_up_clamp, + ) + reference = _reference_backward(config) + grad_output = _grad_output(device, args[3].shape[0], seed=20260817) + + with MoeEp(**config) as op: + _, fc1_c, route_metadata = op(*args) + stash = (fc1_c, route_metadata) + expected = _expected_backward(reference, grad_output, args, stash) + + first = op.backward(grad_output, *args[1:], *stash) + second = op.backward(grad_output, *args[1:], *stash) + torch.cuda.synchronize(device) + + _assert_backward_matches(first, expected, args[3]) + _assert_backward_matches(second, expected, args[3]) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_mxfp8_backward_ep1_uses_explicit_stash_after_reordered_forwards(): + device = _sm107_device() + args_a = make_forward_inputs(device) + args_b = _make_reentrant_case_b(args_a, device) + config = _forward_config(generate_c=True) + reference = _reference_backward(config) + grad_a = _grad_output(device, args_a[3].shape[0], seed=20260818) + grad_b = _grad_output(device, args_b[3].shape[0], seed=20260819) + + with MoeEp(**config) as op: + _, fc1_c_a, metadata_a = op(*args_a) + _, fc1_c_b, metadata_b = op(*args_b) + cases = ( + (grad_b, args_b, (fc1_c_b, metadata_b)), + (grad_a, args_a, (fc1_c_a, metadata_a)), + (grad_b, args_b, (fc1_c_b, metadata_b)), + (grad_a, args_a, (fc1_c_a, metadata_a)), + ) + results = [] + for grad_output, args, stash in cases: + expected = _expected_backward(reference, grad_output, args, stash) + actual = op.backward(grad_output, *args[1:], *stash) + results.append((actual, expected, args[3])) + torch.cuda.synchronize(device) + + for actual, expected, topk_idx in results: + _assert_backward_matches(actual, expected, topk_idx) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +@pytest.mark.parametrize("world_size", [2, 4], ids=["ep2", "ep4"]) +@pytest.mark.parametrize("combine_format", ["bf16", "mxfp8"]) +def test_mxfp8_backward_multi_gpu_matches_reference( + world_size, + combine_format, + tmp_path, +): + _require_distributed_sm107(world_size) + os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") + init_file = ( + tmp_path + / f"{combine_format}_combine_mxfp8_backward_ep{world_size}.init" + ) + mp.spawn( + _distributed_backward_worker, + args=(world_size, str(init_file), combine_format), + nprocs=world_size, + join=True, + ) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_mxfp8_backward_ep2_gate_up_clamp_matches_reference(tmp_path): + world_size = 2 + _require_distributed_sm107(world_size) + os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") + init_file = tmp_path / "bf16_combine_clamped_mxfp8_backward_ep2.init" + mp.spawn( + _distributed_backward_worker, + args=(world_size, str(init_file), "bf16", 0.5), + nprocs=world_size, + join=True, + ) + + +@pytest.mark.L0 +def test_forward_and_backward_mxfp8_combine_are_direct_fp32(): + generator = torch.Generator().manual_seed(20260820) + accumulator = torch.randn(4, 128, generator=generator) * 3.25 + + backward = backward_combine_round_trip( + accumulator, + MoeFormat.MXFP8, + ) + forward = forward_combine_round_trip( + accumulator, + MoeFormat.MXFP8, + ) + + torch.testing.assert_close(backward, forward, rtol=0, atol=0) diff --git a/test/python/moe_ep/test_moe_ep_cutedsl_grad_y2_source.py b/test/python/moe_ep/test_moe_ep_cutedsl_grad_y2_source.py new file mode 100644 index 000000000..221e0cf15 --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_cutedsl_grad_y2_source.py @@ -0,0 +1,63 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Source-only contracts for the upstream grad_y2 and dFC2 scale layout.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + + +_ROOT = Path(__file__).resolve().parents[3] +_CUTEDSL = ( + _ROOT + / "python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin" + / "training/mega" +) +_DGLU = _CUTEDSL / "bwd_dglu/dglu_mxfp8_mega_moe_kernel.py" +_DGLU_EPILOGUE = _CUTEDSL / "bwd_dglu/dglu_mxfp8_fc12_epilogue.py" + + +def _source_and_tree(path: Path) -> tuple[str, ast.Module]: + source = path.read_text(encoding="utf-8") + return source, ast.parse(source, filename=str(path)) + + +@pytest.mark.L0 +def test_dglu_source_exports_upstream_grad_y2_col_quant(): + source, tree = _source_and_tree(_DGLU) + + assert tree is not None + for contract in ( + "enable_grad_y2_col_quant", + "num_ctas_grad_y2_col_quant", + "grad_y2_sizes_region", + "_snapshot_grad_y2_expert_sizes", + "grad_y2_col_quant", + "grad_y2: cute.Tensor", + "grad_y2_sf: cute.Tensor", + ): + assert contract in source + assert source.index( + "self._snapshot_grad_y2_expert_sizes(tidx)" + ) < source.index("self.token_comm.reset_tail()") + assert source.index("self._topk_reduce(") < source.index( + "self.grad_y2_col_quant(" + ) + + +@pytest.mark.L0 +def test_dfc2_scale_source_uses_upstream_mn_major_atoms(): + source, tree = _source_and_tree(_DGLU_EPILOGUE) + + assert tree is not None + assert "def _stg_col_sf_atom_value(" in source + assert "MN-major 128-column × 4-token-block atom" in source + assert "atom_idx * Int64(512)" in source + assert "Int64(hidden_lane) * Int64(16)" in source + assert "Int64(hidden_bank) * Int64(4)" in source + assert "Int64(token_bank)" in source + assert source.count("self._stg_col_sf_atom_value(") >= 2 diff --git a/test/python/moe_ep/test_moe_ep_wgrad_contract.py b/test/python/moe_ep/test_moe_ep_wgrad_contract.py new file mode 100644 index 000000000..4e307b5d6 --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_wgrad_contract.py @@ -0,0 +1,1170 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Focused L0 tests for the public MoeEP wgrad operand contract.""" + +from __future__ import annotations + +import os +from dataclasses import fields, replace +from types import SimpleNamespace + +import pytest +import torch +import torch.multiprocessing as mp + +from cudnn.moe_ep import ( + MoeEp, + MoeEpWgradForwardStash, + MoeEpWgradOperands, +) +from cudnn.moe_ep._megamoe_backend import _capability +from cudnn.moe_ep._megamoe_backend.mxfp8._backward_launch import ( + Mxfp8DgluResult, +) +from cudnn.moe_ep._megamoe_backend.mxfp8._backward_wgrad_export import ( + export_wgrad_operands, +) +from cudnn.moe_ep._megamoe_backend.mxfp8._config import Mxfp8KernelConfig +from cudnn.moe_ep._megamoe_backend.mxfp8._stash import ( + Mxfp8ForwardStash as Mxfp8ForwardStashOwner, +) +from cudnn.moe_ep._megamoe_backend.mxfp8._wgrad_layout import ( + assemble_dfc2_atom_scales, + assemble_discrete_col_requant_scales, + assemble_plain_col_scales, +) +from cudnn.moe_ep._validation import validate_backward +from moe_ep.moe_ep_backward_support import _dense_wgrads_from_operands +from moe_ep.moe_ep_distributed_workers import ( + _distributed_wgrad_worker, + _run_wgrad_operand_case, +) +from moe_ep.moe_ep_forward_support import ( + _require_distributed_sm107, + _sm107_device, +) +from moe_ep.moe_ep_reference import ( + MoeEpReference, + MoeFormat, + WgradOperandsReference, + quantize_blockwise, +) + + +def _inputs(): + activation = torch.randn(2, 128, dtype=torch.bfloat16) + fc1_weight = torch.randn(2, 128, 512, dtype=torch.bfloat16) + fc2_weight = torch.randn(2, 256, 128, dtype=torch.bfloat16) + topk_idx = torch.tensor([[0, -1], [1, 0]], dtype=torch.int32) + topk_weights = torch.randn(2, 2, dtype=torch.float32) + return activation, fc1_weight, fc2_weight, topk_idx, topk_weights + + +def _route_metadata(): + return torch.tensor( + [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], + dtype=torch.int32, + ) + + +def _forward_stash(route_metadata=None): + if route_metadata is None: + route_metadata = _route_metadata() + return MoeEpWgradForwardStash( + fc1_a=torch.empty(128, 512, dtype=torch.float8_e4m3fn), + fc1_sfa=torch.empty(128, 16, dtype=torch.float8_e8m0fnu), + expert_offsets=torch.tensor([256, 512], dtype=torch.int32), + valid_route_counts=torch.tensor([2, 1], dtype=torch.int32), + route_metadata=route_metadata.clone(), + ) + + +def _operator(**overrides): + kwargs = { + "num_experts": 2, + "hidden_size": 128, + "intermediate_size": 256, + "top_k": 2, + "max_tokens_per_rank": 4, + "generate_c": True, + } + kwargs.update(overrides) + if ( + kwargs.get("backward_wgrad_mode") == "operands" + and "token_padding_size" not in overrides + ): + kwargs["token_padding_size"] = 256 + return MoeEp(**kwargs) + + +@pytest.mark.L0 +def test_wgrad_types_are_public_and_have_stable_fields(): + from cudnn import ( + MoeEpWgradForwardStash as TopLevelForwardStash, + MoeEpWgradOperands as TopLevelOperands, + ) + + assert TopLevelForwardStash is MoeEpWgradForwardStash + assert TopLevelOperands is MoeEpWgradOperands + assert [field.name for field in fields(MoeEpWgradForwardStash)] == [ + "fc1_a", + "fc1_sfa", + "expert_offsets", + "valid_route_counts", + "route_metadata", + ] + assert [field.name for field in fields(MoeEpWgradOperands)] == [ + "fc1_a", + "fc1_sfa", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + "expert_offsets", + "valid_route_counts", + "route_metadata", + ] + + +@pytest.mark.L0 +def test_wgrad_mode_is_opt_in_and_requires_generate_c(): + with _operator(generate_c=False) as operator: + assert operator.backward_wgrad_mode == "none" + assert operator._forward_config.backward_wgrad_mode == "none" + + with pytest.raises(ValueError, match="must be 'none' or 'operands'"): + _operator(backward_wgrad_mode="weights") + with pytest.raises(ValueError, match="requires generate_c=True"): + _operator( + generate_c=False, + backward_wgrad_mode="operands", + ) + with pytest.raises(ValueError, match="requires token_padding_size=256"): + _operator( + backward_wgrad_mode="operands", + token_padding_size=128, + ) + with pytest.raises(ValueError, match="requires sf_padding_size=128"): + _operator( + backward_wgrad_mode="operands", + sf_padding_size=256, + ) + + with _operator(backward_wgrad_mode="operands") as operator: + assert operator.backward_wgrad_mode == "operands" + assert operator._forward_config.backward_wgrad_mode == "operands" + + +@pytest.mark.L0 +def test_validate_backward_checks_wgrad_stash_layout_and_route_identity(): + args = _inputs() + route_metadata = _route_metadata() + stash = _forward_stash(route_metadata) + with _operator(backward_wgrad_mode="operands") as operator: + with pytest.raises(TypeError, match="must be a MoeEpWgradForwardStash"): + validate_backward( + operator._forward_config, + torch.randn(2, 128), + *args[1:], + torch.randn(3, 512, dtype=torch.bfloat16), + route_metadata, + ) + + request = validate_backward( + operator._forward_config, + torch.randn(2, 128), + *args[1:], + torch.randn(3, 512, dtype=torch.bfloat16), + route_metadata, + wgrad_forward_stash=stash, + ) + assert request.wgrad_forward_stash is stash + + wrong_scale_shape = replace( + stash, + fc1_sfa=torch.empty(128, 15, dtype=torch.float8_e8m0fnu), + ) + with pytest.raises(ValueError, match="fc1_sfa shape must be"): + validate_backward( + operator._forward_config, + torch.randn(2, 128), + *args[1:], + torch.randn(3, 512, dtype=torch.bfloat16), + route_metadata, + wgrad_forward_stash=wrong_scale_shape, + ) + + wrong_a_stride = replace( + stash, + fc1_a=torch.empty( + 512, + 128, + dtype=torch.float8_e4m3fn, + ).transpose(0, 1), + ) + with pytest.raises(ValueError, match=r"compact \(K, 1\) strides"): + validate_backward( + operator._forward_config, + torch.randn(2, 128), + *args[1:], + torch.randn(3, 512, dtype=torch.bfloat16), + route_metadata, + wgrad_forward_stash=wrong_a_stride, + ) + + wrong_scale_stride = replace( + stash, + fc1_sfa=torch.empty( + 16, + 128, + dtype=torch.float8_e8m0fnu, + ).transpose(0, 1), + ) + with pytest.raises(ValueError, match="fc1_sfa must be contiguous"): + validate_backward( + operator._forward_config, + torch.randn(2, 128), + *args[1:], + torch.randn(3, 512, dtype=torch.bfloat16), + route_metadata, + wgrad_forward_stash=wrong_scale_stride, + ) + + noncanonical_offsets = replace( + stash, + expert_offsets=torch.tensor([512, 768], dtype=torch.int32), + ) + with pytest.raises(ValueError, match="canonical 256-row padding"): + validate_backward( + operator._forward_config, + torch.randn(2, 128), + *args[1:], + torch.randn(3, 512, dtype=torch.bfloat16), + route_metadata, + wgrad_forward_stash=noncanonical_offsets, + ) + + wrong_identity = _forward_stash(route_metadata.flip(0)) + with pytest.raises(ValueError, match="route identity"): + validate_backward( + operator._forward_config, + torch.randn(2, 128), + *args[1:], + torch.randn(3, 512, dtype=torch.bfloat16), + route_metadata, + wgrad_forward_stash=wrong_identity, + ) + + wrong_counts = MoeEpWgradForwardStash( + stash.fc1_a, + stash.fc1_sfa, + stash.expert_offsets, + torch.tensor([1, 2], dtype=torch.int32), + stash.route_metadata, + ) + with pytest.raises(ValueError, match="do not match route_metadata"): + validate_backward( + operator._forward_config, + torch.randn(2, 128), + *args[1:], + torch.randn(3, 512, dtype=torch.bfloat16), + route_metadata, + wgrad_forward_stash=wrong_counts, + ) + + +@pytest.mark.L0 +def test_default_mode_rejects_wgrad_stash_without_changing_default_contract(): + args = _inputs() + route_metadata = _route_metadata() + with _operator() as operator: + with pytest.raises(ValueError, match="only accepted"): + validate_backward( + operator._forward_config, + torch.randn(2, 128), + *args[1:], + torch.randn(3, 512, dtype=torch.bfloat16), + route_metadata, + wgrad_forward_stash=_forward_stash(route_metadata), + ) + + +@pytest.mark.L0 +def test_wgrad_mode_is_backend_capable_and_enables_forward_col_quant( + monkeypatch, +): + with _operator(backward_wgrad_mode="operands") as operator: + monkeypatch.setattr( + _capability, + "_validate_device", + lambda device: pytest.fail("device capability queried"), + ) + _capability.validate_config(operator._forward_config) + kernel_config = Mxfp8KernelConfig.from_forward_config( + operator._forward_config + ) + + assert kernel_config.enable_col_quant is True + assert kernel_config.token_padding_block == 256 + + +@pytest.mark.L0 +def test_forward_col_quant_runtime_uses_static_cute_layout(monkeypatch): + from cudnn.moe_ep._megamoe_backend.mxfp8 import _launch + + tensors = { + name: object() + for name in ( + "activation", + "activation_sf", + "topk_indices", + "topk_scores", + "fc1_weight", + "fc1_weight_sf", + "fc2_weight", + "fc2_weight_sf", + "fc1_c", + "output_data", + "col_quant_data", + "col_quant_sf", + "overflow_flag", + "local_workspace", + "shared_workspace", + ) + } + calls = {} + + def fake_to_cute( + tensor, + assumed_align=16, + *, + dynamic_layout=True, + ): + calls[id(tensor)] = (assumed_align, dynamic_layout) + return tensor + + monkeypatch.setattr(_launch, "_to_cute", fake_to_cute) + monkeypatch.setattr(_launch, "_to_cute_ptr", lambda tensor: tensor) + inputs = SimpleNamespace( + activation=tensors["activation"], + activation_sf=tensors["activation_sf"], + topk_indices=tensors["topk_indices"], + topk_scores=tensors["topk_scores"], + weights=SimpleNamespace( + fc1_weight=tensors["fc1_weight"], + fc1_weight_sf=tensors["fc1_weight_sf"], + fc2_weight=tensors["fc2_weight"], + fc2_weight_sf=tensors["fc2_weight_sf"], + ), + fc1_c=tensors["fc1_c"], + output_data=tensors["output_data"], + col_quant_data=tensors["col_quant_data"], + col_quant_sf=tensors["col_quant_sf"], + overflow_flag=tensors["overflow_flag"], + local_workspace=tensors["local_workspace"], + shared_workspace=tensors["shared_workspace"], + ) + resources = SimpleNamespace( + runtime=SimpleNamespace( + current_stream=lambda: SimpleNamespace(cuda_stream=0) + ), + workspace=SimpleNamespace( + peer_mapping=SimpleNamespace( + to_sym_buffer_host=lambda: object() + ) + ), + ) + + _launch.build_runtime_kwargs(inputs, resources) + + assert calls[id(tensors["col_quant_data"])] == (128, False) + assert calls[id(tensors["col_quant_sf"])] == (16, False) + assert calls[id(tensors["overflow_flag"])] == (4, False) + + +@pytest.mark.L0 +def test_opt_in_forward_and_backward_results_are_backend_representable( + monkeypatch, +): + import cudnn.moe_ep._backend as backend_seam + + args = _inputs() + route_metadata = _route_metadata() + fc1_c = torch.randn(3, 512, dtype=torch.bfloat16) + stash = _forward_stash(route_metadata) + forward_result = (torch.randn(2, 128), fc1_c, route_metadata, stash) + operands = MoeEpWgradOperands( + fc1_a=stash.fc1_a, + fc1_sfa=stash.fc1_sfa, + fc1_b=torch.empty(512, 512, dtype=torch.float8_e4m3fn), + fc1_sfb=torch.empty(512, 16, dtype=torch.float8_e8m0fnu), + fc2_a=torch.empty(256, 512, dtype=torch.float8_e4m3fn), + fc2_sfa=torch.empty(256, 16, dtype=torch.float8_e8m0fnu), + fc2_b=torch.empty(512, 128, dtype=torch.float8_e4m3fn), + fc2_sfb=torch.empty(128, 16, dtype=torch.float8_e8m0fnu), + expert_offsets=stash.expert_offsets, + valid_route_counts=stash.valid_route_counts, + route_metadata=stash.route_metadata, + ) + backward_result = ( + torch.randn(2, 128), + torch.randn(2, 2), + operands, + ) + + class Backend: + backward_request = None + + def forward(self, request): + return forward_result + + def backward(self, request): + self.backward_request = request + return backward_result + + def close(self): + pass + + backend = Backend() + monkeypatch.setattr(backend_seam, "validate_config", lambda config: None) + monkeypatch.setattr(backend_seam, "validate_request", lambda request: None) + monkeypatch.setattr( + backend_seam, + "validate_backward_request", + lambda request: None, + ) + monkeypatch.setattr( + backend_seam, + "create_backend", + lambda config, device: backend, + ) + + with _operator(backward_wgrad_mode="operands") as operator: + assert operator(*args) is forward_result + actual = operator.backward( + torch.randn(2, 128), + *args[1:], + fc1_c, + route_metadata, + wgrad_forward_stash=stash, + ) + + assert actual is backward_result + assert backend.backward_request.wgrad_forward_stash is stash + + +def _blocked_reference(raw: torch.Tensor) -> torch.Tensor: + rows, columns = raw.shape + if columns == 0: + return raw.new_empty((0,)) + padded_rows = (rows + 127) // 128 * 128 + padded_columns = (columns + 3) // 4 * 4 + padded = torch.full( + (padded_rows, padded_columns), + 127, + dtype=torch.uint8, + device=raw.device, + ) + padded[:rows, :columns] = raw + return ( + padded.view(padded_rows // 128, 128, padded_columns // 4, 4) + .permute(0, 2, 1, 3) + .reshape(-1, 4, 32, 4) + .transpose(1, 2) + .reshape(-1) + ) + + +@pytest.mark.L0 +def test_col_requant_scales_accept_upstream_hidden_atom_major_order(): + non_k = 256 + valid_counts = (33, 0, 129) + padded_ends = (256, 256, 512) + source0 = torch.arange(non_k * 4, dtype=torch.int32).to(torch.uint8).reshape( + non_k, + 4, + ) + source2 = ( + torch.arange(non_k * 8, dtype=torch.int32) + 37 + ).to(torch.uint8).reshape(non_k, 8) + target0 = torch.full((non_k, 8), 127, dtype=torch.uint8) + target0[:, :4] = source0 + blocked0 = _blocked_reference(target0) + blocked2 = _blocked_reference(source2) + + # Upstream col requant stores hidden atoms before token atoms inside each + # expert. Its 128-row SF padding is expanded to the 256-row data padding. + packed = torch.cat( + (_blocked_reference(source0), _blocked_reference(source2)) + ) + actual = assemble_discrete_col_requant_scales( + packed, + valid_counts, + padded_ends, + non_k, + 128, + ) + expected = torch.cat((blocked0, blocked2)).reshape(non_k, 16) + + assert actual.shape == (non_k, 16) + assert torch.equal(actual.view(torch.uint8), expected) + + +@pytest.mark.L0 +def test_dfc2_atom_scales_reorder_token_major_atoms_per_expert(): + non_k = 256 + valid_counts = (33, 0, 129) + padded_ends = (256, 256, 512) + source0 = torch.arange(non_k * 4, dtype=torch.int32).to(torch.uint8).reshape( + non_k, + 4, + ) + source2 = ( + torch.arange(non_k * 8, dtype=torch.int32) + 37 + ).to(torch.uint8).reshape(non_k, 8) + target0 = torch.full((non_k, 8), 127, dtype=torch.uint8) + target0[:, :4] = source0 + blocked0 = _blocked_reference(target0) + blocked2 = _blocked_reference(source2) + source_blocked0 = _blocked_reference(source0) + source_blocked2 = _blocked_reference(source2) + + # The dFC2 epilogue writes token atoms before hidden atoms. + physical0 = source_blocked0.reshape(2, 1, 512).permute(1, 0, 2).reshape(-1) + physical2 = source_blocked2.reshape(2, 2, 512).permute(1, 0, 2).reshape(-1) + actual = assemble_dfc2_atom_scales( + torch.cat((physical0, physical2)), + valid_counts, + padded_ends, + non_k, + 128, + ) + expected = torch.cat((blocked0, blocked2)).reshape(non_k, 16) + + assert actual.shape == (non_k, 16) + assert torch.equal(actual.view(torch.uint8), expected) + + +@pytest.mark.L0 +def test_dfc2_atom_scales_deinterleave_gate_up_before_repacking(): + intermediate = 256 + non_k = 2 * intermediate + logical_source = torch.arange( + non_k * 4, + dtype=torch.int32, + ).to(torch.uint8).reshape(non_k, 4) + gate = logical_source[:intermediate].reshape(8, 32, 4) + up = logical_source[intermediate:].reshape(8, 32, 4) + interleaved_source = torch.stack((gate, up), dim=1).reshape(non_k, 4) + target = torch.full((non_k, 8), 127, dtype=torch.uint8) + target[:, :4] = logical_source + + actual = assemble_dfc2_atom_scales( + _blocked_reference(interleaved_source), + (33,), + (256,), + non_k, + 128, + deinterleave_gate_up=intermediate, + ) + + assert torch.equal( + actual.view(torch.uint8), + _blocked_reference(target).reshape(non_k, 8), + ) + + +@pytest.mark.L0 +def test_plain_col_scales_assemble_256_data_padding_with_empty_expert(): + non_k = 256 + counts = (33, 0, 129) + padded_ends = (256, 256, 512) + source = torch.full((12, non_k), 127, dtype=torch.uint8) + source[:2] = torch.arange(2 * non_k, dtype=torch.int32).to( + torch.uint8 + ).reshape(2, non_k) + source[4:9] = ( + torch.arange(5 * non_k, dtype=torch.int32) + 19 + ).to(torch.uint8).reshape(5, non_k) + + raw0 = torch.full((non_k, 8), 127, dtype=torch.uint8) + raw0[:, :2] = source[:2].transpose(0, 1) + raw1 = torch.empty((non_k, 0), dtype=torch.uint8) + raw2 = torch.full((non_k, 8), 127, dtype=torch.uint8) + raw2[:, :5] = source[4:9].transpose(0, 1) + expected = torch.cat( + tuple( + _blocked_reference(raw) + for raw in (raw0, raw1, raw2) + ) + ).reshape(non_k, 16) + + actual = assemble_plain_col_scales( + source.view(torch.float8_e8m0fnu), + counts, + padded_ends, + non_k, + 128, + ) + + assert torch.equal(actual.view(torch.uint8), expected) + empty_plain = assemble_plain_col_scales( + torch.empty(0, non_k, dtype=torch.uint8).view( + torch.float8_e8m0fnu + ), + (0, 0), + (0, 0), + non_k, + 128, + ) + empty_discrete = assemble_discrete_col_requant_scales( + torch.empty(0, dtype=torch.uint8), + (0, 0), + (0, 0), + non_k, + 128, + ) + empty_dfc2 = assemble_dfc2_atom_scales( + torch.empty(0, dtype=torch.uint8), + (0, 0), + (0, 0), + non_k, + 128, + ) + assert empty_plain.shape == (non_k, 0) + assert empty_discrete.shape == (non_k, 0) + assert empty_dfc2.shape == (non_k, 0) + + +@pytest.mark.L0 +def test_forward_materializes_caller_owned_256_padded_operand_stash(): + operator = _operator(backward_wgrad_mode="operands") + config = operator._forward_config + owner = Mxfp8ForwardStashOwner(config, torch.device("cpu")) + request = SimpleNamespace( + topk_idx=torch.tensor([[0, -1], [1, 0]], dtype=torch.int32), + device=torch.device("cpu"), + ) + plan = owner.prepare(request, pool_token_capacity=512) + plan.buffer.zero_() + plan.buffer[0, 0] = 10 + plan.buffer[1, 0] = 20 + plan.buffer[256, 0] = 30 + + def pack(token, slot): + return token | (slot << 32) + + packed_metadata = torch.zeros(512, dtype=torch.int64) + packed_metadata[0] = pack(1, 1) + packed_metadata[1] = pack(0, 0) + packed_metadata[256] = pack(1, 0) + col_data = torch.zeros( + 512, + 128, + dtype=torch.float8_e4m3fn, + ) + col_sf = torch.full( + (512 // 32 * 128,), + 127, + dtype=torch.uint8, + ) + inputs = SimpleNamespace( + fc1_c=plan.buffer, + shared_workspace=packed_metadata.view(torch.uint8), + col_quant_data=col_data, + col_quant_sf=col_sf, + ) + prepared = SimpleNamespace( + token_src_metadata_offset=0, + token_src_metadata_bytes=512 * 8, + pool_token_capacity=512, + ) + + fc1_c, route_metadata, stash = owner.materialize( + plan, + inputs, + prepared, + ) + + assert stash is not None + assert stash.fc1_a.shape == (128, 512) + assert stash.fc1_a.stride(1) == 1 + assert stash.fc1_sfa.shape == (128, 16) + assert stash.expert_offsets.tolist() == [256, 512] + assert stash.valid_route_counts.tolist() == [2, 1] + assert stash.route_metadata is route_metadata + assert route_metadata.tolist() == [ + [0, 0, 0, 0], + [0, 0, 1, 1], + [1, 0, 1, 0], + ] + assert fc1_c[:, 0].tolist() == [20, 10, 30] + preserved = stash.fc1_a.clone() + col_data.fill_(1) + assert torch.equal(stash.fc1_a, preserved) + owner.close() + operator.close() + + +@pytest.mark.L0 +def test_backward_export_owns_outputs_and_uses_grouped_wgrad_strides(): + config = _operator( + backward_wgrad_mode="operands" + )._forward_config + stash = _forward_stash() + request = SimpleNamespace( + config=config, + wgrad_forward_stash=stash, + ) + pool_rows = 512 + sf_rows = 8 + aux = Mxfp8DgluResult( + grad_activation=torch.empty(2, 128), + fc1_recompute=torch.zeros( + pool_rows, + 256, + dtype=torch.float8_e4m3fn, + ), + fc1_recompute_sf=torch.full( + (sf_rows, 256), + 127, + dtype=torch.uint8, + ).view(torch.float8_e8m0fnu), + fc1_col_output=torch.zeros( + pool_rows, + 512, + dtype=torch.float8_e4m3fn, + ), + fc1_col_output_sf=torch.full( + (sf_rows, 512), + 127, + dtype=torch.uint8, + ).view(torch.float8_e8m0fnu), + grad_y2=torch.zeros( + pool_rows, + 128, + dtype=torch.float8_e4m3fn, + ), + grad_y2_sf=torch.full( + (pool_rows // 32 * 128,), + 127, + dtype=torch.uint8, + ), + ) + + operands = export_wgrad_operands(request, aux) + + assert operands.fc1_b.shape == (512, 512) + assert operands.fc1_b.stride(0) == 1 + assert operands.fc2_a.shape == (256, 512) + assert operands.fc2_a.stride(1) == 1 + assert operands.fc2_b.shape == (512, 128) + assert operands.fc2_b.stride(0) == 1 + assert operands.fc1_sfb.shape == (512, 16) + assert operands.fc2_sfa.shape == (256, 16) + assert operands.fc2_sfb.shape == (128, 16) + + preserved_fc1_b = operands.fc1_b.clone() + preserved_fc2_a = operands.fc2_a.clone() + preserved_fc2_b = operands.fc2_b.clone() + aux.fc1_col_output.fill_(1) + aux.fc1_recompute.fill_(1) + aux.grad_y2.fill_(1) + assert torch.equal(operands.fc1_b, preserved_fc1_b) + assert torch.equal(operands.fc2_a, preserved_fc2_a) + assert torch.equal(operands.fc2_b, preserved_fc2_b) + + +def _reference_wgrad_case(topk_idx, topk_weights): + torch.manual_seed(20260821) + token_count = topk_idx.shape[0] + hidden = intermediate = 32 + activation = torch.randn(token_count, hidden) / 4 + fc1_weight = torch.randn(3, hidden, 2 * intermediate) / 8 + fc2_weight = torch.randn(3, intermediate, hidden) / 8 + grad_output = torch.randn(token_count, hidden) / 8 + reference = MoeEpReference( + num_experts=3, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=topk_idx.shape[1], + max_tokens_per_rank=token_count, + generate_c=True, + backward_wgrad_mode="operands", + token_padding_size=256, + ) + _, fc1_c, route_metadata, forward_stash = reference( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ) + _, _, operands = reference.backward( + grad_output, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + wgrad_forward_stash=forward_stash, + ) + return ( + activation, + fc2_weight, + grad_output, + fc1_c, + route_metadata, + operands, + ) + + +@pytest.mark.L0 +def test_reference_wgrad_operands_follow_route_weights_and_dense_formulas(): + topk_idx = torch.tensor( + [[0, 2], [2, 0], [0, 2]], + dtype=torch.int32, + ) + topk_weights = torch.tensor( + [[0.5, 0.25], [0.0, 0.75], [1.0, 0.125]], + dtype=torch.float32, + ) + ( + activation, + fc2_weight, + grad_output, + fc1_c, + route_metadata, + operands, + ) = _reference_wgrad_case(topk_idx, topk_weights) + + assert isinstance(operands, WgradOperandsReference) + assert operands.expert_offsets.tolist() == [256, 256, 512] + assert operands.valid_route_counts.tolist() == [3, 0, 3] + assert operands.fc1_a.shape == (32, 512) + assert operands.fc1_b.shape == (512, 64) + assert operands.fc2_a.shape == (32, 512) + assert operands.fc2_b.shape == (512, 32) + + # Rebuild the staged values independently in compact metadata order. + staged_x = quantize_blockwise( + activation, + MoeFormat.MXFP8, + axis=1, + ).dequantize() + staged_dy = quantize_blockwise( + grad_output, + MoeFormat.MXFP8, + axis=1, + ).dequantize() + staged_w2 = quantize_blockwise( + fc2_weight.transpose(1, 2), + MoeFormat.MXFP8, + axis=1, + ).dequantize().transpose(1, 2) + compact_x = [] + compact_weighted_h = [] + compact_dy = [] + compact_dc = [] + for row, (expert, _, token, slot) in enumerate(route_metadata.tolist()): + c_gate, c_up = fc1_c[row].float().split(32) + sigmoid = torch.sigmoid(c_gate) + silu = c_gate * sigmoid + h = silu * c_up + p = topk_weights[token, slot] + dy = staged_dy[token] + dh = (dy @ staged_w2[expert].transpose(0, 1)) * p + dc = torch.cat( + ( + dh * c_up * sigmoid * (1 + c_gate * (1 - sigmoid)), + dh * silu, + ) + ) + compact_x.append(staged_x[token]) + compact_weighted_h.append(p * h) + compact_dy.append(dy) + compact_dc.append(dc) + + def padded(rows): + result = torch.zeros(512, rows[0].numel()) + result[:3] = torch.stack(rows[:3]) + result[256:259] = torch.stack(rows[3:]) + return result + + expected_x = quantize_blockwise( + padded(compact_x).transpose(0, 1), + MoeFormat.MXFP8, + axis=1, + ) + expected_h = quantize_blockwise( + padded(compact_weighted_h).transpose(0, 1), + MoeFormat.MXFP8, + axis=1, + ) + expected_pdy = quantize_blockwise( + padded(compact_dy), + MoeFormat.MXFP8, + axis=0, + ) + expected_dc = quantize_blockwise( + padded(compact_dc), + MoeFormat.MXFP8, + axis=0, + ) + for actual, expected in ( + (operands.fc1_a, expected_x), + (operands.fc1_b, expected_dc), + (operands.fc2_a, expected_h), + (operands.fc2_b, expected_pdy), + ): + assert torch.equal(actual.data, expected.data) + assert torch.equal(actual.scale, expected.scale) + + # The valid route with p=0 keeps x/dY but contributes zero to both wgrads. + zero_weight_row = 257 + assert operands.fc1_a.dequantize()[:, zero_weight_row].abs().sum() > 0 + assert operands.fc2_a.dequantize()[:, zero_weight_row].eq(0).all() + assert operands.fc1_b.dequantize()[zero_weight_row].eq(0).all() + assert operands.fc2_b.dequantize()[zero_weight_row].abs().sum() > 0 + + dw1, dw2 = operands.dense_wgrads() + a1, b1 = operands.fc1_a.dequantize(), operands.fc1_b.dequantize() + a2, b2 = operands.fc2_a.dequantize(), operands.fc2_b.dequantize() + torch.testing.assert_close(dw1[0], a1[:, :256] @ b1[:256]) + torch.testing.assert_close(dw1[1], torch.zeros_like(dw1[1])) + torch.testing.assert_close(dw1[2], a1[:, 256:] @ b1[256:]) + torch.testing.assert_close(dw2[0], a2[:, :256] @ b2[:256]) + torch.testing.assert_close(dw2[1], torch.zeros_like(dw2[1])) + torch.testing.assert_close(dw2[2], a2[:, 256:] @ b2[256:]) + decoded_dw1, decoded_dw2 = _dense_wgrads_from_operands( + _as_production_operands(operands) + ) + torch.testing.assert_close(decoded_dw1, dw1) + torch.testing.assert_close(decoded_dw2, dw2) + + +@pytest.mark.L0 +def test_reference_wgrad_empty_routes_padding_offsets_and_invalid_routes(): + topk_idx = torch.full((2, 2), -1, dtype=torch.int32) + topk_weights = torch.randn(2, 2) + *_, fc1_c, route_metadata, operands = _reference_wgrad_case( + topk_idx, + topk_weights, + ) + + assert fc1_c.shape == (0, 64) + assert route_metadata.shape == (0, 4) + assert operands.expert_offsets.tolist() == [0, 0, 0] + assert operands.valid_route_counts.tolist() == [0, 0, 0] + assert operands.fc1_a.shape == (32, 0) + assert operands.fc1_b.shape == (0, 64) + assert operands.fc2_a.shape == (32, 0) + assert operands.fc2_b.shape == (0, 32) + for dense in operands.dense_wgrads(): + assert dense.eq(0).all() + + invalid = torch.tensor([[0, -2], [3, -1]], dtype=torch.int32) + with pytest.raises(ValueError, match="out-of-range expert id"): + _reference_wgrad_case(invalid, topk_weights) + + +def _assemble_reference_scale( + tensor, + expert_offsets: torch.Tensor, +) -> torch.Tensor: + k_axis = tensor.axis + scale = tensor.scale + parts = [] + begin = 0 + for end_tensor in expert_offsets: + end = int(end_tensor.item()) + if k_axis == 1: + raw = scale[:, begin // 32 : end // 32] + else: + raw = scale[begin // 32 : end // 32].transpose(0, 1) + parts.append(_blocked_reference(raw.view(torch.uint8))) + begin = end + non_k = tensor.shape[1 - k_axis] + rounded_non_k = (non_k + 127) // 128 * 128 + scale_columns = (tensor.shape[k_axis] // 32 + 3) // 4 * 4 + if parts: + packed = torch.cat(parts) + else: + packed = torch.empty(0, dtype=torch.uint8, device=tensor.device) + return packed.reshape(rounded_non_k, scale_columns).view( + torch.float8_e8m0fnu + ) + + +def _as_production_operands(reference: WgradOperandsReference): + return MoeEpWgradOperands( + fc1_a=reference.fc1_a.data, + fc1_sfa=_assemble_reference_scale( + reference.fc1_a, + reference.expert_offsets, + ), + fc1_b=reference.fc1_b.data.transpose(0, 1).contiguous().transpose(0, 1), + fc1_sfb=_assemble_reference_scale( + reference.fc1_b, + reference.expert_offsets, + ), + fc2_a=reference.fc2_a.data, + fc2_sfa=_assemble_reference_scale( + reference.fc2_a, + reference.expert_offsets, + ), + fc2_b=reference.fc2_b.data.transpose(0, 1).contiguous().transpose(0, 1), + fc2_sfb=_assemble_reference_scale( + reference.fc2_b, + reference.expert_offsets, + ), + expert_offsets=reference.expert_offsets, + valid_route_counts=reference.valid_route_counts, + route_metadata=reference.route_metadata, + ) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_returned_operand_abi_runs_direct_grouped_wgrad(): + if not torch.cuda.is_available(): + pytest.skip("grouped wgrad integration requires CUDA") + device = torch.device("cuda", 0) + major, minor = torch.cuda.get_device_capability(device) + if (major, minor) != (10, 0): + pytest.skip("the in-tree grouped wgrad integration kernel targets SM100") + + import cudnn + + generator = torch.Generator(device=device).manual_seed(20260821) + hidden, intermediate = 128, 256 + activation = torch.randn( + 3, + hidden, + generator=generator, + device=device, + ) / 8 + fc1_weight = torch.randn( + 2, + hidden, + 2 * intermediate, + generator=generator, + device=device, + ) / 8 + fc2_weight = torch.randn( + 2, + intermediate, + hidden, + generator=generator, + device=device, + ) / 8 + topk_idx = torch.tensor( + [[0, 1], [1, 0], [0, 1]], + dtype=torch.int32, + device=device, + ) + topk_weights = torch.tensor( + [[0.5, 0.25], [0.75, 0.5], [1.0, 0.125]], + device=device, + ) + grad_output = torch.randn( + 3, + hidden, + generator=generator, + device=device, + ) / 8 + reference = MoeEpReference( + num_experts=2, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=2, + max_tokens_per_rank=3, + generate_c=True, + backward_wgrad_mode="operands", + token_padding_size=256, + ) + _, fc1_c, metadata, forward_stash = reference( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ) + _, _, logical = reference.backward( + grad_output, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + metadata, + wgrad_forward_stash=forward_stash, + ) + operands = _as_production_operands(logical) + + common = { + "offsets_tensor": operands.expert_offsets, + "output_mode": "dense", + "wgrad_dtype": torch.bfloat16, + "acc_dtype": torch.float32, + "mma_tiler_mn": (128, 128), + "cluster_shape_mn": (1, 1), + "sf_vec_size": 32, + } + fc1 = cudnn.grouped_gemm_wgrad_wrapper_sm100( + a_tensor=operands.fc1_a, + b_tensor=operands.fc1_b, + sfa_tensor=operands.fc1_sfa, + sfb_tensor=operands.fc1_sfb, + **common, + )["wgrad_tensor"] + fc2 = cudnn.grouped_gemm_wgrad_wrapper_sm100( + a_tensor=operands.fc2_a, + b_tensor=operands.fc2_b, + sfa_tensor=operands.fc2_sfa, + sfb_tensor=operands.fc2_sfb, + **common, + )["wgrad_tensor"] + torch.cuda.synchronize(device) + + expected_fc1, expected_fc2 = logical.dense_wgrads() + torch.testing.assert_close( + fc1.float(), + expected_fc1, + rtol=0.15, + atol=0.125, + ) + torch.testing.assert_close( + fc2.float(), + expected_fc2, + rtol=0.15, + atol=0.125, + ) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +@pytest.mark.parametrize("world_size", [1, 2, 4], ids=["ep1", "ep2", "ep4"]) +def test_production_wgrad_operands_run_end_to_end(world_size, tmp_path): + if world_size == 1: + _run_wgrad_operand_case( + device=_sm107_device(), + ep_group=None, + rank=0, + world_size=1, + ) + return + + _require_distributed_sm107(world_size) + os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") + init_file = tmp_path / f"mxfp8_wgrad_operands_ep{world_size}.init" + mp.spawn( + _distributed_wgrad_worker, + args=(world_size, str(init_file)), + nprocs=world_size, + join=True, + ) From 9284bb3052379dc8c8e21a45acf5a77dfb72fceb Mon Sep 17 00:00:00 2001 From: zhibinz Date: Mon, 24 Aug 2026 16:43:00 -0700 Subject: [PATCH 09/31] docs: document MoeEp execution support Describe installation, tensor formats, forward and backward contracts, tuning, lifecycle requirements, and Rubin support boundaries for the new API. --- docs/fe-oss-apis/moe_ep.md | 910 +++++++++++++++++++++++++++++++++++ docs/fe-oss-apis/overview.md | 6 + 2 files changed, 916 insertions(+) create mode 100644 docs/fe-oss-apis/moe_ep.md diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md new file mode 100644 index 000000000..f2ddd1924 --- /dev/null +++ b/docs/fe-oss-apis/moe_ep.md @@ -0,0 +1,910 @@ +# MoE + Expert Parallel API + +Status: public API, validated lazy backend seam, internal runtime/workspace +owners, and executable PyTorch reference. The current device target is the +Rubin training `fwd_glu` kernel plus the restricted `bwd_dglu` path on exactly +SM107 (compute capability 10.7). It accepts MXFP8 E4M3/E8M0 operands or plain +BF16/FP16/FP32 operands staged to MXFP8, supports BF16 or MXFP8 combine with +BF16 output, EP1 through EP16 forward, and EP1/EP2/EP4 backward. Unsupported +combinations fail explicitly rather than returning uncomputed storage. + +The design removes workspace pointers, peer pointer mappers, streams, and +individual scheduler knobs from the semantic runtime interface. Performance +tuning is encapsulated in the optional `MoeEpTuningConfig`. + +## 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`. +- In the public contract, `output_format` describes the + post-top-k-reduction `(T, H)` result and `combine_format` independently + describes each per-route FC2 contribution on the EP return path. +- The public API reserves BF16, MXFP8, and NVFP4 spellings for both choices. + Quantized output is a data-plus-scale object, never a scale-free PyTorch + tensor. This is a future contract, not the current device capability. +- The current Rubin backend accepts `combine_format="bf16"` or `"mxfp8"` and + requires `output_format="bf16"`. NVFP4 combine and quantized public output + are rejected before runtime initialization. + +The distinction between public output and combine traffic is intentional. A +future transport-only MXFP8/NVFP4 mode would set `combine_format` to that +format and keep `output_format="bf16"`. The current SM107 MXFP8 combine path is +the training kernel's direct `MXFP8(FP32 accumulator)` conversion. + +## Initial device-backend implementation scope + +The public contract below remains the target. Device support will be enabled +incrementally, and unsupported combinations must fail explicitly rather than +return uninitialized storage. + +The current implementation connects the complete EP-subgroup path through the +executable CuTeDSL backend. Its deliberately narrow capability is: + +- the device must report exactly compute capability `(10, 7)` (SM107). If + `CUTE_DSL_ARCH` is unset, the compile runner sets it to `sm_107a`; an + existing value must target `sm_107` or `sm_107a`; +- `activation`, `fc1_weight`, and `fc2_weight` may be MXFP8 + `BlockScaledTensor` objects using logical, unswizzled scales and `axis=1`, + or plain BF16/FP16/FP32 tensors that staging quantizes to MXFP8; +- MXFP8 payloads use FP8 E4M3 and scales use FP8 E8M0, with the shapes specified + in "Block-scaled representation"; +- `combine_format` may be `"bf16"` or `"mxfp8"` and `output_format` must be + `"bf16"`; both + `generate_c=False` and eager-only `generate_c=True` forward are supported; +- `generate_c=True` performs one lockstep route-count collective and returns + fresh compact `fc1_c`/`route_metadata` tensors. The default path uses + 128-row internal expert alignment. Opt-in + `backward_wgrad_mode="operands"` requires 256-row alignment and additionally + returns caller-owned MXFP8 grouped-wgrad operands. Training execution does + not support CUDA Graph capture; the restricted device backward returns + activation and router-weight gradients for EP1/EP2/EP4; +- forward and backward support optional `gate_up_clamp`; Rubin training + execution requires `apply_topk_in_fc1=True`; +- `max_tokens_per_rank` must be explicitly positive; `top_k <= 32`, + `H % 128 == 0`, and `I % 256 == 0`; +- `ep_group=None` remains explicit single-rank execution. Distributed execution + accepts any initialized `torch.distributed.ProcessGroup`, including + non-contiguous global-rank membership, with EP2 through EP16. The public + contract requires `top_k <= num_experts` and the device path additionally + requires `top_k <= 32`; `top_k` may exceed `experts_per_rank`. Expert + ownership, peer tables, and route metadata use dense group-relative EP ranks; + experts remain contiguous and equally partitioned; +- every subgroup rank must call forward, backward, warmup, and graph replay + (where supported) in lockstep, including zero-token and zero-valid-route + ranks. Teardown must also be coordinated, although `close()` does not insert + a process-group barrier. Validation is rank-local, and collective + participation is a caller contract rather than an extra host-synchronized + control collective; +- the first lazy subgroup forward launch performs a one-time readiness + rendezvous after staging and JIT: each rank synchronizes its current stream, + all-gathers the effective tuning signature and rejects a mismatch, then + enters a process-group barrier before peer metadata writes can begin. The + first distributed backward launch likewise synchronizes and barriers without + a second tuning gather; EP1 skips both rendezvous. Subsequent eager launches + and graph replays do not add this control collective; +- while the process-global NVSHMEM runtime is active, all backend instances in + that process share one CUDA device binding, one ordered EP-subgroup + membership, and one reference count; distributed deployment uses one process + per GPU. A different subgroup cannot become active in the same process until + the first is fully released. A non-WORLD subgroup does not attach to an + externally initialized NVSHMEM runtime whose membership cannot be verified; + a matching full-WORLD external runtime may be attached but is never finalized + by this backend; +- each operator owns its local workspace and its NVSHMEM-symmetric workspace, + plus transformed-weight and reduction scratch. Allocations are plan-scoped + and stable across eager calls; `generate_c=True` additionally owns the BF16 + high-watermark C buffer described above; +- the forward and backward compile paths instantiate the vendored Rubin + training `Sm107MegaMoEMxfp8GluKernel` and + `Sm107MegaMoEMxfp8DgluKernel`, respectively. They do not select the + Blackwell inference MegaMoE implementation. + +Plain BF16/FP16/FP32 operands do not select a separate floating-point kernel +specialization: staging quantizes them to MXFP8 E4M3/E8M0 before launch. +Pre-quantized MXFP8 payloads and logical scales are preserved rather than +dequantized and requantized. + +The private CuTeDSL source snapshot lives under +`python/cudnn/moe_ep/_megamoe_backend/cutedsl_src`. Its +`VENDOR_INFO.md` records the upstream source revisions, vendoring dates, +copied-file manifest, local import adaptations, and update procedure. The +Apache-2.0 headers and redistribution terms are packaged beside the source. +The current kernel entry points are vendored under +`kernel_src/rubin/training/mega/fwd_glu` and +`kernel_src/rubin/training/mega/bwd_dglu`; those packages must use +package-relative imports, must not depend on a sibling `cutedsl_megamoe` +checkout, and must not import `kernel_src.blackwell`. + +After installing the `moe_ep` optional dependencies, run the L0 device +validation with: + +```bash +python -m pytest \ + test/python/fe_api/moe_ep/test_moe_ep_forward.py \ + test/python/fe_api/moe_ep/test_moe_ep_backward.py \ + -m L0 +``` + +These forward and backward core suites include public-contract and capability +checks, staging, runtime, reference, and supported single-rank numerical +coverage. On SM107, the L0 forward suite exercises the production compile and +launch path. Full numerical parity, stress, CUDA Graph replay, multi-rank +forward, and device backward acceptance are covered by L1 tests and the +hardware/container runner rather than by this L0 command alone. + +The ordinary L1 hardware/container runner is single-node: its distributed +forward matrix uses `mp.spawn` for EP2, EP3, and EP4. EP7, EP12, EP15, and +EP16 use the torchrun-native multi-node suite instead. From the first node of +an existing Slurm allocation, start one torchrun agent per selected node with: + +```bash +torchrun \ + --nnodes="${NNODES}" \ + --node-rank="${NODE_RANK}" \ + --nproc-per-node="${NPROC_PER_NODE}" \ + --master-addr="${MASTER_ADDR}" \ + --master-port="${MASTER_PORT}" \ + -m pytest \ + test/python/fe_api/moe_ep/test_moe_ep_forward_multinode.py \ + -m "L1 and moe_ep_multinode" \ + -k "${PYTEST_FILTER}" +``` + +NVSHMEM requires the same number of EP PEs on every participating node. Use +`NNODES=7`, `NPROC_PER_NODE=2`, and `PYTEST_FILTER=ep7-world14` for EP7; its +subgroup selects local rank zero on every node. Use `NNODES=5`, +`NPROC_PER_NODE=4`, and `PYTEST_FILTER=ep15-world20` for EP15; its subgroup +selects local ranks zero through two on every node. EP12 uses `NNODES=3`, +`NPROC_PER_NODE=4`, and `PYTEST_FILTER=ep12-world12`; EP16 uses `NNODES=4`, +`NPROC_PER_NODE=4`, and `PYTEST_FILTER=ep16-world16`. The remaining WORLD +ranks synchronize without constructing `MoeEp`. The Slurm harness must provide +a distinct `NODE_RANK` to each node and shared `MASTER_ADDR`/`MASTER_PORT` +values. Each case checks BF16 and MXFP8 combine against the executable +reference plus all-`-1` route behavior. Capability support alone is not a +hardware PASS: preserve the torchrun logs for acceptance evidence. + +PyTorch is a prerequisite of the `MoeEp` API and remains in the `moe_ep` +optional extra rather than becoming a base dependency of the entire +`nvidia-cudnn-frontend` distribution. The same extra selects the CUDA-13 +CuTeDSL stack and `nvshmem4py-cu13>=0.3.1`. Given an installed PyTorch, +ordinary `import cudnn` and `import cudnn.moe_ep` do not import CuTeDSL, +NVSHMEM4Py, or CUDA Python and do not initialize CUDA. Missing CuTeDSL or +NVSHMEM runtime components are reported as `BackendUnavailableError` only +when a supported device forward first needs the backend. + +The hardware, API, stress, and packaging runner defaults to +`MOE_EP_DEPENDENCY_MODE=rubin-internal`. In this mode it first removes every +installed `nvidia-cutlass-dsl*` distribution, then installs the latest +pre-release `nvidia-cutlass-dsl` from NVIDIA's Rubin-capable CUTLASS DSL master +index with PyPI as a dependency fallback. The runner verifies that +`cutlass.utils.rubin_helpers` is importable before compiling the kernel. +`MOE_EP_DEPENDENCY_MODE=latest` and `minimum` retain public-release acceptance; +the latter pins CUTLASS DSL, NVSHMEM4Py, and Apache TVM FFI to 4.8.0, 0.3.1, +and 0.1.11 respectively. The public 4.7.0 wheel does not contain +`rubin_helpers`. The packaging runner also validates an isolated wheel import +and the vendored Rubin import closure. + +## Public API + +The production class is `cudnn.moe_ep.MoeEp`. The test-only +`MoeEpReference` is the executable semantic and numerical oracle; it is not the +device backend. + +```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, + backward_wgrad_mode: Literal["none", "operands"] = "none", + token_padding_size: int = 128, + sf_padding_size: int = 128, + tuning: Optional[MoeEpTuningConfig] = None, + ) -> 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] + | tuple[ + Tensor | BlockScaledTensor, + Tensor, + Tensor, + MoeEpWgradForwardStash, + ] + ): ... + + def warmup( + self, + activation: Tensor | BlockScaledTensor, + fc1_weight: Tensor | BlockScaledTensor, + fc2_weight: Tensor | BlockScaledTensor, + topk_idx: Tensor, + topk_weights: Tensor, + ) -> None: ... + + def backward( + self, + grad_output: Tensor, + fc1_weight: Tensor | BlockScaledTensor, + fc2_weight: Tensor | BlockScaledTensor, + topk_idx: Tensor, + topk_weights: Tensor, + fc1_c: Tensor, + route_metadata: Tensor, + *, + wgrad_forward_stash: Optional[MoeEpWgradForwardStash] = None, + ) -> ( + tuple[Tensor, Tensor] + | tuple[Tensor, Tensor, MoeEpWgradOperands] + ): ... + + def close(self) -> None: ... + + def __enter__(self) -> "MoeEp": ... + + def __exit__(self, exc_type, exc_value, traceback) -> bool: ... +``` + +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. + +`close()` is terminal and idempotent. Once device workspaces have been created, +the backend synchronizes outstanding CUDA work before releasing transformed +weights, workspace, and runtime ownership. It does not issue a process-group +barrier; distributed callers must coordinate teardown so one rank cannot +release symmetric storage while a peer is still launching or replaying. +The context-manager form is preferred when deterministic release matters. + +### 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 constructor, but the current device capability gate requires an explicit positive value on first execution. | +| `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 the combine-rounded FC2 route contribution in the standalone top-k reducer. | +| `gate_up_clamp` | If set, use `gate = clamp(gate, max=abs(limit))` and `up = clamp(up, min=-abs(limit), max=abs(limit))`. | +| `generate_c` | Training integration: additionally return `fc1_c` (the BF16-rounded pre-SwiGLU FC1 accumulator of every route this rank's experts processed) and its row-aligned `route_metadata`. | +| `backward_wgrad_mode` | `"none"` preserves the default API. `"operands"` opts into caller-owned MXFP8 operands for external grouped FC1/FC2 wgrad GEMMs; it requires `generate_c=True`, `token_padding_size=256`, and `sf_padding_size=128`. | +| `token_padding_size` | Token-dimension padding used by the Rubin execution plan; positive integer, default 128. | +| `sf_padding_size` | Scale-factor padding; positive multiple of 128, default 128. Operand mode currently requires exactly 128. | +| `tuning` | Optional `MoeEpTuningConfig`; every rank in an EP group must use the same effective configuration. | + +The constructor validates public static dimensions, format alignment, padding, +and tuning. MXFP8 format spellings require `H % 32 == 0`; NVFP4 format +spellings require `H % 16 == 0`. Device-specific requirements such as SM107, +`H % 128 == 0`, `I % 256 == 0`, `top_k <= 32`, and explicit positive capacity +are checked lazily before backend initialization. + +The wgrad mode is strictly opt-in. With the default +`backward_wgrad_mode="none"`, the constructor default, forward return, backward +signature, backward return, padding default, allocation behavior, and +documented numerical semantics are unchanged. + +Scheduler settings are exposed through the optional `MoeEpTuningConfig` +object. The current public knobs are `token_back_mode`, `epi_flag_batch`, +`token_in_flag_batch`, `group_hint`, and `reduce_topk_in_kernel`; they are +keyword configuration rather than positional runtime arguments. + +`reduce_topk_in_kernel=True` enables the in-kernel top-k reduction path and is +restricted to BF16 combine/output, `apply_topk_in_fc1=True`, and +`token_back_mode="epi_warps"`. BF16 reduction order can change rounding, so +this path is accepted against the documented numerical tolerance rather than +bitwise Form A equality. + +### 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/FP16/FP32 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. + +The first successful backend creation binds a `MoeEp` instance and its stable +workspace allocations to that device. A later call on another device raises +`ValueError`; callers must create one `MoeEp` instance per device. + +### 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. + +With `generate_c=True`, the call instead returns +`(output, fc1_c, route_metadata)`; see the training-integration sections +below. + +With both `generate_c=True` and `backward_wgrad_mode="operands"`, it returns +`(output, fc1_c, route_metadata, wgrad_forward_stash)`. The fourth item belongs +to that exact routed forward call and must be passed to its corresponding +backward call. + +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`; the +production entry point is `MoeEp.backward`, while `MoeEpReference.backward` +defines its executable semantic oracle. The restricted Rubin MXFP8 device path +supports BF16/MXFP8 combine, BF16 output, EP1/EP2/EP4, +`apply_topk_in_fc1=True`, and optional `gate_up_clamp`. 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)` | Any floating dtype on the request device; incoming gradient of the *dequantized* public output (all encodes are straight-through). | +| `fc1_weight`, `fc2_weight`, `topk_idx`, `topk_weights` | as in forward | the framework re-supplies the forward weights and routing 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. | +| `wgrad_forward_stash` | `MoeEpWgradForwardStash` | Keyword-only and required only in operand mode. It contains the forward `x.T` operand plus padded expert offsets/counts and exact route identity. Passing it in default mode is an error. | + +Default mode returns two 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_topk_weights` | `(T, K)` | per-route router-weight gradient; exact zero at `-1` slots. | + +Operand mode returns +`(grad_activation, grad_topk_weights, wgrad_operands)`. The first two values +retain the default meanings and dtypes; `wgrad_operands` is described below. + +On the Rubin device path, dGLU materializes `grad_activation` in BF16 and the +public wrapper widens it to FP32, so its numerical granularity remains BF16. +The semantic router-weight gradient is recomputed from the unquantized +floating `grad_output` and returned in FP32. + +Internally, the compact public `fc1_c` rows are lowered into an external +pool-layout BF16 `fc1_preact` tensor. The kernel also writes a pre-zeroed, +symmetric source-domain FP32 dprob plane with shape +`(max_tokens_per_rank, top_k)` for its peer-atomic ABI; that internal plane is +not the public return because the public straight-through semantics use the +FP32 recomputation described above. + +The save-set, recompute rules, and numerical conventions behind this +signature are specified in "Saved tensors for backward" below. + +### Grouped-wgrad operand contract + +`backward_wgrad_mode="operands"` exposes the operands needed by the existing +`GroupedGemmWgradSm100` / `grouped_gemm_wgrad_wrapper_sm100` Tensor2D ABI. It +does not launch those GEMMs and does not return dense weight gradients. + +For local expert `e`, let `R_e` be its valid route count, +`P_e = ceil(R_e / 256) * 256`, and `Kp = sum_e P_e`. Every expert occupies one +contiguous range in the shared K dimension. Valid rows precede zero padding; +an empty expert contributes zero extent, so cumulative offsets may repeat. +`expert_offsets[e] = sum_{j <= e} P_j`, and `valid_route_counts[e] = R_e`. + +Forward returns `MoeEpWgradForwardStash` with: + +| Field | Logical shape | Meaning | +|---|---:|---| +| `fc1_a`, `fc1_sfa` | `(H, Kp)`, `(round_up(H,128), round_up(Kp/32,4))` | MXFP8 `x.T` data and assembled E8M0 scales. | +| `expert_offsets` | `(E_local,)` | Int32 cumulative 256-padded expert ends. | +| `valid_route_counts` | `(E_local,)` | Int32 unpadded route counts. | +| `route_metadata` | `(local_routes, 4)` | The same compact route identity returned beside `fc1_c`. | + +Backward returns `MoeEpWgradOperands`. It carries those five fields plus: + +| Field | Logical shape | Meaning | +|---|---:|---| +| `fc1_b`, `fc1_sfb` | `(Kp, 2I)`, `(round_up(2I,128), round_up(Kp/32,4))` | MXFP8 `dC`, where `C = x @ W1`. | +| `fc2_a`, `fc2_sfa` | `(I, Kp)`, `(round_up(I,128), round_up(Kp/32,4))` | MXFP8 route-weighted recomputed `(p * h).T`, `h = silu(gate) * up`. | +| `fc2_b`, `fc2_sfb` | `(Kp, H)`, `(round_up(H,128), round_up(Kp/32,4))` | MXFP8 unweighted routed `dY`. | + +For each expert range, including its zero padding, the represented dense +operations are: + +```text +dW1[e] = fc1_a[e] @ fc1_b[e] = x[e].T @ dC[e] +dW2[e] = fc2_a[e] @ fc2_b[e] = (p[e] * h[e]).T @ dY[e] +``` + +The data operands are E4M3. Scale factors are E8M0 with logical 1x32 scaling +and grouped-wgrad's assembled physical 128x4 scale tiles. A operands have +unit K stride; B operands use the grouped-wgrad K-major view with unit K +stride. + +The Rubin staging/quantization order is part of the operand model: + +1. `x` is first staged to MXFP8 along H (plain inputs only), routed and padded, + then column-requantized along expert K to form `fc1_a`. +2. BF16 `fc1_c` is recomputed through clamp/SwiGLU without a router weight; + that `h` is column-requantized along K to form `fc2_a`. +3. `grad_output` is staged to MXFP8 along H before re-dispatch. The route + weight is then applied exactly once, and the result is + column-requantized along K to form `fc2_b`. +4. Staged `dY` and backward-staged `W2.T` produce `dH`; the route weight is + applied before the SwiGLU derivative, and the resulting `dC` is directly + column-requantized along K to form `fc1_b`. + +All returned stash and operand tensors are caller-owned allocations. They do +not alias reusable execution-plan workspace and are not overwritten by later +forward/backward calls. The caller must retain the forward stash through its +matching backward call and retain the returned operands until every external +grouped-wgrad consumer has completed. Route metadata and counts are validated +against the matching call; stashes are not interchangeable between different +routing inputs. + +### Example + +The following illustrates the reserved future quantized-output contract. It is +valid reference-level API semantics, but the current SM107 device backend +rejects `output_format="nvfp4"`. The shown `combine_format="mxfp8"` is supported +by the device backend when paired with `output_format="bf16"`. + +```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]) +combine[t,k] = dequantize(quantize(expert[t,k], combine_format)) + +if not apply_topk_in_fc1: + combine[t,k] *= p[t,k] + +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. + +For `combine_format="mxfp8"`, the current Rubin forward and backward paths +directly convert each FP32 route accumulator to MXFP8 before top-k reduction. + +Moving the router weight from the post-SwiGLU intermediate to the +combine-rounded FC2 contribution is algebraically equivalent only without the +intervening low-precision conversions. The option is therefore semantic and is +fixed in the constructor, matching the MegaMoE 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 | public `axis=1`: contraction axis for weights, feature/output axis otherwise | +| NVFP4 | packed FP4 E2M1, low nibble first | FP8 E4M3 | 16 | public `axis=1`: contraction axis for weights, feature/output axis otherwise | + +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 interface + +| concept | Public 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 | +| forward column-requant `x.T` | `MoeEpWgradForwardStash.fc1_a/fc1_sfa` in operand mode | +| dGLU column-requant `dC`, `(p*h).T`, `dY` | `MoeEpWgradOperands` returned by backward | + +### `fc1_c` and `route_metadata` (training integration) + +With `generate_c=True`, `__call__` returns `(output, fc1_c, route_metadata)`. +Operand mode appends `wgrad_forward_stash` as a fourth item. +`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 output gradients to the expert ranks, which is where the +stashed preactivations 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 +FC1 FP32 accumulator rounded to BF16 and captured **before** SwiGLU, before the +gate/up clamp, and without the router weight (which applies after SwiGLU when +`apply_topk_in_fc1=True`). Columns `[0:I)` are gate and `[I:2I)` are up. The +kernel's mode-dependent per-expert padding (128 rows by default, 256 in operand +mode) and internal gate/up-interleaved layout are absent from the logical +`fc1_c` 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 + +By default, `MoeEpReference.backward(grad_output, fc1_weight, fc2_weight, +topk_idx, topk_weights, fc1_c, route_metadata)` returns +`(grad_activation, grad_topk_weights)` and is the executable statement of the +save-set. Operand mode also takes `wgrad_forward_stash=` and appends a +`WgradOperandsReference` result. 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; it yields `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ᵀ`. | +| `topk_idx`, `topk_weights` | source rank (framework inputs) | `topk_idx` regenerates the exact dispatch plan; `topk_weights` scales the activation-gradient path 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). | +| `wgrad_forward_stash` | expert rank (operand mode only) | Caller-owned MXFP8 `x.T` plus 256-padded offsets/counts and route identity; supplies FC1 A and fixes the grouped K segmentation. | + +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. + +In the reference backward, input decode, `combine_format`, and `output_format` +round trips 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; 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 it is +`d_w = ⟨d_y, y_pre⟩`. The Rubin device dGLU path additionally MXFP8-stages +`grad_output` and transposed weights, modeled in parity tests by +`backward_operand_format="mxfp8"`, while recomputing `grad_topk_weights` from +the original floating `grad_output`. + +## 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. + - Bootstrap and allocate once outside CUDA graph capture; capture requires a + prior warmup whose stream has completed before capture begins. + - `MoeEp.warmup(...)` runs one complete eager forward and synchronizes its + CUDA device. For an EP subgroup, all member ranks must call it concurrently + and the caller must align ranks before capture; the method deliberately + does not issue a process-group barrier. + - Distributed graph capture is per-rank and replay is lockstep. Capture + records the cross-rank kernel without executing it; every replay must be + issued by every EP rank in the same iteration. + - Cache successful eager route-value validation by tensor identity and + version so unchanged routing avoids repeated host synchronization. + - 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. + - A captured graph has a static-storage contract: the `MoeEp` instance, + workspace, captured input/output storages, and transformed weights must + outlive every replay. Weights must not be modified after capture; replace + or modify them only after retiring the graph, then warm up and recapture. + - The same `MoeEp` instance must not be used concurrently by graph replay + and eager execution. Eager calls on different streams are serialized by a + completion event before shared workspace staging. + - Warmup validates route values. Capture/replay requires every route to + remain `-1` or a valid global expert ID; data-dependent host validation is + not capturable. + - Inference tensors without a PyTorch version counter are repacked on every + eager call and are not accepted as weights during graph capture. + - `generate_c=True` is eager-only because its exact `(local_routes, 2I)` + result requires data-dependent route counting and compaction. Capture is + rejected before the count collective, allocation, or kernel launch. + - Distributed callers must coordinate `close()` because symmetric allocation + release and an owned NVSHMEM finalization must not race a peer that is + still launching or replaying the operator. `close()` itself does not issue + a process-group barrier. + +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. +- 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. +- `WgradForwardStashReference` and `WgradOperandsReference` model logical + MXFP8 values, 256-route expert padding, and dense + `x.T @ dC` / `(p*h).T @ dY` results. Their scales are logical rather than + the production 128x4 physical assembly. +- It has no expert-capacity drop policy beyond explicit `-1` routes. + +Device parity tests configure the same `MoeEpReference` with +`intermediate_format="mxfp8"` to model the internal post-SwiGLU MXFP8 round +trip before FC2, and with `backward_operand_format="mxfp8"` to model dGLU +operand staging. These options are diagnostic backend approximations and are +deliberately not part of the public mathematical contract above. + +## Validation and test matrix + +Public-contract and reference tests remain broader than the current device +backend. They cover MXFP8/NVFP4 representation and quantization semantics, +both router-weight locations, quantized combine/output round trips, and the +explicit reference backward. Those checks do not imply kernel capability. + +The status below applies specifically to the production Rubin device backend. +“Hardware-validated” means that the case passed on SM107 (compute capability +10.7); CPU/reference tests, collection, and skips do not establish that status. + +### Hardware-validated on SM107 + +- EP1 forward has passed with MXFP8 E4M3/E8M0 operands and with mixed plain + BF16/FP16/FP32 activation/weights, BF16 or MXFP8 combine, BF16 output, + `apply_topk_in_fc1=True`, gate/up clamp, all-`-1` routes, and fresh output + ownership. INT32 and INT64 routing indices and supported shape/top-k cases + through `top_k=32` have also passed. +- Eager EP1 `generate_c=True` has passed with the Rubin-required + `apply_topk_in_fc1=True`. The checks cover compact BF16 `fc1_c`, INT32 + `route_metadata`, pre-clamp/unweighted values, repeated calls, and fresh + tensor ownership. +- EP1 warmup, non-default tuning, 100-call eager stress, and CUDA Graph replay + have passed. CUDA Graph replay with post-FC2 router weighting has also passed. +- Single-node WORLD-group eager forward parity has passed for EP2, EP3, and + EP4 with BF16 and MXFP8 combine, including all-`-1` route behavior. +- A WORLD4 eager test has passed for disjoint non-contiguous EP2 subgroups + `[0,2]` and `[1,3]`, including the case where subgroup rank zero is not global + rank zero. +- Multi-node eager forward parity has passed for EP12/WORLD12 (three nodes, + four PEs per node) and EP16/WORLD16 (four nodes, four PEs per node), with + BF16 and MXFP8 combine and all-`-1` route behavior. +- Restricted MXFP8 backward has passed for EP1, EP2, and EP4 with BF16 and + MXFP8 combine. It checks activation and router-weight gradients against the + executable reference; EP1 additionally covers repeated calls and explicitly + reordered forward stashes. +- The returned operand field/stride/scale ABI has passed direct execution + through both FC1 and FC2 grouped-wgrad GEMMs on SM100 using + reference-generated operands. This establishes consumer ABI integration, not + end-to-end Rubin operand production. + +### Supported but awaiting hardware validation + +- Forward capability accepts every EP size from EP1 through EP16. EP5, EP6, + EP8 through EP11, EP13, and EP14 have no current hardware acceptance case. +- EP7/WORLD14 is defined as a seven-node subgroup with one EP PE per node, and + EP15/WORLD20 is defined as a five-node subgroup with three EP PEs per node. + Both acceptance cases remain pending on allocations with the required node + counts. +- Forward CUDA Graph capture/replay and lifecycle contracts apply to + distributed EP groups, but current EP2+ acceptance runs cover eager parity + only; distributed stress and graph replay remain pending. +- Plain/mixed operands, gate/up clamp, post-FC2 router weighting, and explicit + `generate_c=True` output semantics are hardware-validated at EP1, but are not + separately validated across every supported distributed EP size. +- End-to-end Rubin forward/backward production of wgrad operands remains + awaiting SM107 hardware validation. The available SM100 test can execute the + grouped-wgrad consumer but cannot execute the SM107-only MegaMoE producer. + +### Currently unsupported by the device backend + +- Devices other than SM107, EP sizes above 16, non-positive or unspecified + `max_tokens_per_rank`, `hidden_size` not divisible by 128, + `intermediate_size` not divisible by 256, and `top_k > 32`. +- Native NVFP4 operands or combine, any non-BF16 public output, and plain + operand dtypes other than BF16/FP16/FP32. MXFP8 operands and BF16/MXFP8 + combine with BF16 output are supported. +- Backward outside EP1/EP2/EP4, execution with + `apply_topk_in_fc1=False`, non-BF16 output, or backward CUDA Graph capture. +- Wgrad operand mode with padding other than 256, without `generate_c=True`, + or outside the restricted Rubin MXFP8 backward configuration. The mode + returns operands only; dense `dW1`/`dW2` computation remains an explicit + grouped-wgrad call by the integration layer. +- CUDA Graph capture with `generate_c=True`, same-process concurrent EP + subgroups, expert bias, shared/dense experts inside this operator, implicit + top-k normalization, capacity-factor routing, and implicit route drop. +- Dense weight-gradient returns and an integrated `torch.autograd.Function` + wrapper. + +Source/packaging tests separately require the vendored Rubin +`training/mega/fwd_glu` and `training/mega/bwd_dglu` packages, reject sibling +`cutedsl_megamoe` dependencies and `kernel_src.blackwell` imports, and validate +isolated-wheel imports. These checks validate packaging rather than additional +device capabilities. + +Future format families must add the same CUDA/NVSHMEM, stress, graph, and +isolated-package matrix before their capability gates are removed. + +## Current 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; +- backward semantics are fixed by `MoeEpReference.backward` (consuming the + `generate_c=True` stash); opt-in wgrad operands are exposed, while dense + weight-gradient returns and a `torch.autograd` wrapper are not part of this + API; +- 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 2c5439931..5edf5a3bf 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -36,6 +36,7 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d - [RMSNorm + RHT + Amax](rmsnorm_rht_amax.md) - [SDPA Backward (SM120)](attention/sdpa_bwd_sm120.md) - [RMSNorm + SiLU](rmsnorm_silu.md) +- [MoE + Expert Parallel API](moe_ep.md) ## Installation and setup @@ -51,6 +52,11 @@ pip install --group jax # jax >= 0.5 (XLA entry points via cutlass.jax, ship ``` (For the published wheel, `pip install torch torch-c-dlpack-ext` or `pip install "jax>=0.5"` directly.) +MoE + Expert Parallel requires its dedicated optional dependencies: +```bash +pip install nvidia-cudnn-frontend[moe_ep] +``` + After installation, you can import the APIs directly from the `cudnn` package, i.e. `from cudnn import {your_operation}` ## API Usage From b562effbf8d1e9fb2b591c492717284b27eccf9e Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 26 Aug 2026 19:04:23 +0000 Subject: [PATCH 10/31] ep size restriction remove Signed-off-by: Varun Thumbe --- docs/fe-oss-apis/moe_ep.md | 27 ++++++++++--------- .../cudnn/moe_ep/_megamoe_backend/README.md | 19 ++++++------- .../moe_ep/_megamoe_backend/_capability.py | 9 ------- .../moe_ep/_megamoe_backend/mxfp8/_config.py | 6 ++--- python/cudnn/moe_ep/api.py | 3 ++- test/python/moe_ep/test_moe_ep_backward.py | 22 ++++++++++++++- test/python/moe_ep/test_moe_ep_forward.py | 25 +++++++++++++++++ 7 files changed, 75 insertions(+), 36 deletions(-) diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md index f2ddd1924..ba542ca2f 100644 --- a/docs/fe-oss-apis/moe_ep.md +++ b/docs/fe-oss-apis/moe_ep.md @@ -5,7 +5,7 @@ owners, and executable PyTorch reference. The current device target is the Rubin training `fwd_glu` kernel plus the restricted `bwd_dglu` path on exactly SM107 (compute capability 10.7). It accepts MXFP8 E4M3/E8M0 operands or plain BF16/FP16/FP32 operands staged to MXFP8, supports BF16 or MXFP8 combine with -BF16 output, EP1 through EP16 forward, and EP1/EP2/EP4 backward. Unsupported +BF16 output, and any positive EP size for forward and backward. Unsupported combinations fail explicitly rather than returning uncomputed storage. The design removes workspace pointers, peer pointer mappers, streams, and @@ -65,14 +65,15 @@ executable CuTeDSL backend. Its deliberately narrow capability is: `backward_wgrad_mode="operands"` requires 256-row alignment and additionally returns caller-owned MXFP8 grouped-wgrad operands. Training execution does not support CUDA Graph capture; the restricted device backward returns - activation and router-weight gradients for EP1/EP2/EP4; + activation and router-weight gradients; backward hardware acceptance + currently covers EP1/EP2/EP4; - forward and backward support optional `gate_up_clamp`; Rubin training execution requires `apply_topk_in_fc1=True`; - `max_tokens_per_rank` must be explicitly positive; `top_k <= 32`, `H % 128 == 0`, and `I % 256 == 0`; - `ep_group=None` remains explicit single-rank execution. Distributed execution accepts any initialized `torch.distributed.ProcessGroup`, including - non-contiguous global-rank membership, with EP2 through EP16. The public + non-contiguous global-rank membership, with any positive EP size. The public contract requires `top_k <= num_experts` and the device path additionally requires `top_k <= 32`; `top_k` may exceed `experts_per_rank`. Expert ownership, peer tables, and route metadata use dense group-relative EP ranks; @@ -367,8 +368,8 @@ does not change its output structure across calls. `backward` requires the operator to be constructed with `generate_c=True`; the production entry point is `MoeEp.backward`, while `MoeEpReference.backward` -defines its executable semantic oracle. The restricted Rubin MXFP8 device path -supports BF16/MXFP8 combine, BF16 output, EP1/EP2/EP4, +defines its executable semantic oracle. The Rubin MXFP8 device path supports +BF16/MXFP8 combine, BF16 output, any positive EP size, `apply_topk_in_fc1=True`, and optional `gate_up_clamp`. It is a collective: every rank in `ep_group` must call it because gradients re-dispatch along the identical forward routes. @@ -837,7 +838,7 @@ The status below applies specifically to the production Rubin device backend. - Multi-node eager forward parity has passed for EP12/WORLD12 (three nodes, four PEs per node) and EP16/WORLD16 (four nodes, four PEs per node), with BF16 and MXFP8 combine and all-`-1` route behavior. -- Restricted MXFP8 backward has passed for EP1, EP2, and EP4 with BF16 and +- MXFP8 backward has passed for EP1, EP2, and EP4 with BF16 and MXFP8 combine. It checks activation and router-weight gradients against the executable reference; EP1 additionally covers repeated calls and explicitly reordered forward stashes. @@ -848,8 +849,10 @@ The status below applies specifically to the production Rubin device backend. ### Supported but awaiting hardware validation -- Forward capability accepts every EP size from EP1 through EP16. EP5, EP6, - EP8 through EP11, EP13, and EP14 have no current hardware acceptance case. +- Forward capability accepts every positive EP size. EP5, EP6, EP8 through + EP11, EP13, EP14, and sizes above EP16 have no current hardware acceptance + case. Sizes above EP16 use the generated vector peer-offset path instead of + the fixed 128-byte by-value table used through EP16. - EP7/WORLD14 is defined as a seven-node subgroup with one EP PE per node, and EP15/WORLD20 is defined as a five-node subgroup with three EP PEs per node. Both acceptance cases remain pending on allocations with the required node @@ -866,14 +869,14 @@ The status below applies specifically to the production Rubin device backend. ### Currently unsupported by the device backend -- Devices other than SM107, EP sizes above 16, non-positive or unspecified - `max_tokens_per_rank`, `hidden_size` not divisible by 128, +- Devices other than SM107, non-positive or unspecified `max_tokens_per_rank`, + `hidden_size` not divisible by 128, `intermediate_size` not divisible by 256, and `top_k > 32`. - Native NVFP4 operands or combine, any non-BF16 public output, and plain operand dtypes other than BF16/FP16/FP32. MXFP8 operands and BF16/MXFP8 combine with BF16 output are supported. -- Backward outside EP1/EP2/EP4, execution with - `apply_topk_in_fc1=False`, non-BF16 output, or backward CUDA Graph capture. +- Backward execution with `apply_topk_in_fc1=False`, non-BF16 output, or + backward CUDA Graph capture. - Wgrad operand mode with padding other than 256, without `generate_c=True`, or outside the restricted Rubin MXFP8 backward configuration. The mode returns operands only; dense `dW1`/`dW2` computation remains an explicit diff --git a/python/cudnn/moe_ep/_megamoe_backend/README.md b/python/cudnn/moe_ep/_megamoe_backend/README.md index 8fddc1a52..29539fda3 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/README.md +++ b/python/cudnn/moe_ep/_megamoe_backend/README.md @@ -25,8 +25,8 @@ separate: a request may be valid for `MoeEp` but unavailable in this backend. - `hidden_size` must be divisible by 128, and `intermediate_size` must be divisible by 256. - `top_k` must not exceed 32. -- Forward EP size must not exceed 16 because the validated peer-mapper ABI - carries a fixed 128-byte by-value offset table. +- EP sizes above 16 use a generated vector peer-offset table; EP sizes through + 16 use the fixed 128-byte by-value table. These are backend limits, not additional public `MoeEp` semantics. They remain precise, product-specific capability gates rather than hidden padding or a @@ -37,9 +37,9 @@ silent numerical fallback. `MoeEp.backward` has a validated backend seam and requires a forward stash from `generate_c=True`. In the default `backward_wgrad_mode="none"`, the restricted Rubin MXFP8 path returns -`(grad_activation, grad_topk_weights)` for EP1/EP2/EP4 with BF16 or MXFP8 -combine, BF16 output, `apply_topk_in_fc1=True`, optional `gate_up_clamp`, and -eager execution. It uses `fc1_c` and `route_metadata` to reconstruct an +`(grad_activation, grad_topk_weights)` for any positive EP size with BF16 or +MXFP8 combine, BF16 output, `apply_topk_in_fc1=True`, optional +`gate_up_clamp`, and eager execution. It uses `fc1_c` and `route_metadata` to reconstruct an external pool-layout `fc1_preact` tensor and converts `grad_output` to FP32 before re-dispatching it for semantic dprob. The kernel's source-domain dprob plane is symmetric and reset before every launch; the public router-weight @@ -103,7 +103,7 @@ stashes from different forwards. This mode only produces operands; it does not launch grouped wgrad or return dense `dW1`/`dW2`. It remains eager-only and inherits the restricted Rubin -MXFP8 backward gates (SM107, EP1/EP2/EP4, BF16 output, +MXFP8 backward gates (SM107, BF16 output, `apply_topk_in_fc1=True`, and BF16/MXFP8 combine). End-to-end operand production still requires SM107 acceptance. Direct FC1/FC2 consumer execution has been validated separately on SM100 with reference-generated operands. @@ -111,8 +111,8 @@ has been validated separately on SM100 with reference-generated operands. The dGLU product emits BF16 `grad_activation`; the backend converts it to FP32 for the public return. This is a documented BF16-rounded numerical limitation, not strict FP32 dgrad parity. `apply_topk_in_fc1=False`, NVFP4 operands or -combine, non-BF16 output, backward CUDA Graph capture, and EP sizes outside -1/2/4 remain capability-gated. +combine, non-BF16 output, and backward CUDA Graph capture remain +capability-gated. ## Validation boundary @@ -132,4 +132,5 @@ PASS for EP12 and EP16; EP7 and EP15 remain pending. End-to-end device forward/backward parity requires SM107 hardware and the `moe_ep` optional runtime dependencies, including a CuTeDSL installation that provides `cutlass.utils.rubin_helpers`. Backward acceptance remains limited to -EP1/EP2/EP4 and is not expanded by the multi-node forward suite. +EP1/EP2/EP4; larger EP sizes are enabled but not covered by current backward +hardware acceptance. diff --git a/python/cudnn/moe_ep/_megamoe_backend/_capability.py b/python/cudnn/moe_ep/_megamoe_backend/_capability.py index d8d0124cf..a2d949c48 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/_capability.py +++ b/python/cudnn/moe_ep/_megamoe_backend/_capability.py @@ -125,11 +125,6 @@ def validate_config(config: ForwardConfig) -> None: raise NotImplementedError( "MoeEp SM107 MXFP8 dispatch currently requires top_k <= 32" ) - if config.ep_size > 16: - raise NotImplementedError( - "MoeEp SM107 MXFP8 execution supports at most EP16 because the " - "validated peer-mapper ABI uses a 128-byte by-value offset table" - ) if not config.apply_topk_in_fc1: raise NotImplementedError( "MoeEp Rubin training MegaMoE requires apply_topk_in_fc1=True" @@ -164,10 +159,6 @@ def validate_backward_request(request: ValidatedBackwardRequest) -> None: raise NotImplementedError( "MoeEp MXFP8 backward currently requires output_format='bf16'" ) - if config.ep_size not in (1, 2, 4): - raise NotImplementedError( - "MoeEp MXFP8 backward currently supports EP1/EP2/EP4" - ) if not config.apply_topk_in_fc1: raise NotImplementedError( "MoeEp MXFP8 backward currently requires apply_topk_in_fc1=True" diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py index 647ef07f0..a26a17971 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py @@ -62,10 +62,8 @@ def __post_init__(self) -> None: @classmethod def from_forward_config(cls, config: ForwardConfig) -> "Mxfp8KernelConfig": - if config.ep_size < 1 or config.ep_size > 16: - raise NotImplementedError( - "MXFP8 execution supports EP subgroup sizes from 1 through 16" - ) + if config.ep_size < 1: + raise ValueError("MXFP8 execution requires a positive EP size") if config.ep_rank < 0 or config.ep_rank >= config.ep_size: raise ValueError( f"ep_rank {config.ep_rank} is outside EP size {config.ep_size}" diff --git a/python/cudnn/moe_ep/api.py b/python/cudnn/moe_ep/api.py index 900ab4b65..4fcdcdcae 100644 --- a/python/cudnn/moe_ep/api.py +++ b/python/cudnn/moe_ep/api.py @@ -440,7 +440,8 @@ def backward( ``backward_wgrad_mode="operands"``, ``wgrad_forward_stash`` is required and the return tuple has a third ``MoeEpWgradOperands`` item. The Rubin MXFP8 device path supports BF16/MXFP8 combine and BF16 output - on EP1/EP2/EP4 under its documented capability gates. Forward and + for any positive EP size under its documented capability gates. + Backward has hardware acceptance coverage at EP1/EP2/EP4. Forward and backward both quantize each FP32 route accumulator directly to MXFP8 before top-k reduction. """ diff --git a/test/python/moe_ep/test_moe_ep_backward.py b/test/python/moe_ep/test_moe_ep_backward.py index 804c87cee..b58d3e9ae 100644 --- a/test/python/moe_ep/test_moe_ep_backward.py +++ b/test/python/moe_ep/test_moe_ep_backward.py @@ -889,7 +889,6 @@ def test_moe_ep_backward_requires_generate_c(): ("overrides", "message"), [ ({"output_format": "mxfp8"}, "output_format='bf16'"), - ({"ep_size": 8}, "supports EP1/EP2/EP4"), ({"apply_topk_in_fc1": False}, "apply_topk_in_fc1=True"), ], ) @@ -942,6 +941,27 @@ def test_mxfp8_backward_capability_accepts_gate_up_clamp(monkeypatch): _capability.validate_backward_request(request) +@pytest.mark.L0 +def test_mxfp8_backward_capability_accepts_ep_above_16(monkeypatch): + config = _config(ep_size=32, ep_rank=31) + args = _inputs() + request = _validate_backward( + config, + torch.randn(2, 128), + args, + torch.randn(3, 512, dtype=torch.bfloat16), + torch.zeros(3, 4, dtype=torch.int32), + ) + monkeypatch.setattr(_capability, "_validate_device", lambda device: None) + monkeypatch.setattr( + _capability, + "_is_cuda_stream_capturing", + lambda device: False, + ) + + _capability.validate_backward_request(request) + + @pytest.mark.L0 def test_mxfp8_backward_capability_rejects_cuda_graph_capture(monkeypatch): config = _config() diff --git a/test/python/moe_ep/test_moe_ep_forward.py b/test/python/moe_ep/test_moe_ep_forward.py index 8ca620596..fbddf6121 100644 --- a/test/python/moe_ep/test_moe_ep_forward.py +++ b/test/python/moe_ep/test_moe_ep_forward.py @@ -1411,6 +1411,31 @@ def test_resolve_runtime_world_revalidates_ordered_membership(monkeypatch): _resolve_world(config) +@pytest.mark.L0 +def test_megamoe_capability_and_kernel_config_accept_ep_above_16(): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend._capability import validate_config + from cudnn.moe_ep._megamoe_backend.mxfp8._config import ( + Mxfp8KernelConfig, + ) + + with MoeEp(**_forward_config()) as op: + config = replace( + op._forward_config, + num_experts=32, + experts_per_rank=1, + ep_size=32, + ep_rank=31, + ep_group=object(), + ep_global_ranks=tuple(range(32)), + ) + + validate_config(config) + kernel_config = Mxfp8KernelConfig.from_forward_config(config) + assert kernel_config.world_size == 32 + assert kernel_config.local_rank == 31 + + @pytest.mark.L0 def test_megamoe_capability_accepts_nonworld_subgroup_config(): from cudnn.moe_ep._contracts import ForwardConfig From cd35ed5ad123f635377ee299ad75a07ce0827428 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 26 Aug 2026 20:01:01 +0000 Subject: [PATCH 11/31] add drop_on_overflow and max recv size per rank configuration Signed-off-by: Varun Thumbe --- docs/fe-oss-apis/moe_ep.md | 10 ++++++-- python/cudnn/moe_ep/_contracts.py | 2 ++ .../cudnn/moe_ep/_megamoe_backend/README.md | 3 +++ .../mxfp8/_backward_compile.py | 8 ++---- .../moe_ep/_megamoe_backend/mxfp8/_config.py | 8 +++++- python/cudnn/moe_ep/api.py | 16 ++++++++++++ test/python/moe_ep/test_moe_ep_forward.py | 25 ++++++++++++++++++- 7 files changed, 62 insertions(+), 10 deletions(-) diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md index ba542ca2f..59df7de1a 100644 --- a/docs/fe-oss-apis/moe_ep.md +++ b/docs/fe-oss-apis/moe_ep.md @@ -210,6 +210,8 @@ class MoeEp: top_k: int, ep_group: Optional[torch.distributed.ProcessGroup] = None, max_tokens_per_rank: Optional[int] = None, + max_recv_size_per_rank: Optional[int] = None, + drop_on_overflow: bool = False, output_format: Literal["bf16", "mxfp8", "nvfp4"] = "bf16", combine_format: Literal["bf16", "mxfp8", "nvfp4"] = "bf16", apply_topk_in_fc1: bool = True, @@ -293,6 +295,8 @@ The context-manager form is preferred when deterministic release matters. | `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 constructor, but the current device capability gate requires an explicit positive value on first execution. | +| `max_recv_size_per_rank` | Optional bound on routed rows received by one EP rank. The default is the conservative `ep_size * max_tokens_per_rank * top_k`; a smaller bound reduces workspace size. | +| `drop_on_overflow` | If `True`, truncate routed rows beyond `max_recv_size_per_rank`; if `False` (default), overflow traps instead of silently changing results. | | `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 the combine-rounded FC2 route contribution in the standalone top-k reducer. | @@ -611,8 +615,10 @@ 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. +is `ep_size * max_tokens_per_rank * top_k`. `max_recv_size_per_rank` selects a +smaller static workspace bound; callers must provide a router capacity contract +or accept `drop_on_overflow=True`. With `drop_on_overflow=False`, exceeding the +bound fails explicitly. ## Mapping to the MegaMoE interface diff --git a/python/cudnn/moe_ep/_contracts.py b/python/cudnn/moe_ep/_contracts.py index bad8a0c60..b0c46d57b 100644 --- a/python/cudnn/moe_ep/_contracts.py +++ b/python/cudnn/moe_ep/_contracts.py @@ -37,6 +37,8 @@ class ForwardConfig: sf_padding_size: int tuning: MoeEpTuningConfig backward_wgrad_mode: Literal["none", "operands"] = "none" + max_recv_size_per_rank: Optional[int] = None + drop_on_overflow: bool = False @dataclass(frozen=True) diff --git a/python/cudnn/moe_ep/_megamoe_backend/README.md b/python/cudnn/moe_ep/_megamoe_backend/README.md index 29539fda3..0b622fb14 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/README.md +++ b/python/cudnn/moe_ep/_megamoe_backend/README.md @@ -27,6 +27,9 @@ separate: a request may be valid for `MoeEp` but unavailable in this backend. - `top_k` must not exceed 32. - EP sizes above 16 use a generated vector peer-offset table; EP sizes through 16 use the fixed 128-byte by-value table. +- `max_recv_size_per_rank` may bound per-rank routed rows below the conservative + `EP * max_tokens_per_rank * top_k` workspace size. Overflow traps by default + or truncates only when `drop_on_overflow=True`. These are backend limits, not additional public `MoeEp` semantics. They remain precise, product-specific capability gates rather than hidden padding or a diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py index 77530ae39..931f1d55b 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py @@ -148,14 +148,10 @@ def prepare_backward_kernel( local_rank=0, num_topk=config.top_k, max_tokens_per_rank=config.max_tokens_per_rank, - max_recv_size_per_rank=( - config.world_size - * config.max_tokens_per_rank - * config.top_k - ), + max_recv_size_per_rank=config.max_recv_size_per_rank, hidden=config.hidden, launch_cluster_count=launch_cluster_count, - drop_on_overflow=True, + drop_on_overflow=config.drop_on_overflow, fc2_in_kernel_topk_reduce=False, token_back_mode="epi_warps", epi_flag_batch=config.epi_flag_batch, diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py index a26a17971..e6a6040df 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py @@ -70,9 +70,14 @@ def from_forward_config(cls, config: ForwardConfig) -> "Mxfp8KernelConfig": ) if config.max_tokens_per_rank is None: raise ValueError("MXFP8 execution requires max_tokens_per_rank") - max_recv_size_per_rank = ( + worst_case_recv_size = ( config.ep_size * config.max_tokens_per_rank * config.top_k ) + max_recv_size_per_rank = ( + worst_case_recv_size + if config.max_recv_size_per_rank is None + else min(config.max_recv_size_per_rank, worst_case_recv_size) + ) if max_recv_size_per_rank <= 0: raise ValueError("max_recv_size_per_rank must be positive") return cls( @@ -87,6 +92,7 @@ def from_forward_config(cls, config: ForwardConfig) -> "Mxfp8KernelConfig": gate_up_clamp=config.gate_up_clamp, generate_c=config.generate_c, max_recv_size_per_rank=max_recv_size_per_rank, + drop_on_overflow=config.drop_on_overflow, combine_format=combine_wire_format(config.combine_format), enable_col_quant=( config.backward_wgrad_mode == "operands" diff --git a/python/cudnn/moe_ep/api.py b/python/cudnn/moe_ep/api.py index 4fcdcdcae..cb6f5674d 100644 --- a/python/cudnn/moe_ep/api.py +++ b/python/cudnn/moe_ep/api.py @@ -112,6 +112,8 @@ def __init__( top_k: int, ep_group: Optional[dist.ProcessGroup] = None, max_tokens_per_rank: Optional[int] = None, + max_recv_size_per_rank: Optional[int] = None, + drop_on_overflow: bool = False, output_format: Union[MoeFormat, str] = MoeFormat.BF16, combine_format: Union[MoeFormat, str] = MoeFormat.BF16, apply_topk_in_fc1: bool = True, @@ -141,6 +143,16 @@ def __init__( raise ValueError( "max_tokens_per_rank must be a non-negative integer or None" ) + if max_recv_size_per_rank is not None and ( + isinstance(max_recv_size_per_rank, bool) + or not isinstance(max_recv_size_per_rank, int) + or max_recv_size_per_rank <= 0 + ): + raise ValueError( + "max_recv_size_per_rank must be a positive integer or None" + ) + if not isinstance(drop_on_overflow, bool): + raise ValueError("drop_on_overflow must be a bool") if not isinstance(apply_topk_in_fc1, bool): raise ValueError("apply_topk_in_fc1 must be a bool") if not isinstance(generate_c, bool): @@ -208,6 +220,8 @@ def __init__( self.ep_global_ranks = ep_global_ranks self.experts_per_rank = num_experts // ep_size self.max_tokens_per_rank = max_tokens_per_rank + self.max_recv_size_per_rank = max_recv_size_per_rank + self.drop_on_overflow = drop_on_overflow self.output_format = _parse_format(output_format) self.combine_format = _parse_format(combine_format) self.apply_topk_in_fc1 = apply_topk_in_fc1 @@ -246,6 +260,8 @@ def __init__( ep_group=self.ep_group, ep_global_ranks=self.ep_global_ranks, max_tokens_per_rank=self.max_tokens_per_rank, + max_recv_size_per_rank=self.max_recv_size_per_rank, + drop_on_overflow=self.drop_on_overflow, output_format=self.output_format.value, combine_format=self.combine_format.value, apply_topk_in_fc1=self.apply_topk_in_fc1, diff --git a/test/python/moe_ep/test_moe_ep_forward.py b/test/python/moe_ep/test_moe_ep_forward.py index fbddf6121..ff5cc2092 100644 --- a/test/python/moe_ep/test_moe_ep_forward.py +++ b/test/python/moe_ep/test_moe_ep_forward.py @@ -212,7 +212,7 @@ def test_moe_ep_tuning_public_contract_mapping_and_cache_key(): assert effective["effective_group_hint"] == 768 assert effective["fc2_in_kernel_topk_reduce"] is True assert effective["launch_cluster_count"] == 123 - assert effective["drop_on_overflow"] is True + assert effective["drop_on_overflow"] is False assert effective["enable_col_quant"] is False assert "output_format" not in effective @@ -272,6 +272,29 @@ def test_internal_column_requant_config_is_disabled_by_default_and_cache_distinc ) +@pytest.mark.L0 +def test_bounded_receive_capacity_propagates_to_kernel_config(): + from cudnn import MoeEp + from cudnn.moe_ep._megamoe_backend.mxfp8._config import ( + Mxfp8KernelConfig, + ) + + with MoeEp( + **_forward_config(), + max_recv_size_per_rank=7, + drop_on_overflow=False, + ) as op: + config = Mxfp8KernelConfig.from_forward_config(op._forward_config) + + assert config.max_recv_size_per_rank == 7 + assert config.drop_on_overflow is False + + with pytest.raises(ValueError, match="max_recv_size_per_rank"): + MoeEp(**_forward_config(), max_recv_size_per_rank=0) + with pytest.raises(ValueError, match="drop_on_overflow"): + MoeEp(**_forward_config(), drop_on_overflow=1) + + @pytest.mark.L0 @pytest.mark.parametrize( ("public_format", "wire_format"), From c7cf20506a11a812140551d9a14764e5b098360b Mon Sep 17 00:00:00 2001 From: zhibinz Date: Fri, 28 Aug 2026 07:58:05 -0700 Subject: [PATCH 12/31] Add graph-stable fixed-resource MoeEP training Replace the dynamic stash path with slot/lane resources, integrate the Rubin WGrad kernels, and harden distributed runtime lifecycle and coverage. --- python/cudnn/__init__.py | 17 +- python/cudnn/moe_ep/__init__.py | 14 +- python/cudnn/moe_ep/_backend.py | 44 +- python/cudnn/moe_ep/_contracts.py | 21 +- .../moe_ep/_megamoe_backend/_capability.py | 31 +- python/cudnn/moe_ep/_megamoe_backend/_comm.py | 81 +- python/cudnn/moe_ep/_megamoe_backend/_plan.py | 8 +- .../cudnn/moe_ep/_megamoe_backend/_runtime.py | 253 +- .../moe_ep/_megamoe_backend/_workspace.py | 19 +- .../cutedsl_src/VENDOR_INFO.md | 9 +- .../cutedsl_src/helpers/smem_workspace.py | 1 + .../mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py | 268 +- .../bwd_dglu/dglu_mxfp8_fc12_extension.py | 136 +- .../mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py | 73 +- .../bwd_dglu/dglu_mxfp8_mega_moe_kernel.py | 48 +- .../mega/fwd_glu/glu_mxfp8_col_requant.py | 98 +- .../mega/fwd_glu/glu_mxfp8_fc12_epilogue.py | 1 + .../mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py | 19 +- .../moe_ep/_megamoe_backend/mxfp8/_adapter.py | 21 +- .../moe_ep/_megamoe_backend/mxfp8/_backend.py | 160 +- .../_megamoe_backend/mxfp8/_backward.py | 112 - .../mxfp8/_backward_dispatch.py | 204 -- .../_megamoe_backend/mxfp8/_backward_dprob.py | 87 - .../mxfp8/_backward_launch.py | 5 + .../mxfp8/_backward_layout.py | 82 - .../mxfp8/_backward_staging.py | 516 ---- .../mxfp8/_backward_wgrad_export.py | 226 -- .../moe_ep/_megamoe_backend/mxfp8/_compile.py | 36 + .../moe_ep/_megamoe_backend/mxfp8/_stash.py | 302 --- .../mxfp8/_training_execute.py | 278 ++ .../mxfp8/_training_resources.py | 1336 ++++++++++ .../_megamoe_backend/mxfp8/_training_stage.py | 199 ++ .../mxfp8/_training_stage_kernel.py | 134 + .../mxfp8/_training_weights.py | 356 +++ .../_megamoe_backend/mxfp8/_training_wgrad.py | 176 ++ .../mxfp8/_training_wgrad_kernel.py | 177 ++ .../_megamoe_backend/mxfp8/_wgrad_layout.py | 410 --- python/cudnn/moe_ep/_types.py | 272 +- python/cudnn/moe_ep/_validation.py | 454 +--- python/cudnn/moe_ep/api.py | 245 +- test/python/moe_ep/moe_ep_backward_support.py | 212 -- .../moe_ep/moe_ep_distributed_workers.py | 591 ++--- test/python/moe_ep/moe_ep_forward_support.py | 319 --- test/python/moe_ep/moe_ep_test_data.py | 168 -- test/python/moe_ep/moe_ep_test_support.py | 1402 ++++++++++ .../moe_ep/probe_moe_ep_training_graph.py | 994 ++++++++ test/python/moe_ep/test_moe_ep_backward.py | 2256 ++++++++++------- .../test_moe_ep_cutedsl_grad_y2_source.py | 63 - test/python/moe_ep/test_moe_ep_forward.py | 322 +-- .../moe_ep/test_moe_ep_forward_multinode.py | 209 -- test/python/moe_ep/test_moe_ep_multinode.py | 365 +++ .../moe_ep/test_moe_ep_wgrad_contract.py | 1170 --------- 52 files changed, 8405 insertions(+), 6595 deletions(-) delete mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward.py delete mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dispatch.py delete mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dprob.py delete mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_layout.py delete mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_staging.py delete mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_wgrad_export.py delete mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_stash.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py delete mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_wgrad_layout.py delete mode 100644 test/python/moe_ep/moe_ep_backward_support.py delete mode 100644 test/python/moe_ep/moe_ep_forward_support.py delete mode 100644 test/python/moe_ep/moe_ep_test_data.py create mode 100644 test/python/moe_ep/moe_ep_test_support.py create mode 100644 test/python/moe_ep/probe_moe_ep_training_graph.py delete mode 100644 test/python/moe_ep/test_moe_ep_cutedsl_grad_y2_source.py delete mode 100644 test/python/moe_ep/test_moe_ep_forward_multinode.py create mode 100644 test/python/moe_ep/test_moe_ep_multinode.py delete mode 100644 test/python/moe_ep/test_moe_ep_wgrad_contract.py diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 065fd7767..395228ae0 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -315,9 +315,12 @@ def _dlopen_cudnn(): "moe_ep", "BlockScaledTensor", "MoeEp", + "MoeEpExecutionLane", + "MoeEpTrainingResources", + "MoeEpTrainingSlot", + "MoeEpTrainingWeights", + "MoeEpTrainingWgradOperands", "MoeEpTuningConfig", - "MoeEpWgradForwardStash", - "MoeEpWgradOperands", "MoeFormat", "MoeTensor", } @@ -327,9 +330,15 @@ def _dlopen_cudnn(): "moe_ep": (".moe_ep", None), "BlockScaledTensor": (".moe_ep", "BlockScaledTensor"), "MoeEp": (".moe_ep", "MoeEp"), + "MoeEpExecutionLane": (".moe_ep", "MoeEpExecutionLane"), + "MoeEpTrainingResources": (".moe_ep", "MoeEpTrainingResources"), + "MoeEpTrainingSlot": (".moe_ep", "MoeEpTrainingSlot"), + "MoeEpTrainingWeights": (".moe_ep", "MoeEpTrainingWeights"), + "MoeEpTrainingWgradOperands": ( + ".moe_ep", + "MoeEpTrainingWgradOperands", + ), "MoeEpTuningConfig": (".moe_ep", "MoeEpTuningConfig"), - "MoeEpWgradForwardStash": (".moe_ep", "MoeEpWgradForwardStash"), - "MoeEpWgradOperands": (".moe_ep", "MoeEpWgradOperands"), "MoeFormat": (".moe_ep", "MoeFormat"), "MoeTensor": (".moe_ep", "MoeTensor"), "BSA": (".block_sparse_attention", "BSA"), diff --git a/python/cudnn/moe_ep/__init__.py b/python/cudnn/moe_ep/__init__.py index ed160766f..8557cff2d 100644 --- a/python/cudnn/moe_ep/__init__.py +++ b/python/cudnn/moe_ep/__init__.py @@ -4,8 +4,11 @@ from ._tuning import MoeEpTuningConfig from ._types import ( BlockScaledTensor, - MoeEpWgradForwardStash, - MoeEpWgradOperands, + MoeEpExecutionLane, + MoeEpTrainingResources, + MoeEpTrainingSlot, + MoeEpTrainingWeights, + MoeEpTrainingWgradOperands, MoeFormat, MoeTensor, ) @@ -14,9 +17,12 @@ __all__ = [ "BlockScaledTensor", "MoeEp", + "MoeEpExecutionLane", + "MoeEpTrainingResources", + "MoeEpTrainingSlot", + "MoeEpTrainingWeights", + "MoeEpTrainingWgradOperands", "MoeEpTuningConfig", - "MoeEpWgradForwardStash", - "MoeEpWgradOperands", "MoeFormat", "MoeTensor", ] diff --git a/python/cudnn/moe_ep/_backend.py b/python/cudnn/moe_ep/_backend.py index 4efa856ab..11e84a977 100644 --- a/python/cudnn/moe_ep/_backend.py +++ b/python/cudnn/moe_ep/_backend.py @@ -10,43 +10,20 @@ from __future__ import annotations -from typing import Protocol, Tuple, Union +from typing import Protocol import torch -from ._contracts import ( - ForwardConfig, - ValidatedBackwardRequest, - ValidatedForwardRequest, -) -from ._types import MoeEpWgradForwardStash, MoeEpWgradOperands, MoeTensor - - -ForwardResult = Union[ - MoeTensor, - Tuple[MoeTensor, torch.Tensor, torch.Tensor], - Tuple[ - MoeTensor, - torch.Tensor, - torch.Tensor, - MoeEpWgradForwardStash, - ], -] -BackwardResult = Union[ - Tuple[torch.Tensor, torch.Tensor], - Tuple[torch.Tensor, torch.Tensor, MoeEpWgradOperands], -] +from ._contracts import ForwardConfig, ValidatedForwardRequest +from ._types import MoeTensor class MoeEpBackend(Protocol): """Instance-local backend created lazily for one static ``MoeEp`` config.""" - def forward(self, request: ValidatedForwardRequest) -> ForwardResult: + def forward(self, request: ValidatedForwardRequest) -> MoeTensor: """Execute one already-validated forward request.""" - def backward(self, request: ValidatedBackwardRequest) -> BackwardResult: - """Execute one already-validated backward request.""" - def close(self) -> None: """Release backend-owned resources.""" @@ -71,16 +48,6 @@ def validate_request(request: ValidatedForwardRequest) -> None: validate(request) -def validate_backward_request(request: ValidatedBackwardRequest) -> None: - """Run the selected backend's backward capability gate lazily.""" - - from ._megamoe_backend._capability import ( - validate_backward_request as validate, - ) - - validate(request) - - def create_backend( config: ForwardConfig, device: torch.device, @@ -93,12 +60,9 @@ def create_backend( __all__ = [ - "BackwardResult", "BackendUnavailableError", "MoeEpBackend", - "ForwardResult", "create_backend", "validate_config", - "validate_backward_request", "validate_request", ] diff --git a/python/cudnn/moe_ep/_contracts.py b/python/cudnn/moe_ep/_contracts.py index b0c46d57b..9acb907fb 100644 --- a/python/cudnn/moe_ep/_contracts.py +++ b/python/cudnn/moe_ep/_contracts.py @@ -11,7 +11,7 @@ import torch from ._tuning import MoeEpTuningConfig -from ._types import MoeEpWgradForwardStash, MoeTensor +from ._types import MoeTensor @dataclass(frozen=True) @@ -55,26 +55,7 @@ class ValidatedForwardRequest: device: torch.device -@dataclass(frozen=True) -class ValidatedBackwardRequest: - """Runtime inputs that have passed the public backward contract.""" - - config: ForwardConfig - grad_output: torch.Tensor - fc1_weight: MoeTensor - fc2_weight: MoeTensor - topk_idx: torch.Tensor - topk_weights: torch.Tensor - fc1_c: torch.Tensor - route_metadata: torch.Tensor - token_count: int - local_routes: int - device: torch.device - wgrad_forward_stash: Optional[MoeEpWgradForwardStash] = None - - __all__ = [ "ForwardConfig", - "ValidatedBackwardRequest", "ValidatedForwardRequest", ] diff --git a/python/cudnn/moe_ep/_megamoe_backend/_capability.py b/python/cudnn/moe_ep/_megamoe_backend/_capability.py index a2d949c48..d34f74038 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/_capability.py +++ b/python/cudnn/moe_ep/_megamoe_backend/_capability.py @@ -15,7 +15,6 @@ from .._contracts import ( ForwardConfig, - ValidatedBackwardRequest, ValidatedForwardRequest, ) from .._types import BlockScaledTensor, MoeFormat @@ -73,10 +72,10 @@ def _validate_wgrad_config(config: ForwardConfig) -> None: raise ValueError( "backward_wgrad_mode='operands' requires generate_c=True" ) - if config.token_padding_size != 256: + if config.token_padding_size != 128: raise ValueError( "backward_wgrad_mode='operands' requires " - "token_padding_size=256" + "token_padding_size=128" ) if config.sf_padding_size != 128: raise ValueError( @@ -144,33 +143,7 @@ def validate_request(request: ValidatedForwardRequest) -> None: _validate_device(request.device) -def validate_backward_request(request: ValidatedBackwardRequest) -> None: - """Reject backward requests outside the Rubin MXFP8 training path.""" - - _validate_wgrad_config(request.config) - for name, tensor in ( - ("fc1_weight", request.fc1_weight), - ("fc2_weight", request.fc2_weight), - ): - _validate_operand(name, tensor) - _validate_device(request.device) - config = request.config - if config.output_format != MoeFormat.BF16.value: - raise NotImplementedError( - "MoeEp MXFP8 backward currently requires output_format='bf16'" - ) - if not config.apply_topk_in_fc1: - raise NotImplementedError( - "MoeEp MXFP8 backward currently requires apply_topk_in_fc1=True" - ) - if _is_cuda_stream_capturing(request.device): - raise NotImplementedError( - "MoeEp MXFP8 backward does not support CUDA graph capture" - ) - - __all__ = [ - "validate_backward_request", "validate_config", "validate_request", ] diff --git a/python/cudnn/moe_ep/_megamoe_backend/_comm.py b/python/cudnn/moe_ep/_megamoe_backend/_comm.py index b1379e4cd..d837558e7 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/_comm.py +++ b/python/cudnn/moe_ep/_megamoe_backend/_comm.py @@ -5,12 +5,17 @@ from __future__ import annotations +import time from dataclasses import dataclass from typing import Optional, Protocol import torch -from ._runtime import RuntimeHandle, RuntimeUnavailableError +from ._runtime import ( + RuntimeHandle, + RuntimeUnavailableError, + _runtime_debug, +) class SymmetricMemoryProvider(Protocol): @@ -53,34 +58,86 @@ def _core(): def allocate(self, nbytes: int, device: torch.device) -> torch.Tensor: del device # NVSHMEM allocates on the device bound during runtime init. + started_at = time.monotonic() + _runtime_debug("symmetric.allocate.begin", nbytes=nbytes) try: - return self._core().tensor( + tensor = self._core().tensor( (nbytes,), dtype=torch.uint8, release=False, except_on_del=True, ) except Exception as exc: + _runtime_debug( + "symmetric.allocate.error", + nbytes=nbytes, + error_type=type(exc).__name__, + error=repr(exc), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) raise RuntimeUnavailableError( f"failed to allocate {nbytes} bytes from the NVSHMEM symmetric heap" ) from exc + _runtime_debug( + "symmetric.allocate.end", + nbytes=nbytes, + data_ptr=hex(tensor.data_ptr()), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) + return tensor def free(self, tensor: torch.Tensor) -> None: + started_at = time.monotonic() + _runtime_debug( + "symmetric.free.begin", + nbytes=tensor.numel() * tensor.element_size(), + data_ptr=hex(tensor.data_ptr()), + ) try: self._core().free_tensor(tensor) except Exception as exc: + _runtime_debug( + "symmetric.free.error", + error_type=type(exc).__name__, + error=repr(exc), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) raise RuntimeUnavailableError( "failed to free the NVSHMEM symmetric root slab" ) from exc + _runtime_debug( + "symmetric.free.end", + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) def peer_address(self, tensor: torch.Tensor, peer: int) -> int: + started_at = time.monotonic() + _runtime_debug( + "symmetric.peer-map.begin", + peer=peer, + data_ptr=hex(tensor.data_ptr()), + ) try: peer_tensor = self._core().get_peer_tensor(tensor, peer) except Exception as exc: + _runtime_debug( + "symmetric.peer-map.error", + peer=peer, + error_type=type(exc).__name__, + error=repr(exc), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) raise RuntimeUnavailableError( f"failed to map symmetric root slab for peer {peer}" ) from exc - return int(peer_tensor.data_ptr()) + peer_pointer = int(peer_tensor.data_ptr()) + _runtime_debug( + "symmetric.peer-map.end", + peer=peer, + peer_data_ptr=hex(peer_pointer), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) + return peer_pointer @dataclass(frozen=True) @@ -159,6 +216,12 @@ def ensure_allocated(self) -> None: "symmetric slab has an allocation pending cleanup" ) + _runtime_debug( + "symmetric-slab.ensure.begin", + nbytes=self._nbytes, + world_size=self._runtime.world_size, + ep_rank=self._runtime.rank, + ) root = self._provider.allocate(self._nbytes, self._runtime.device) if not isinstance(root, torch.Tensor): raise TypeError( @@ -183,7 +246,13 @@ def ensure_allocated(self) -> None: raise try: + _runtime_debug( + "symmetric-slab.zero.begin", + nbytes=self._nbytes, + data_ptr=hex(root.data_ptr()), + ) root.zero_() + _runtime_debug("symmetric-slab.zero.enqueued") base_address = int(root.data_ptr()) offsets = [] @@ -211,6 +280,12 @@ def ensure_allocated(self) -> None: raise self._mapping = mapping + _runtime_debug( + "symmetric-slab.ensure.end", + nbytes=self._nbytes, + base_address=hex(mapping.base_address), + offsets=mapping.offsets, + ) @property def nbytes(self) -> int: diff --git a/python/cudnn/moe_ep/_megamoe_backend/_plan.py b/python/cudnn/moe_ep/_megamoe_backend/_plan.py index a2b72b5fa..c78b880cb 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/_plan.py +++ b/python/cudnn/moe_ep/_megamoe_backend/_plan.py @@ -11,11 +11,7 @@ import torch -from .._contracts import ( - ForwardConfig, - ValidatedBackwardRequest, - ValidatedForwardRequest, -) +from .._contracts import ForwardConfig, ValidatedForwardRequest from ._comm import SymmetricMemoryProvider from ._runtime import RuntimeHandle, RuntimeManager, get_runtime_manager from ._workspace import ( @@ -82,7 +78,7 @@ def closed(self) -> bool: def prepare( self, - request: ValidatedForwardRequest | ValidatedBackwardRequest, + request: ValidatedForwardRequest, ) -> PreparedResources: with self._lock: if self._closed: diff --git a/python/cudnn/moe_ep/_megamoe_backend/_runtime.py b/python/cudnn/moe_ep/_megamoe_backend/_runtime.py index 9976acd11..126431d85 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/_runtime.py +++ b/python/cudnn/moe_ep/_megamoe_backend/_runtime.py @@ -10,10 +10,16 @@ from __future__ import annotations +import faulthandler import logging +import os +import socket +import sys import threading +import time from dataclasses import dataclass from enum import Enum +from pathlib import Path from typing import Callable, Optional, Protocol import torch @@ -24,6 +30,104 @@ _logger = logging.getLogger(__name__) +def _runtime_debug_enabled() -> bool: + return os.environ.get("MOE_EP_DEBUG_RUNTIME", "0") == "1" + + +def _runtime_debug(event: str, **details: object) -> None: + if not _runtime_debug_enabled(): + return + fields = { + "time": f"{time.monotonic():.6f}", + "host": socket.gethostname(), + "pid": os.getpid(), + "rank": os.environ.get("RANK", "?"), + "local_rank": os.environ.get("LOCAL_RANK", "?"), + "event": event, + **details, + } + print( + "[moe-ep-runtime] " + + " ".join(f"{name}={value}" for name, value in fields.items()), + file=sys.stderr, + flush=True, + ) + + +def _runtime_debug_init_status(core) -> object: + if not _runtime_debug_enabled(): + return "debug-disabled" + try: + status = core.init_status() + except (AttributeError, RuntimeError): + return "unavailable" + return getattr(status, "name", status) + + +class _RuntimeWatchdog: + """Emit Python stacks and kernel wait channels while NVSHMEM init is blocked.""" + + def __init__(self, event: str) -> None: + self._event = event + self._stopped = threading.Event() + try: + self._interval = float( + os.environ.get("MOE_EP_RUNTIME_WATCHDOG_SECONDS", "30") + ) + except ValueError: + self._interval = 30.0 + self._thread: Optional[threading.Thread] = None + + def start(self) -> None: + if not _runtime_debug_enabled() or self._interval <= 0: + return + # faulthandler's timer is implemented outside the Python interpreter + # lock, so it still emits the main-thread stack if a native NVSHMEM + # call holds the GIL. The Python thread adds /proc wait-channel data + # whenever the GIL remains schedulable. + faulthandler.dump_traceback_later( + self._interval, + repeat=True, + file=sys.stderr, + ) + self._thread = threading.Thread( + target=self._run, + name="moe-ep-runtime-watchdog", + daemon=True, + ) + self._thread.start() + + def close(self) -> None: + self._stopped.set() + if _runtime_debug_enabled() and self._interval > 0: + faulthandler.cancel_dump_traceback_later() + if self._thread is not None: + self._thread.join(timeout=1) + + def _run(self) -> None: + sample = 0 + while not self._stopped.wait(self._interval): + sample += 1 + wait_channels: list[str] = [] + for task_dir in sorted(Path("/proc/self/task").glob("[0-9]*")): + try: + thread_name = (task_dir / "comm").read_text().strip() + wait_channel = (task_dir / "wchan").read_text().strip() + except OSError as exc: + wait_channels.append(f"{task_dir.name}:unavailable({exc.errno})") + else: + wait_channels.append( + f"{task_dir.name}:{thread_name}:{wait_channel or '-'}" + ) + _runtime_debug( + "watchdog", + blocked_event=self._event, + sample=sample, + threads=";".join(wait_channels), + ) + faulthandler.dump_traceback(file=sys.stderr, all_threads=True) + + class RuntimeUnavailableError(RuntimeError): """The requested runtime cannot be loaded or initialized.""" @@ -183,6 +287,14 @@ def initialize(self, device: torch.device, world: RuntimeWorld) -> None: raise ValueError("NVSHMEM initialization requires a process group") core = _load_nvshmem_core() + started_at = time.monotonic() + _runtime_debug( + "initialize.begin", + device=device, + ep_rank=world.rank, + ep_size=world.size, + global_ranks=world.global_ranks, + ) try: import numpy as np @@ -194,8 +306,10 @@ def initialize(self, device: torch.device, world: RuntimeWorld) -> None: torch.cuda.set_device(device) cuda_device = Device(device.index) cuda_device.set_current() + _runtime_debug("initialize.cuda-current", device=device) uid = core.get_unique_id(empty=(world.rank != 0)) + _runtime_debug("initialize.uid-created") uid_bytes = uid._data.view(np.uint8).copy() uid_tensor = torch.from_numpy(uid_bytes) group_backend = dist.get_backend(world.group) @@ -209,24 +323,51 @@ def initialize(self, device: torch.device, world: RuntimeWorld) -> None: raise RuntimeError( "EP subgroup root changed during NVSHMEM bootstrap" ) + _runtime_debug( + "initialize.uid-broadcast.begin", + backend=group_backend, + root_global_rank=root_global_rank, + tensor_device=uid_tensor.device, + tensor_bytes=uid_tensor.numel() * uid_tensor.element_size(), + ) dist.broadcast( uid_tensor, src=root_global_rank, group=world.group, ) + _runtime_debug("initialize.uid-broadcast.end") + _runtime_debug("initialize.torch-barrier.begin") dist.barrier(group=world.group) + _runtime_debug("initialize.torch-barrier.end") uid._data[:] = uid_tensor.cpu().numpy().view(uid._data.dtype) - core.init( - device=cuda_device, - uid=uid, - rank=world.rank, - nranks=world.size, - initializer_method="uid", + watchdog = _RuntimeWatchdog("core.init") + _runtime_debug("initialize.core-init.begin") + watchdog.start() + try: + core.init( + device=cuda_device, + uid=uid, + rank=world.rank, + nranks=world.size, + initializer_method="uid", + ) + finally: + watchdog.close() + _runtime_debug( + "initialize.core-init.end", + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + init_status=_runtime_debug_init_status(core), ) except RuntimeUnavailableError: raise except Exception as exc: + _runtime_debug( + "initialize.error", + error_type=type(exc).__name__, + error=repr(exc), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) raise RuntimeUnavailableError( "failed to initialize the NVSHMEM EP subgroup runtime" ) from exc @@ -259,10 +400,31 @@ def device(self) -> torch.device: ) from exc def finalize(self) -> None: + core = _load_nvshmem_core() + started_at = time.monotonic() + _runtime_debug( + "finalize.begin", + init_status=_runtime_debug_init_status(core), + ) + watchdog = _RuntimeWatchdog("core.finalize") + watchdog.start() try: - _load_nvshmem_core().finalize() + core.finalize() except Exception as exc: + _runtime_debug( + "finalize.error", + error_type=type(exc).__name__, + error=repr(exc), + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + ) raise RuntimeUnavailableError("failed to finalize NVSHMEM") from exc + finally: + watchdog.close() + _runtime_debug( + "finalize.end", + elapsed_seconds=f"{time.monotonic() - started_at:.3f}", + init_status=_runtime_debug_init_status(core), + ) @dataclass @@ -343,9 +505,11 @@ def __init__( _DefaultNvshmemRuntimeProvider ), world_resolver: Callable[[ForwardConfig], RuntimeWorld] = _resolve_world, + keep_alive: bool = False, ) -> None: self._provider_factory = provider_factory self._world_resolver = world_resolver + self._keep_alive = keep_alive @property def ref_count(self) -> int: @@ -370,6 +534,12 @@ def acquire( with _PROCESS_RUNTIME_REGISTRY.lock: if _PROCESS_RUNTIME_REGISTRY.active is not None: active = _PROCESS_RUNTIME_REGISTRY.active + _runtime_debug( + "manager.acquire-reuse.begin", + ref_count=active.ref_count, + owns_runtime=active.owns_runtime, + cleanup_required=active.cleanup_required, + ) if active.cleanup_required: raise RuntimeError( "MegaMoE process runtime requires cleanup before reacquire" @@ -385,6 +555,10 @@ def acquire( "EP subgroup" ) active.ref_count += 1 + _runtime_debug( + "manager.acquire-reuse.end", + ref_count=active.ref_count, + ) return RuntimeHandle( self, active.token, @@ -395,9 +569,19 @@ def acquire( provider: Optional[NvshmemRuntimeProvider] = None owns_runtime = False + _runtime_debug( + "manager.acquire-new.begin", + device=device, + ep_rank=world.rank, + ep_size=world.size, + ) if world.size > 1: provider = self._provider_factory() status = provider.initialization_state() + _runtime_debug( + "manager.acquire-new.state", + init_status=status.value, + ) if status is RuntimeInitState.PARTIAL: raise RuntimeError( "cannot attach to a partially initialized NVSHMEM runtime" @@ -481,6 +665,11 @@ def acquire( provider=provider, owns_runtime=owns_runtime, ) + _runtime_debug( + "manager.acquire-new.end", + owns_runtime=owns_runtime, + ref_count=1, + ) return RuntimeHandle(self, token, device, world, owns_runtime) @staticmethod @@ -559,16 +748,47 @@ def retry_cleanup(self) -> None: active.provider.finalize() _PROCESS_RUNTIME_REGISTRY.active = None + def shutdown(self) -> None: + """Finalize an idle process runtime at a caller-controlled collective point.""" + + with _PROCESS_RUNTIME_REGISTRY.lock: + active = _PROCESS_RUNTIME_REGISTRY.active + if active is None: + return + if active.ref_count != 0: + raise RuntimeError( + "cannot shut down the MegaMoE process runtime while " + f"{active.ref_count} handles remain active" + ) + if active.provider is not None and ( + active.owns_runtime or active.cleanup_required + ): + _runtime_debug( + "manager.shutdown-finalize.begin", + cleanup_required=active.cleanup_required, + ) + active.provider.finalize() + _runtime_debug("manager.shutdown-finalize.end") + _PROCESS_RUNTIME_REGISTRY.active = None + _runtime_debug("manager.shutdown.end") + def _release(self, token: object) -> None: with _PROCESS_RUNTIME_REGISTRY.lock: active = _PROCESS_RUNTIME_REGISTRY.active if active is None or active.token is not token: + _runtime_debug("manager.release-stale") return if active.ref_count <= 0: raise RuntimeError( "MegaMoE process runtime has invalid release state" ) + _runtime_debug( + "manager.release.begin", + ref_count=active.ref_count, + owns_runtime=active.owns_runtime, + cleanup_required=active.cleanup_required, + ) if active.cleanup_required: if active.ref_count != 1 or active.provider is None: raise RuntimeError( @@ -576,22 +796,39 @@ def _release(self, token: object) -> None: ) active.provider.finalize() _PROCESS_RUNTIME_REGISTRY.active = None + _runtime_debug("manager.release.cleanup-retry.end") return if active.ref_count > 1: active.ref_count -= 1 + _runtime_debug( + "manager.release-retained", + ref_count=active.ref_count, + ) + return + + if self._keep_alive and not active.cleanup_required: + active.ref_count = 0 + _runtime_debug( + "manager.release-idle", + owns_runtime=active.owns_runtime, + ) return if active.owns_runtime and active.provider is not None: + _runtime_debug("manager.release-finalize.begin") try: active.provider.finalize() except Exception: active.cleanup_required = True + _runtime_debug("manager.release-finalize.error") raise + _runtime_debug("manager.release-finalize.end") _PROCESS_RUNTIME_REGISTRY.active = None + _runtime_debug("manager.release.end", ref_count=0) -_DEFAULT_RUNTIME_MANAGER = RuntimeManager() +_DEFAULT_RUNTIME_MANAGER = RuntimeManager(keep_alive=True) def get_runtime_manager() -> RuntimeManager: diff --git a/python/cudnn/moe_ep/_megamoe_backend/_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/_workspace.py index 97f9ba793..ad65394e9 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/_workspace.py +++ b/python/cudnn/moe_ep/_megamoe_backend/_workspace.py @@ -19,7 +19,7 @@ SymmetricSlab, _TorchMemoryProvider, ) -from ._runtime import RuntimeHandle +from ._runtime import RuntimeHandle, _runtime_debug def _align_up(value: int, alignment: int) -> int: @@ -257,7 +257,13 @@ def __init__( raise ValueError(f"local slab size must be positive, got {nbytes}") self._provider = provider self._nbytes = nbytes + _runtime_debug("local-slab.allocate.begin", nbytes=nbytes, device=device) root = provider.allocate(nbytes, device) + _runtime_debug( + "local-slab.allocate.end", + nbytes=nbytes, + data_ptr=hex(root.data_ptr()) if isinstance(root, torch.Tensor) else "?", + ) self._root: Optional[torch.Tensor] = None try: if not isinstance(root, torch.Tensor): @@ -274,7 +280,9 @@ def __init__( ) if not root.is_contiguous(): raise ValueError("local root tensor must be contiguous") + _runtime_debug("local-slab.zero.begin", nbytes=nbytes) root.zero_() + _runtime_debug("local-slab.zero.enqueued", nbytes=nbytes) except Exception: if isinstance(root, torch.Tensor): provider.free(root) @@ -372,6 +380,11 @@ def ensure_allocated(self) -> None: return self.runtime.ensure_open() + _runtime_debug( + "workspace.allocate.begin", + local_bytes=self.local_layout.total_bytes, + symmetric_bytes=self.symmetric_layout.total_bytes, + ) local = _LocalSlab( self.local_layout.total_bytes, self.runtime.device, @@ -379,13 +392,16 @@ def ensure_allocated(self) -> None: ) self._local = local try: + _runtime_debug("workspace.symmetric-slab.create.begin") symmetric = SymmetricSlab( self.runtime, self.symmetric_layout.total_bytes, provider=self._symmetric_provider, ) self._symmetric = symmetric + _runtime_debug("workspace.symmetric-slab.ensure.begin") symmetric.ensure_allocated() + _runtime_debug("workspace.symmetric-slab.ensure.end") except Exception: try: if self._symmetric is not None: @@ -398,6 +414,7 @@ def ensure_allocated(self) -> None: self._cleanup_required = True raise raise + _runtime_debug("workspace.allocate.end") def views(self, token_count: int) -> WorkspaceViews: with self._lock: diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR_INFO.md b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR_INFO.md index 7aab69c5a..6c293ef91 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR_INFO.md +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR_INFO.md @@ -9,11 +9,10 @@ - Selected forward updates and backward dGLU upstream revision: `92dd334af2eeedb36087834354b58ace08e880c6` - Latest synchronized upstream revision: - `5a43c8523ea5215923c2fc8d0abae75bd6762011` (merge of source revision - `dc05bbdf38350a0eb67e9d9440e3c7c0e21e99fc`) + `5b89819cb16069dfe20a1a0ba0778d35cb428352` - Vendoring dates: 2026-08-11 (base), 2026-08-17 (selected updates), and - 2026-08-20 and 2026-08-24 (latest synchronizations). -- On 2026-08-24 every vendored Python source except the intentionally minimal + 2026-08-20, 2026-08-24, and 2026-08-28 (latest synchronizations). +- On 2026-08-28 every vendored Python source except the intentionally minimal `kernel_src/rubin/training/__init__.py` was synchronized byte-for-byte with the revision above. Other integration-specific behavior lives outside this directory. @@ -110,7 +109,7 @@ quant_def.py - Every `.py` file listed in the manifest except `kernel_src/rubin/training/__init__.py` is a byte-for-byte copy of the same - relative path at revision `5a43c8523ea5215923c2fc8d0abae75bd6762011`. + relative path at revision `5b89819cb16069dfe20a1a0ba0778d35cb428352`. - `kernel_src/rubin/training/__init__.py` is intentionally reduced to a package marker. This avoids vendoring and eagerly importing the unused Rubin traditional-wgrad product. diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py index ce412673e..e6fe04c53 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py @@ -425,3 +425,4 @@ def _require_finalized(self) -> None: "SmemWorkspace", "SwizzleSpec", ] + diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py index 46b57451d..55d78e996 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py @@ -148,6 +148,12 @@ def __init__( self._dfc2_recompute = dfc2_recompute self._dfc2_col_output = dfc2_col_output + # One PipelineTmaStore stage holds every data plane produced by a dFC2 subtile + self._d_output_slots = ( + 2 + + (2 if dfc2_col_output else 0) + + (1 if dfc2_recompute else 0) + ) # combine_format determines the dfc1 (final grad_x) combine encoding. if combine_format is None: @@ -172,7 +178,7 @@ def __init__( ) pass - + # -- Codegen-time queries -- @property @@ -187,6 +193,10 @@ def num_acc_pipeline_stages(self) -> int: def num_acc_stage(self) -> int: return self._num_acc_stage + @property + def d_output_slots(self) -> int: + return self._d_output_slots + @property def subtile_cnt(self) -> int: return self._subtile_cnt @@ -280,6 +290,24 @@ def preact_smem_layout_one_stage(self) -> Union[cute.Layout, cute.ComposedLayout def preact_bytes_per_stage(self) -> int: return cute.size_in_bytes(cutlass.BFloat16, self.preact_smem_layout_one_stage) + @cute.jit + def _store_aux_row_smem(self, r_data: cute.Tensor, s_data: cute.Tensor) -> None: + """Store one 32-byte token row as two swizzle-safe 128-bit segments.""" + segment_elements = 16 + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.fc1_output_dtype, + num_bits_per_copy=128, + ) + r_segments = cute.zipped_divide(r_data, (segment_elements,)) + s_segments = cute.zipped_divide(s_data, (segment_elements,)) + for segment in cutlass.range_constexpr(EpilogueTileN // segment_elements): + cute.copy( + store_atom, + cute.coalesce(r_segments[None, segment]), + cute.coalesce(s_segments[None, segment]), + ) + @cute.jit def _run_dfc2_task_tile( self, @@ -290,8 +318,10 @@ def _run_dfc2_task_tile( sched_ext, gmem_fc1_output: cute.Tensor, gmem_fc1_output_sf: cute.Tensor, + tma_atom_fc1_recompute: cute.CopyAtom, gmem_fc1_recompute: cute.Tensor, gmem_fc1_recompute_sf: cute.Tensor, + tma_atom_fc1_col_output: cute.CopyAtom, gmem_fc1_col_output: cute.Tensor, gmem_fc1_col_output_sf: cute.Tensor, c_pipeline, @@ -461,15 +491,48 @@ def _run_dfc2_task_tile( (self._cta_tile_m, EpilogueTileN, 1), (base_token_tile, gate_col_idx + cutlass.Int32(1), cutlass.Int32(0)), )[(None, None, 0)] + g_col_gate = None + g_col_up = None + if cutlass.const_expr(self._dfc2_col_output): + g_col_gate = cute.local_tile( + real_fc1_col_output, + (self._cta_tile_m, EpilogueTileN, 1), + (base_token_tile, gate_col_idx, cutlass.Int32(0)), + )[(None, None, 0)] + g_col_up = cute.local_tile( + real_fc1_col_output, + (self._cta_tile_m, EpilogueTileN, 1), + (base_token_tile, gate_col_idx + cutlass.Int32(1), cutlass.Int32(0)), + )[(None, None, 0)] + g_recompute = None + if cutlass.const_expr(self._dfc2_recompute): + recompute_col_idx = ( + work_tile_info.tile_n_idx + * cutlass.Int32(self._cta_tile_n // EpilogueTileN) + + subtile_idx + ) + g_recompute = cute.local_tile( + real_fc1_recompute, + (self._cta_tile_m, EpilogueTileN, 1), + (base_token_tile, recompute_col_idx, cutlass.Int32(0)), + )[(None, None, 0)] # TMA issue (warp 0 only). - d_n_slots = cutlass.const_expr(d_num_stage // 2) - d_slot = cutlass.Int32(2) * (cutlass.Int32(i) % cutlass.Int32(d_n_slots)) + d_outputs_per_stage = cutlass.const_expr(self._d_output_slots) + d_n_stages = cutlass.const_expr(d_num_stage // d_outputs_per_stage) + d_slot = cutlass.Int32(d_outputs_per_stage) * ( + cutlass.Int32(i) % cutlass.Int32(d_n_stages) + ) if warp_idx == self._epilogue_warp_ids[0]: - self.tma_store_dfc2_output( + self.tma_store_dfc2_outputs( smem_d_buffer, tma_atom_grad_y1, g_gate, g_up, + tma_atom_fc1_col_output, + g_col_gate, + g_col_up, + tma_atom_fc1_recompute, + g_recompute, valid_tokens, d_pipeline, d_slot, @@ -479,7 +542,7 @@ def _run_dfc2_task_tile( valid_inter = real_fc1_output.shape[1] self._stg_sf_dfc2(rmem_sf, real_fc1_output_sf, work_tile_info, tidx, valid_inter) - + # fc1_recompute SFs if cutlass.const_expr(self._dfc2_recompute): valid_inter_recompute = real_fc1_recompute.shape[1] @@ -618,7 +681,7 @@ def _run_dfc2_subtile( c_shape = cute.make_layout(((1, EN,), 1, 1), stride=((0, 1,), 0, 0)).shape c_gate = cute.make_rmem_tensor(c_shape, self.fc1_output_dtype) c_up = cute.make_rmem_tensor(c_shape, self.fc1_output_dtype) - # c_recompute: flat MXFP8 row for per-thread STG + # c_recompute: flat MXFP8 row for token-major output staging c_recompute = cute.make_rmem_tensor(cute.make_layout(EN).shape, self.fc1_output_dtype) is_valid_row = token_row_in_cta < valid_tokens @@ -667,32 +730,6 @@ def _run_dfc2_subtile( rmem_sf_col_output[2 * _k] = qg_col rmem_sf_col_output[2 * _k + 1] = qu_col - # STG gate + up MXFP8 cols -- only valid rows and in-range N strips. - if is_valid_row: - n_col_strips = 2 * (self._cta_tile_n // EN) - gate_strip_idx = ( - work_tile_info.tile_n_idx * cutlass.Int32(n_col_strips) - + subtile_idx * cutlass.Int32(2) - ) - up_strip_idx = gate_strip_idx + cutlass.Int32(1) - token_idx = ( - work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + token_row_in_cta - ) - for strip_idx, r_data in ((gate_strip_idx, c_gate_col), (up_strip_idx, c_up_col)): - if strip_idx * cutlass.Int32(EN) < real_fc1_col_output.shape[1]: - strip_base = cute.local_tile( - real_fc1_col_output, (1, EN, 1), - (token_idx, strip_idx, cutlass.Int32(0)), - ) - strip_ptr = cute.make_ptr( - self.fc1_output_dtype, - strip_base.iterator.toint(), - cute.AddressSpace.gmem, - assumed_align=EN, - ) - gmem_strip = cute.make_tensor(strip_ptr, cute.make_layout(EN)) - cute.autovec_copy(r_data, gmem_strip) - # quantize each half to MXFP8 + E8M0 row SF (per-thread, no warp reduction) ---- qg = quant_sfd_row( d_gate, c_gate, norm_const, self._sf_vec_size, self.sf_dtype, self.fc1_output_dtype, @@ -707,7 +744,7 @@ def _run_dfc2_subtile( rmem_sf[2 * _k] = qg rmem_sf[2 * _k + 1] = qu - # dfc2_recompute: forward swiglu + col quant + per-thread STG + # dfc2_recompute: forward swiglu + column quantization if cutlass.const_expr(self._dfc2_recompute): c_recompute_f32 = cute.make_rmem_tensor(r_layout.shape, self.acc_dtype) if cutlass.const_expr(self._act_func == "swiglu"): @@ -726,42 +763,36 @@ def _run_dfc2_subtile( if subtile_idx == cutlass.Int32(_k): rmem_sf_recompute[_k] = qc - # STG c_recompute -- only valid rows and in-range N strips. - if is_valid_row: - c_col_idx = ( - work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n // EN) - + subtile_idx - ) - expert_local_token_idx = ( - work_tile_info.tile_m_idx * cutlass.Int32(self._cta_tile_m) + token_row_in_cta - ) - if c_col_idx * cutlass.Int32(EN) < real_fc1_recompute.shape[1]: - c_base = cute.local_tile( - real_fc1_recompute, (1, EN, 1), - (expert_local_token_idx, c_col_idx, cutlass.Int32(0)), - ) - c_ptr = cute.make_ptr( - self.fc1_output_dtype, - c_base.iterator.toint(), - cute.AddressSpace.gmem, - assumed_align=EN, - ) - gmem_c = cute.make_tensor(c_ptr, cute.make_layout(EN)) - cute.autovec_copy(c_recompute, gmem_c) - # BARRIER: drain PREVIOUS subtile's TMA BEFORE R2S. if warp_idx == self._epilogue_warp_ids[0]: d_pipeline.producer_acquire() epilog_sync.arrive_and_wait() # Write d to smem. - d_n_slots = cutlass.const_expr(d_num_stage // 2) - d_slot = cutlass.Int32(2) * (subtile_i % cutlass.Int32(d_n_slots)) + d_outputs_per_stage = cutlass.const_expr(self._d_output_slots) + d_n_stages = cutlass.const_expr(d_num_stage // d_outputs_per_stage) + d_slot = cutlass.Int32(d_outputs_per_stage) * ( + subtile_i % cutlass.Int32(d_n_stages) + ) thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) sd = thr_copy_r2s.partition_D(smem_d) cute.copy(tiled_copy_r2s, c_gate, sd[(None, None, None, d_slot)]) cute.copy(tiled_copy_r2s, c_up, sd[(None, None, None, d_slot + cutlass.Int32(1))]) + # Auxiliary data planes use the public token-major ABI. + next_slot = d_slot + cutlass.Int32(2) + if cutlass.const_expr(self._dfc2_col_output): + s_col_gate = cute.slice_(smem_d, (token_row_in_cta, None, next_slot)) + s_col_up = cute.slice_( + smem_d, (token_row_in_cta, None, next_slot + cutlass.Int32(1)) + ) + self._store_aux_row_smem(c_gate_col, s_col_gate) + self._store_aux_row_smem(c_up_col, s_col_up) + next_slot = next_slot + cutlass.Int32(2) + if cutlass.const_expr(self._dfc2_recompute): + s_recompute = cute.slice_(smem_d, (token_row_in_cta, None, next_slot)) + self._store_aux_row_smem(c_recompute, s_recompute) + iket.range_pop() return c_consumer_state, subtile_dprob @@ -809,33 +840,23 @@ def _stg_col_sf_atom_value( self, real_sf: cute.Tensor, row_block, - col, - hidden_atoms, + feature, + _feature_atoms, sf_value, ) -> None: - """Store one col-SF in the MN-major 128-column × 4-token-block atom.""" + """Store one SF in a 128-feature × 4-token-block atom.""" token_atom = row_block // cutlass.Int32(4) token_bank = row_block % cutlass.Int32(4) - hidden_atom = col // cutlass.Int32(128) - hidden_bank = (col // cutlass.Int32(32)) % cutlass.Int32(4) - hidden_lane = col % cutlass.Int32(32) - atom_idx = Int64(token_atom) * Int64(hidden_atoms) + Int64(hidden_atom) - byte_offset = ( - atom_idx * Int64(512) - + Int64(hidden_lane) * Int64(16) - + Int64(hidden_bank) * Int64(4) - + Int64(token_bank) - ) - sf_ptr = cute.make_ptr( - self.sf_dtype, - real_sf.iterator.toint() + byte_offset, - cute.AddressSpace.gmem, - assumed_align=1, + feature_atom = feature // cutlass.Int32(128) + feature_bank = (feature // cutlass.Int32(32)) % cutlass.Int32(4) + feature_lane = feature % cutlass.Int32(32) + atom_byte = ( + feature_lane * cutlass.Int32(16) + + feature_bank * cutlass.Int32(4) + + token_bank ) - gmem_sf1 = cute.make_tensor(sf_ptr, cute.make_layout(1)) - r_sf1 = cute.make_rmem_tensor(cute.make_layout(1).shape, self.sf_dtype) - r_sf1[0] = sf_value.to(self.sf_dtype) - cute.autovec_copy(r_sf1, gmem_sf1) + if feature_atom < real_sf.shape[0] and token_atom < real_sf.shape[1]: + real_sf[feature_atom, token_atom, atom_byte] = sf_value.to(self.sf_dtype) @cute.jit def _stg_sf_recompute( @@ -863,12 +884,9 @@ def _stg_sf_recompute( col_base = ( work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n) ) - # Row predicate: a warp stores its SF only if its 32-row block overlaps - # the CTA tile's valid rows. - warp_rows_valid = (warp_idx_local * cutlass.Int32(sf_vec_size)) < valid_tokens for s in cutlass.range_constexpr(self._cta_tile_n // EN): col = col_base + cutlass.Int32(s * EN) + warp_lane_idx - if warp_rows_valid and col < valid_inter: + if col < valid_inter: self._stg_col_sf_atom_value( real_fc1_recompute_sf, row_block, @@ -903,11 +921,10 @@ def _stg_sf_col_output( col_base = ( work_tile_info.tile_n_idx * cutlass.Int32(self._cta_tile_n * 2) ) - warp_rows_valid = (warp_idx_local * cutlass.Int32(sf_vec_size)) < valid_tokens for s in cutlass.range_constexpr(self._cta_tile_n // EN): for gu in cutlass.range_constexpr(2): col = col_base + cutlass.Int32((2 * s + gu) * EN) + warp_lane_idx - if warp_rows_valid and col < valid_inter: + if col < valid_inter: self._stg_col_sf_atom_value( real_fc1_col_output_sf, row_block, @@ -917,33 +934,74 @@ def _stg_sf_col_output( ) @cute.jit - def tma_store_dfc2_output( + def _tma_store_tile( self, - smem_d_buffer: cute.Tensor, + smem_tile: cute.Tensor, tma_atom: cute.CopyAtom, + gmem_tile: cute.Tensor, + ) -> None: + tma_smem_src, tma_gmem_dst = cpasync.tma_partition( + tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(smem_tile, 0, 2), + cute.group_modes(gmem_tile, 0, 2), + ) + cute.copy(tma_atom, tma_smem_src, tma_gmem_dst) + + @cute.jit + def tma_store_dfc2_outputs( + self, + smem_d_buffer: cute.Tensor, + tma_atom_grad_y1: cute.CopyAtom, g_gate_2d: cute.Tensor, g_up_2d: cute.Tensor, + tma_atom_col_output: cute.CopyAtom, + g_col_gate_2d, + g_col_up_2d, + tma_atom_recompute: cute.CopyAtom, + g_recompute_2d, valid_tokens, d_pipeline, d_slot, ) -> None: - """TMA-store one subtile's gate+up from shared sD stages to grad_y1 GMEM.""" - sD_gate = cute.slice_(smem_d_buffer, (None, None, d_slot)) - sD_up = cute.slice_(smem_d_buffer, (None, None, d_slot + cutlass.Int32(1))) - bSG_sD_gate, bSG_g_gate = cpasync.tma_partition( - tma_atom, 0, cute.make_layout(1), - cute.group_modes(sD_gate, 0, 2), - cute.group_modes(g_gate_2d, 0, 2), - ) - bSG_sD_up, bSG_g_up = cpasync.tma_partition( - tma_atom, 0, cute.make_layout(1), - cute.group_modes(sD_up, 0, 2), - cute.group_modes(g_up_2d, 0, 2), - ) + """Issue one TMA store group for every dFC2 data plane.""" tile_is_valid = valid_tokens > cutlass.Int32(0) if tile_is_valid: - cute.copy(tma_atom, bSG_sD_gate, bSG_g_gate) - cute.copy(tma_atom, bSG_sD_up, bSG_g_up) + self._tma_store_tile( + cute.slice_(smem_d_buffer, (None, None, d_slot)), + tma_atom_grad_y1, + g_gate_2d, + ) + self._tma_store_tile( + cute.slice_( + smem_d_buffer, (None, None, d_slot + cutlass.Int32(1)) + ), + tma_atom_grad_y1, + g_up_2d, + ) + next_slot = d_slot + cutlass.Int32(2) + if cutlass.const_expr(self._dfc2_col_output): + self._tma_store_tile( + cute.slice_(smem_d_buffer, (None, None, next_slot)), + tma_atom_col_output, + g_col_gate_2d, + ) + self._tma_store_tile( + cute.slice_( + smem_d_buffer, + (None, None, next_slot + cutlass.Int32(1)), + ), + tma_atom_col_output, + g_col_up_2d, + ) + next_slot = next_slot + cutlass.Int32(2) + if cutlass.const_expr(self._dfc2_recompute): + self._tma_store_tile( + cute.slice_(smem_d_buffer, (None, None, next_slot)), + tma_atom_recompute, + g_recompute_2d, + ) d_pipeline.producer_commit() @@ -1333,8 +1391,10 @@ def run( sched_ext, gmem_fc1_output: cute.Tensor, gmem_fc1_output_sf: cute.Tensor, + tma_atom_fc1_recompute: cute.CopyAtom, gmem_fc1_recompute: Optional[cute.Tensor], gmem_fc1_recompute_sf: Optional[cute.Tensor], + tma_atom_fc1_col_output: cute.CopyAtom, gmem_fc1_col_output: Optional[cute.Tensor], gmem_fc1_col_output_sf: Optional[cute.Tensor], smem_preact_buffer: cute.Tensor, @@ -1395,8 +1455,10 @@ def run( sched_ext=sched_ext, gmem_fc1_output=gmem_fc1_output, gmem_fc1_output_sf=gmem_fc1_output_sf, + tma_atom_fc1_recompute=tma_atom_fc1_recompute, gmem_fc1_recompute=gmem_fc1_recompute, gmem_fc1_recompute_sf=gmem_fc1_recompute_sf, + tma_atom_fc1_col_output=tma_atom_fc1_col_output, gmem_fc1_col_output=gmem_fc1_col_output, gmem_fc1_col_output_sf=gmem_fc1_col_output_sf, c_pipeline=c_pipeline, diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py index 11aafe176..c6746baef 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py @@ -2,24 +2,115 @@ # SPDX-License-Identifier: BSD-3-Clause """Sched extension for the fused fc1+fc2 dGLU-backward MXFP8 kernel.""" +import dataclasses from typing import Optional, Tuple import cutlass import cutlass.cute as cute from cutlass.cute.typing import Pointer +from cutlass.cutlass_dsl import Int32, Int64, extract_mlir_values, new_from_mlir_values from ..fwd_glu.glu_mxfp8_fc12_extension import GluMxFp8Fc12SchedExtension from .....schedulers.fc12_mapping import NonSwapAbFc12WorkTileInfo -def _rewrite_tensor_shape(tensor: cute.Tensor, new_shape: Tuple) -> cute.Tensor: - return cute.make_tensor(tensor.iterator, cute.make_layout(new_shape, stride=tensor.stride)) +@dataclasses.dataclass(frozen=True) +class DgluMxFp8Fc12SchedExtension(GluMxFp8Fc12SchedExtension): + """dGLU adapter for token-major auxiliary data and per-expert blocked SF.""" + expert_token_sizes: Optional[cute.Tensor] = None + token_padding_block: int = 128 + sf_padding_block: int = 128 -class DgluMxFp8Fc12SchedExtension(GluMxFp8Fc12SchedExtension): - """ - Sched extension for the fused fc1+fc2 dGLU-backward MXFP8 kernel. - """ + def __post_init__(self) -> None: + super().__post_init__() + if self.expert_token_sizes is None: + raise ValueError("dGLU auxiliaries require expert_token_sizes.") + if self.token_padding_block != self.sf_padding_block or self.token_padding_block % 128 != 0: + raise ValueError("dGLU auxiliaries require equal token/SF padding divisible by 128.") + + def __extract_mlir_values__(self) -> list: + values = super().__extract_mlir_values__() + values.extend(extract_mlir_values(self.expert_token_sizes)) + return values + + def __new_from_mlir_values__(self, values: list) -> "DgluMxFp8Fc12SchedExtension": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + fc1_done_counter_pointer = rebuild(self.fc1_done_counter_pointer) + fc2_spin_threshold = rebuild(self.fc2_spin_threshold) + fc1_ready_counter_pointer = ( + rebuild(self.fc1_ready_counter_pointer) if self.fc1_ready_counter_pointer is not None else None + ) + expert_token_sizes = rebuild(self.expert_token_sizes) + if value_index != len(values): + raise ValueError( + f"DgluMxFp8Fc12SchedExtension MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return type(self)( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter_pointer, + fc2_spin_threshold=fc2_spin_threshold, + fc1_ready_counter_pointer=fc1_ready_counter_pointer, + cluster_m=self.cluster_m, + expert_token_sizes=expert_token_sizes, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, + ) + + @cute.jit + def _physical_token_count(self, work_tile_info: NonSwapAbFc12WorkTileInfo, padding_block: int): + expert_idx = work_tile_info.expert_idx + valid_tokens = Int32(self.expert_token_sizes[expert_idx]) + return ((valid_tokens + Int32(padding_block - 1)) // Int32(padding_block)) * Int32(padding_block) + + @cute.jit + def _aux_data_tensor( + self, + tensor: cute.Tensor, + work_tile_info: NonSwapAbFc12WorkTileInfo, + feature_extent, + ) -> cute.Tensor: + physical_tokens = self._physical_token_count(work_tile_info, self.token_padding_block) + real = cute.domain_offset( + (work_tile_info.cumulative_data_physical_row, 0, 0), + tensor, + ) + return cute.make_tensor( + real.iterator, + cute.make_layout( + (physical_tokens, feature_extent, Int32(1)), + stride=real.stride, + ), + ) + + @cute.jit + def _aux_sf_tensor( + self, + tensor: cute.Tensor, + work_tile_info: NonSwapAbFc12WorkTileInfo, + feature_padded, + ) -> cute.Tensor: + physical_tokens = self._physical_token_count(work_tile_info, self.sf_padding_block) + feature_atoms = Int32(feature_padded) // Int32(128) + token_atoms = physical_tokens // Int32(128) + element_offset = Int64(feature_padded) * ( + Int64(work_tile_info.cumulative_sf_physical_row) // Int64(self.sf_vec_size) + ) + return cute.make_tensor( + tensor.iterator + element_offset, + cute.make_layout( + (feature_atoms, token_atoms, Int32(512)), + stride=(token_atoms * Int32(512), Int32(512), Int32(1)), + ), + ) @cute.jit def get_gmem_tensor( @@ -29,44 +120,19 @@ def get_gmem_tensor( work_tile_info: NonSwapAbFc12WorkTileInfo, ) -> Tuple[cute.Tensor, Optional[Pointer]]: """dGLU-backward operand views; every other name delegates to the base.""" - data_token_offset = work_tile_info.cumulative_data_physical_row - sf_token_offset = work_tile_info.cumulative_sf_physical_row - shape = gmem_tensor_in_moe_view.shape - c1 = cutlass.Int32(1) - sf_vec_size = self.sf_vec_size if cutlass.const_expr(tensor_name == "recompute"): - # Forward-swiglu recompute data tensor. - real = cute.domain_offset( - (data_token_offset, 0, 0), gmem_tensor_in_moe_view - ) - real = _rewrite_tensor_shape(real, (shape[0], shape[1], c1)) # type: ignore[index] - return (real, None) + return (self._aux_data_tensor(gmem_tensor_in_moe_view, work_tile_info, shape[1]), None) elif cutlass.const_expr(tensor_name == "sfrecompute"): - # Per-expert base for atom-packed col-SF of the forward-swiglu recompute. - real = cute.domain_offset( - (sf_token_offset // sf_vec_size, 0, 0), gmem_tensor_in_moe_view - ) - real = _rewrite_tensor_shape(real, (shape[0], shape[1], c1)) # type: ignore[index] - return (real, None) + return (self._aux_sf_tensor(gmem_tensor_in_moe_view, work_tile_info, shape[0]), None) elif cutlass.const_expr(tensor_name == "col_output"): - # Col-quant grad_y1 data tensor (alongside row-quant "d"). - real = cute.domain_offset( - (data_token_offset, 0, 0), gmem_tensor_in_moe_view - ) - real = _rewrite_tensor_shape(real, (shape[0], shape[1], c1)) # type: ignore[index] - return (real, None) + return (self._aux_data_tensor(gmem_tensor_in_moe_view, work_tile_info, shape[1]), None) elif cutlass.const_expr(tensor_name == "sfcol_output"): - # Per-expert base for atom-packed col-SF of the grad_y1 col output. - real = cute.domain_offset( - (sf_token_offset // sf_vec_size, 0, 0), gmem_tensor_in_moe_view - ) - real = _rewrite_tensor_shape(real, (shape[0], shape[1], c1)) # type: ignore[index] - return (real, None) + return (self._aux_sf_tensor(gmem_tensor_in_moe_view, work_tile_info, shape[0]), None) return GluMxFp8Fc12SchedExtension.get_gmem_tensor( self, tensor_name, gmem_tensor_in_moe_view, work_tile_info diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py index c1f51a930..cff668b4a 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py @@ -418,7 +418,8 @@ def _setup_attributes(self) -> None: self.num_c_stage = 2 assert self.num_c_stage % 2 == 0, f"num_c_stage must be even, got {self.num_c_stage}" self.num_c_pipe_stage = self.num_c_stage // 2 - self.num_d_stage = 2 + # One PipelineTmaStore stage contains every dFC2 data output tile. + self.num_d_stage = self.epilogue.d_output_slots c_bytes_total = self.num_c_stage * self.epilogue.preact_bytes_per_stage d_bytes_total = self.num_d_stage * self.epilogue.d_bytes_per_stage self.c_bytes_total = c_bytes_total @@ -876,21 +877,19 @@ def __call__( ), ) - # fc1_recompute: forward swiglu recomputed from the fc1 c-tensor. - intermediate_downproj_half = fc1_recompute.shape[1] + # dFC2 auxiliary data planes use the public token-major ABI. fc1_recompute_gemm = cute.make_tensor( fc1_recompute.iterator, cute.make_layout( - (tokens_sum, intermediate_downproj_half, 1), + (fc1_recompute.shape[0], fc1_recompute.shape[1], 1), stride=(fc1_recompute.stride[0], fc1_recompute.stride[1], 0), ), ) - # fc1_col_output: col-quantized grad_y1 alongside row-quant fc1_output. fc1_col_output_gemm = cute.make_tensor( fc1_col_output.iterator, cute.make_layout( - (tokens_sum, intermediate_downproj, 1), + (fc1_col_output.shape[0], fc1_col_output.shape[1], 1), stride=(fc1_col_output.stride[0], fc1_col_output.stride[1], 0), ), ) @@ -1071,6 +1070,18 @@ def __call__( self.epilogue.d_smem_layout_one_stage, self.epilogue.d_epi_tile, ) + tma_atom_fc1_recompute, tma_tensor_fc1_recompute = cpasync.make_tiled_tma_atom( + grad_y1_tma_op, + fc1_recompute_gemm, + self.epilogue.d_smem_layout_one_stage, + self.epilogue.d_epi_tile, + ) + tma_atom_fc1_col_output, tma_tensor_fc1_col_output = cpasync.make_tiled_tma_atom( + grad_y1_tma_op, + fc1_col_output_gemm, + self.epilogue.d_smem_layout_one_stage, + self.epilogue.d_epi_tile, + ) # fc1 SFC GMEM tensor (= fc1_output_sf user view). No TMA atom; it is # per-thread STG. @@ -1082,23 +1093,19 @@ def __call__( ), ) - # fc1_recompute SFC GMEM storage. The epilogue uses the iterator as the - # base of per-expert MN-major 128-column × 4-token-block atoms. - fc1_recompute_sf_row_blocks = fc1_recompute_sf.shape[0] + # Token-major blocked SF carriers. fc1_recompute_sf_gemm = cute.make_tensor( fc1_recompute_sf.iterator, cute.make_layout( - (fc1_recompute_sf_row_blocks, intermediate_downproj_half, 1), + (fc1_recompute_sf.shape[0], fc1_recompute_sf.shape[1], 1), stride=(fc1_recompute_sf.stride[0], fc1_recompute_sf.stride[1], 0), ), ) - # fc1_col_output SFC GMEM storage, likewise atom-packed by the epilogue. - fc1_col_output_sf_row_blocks = fc1_col_output_sf.shape[0] fc1_col_output_sf_gemm = cute.make_tensor( fc1_col_output_sf.iterator, cute.make_layout( - (fc1_col_output_sf_row_blocks, intermediate_downproj, 1), + (fc1_col_output_sf.shape[0], fc1_col_output_sf.shape[1], 1), stride=(fc1_col_output_sf.stride[0], fc1_col_output_sf.stride[1], 0), ), ) @@ -1216,11 +1223,12 @@ def __call__( tma_tensor_fc1_preact, tma_atom_grad_y1, tma_tensor_grad_y1, - # fc1_recompute (forward swiglu) — per-thread STG, N = inter_half - fc1_recompute_gemm, + # token-major auxiliary data — TMA S2G stores + tma_atom_fc1_recompute, + tma_tensor_fc1_recompute, fc1_recompute_sf_gemm, - # fc1_col_output (col-quant grad_y1) — per-thread STG, N = intermediate - fc1_col_output_gemm, + tma_atom_fc1_col_output, + tma_tensor_fc1_col_output, fc1_col_output_sf_gemm, # topk / beta / dprob + cross-phase sync workspace topk_scores, @@ -1298,11 +1306,12 @@ def kernel( # grad_y1 (dfc2 output) — TMA S2G store tma_atom_grad_y1: cute.CopyAtom, tma_tensor_grad_y1: cute.Tensor, - # fc1_recompute (forward swiglu) — per-thread STG, N = inter_half - fc1_recompute_gemm: cute.Tensor, + # token-major auxiliary data — TMA S2G stores + tma_atom_fc1_recompute: cute.CopyAtom, + tma_tensor_fc1_recompute: cute.Tensor, fc1_recompute_sf_gemm: cute.Tensor, - # fc1_col_output (col-quant grad_y1) — per-thread STG, N = intermediate - fc1_col_output_gemm: cute.Tensor, + tma_atom_fc1_col_output: cute.CopyAtom, + tma_tensor_fc1_col_output: cute.Tensor, fc1_col_output_sf_gemm: cute.Tensor, # topk / beta / dprob + cross-phase sync workspace topk_scores: cute.Tensor, @@ -1349,6 +1358,13 @@ def kernel( fc1_weight_gemm.shape[0] + self.cta_tile_shape_mnk[1] - 1 ) // self.cta_tile_shape_mnk[1] * self.epilogue._atom_thr_size + if cutlass.const_expr(self.enable_token_comm): + _aux_expert_sizes = self.token_comm.local_expert_sizes( + self._mega_device_workspace, mega_local_rank + ) + else: + _aux_expert_sizes = expert_token_sizes + ext = DgluMxFp8Fc12SchedExtension( sf_vec_size=self.sf_vec_size, fc1_done_counter_pointer=fc1_done_counter.iterator, @@ -1360,6 +1376,9 @@ def kernel( ), # Fold the 2 CTAs of a cluster onto one fc1_ready slot cluster_m=self.epilogue._atom_thr_size, + expert_token_sizes=_aux_expert_sizes, + token_padding_block=self.token_padding_block, + sf_padding_block=self.sf_padding_block, ) warp_idx = cute.arch.warp_idx() @@ -1423,7 +1442,7 @@ class SharedStorage: ], 1024, ] - # grad_y1 (dfc2 output) store staging — stage 0 = gate, stage 1 = up. + # Unified dFC2 data-output staging; slot count is compile-time gated. sD: cute.struct.Align[ cute.struct.MemRange[ self.fc1_output_dtype, @@ -1502,7 +1521,7 @@ class SharedStorage: 32 * len(self.epilogue_warp_id), ) d_pipeline = pipeline.PipelineTmaStore.create( - num_stages=num_d_stage // 2, + num_stages=num_d_stage // self.epilogue.d_output_slots, producer_group=d_producer_group, ) @@ -1588,7 +1607,7 @@ class SharedStorage: swizzle=preact_smem_layout_staged.inner, ) - # grad_y1 (dfc2 output) store staging tensor (stage 0 = gate, stage 1 = up). + # Unified dFC2 data-output store staging tensor. d_smem_layout_staged = self.epilogue.d_staged_smem_layout(num_d_stage) sD = storage.sD.get_tensor( d_smem_layout_staged.outer, @@ -2249,9 +2268,11 @@ class SharedStorage: sched_ext=ext, gmem_fc1_output=tma_tensor_grad_y1, gmem_fc1_output_sf=fc1_output_sf_gemm, - gmem_fc1_recompute=fc1_recompute_gemm, + tma_atom_fc1_recompute=tma_atom_fc1_recompute, + gmem_fc1_recompute=tma_tensor_fc1_recompute, gmem_fc1_recompute_sf=fc1_recompute_sf_gemm, - gmem_fc1_col_output=fc1_col_output_gemm, + tma_atom_fc1_col_output=tma_atom_fc1_col_output, + gmem_fc1_col_output=tma_tensor_fc1_col_output, gmem_fc1_col_output_sf=fc1_col_output_sf_gemm, smem_preact_buffer=sPre, c_pipeline=c_pipeline, diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py index b7844237a..a53770617 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py @@ -10,7 +10,7 @@ import cutlass import cutlass.cute as cute from cutlass.cute.typing import AddressSpace -from cutlass.cutlass_dsl import Int64 +from cutlass.cutlass_dsl import Int32, Int64 from ......api import ImplDesc, KernelClass, ProblemDesc, StaticOrRuntimeIntegerType from ......helpers.device_workspace import DeviceWorkspace @@ -167,7 +167,7 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): stream=make_fake_stream(), ) fake_arguments["grad_y2"] = fake_tensor( - self.ab_dtype, aux_shapes["grad_y2"], (1, 0), {0}, 16 + self.ab_dtype, aux_shapes["grad_y2"], (0, 1), set(), 128 ) fake_arguments["grad_y2_sf"] = fake_tensor( cutlass.Uint8, aux_shapes["grad_y2_sf"], (0,), set(), 16 @@ -525,6 +525,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: num_persistent_ctas=self.num_ctas_grad_y2_col_quant, token_padding_block=self.token_padding_block, sf_padding_block=self.sf_padding_block, + dst_k_major=True, ) def _smem_misc_budget_bytes(self) -> int: @@ -533,16 +534,20 @@ def _smem_misc_budget_bytes(self) -> int: return super()._smem_misc_budget_bytes() + self._token_comm_smem_bytes + _sched + self._SMEM_ALLOC_MARGIN def get_aux_output_shapes(self) -> dict: - """Shapes of the fixed-ABI dFC2 auxiliary outputs.""" + """Shapes of the fixed-ABI dFC2 auxiliary outputs. + + Data planes are compact token-major matrices. Column-quantized scale + planes retain the WGrad 128x4 atom layout. + """ data_token_capacity = self.token_comm.worst_case_token_count sf_token_capacity = self.token_comm.worst_case_sf_token_count column_sf_row_count = sf_token_capacity // self.sf_vec_size return { "dprob": (self.max_tokens_per_rank, self.num_topk), "fc1_recompute": (data_token_capacity, self.intermediate_downproj), - "fc1_recompute_sf": (column_sf_row_count, self.intermediate_downproj), + "fc1_recompute_sf": (round_up(self.intermediate_downproj, 128), column_sf_row_count), "fc1_col_output": (data_token_capacity, self.intermediate_gateup), - "fc1_col_output_sf": (column_sf_row_count, self.intermediate_gateup), + "fc1_col_output_sf": (round_up(self.intermediate_gateup, 128), column_sf_row_count), "grad_y2": (data_token_capacity, self.hidden), "grad_y2_sf": (sf_token_capacity * (self.hidden // self.sf_vec_size),), } @@ -552,7 +557,7 @@ def get_fc1_preact_shape(self) -> Tuple[int, int]: return (self.token_comm.worst_case_token_count, self.intermediate_gateup) @cute.jit - def _validate_fixed_pool_tensor(self, tensor: cute.Tensor, dtype, expected_shape) -> None: + def _validate_fixed_pool_tensor(self, tensor: cute.Tensor, dtype, expected_shape, expected_stride=None) -> None: if cutlass.const_expr(tensor.element_type is not dtype): raise TypeError("pool-domain tensor has an unexpected element type.") if cutlass.const_expr(cute.rank(tensor.layout) != 2): @@ -564,8 +569,9 @@ def _validate_fixed_pool_tensor(self, tensor: cute.Tensor, dtype, expected_shape or tensor.shape[1] != expected_shape[1] ): raise ValueError(f"pool-domain tensor must have static shape {expected_shape}.") - if cutlass.const_expr(tensor.stride[0] != expected_shape[1] or tensor.stride[1] != 1): - raise ValueError("pool-domain tensor must be compact row-major.") + stride = (expected_shape[1], 1) if expected_stride is None else expected_stride + if cutlass.const_expr(tensor.stride[0] != stride[0] or tensor.stride[1] != stride[1]): + raise ValueError("pool-domain tensor has an unexpected stride.") def _build_megamoe_device_workspace(self) -> DeviceWorkspace: """Register internal dGLU pools, counters, and token-comm regions.""" @@ -730,8 +736,6 @@ def token_comm_hook_kernel_tail(self, token_comm_args, *, warp_idx, lane_idx, ti @cute.jit def _snapshot_grad_y2_expert_sizes(self, tidx) -> None: """Preserve local expert counts before token_comm tail reset.""" - from cutlass.cutlass_dsl import Int32 - dw = self._mega_device_workspace if self.token_comm._linear_cta_idx == Int32(0): sizes = self.token_comm.local_expert_sizes(dw, self.token_comm._local_rank) @@ -762,10 +766,10 @@ def __call__( output_activation: cute.Tensor, # (max_tokens_per_rank, topk, hidden) BF16 overflow_flag: cute.Tensor, # (1,) Int32, per-rank FC12 overflow output dprob: cute.Tensor, # (max_tokens_per_rank, topk) Float32; symmetric, pre-zeroed - fc1_recompute: cute.Tensor, # (pool_token_capacity, inter_downproj) - fc1_recompute_sf: cute.Tensor, # (col_sf_rows, inter_downproj) E8M0 - fc1_col_output: cute.Tensor, # (pool_token_capacity, intermediate_gateup) - fc1_col_output_sf: cute.Tensor, # (col_sf_rows, intermediate_gateup) E8M0 + fc1_recompute: cute.Tensor, # (pool_token_capacity, inter_downproj), token-major + fc1_recompute_sf: cute.Tensor, # WGrad2 SFA: (inter_padded, col_sf_rows) + fc1_col_output: cute.Tensor, # (pool_token_capacity, gateup), token-major + fc1_col_output_sf: cute.Tensor, # WGrad1 SFB: (gateup_padded, col_sf_rows) grad_y2: cute.Tensor, # (pool_token_capacity, hidden) token-axis MXFP8 grad_y2_sf: cute.Tensor, # flat MN-major E8M0 bytes local_workspace: cute.Pointer, @@ -790,13 +794,27 @@ def __call__( aux_shapes["fc1_recompute_sf"], ) self._validate_fixed_pool_tensor( - fc1_col_output, self.ab_dtype, aux_shapes["fc1_col_output"] + fc1_col_output, + self.ab_dtype, + aux_shapes["fc1_col_output"], ) self._validate_fixed_pool_tensor( fc1_col_output_sf, self.token_comm.activation_sf_dtype, aux_shapes["fc1_col_output_sf"], ) + self._validate_fixed_pool_tensor( + grad_y2, + self.ab_dtype, + aux_shapes["grad_y2"], + (1, aux_shapes["grad_y2"][0]), + ) + if cutlass.const_expr( + grad_y2_sf.element_type is not cutlass.Uint8 + or cute.rank(grad_y2_sf.layout) != 1 + or grad_y2_sf.shape[0] != aux_shapes["grad_y2_sf"][0] + ): + raise ValueError("grad_y2_sf must be the fixed-size flat Uint8 carrier.") self.token_comm.launch_router( topk_indices=topk_idx, topk_scores=topk_weights, diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py index bd54c92f3..dbf7b9eb2 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py @@ -399,6 +399,7 @@ def __init__( sf_padding_block: int = SfPaddingBlock, *, scaled_cvt: "bool | None" = None, + dst_k_major: bool = False, ) -> None: """``scaled_cvt`` selects the requant path: ``None`` asks the compilation target, ``True`` and ``False`` force the block-scaled and @@ -410,6 +411,7 @@ def __init__( self.quant_type = quant_type self.token_padding_block = int(token_padding_block) self.sf_padding_block = int(sf_padding_block) + self.dst_k_major = bool(dst_k_major) self._require_token_padding_block(self.token_padding_block) if self.sf_padding_block != self.SfAtomNonK: @@ -627,16 +629,20 @@ def __call__( (BOX_T, BOX_H), ) - dst_u32 = cute.make_tensor( - cute.recast_ptr(dst_data.iterator, dtype=cutlass.Uint32), - cute.make_layout((dst_data.shape[0], HID_U32), stride=(HID_U32, 1)), - ) - tma_atom_st, tma_tensor_st = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), - dst_u32, - cute.make_layout((BOX_T, BOX_H), stride=(BOX_H, 1)), - (BOX_T, BOX_H), - ) + if cutlass.const_expr(self.dst_k_major): + tma_atom_st = None + tma_tensor_st = None + else: + dst_u32 = cute.make_tensor( + cute.recast_ptr(dst_data.iterator, dtype=cutlass.Uint32), + cute.make_layout((dst_data.shape[0], HID_U32), stride=(HID_U32, 1)), + ) + tma_atom_st, tma_tensor_st = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + dst_u32, + cute.make_layout((BOX_T, BOX_H), stride=(BOX_H, 1)), + (BOX_T, BOX_H), + ) k = self.ws_kernel( src_data, src_sf_u8, expert_token_sizes, dst_data, dst_sf_u8, tma_atom, tma_tensor, tma_atom_st, tma_tensor_st, TOKPAD, @@ -803,7 +809,7 @@ def ws_kernel( else: self.consume_scaled( smem_data_base, smem_sf_in_base, smem_sf_out_base, mbar_full, mbar_empty, - tbl_vend, tbl_data, tbl_sf, dst_sf_base, + tbl_vend, tbl_data, tbl_sf, dst_data, dst_sf_base, bidx, grid_dim_x, total_tiles, warp_idx - Int32(self.ProducerWarps), lane_idx, tma_atom_st, tma_tensor_st, TOKPAD, @@ -904,7 +910,7 @@ def produce( @cute.jit def consume_scaled( self, smem_data_base, smem_sf_in_base, smem_sf_out_base, mbar_full, mbar_empty, - tbl_vend, tbl_data, tbl_sf, dst_sf_base, + tbl_vend, tbl_data, tbl_sf, dst_data, dst_sf_base, bidx, grid_dim_x, total_tiles, cw, lane_idx, tma_atom_st=None, tma_tensor_st=None, token_padding_block: cutlass.Constexpr = None, @@ -938,20 +944,29 @@ def consume_scaled( hb0 = seg * Int32(C) + (lane_idx * Int32(LW)) // Int32(32) sf_lane_off = tb * Int32(4) - BOX_HS = cutlass.const_expr(self.TmaBoxHidU32) - sDo = cute.make_tensor( - cute.make_ptr( - cutlass.Uint32, smem_data_base, AddressSpace.smem, assumed_align=128, - ), - cute.make_layout((TOK, BOX_HS, S), stride=(BOX_HS, 1, TOK * BOX_HS)), - ) - gDo = cute.group_modes( - cute.local_tile(tma_tensor_st, (TOK, BOX_HS), (None, None)), 0, 2 - ) - tDsDo, tDgDo = cpasync.tma_partition( - tma_atom_st, 0, cute.make_layout(1), cute.group_modes(sDo, 0, 2), gDo, - ) - cpasync.prefetch_descriptor(tma_atom_st) + if cutlass.const_expr(self.dst_k_major): + sDo_u8 = cute.make_tensor( + cute.make_ptr( + cutlass.Uint8, smem_data_base, AddressSpace.smem, assumed_align=128, + ), + cute.make_layout((TOK, W, S), stride=(W, 1, TOK * W)), + ) + dst_u8_pointer = cute.recast_ptr(dst_data.iterator, dtype=cutlass.Uint8) + else: + BOX_HS = cutlass.const_expr(self.TmaBoxHidU32) + sDo = cute.make_tensor( + cute.make_ptr( + cutlass.Uint32, smem_data_base, AddressSpace.smem, assumed_align=128, + ), + cute.make_layout((TOK, BOX_HS, S), stride=(BOX_HS, 1, TOK * BOX_HS)), + ) + gDo = cute.group_modes( + cute.local_tile(tma_tensor_st, (TOK, BOX_HS), (None, None)), 0, 2 + ) + tDsDo, tDgDo = cpasync.tma_partition( + tma_atom_st, 0, cute.make_layout(1), cute.group_modes(sDo, 0, 2), gDo, + ) + cpasync.prefetch_descriptor(tma_atom_st) t = Int32(0) work_idx = Int32(bidx) @@ -1001,12 +1016,31 @@ def consume_scaled( # The whole tile goes out, padding rows included; those were # neutralised to zero in shared memory, which is what the pool # expects to find there. - if cw == Int32(0): - cute.copy( - tma_atom_st, - tDsDo[(None, stage)], - tDgDo[(None, token_tile, hid_begin // Int32(W))], - ) + if cutlass.const_expr(self.dst_k_major): + linear = cw * Int32(32) + lane_idx + while linear < Int32(TOK * W): + token_in_tile = linear // Int32(W) + feature_in_tile = linear % Int32(W) + token = data_row0 + token_in_tile + feature = hid_begin + feature_in_tile + if token < dst_data.shape[0] and feature < dst_data.shape[1]: + dst_offset = ( + Int64(feature) * Int64(dst_data.shape[0]) + + Int64(token) + ) + dst_slot = cute.make_tensor( + dst_u8_pointer + dst_offset, + cute.make_layout(1), + ) + dst_slot[0] = sDo_u8[token_in_tile, feature_in_tile, stage] + linear = linear + Int32(CONS_THREADS) + else: + if cw == Int32(0): + cute.copy( + tma_atom_st, + tDsDo[(None, stage)], + tDgDo[(None, token_tile, hid_begin // Int32(W))], + ) if cw == Int32(0) and lane_idx >= Int32(16) and lane_idx < Int32(16) + Int32(HATOMS) * sf_live: atom = lane_idx - Int32(16) cp_async_bulk_s2g( diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py index 552965c7b..bca96e800 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py @@ -1645,3 +1645,4 @@ def run( ) flag_tracker.fire() + diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py index 2d39dc12f..646161a13 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py @@ -164,10 +164,10 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): fake_arguments["fc1_c"] = None if self.enable_col_quant: - # col_quant_data shares the dispatch pool's row-major (token, hidden) layout; + # col_quant_data is logically (token, hidden) with token unit stride; # col_quant_sf is flat concat_e [hidden_atom][token_atom] E8M0 bytes. fake_arguments["col_quant_data"] = fake_tensor( - self.ab_dtype, aux_shapes["col_quant_data"], (1, 0), set(), 16 + self.ab_dtype, aux_shapes["col_quant_data"], (0, 1), set(), 16 ) fake_arguments["col_quant_sf"] = fake_tensor( cutlass.Uint8, aux_shapes["col_quant_sf"], (0,), set(), 16 @@ -515,6 +515,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: num_persistent_ctas=self._col_quant_num_ctas, token_padding_block=self.token_padding_block, sf_padding_block=self.sf_padding_block, + dst_k_major=True, ) def get_aux_output_shapes(self) -> dict: @@ -529,7 +530,9 @@ def get_aux_output_shapes(self) -> dict: } @cute.jit - def _validate_fixed_matrix(self, tensor: cute.Tensor, dtype, expected_shape) -> None: + def _validate_fixed_matrix( + self, tensor: cute.Tensor, dtype, expected_shape, expected_stride=None + ) -> None: if cutlass.const_expr(tensor.element_type is not dtype): raise TypeError("pool-domain matrix has an unexpected element type.") if cutlass.const_expr(cute.rank(tensor.layout) != 2): @@ -541,8 +544,9 @@ def _validate_fixed_matrix(self, tensor: cute.Tensor, dtype, expected_shape) -> or tensor.shape[1] != expected_shape[1] ): raise ValueError(f"pool-domain matrix must have static shape {expected_shape}.") - if cutlass.const_expr(tensor.stride[0] != expected_shape[1] or tensor.stride[1] != 1): - raise ValueError("pool-domain matrix must be compact row-major.") + stride = (expected_shape[1], 1) if expected_stride is None else expected_stride + if cutlass.const_expr(tensor.stride[0] != stride[0] or tensor.stride[1] != stride[1]): + raise ValueError("pool-domain matrix has an unexpected stride.") @cute.jit def _validate_fixed_vector(self, tensor: cute.Tensor, dtype, expected_size: int) -> None: @@ -739,7 +743,10 @@ def __call__( if cutlass.const_expr(col_quant_data is None or col_quant_sf is None): raise ValueError("enable_col_quant=True requires data and scale outputs.") self._validate_fixed_matrix( - col_quant_data, self.ab_dtype, aux_shapes["col_quant_data"] + col_quant_data, + self.ab_dtype, + aux_shapes["col_quant_data"], + (1, aux_shapes["col_quant_data"][0]), ) self._validate_fixed_vector( col_quant_sf, cutlass.Uint8, aux_shapes["col_quant_sf"][0] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py index ef5d4ee75..34ec485f1 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py @@ -108,6 +108,23 @@ def _typed_view( return byte_tensor.view(dtype).reshape(shape) +def _typed_k_major_view( + byte_tensor: torch.Tensor, + dtype: torch.dtype, + shape: tuple[int, int], +) -> torch.Tensor: + """Return a rank-2 ``(K,N)`` view with K as the unit-stride mode.""" + + if len(shape) != 2: + raise ValueError(f"K-major view requires rank-2 shape, got {shape}") + rows, columns = shape + return _typed_view( + byte_tensor, + dtype, + (columns, rows), + ).transpose(0, 1) + + def _as_bytes(tensor: torch.Tensor) -> torch.Tensor: return tensor.view(torch.uint8) @@ -500,7 +517,7 @@ def stage( raise ValueError( "enabled column requant requires positive output capacities" ) - col_quant_data = _typed_view( + col_quant_data = _typed_k_major_view( local["col_quant_data"], _MXFP8_DATA_DTYPE, (col_quant_data_rows, config.hidden), @@ -536,7 +553,7 @@ def stage( topk_weights[:token_count].copy_(request.topk_weights) _as_bytes(output_data).zero_() if col_quant_data is not None: - _as_bytes(col_quant_data).zero_() + local["col_quant_data"].zero_() if col_quant_sf is not None: col_quant_sf.zero_() overflow_flag.zero_() diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py index b95f61415..c3e67cf87 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py @@ -6,19 +6,17 @@ from __future__ import annotations import threading +from dataclasses import replace import torch import torch.distributed as dist from ..._backend import BackendUnavailableError -from ..._contracts import ( - ForwardConfig, - ValidatedBackwardRequest, - ValidatedForwardRequest, -) +from ..._contracts import ForwardConfig, ValidatedForwardRequest +from ..._types import MoeEpTrainingWeights from .._plan import ExecutionPlanOwner from ._adapter import Mxfp8InputAdapter -from ._backward import Mxfp8BackwardExecutor +from ._backward_compile import prepare_backward_kernel from ._compile import ( CompiledMxfp8Kernel, PreparedMxfp8Kernel, @@ -27,7 +25,6 @@ ) from ._config import Mxfp8KernelConfig from ._launch import launch_forward -from ._stash import Mxfp8ForwardStash class Mxfp8Backend: @@ -38,11 +35,6 @@ def __init__(self, config: ForwardConfig, device: torch.device) -> None: self.device = torch.device(device) self.kernel_config = Mxfp8KernelConfig.from_forward_config(config) self._adapter = Mxfp8InputAdapter() - self._stash = ( - Mxfp8ForwardStash(config, self.device) - if config.generate_c - else None - ) self._prepared_kernel: PreparedMxfp8Kernel | None = None self._compiled: CompiledMxfp8Kernel | None = None self._plan: ExecutionPlanOwner | None = None @@ -52,7 +44,7 @@ def __init__(self, config: ForwardConfig, device: torch.device) -> None: self._completion_recorded = False self._device_work_may_be_pending = False self._ep_launch_ready = config.ep_size == 1 - self._backward_executor: Mxfp8BackwardExecutor | None = None + self._training_resource_owner = None self._lock = threading.RLock() @property @@ -127,11 +119,6 @@ def forward(self, request: ValidatedForwardRequest): with torch.cuda.device(self.device): capturing = torch.cuda.is_current_stream_capturing() - if capturing and self._stash is not None: - raise NotImplementedError( - "MoeEp generate_c=True is eager-only and does not " - "support CUDA graph capture" - ) if ( capturing and not self._adapter.weights_have_version_counters( @@ -176,14 +163,6 @@ def forward(self, request: ValidatedForwardRequest): # a retry on another stream cannot race those writes. device_work_attempted = True resources = self._plan.prepare(request) - stash_plan = ( - None - if self._stash is None - else self._stash.prepare( - request, - pool_token_capacity=prepared.pool_token_capacity, - ) - ) inputs = self._adapter.stage( request, resources, @@ -208,11 +187,7 @@ def forward(self, request: ValidatedForwardRequest): ), col_quant_data_rows=prepared.col_quant_data_rows, col_quant_sf_elements=prepared.col_quant_sf_elements, - fc1_c=( - None - if stash_plan is None - else stash_plan.buffer - ), + fc1_c=None, ) self._compiled = compile_or_get( prepared, @@ -225,26 +200,6 @@ def forward(self, request: ValidatedForwardRequest): inputs, resources, ) - if self._stash is not None: - assert stash_plan is not None - ( - fc1_c, - route_metadata, - wgrad_stash, - ) = self._stash.materialize( - stash_plan, - inputs, - prepared, - ) - if wgrad_stash is None: - output = (output, fc1_c, route_metadata) - else: - output = ( - output, - fc1_c, - route_metadata, - wgrad_stash, - ) except (ImportError, OSError) as exc: raise BackendUnavailableError( "MoeEp MXFP8 backend requires the 'moe_ep' optional " @@ -264,42 +219,70 @@ def forward(self, request: ValidatedForwardRequest): self._warmed_up = True return output - def backward(self, request: ValidatedBackwardRequest): - """Run the restricted explicit dgrad/dprob Rubin MXFP8 path.""" + def prepare_training_resources( + self, + weights: MoeEpTrainingWeights, + *, + slot_count: int, + lane_count: int, + ): + """Allocate the fixed slot/lane roots used by the training graph path.""" with self._lock: if self._closed: raise RuntimeError("MoeEp MXFP8 backend is closed") - if request.device != self.device: - raise ValueError( - f"MoeEp MXFP8 backend is bound to {self.device}, " - f"got {request.device}" - ) - stream = torch.cuda.current_stream(self.device) - if self._device_work_may_be_pending: - torch.cuda.synchronize(self.device) - self._device_work_may_be_pending = False - if self._completion_event is None: - self._completion_event = torch.cuda.Event() - elif self._completion_recorded: - stream.wait_event(self._completion_event) - if self._backward_executor is None: - self._backward_executor = Mxfp8BackwardExecutor( - self.config, - self.device, - ) + if self._training_resource_owner is not None: + raise RuntimeError("MoeEp training resources already exist") + training_config = replace( + self.config, + generate_c=True, + backward_wgrad_mode="operands", + token_padding_size=128, + sf_padding_size=128, + ) + training_kernel_config = Mxfp8KernelConfig.from_forward_config( + training_config + ) + # Graph transport must complete its cross-rank protocol before the + # frontend applies the public trap/drop policy at graph tail. + graph_kernel_config = replace( + training_kernel_config, + drop_on_overflow=True, + # Upstream 5b89819's forward col-requant accepts token + # padding 128/256 but fixes SF atoms at 128; its dGLU + # auxiliaries require token and SF padding to match. The + # graph-only fixed-capacity intersection is therefore 128. + token_padding_block=128, + sf_padding_block=128, + ) + forward = prepare_kernel( + training_config, + graph_kernel_config, + self.device, + ) + backward = prepare_backward_kernel( + training_config, + graph_kernel_config, + self.device, + ) + from ._training_resources import Mxfp8TrainingResourceOwner + + owner = Mxfp8TrainingResourceOwner( + training_config, + self.device, + forward, + backward, + weights, + slot_count=slot_count, + lane_count=lane_count, + ) try: - result = self._backward_executor.run(request) - finally: - try: - self._completion_event.record(stream) - self._completion_recorded = True - self._device_work_may_be_pending = False - except Exception: - self._completion_recorded = False - self._device_work_may_be_pending = True - raise - return result + owner.prepare() + except Exception: + owner.close() + raise + self._training_resource_owner = owner + return owner def close(self) -> None: with self._lock: @@ -311,14 +294,15 @@ def close(self) -> None: "MoeEp MXFP8 backend cannot be closed during " "CUDA graph capture" ) - if self._plan is not None or self._backward_executor is not None: + if ( + self._plan is not None + or self._training_resource_owner is not None + ): torch.cuda.synchronize(self.device) self._adapter.close() - if self._backward_executor is not None: - self._backward_executor.close() - self._backward_executor = None - if self._stash is not None: - self._stash.close() + if self._training_resource_owner is not None: + self._training_resource_owner.close() + self._training_resource_owner = None if self._plan is not None: self._plan.close() self._plan = None diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward.py deleted file mode 100644 index 2d642f373..000000000 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward.py +++ /dev/null @@ -1,112 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Explicit dgrad/dprob Rubin MXFP8 backward orchestration.""" - -from __future__ import annotations - -import torch -import torch.distributed as dist - -from ..._backend import BackendUnavailableError -from ..._contracts import ForwardConfig, ValidatedBackwardRequest -from .._plan import ExecutionPlanOwner -from ._backward_compile import ( - CompiledMxfp8BackwardKernel, - PreparedMxfp8BackwardKernel, - compile_backward_or_get, - prepare_backward_kernel, -) -from ._backward_dispatch import Mxfp8BackwardRedispatch -from ._backward_dprob import return_grad_topk_weights -from ._backward_launch import launch_backward_dglu -from ._backward_layout import Mxfp8BackwardLayout -from ._backward_staging import stage_backward -from ._backward_wgrad_export import export_wgrad_operands -from ._config import Mxfp8KernelConfig - - -class Mxfp8BackwardExecutor: - """Own only compiled products and reusable capacity workspaces.""" - - def __init__(self, config: ForwardConfig, device: torch.device) -> None: - self.config = config - self.device = torch.device(device) - self.kernel_config = Mxfp8KernelConfig.from_forward_config(config) - self._prepared: PreparedMxfp8BackwardKernel | None = None - self._compiled: CompiledMxfp8BackwardKernel | None = None - self._plan: ExecutionPlanOwner | None = None - self._ep_launch_ready = config.ep_size == 1 - - def _ensure_prepared(self) -> PreparedMxfp8BackwardKernel: - if self._prepared is None: - try: - self._prepared = prepare_backward_kernel( - self.config, - self.kernel_config, - self.device, - ) - except (ImportError, OSError) as exc: - raise BackendUnavailableError( - "MoeEp MXFP8 backward requires the 'moe_ep' optional " - "dependencies and their shared libraries" - ) from exc - return self._prepared - - def _ensure_ep_launch_ready(self) -> None: - if self._ep_launch_ready: - return - if self.config.ep_group is None: - raise RuntimeError( - "distributed MXFP8 backward requires an EP process group" - ) - torch.cuda.current_stream(self.device).synchronize() - dist.barrier(group=self.config.ep_group) - self._ep_launch_ready = True - - def run( - self, - request: ValidatedBackwardRequest, - ): - prepared = self._ensure_prepared() - if self._plan is None: - self._plan = ExecutionPlanOwner( - self.config, - self.device, - prepared.workspace_requirements, - ) - - layout = Mxfp8BackwardLayout.from_request(request) - redispatched = Mxfp8BackwardRedispatch(request).run() - resources = self._plan.prepare(request) - inputs = stage_backward(request, layout, prepared, resources) - self._compiled = compile_backward_or_get( - prepared, - inputs, - resources, - ) - self._ensure_ep_launch_ready() - dglu = launch_backward_dglu( - self._compiled, - inputs, - resources, - ) - grad_topk_weights = return_grad_topk_weights( - request, - redispatched.grad_output, - ) - if request.config.backward_wgrad_mode == "operands": - operands = export_wgrad_operands(request, dglu) - return dglu.grad_activation, grad_topk_weights, operands - return dglu.grad_activation, grad_topk_weights - - def close(self) -> None: - if self._plan is not None: - self._plan.close() - self._plan = None - self._prepared = None - self._compiled = None - self._ep_launch_ready = self.config.ep_size == 1 - - -__all__ = ["Mxfp8BackwardExecutor"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dispatch.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dispatch.py deleted file mode 100644 index f073e17c1..000000000 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dispatch.py +++ /dev/null @@ -1,204 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Explicit grad-output re-dispatch for semantic router gradients.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Sequence - -import torch -import torch.distributed as dist - -from ..._contracts import ValidatedBackwardRequest - - -@dataclass(frozen=True) -class _DispatchPlan: - send_token: torch.Tensor - send_slot: torch.Tensor - send_local_expert: torch.Tensor - send_counts: tuple[int, ...] - recv_counts: tuple[int, ...] - - -@dataclass(frozen=True) -class RedispatchedGradOutput: - """Route rows in the compact public ``route_metadata`` order.""" - - grad_output: torch.Tensor - - -class Mxfp8BackwardRedispatch: - """Recreate the identical-route grad-output exchange for dprob.""" - - def __init__(self, request: ValidatedBackwardRequest) -> None: - self.request = request - self.config = request.config - - def _collective_device(self, device: torch.device) -> torch.device: - if ( - device.type != "cpu" - and self.config.ep_size > 1 - and dist.get_backend(self.config.ep_group) == "gloo" - ): - return torch.device("cpu") - return device - - def _exchange_counts(self, send_counts: torch.Tensor) -> torch.Tensor: - if self.config.ep_size == 1: - return send_counts.clone() - staged = send_counts.to(self._collective_device(send_counts.device)) - recv_counts = torch.empty_like(staged) - dist.all_to_all_single( - recv_counts, - staged, - group=self.config.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.config.ep_size == 1: - return send.clone() - comm_device = self._collective_device(send.device) - staged = send.contiguous().to(comm_device) - recv = torch.empty( - (sum(recv_counts), *send.shape[1:]), - 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.config.ep_group, - ) - return recv.to(send.device) - - def _plan(self) -> _DispatchPlan: - config = self.config - flat_expert = self.request.topk_idx.reshape(-1).to(torch.int64) - valid = flat_expert != -1 - token = torch.arange( - self.request.token_count, - dtype=torch.int64, - device=self.request.device, - ).repeat_interleave(config.top_k) - slot = torch.arange( - config.top_k, - dtype=torch.int64, - device=self.request.device, - ).repeat(self.request.token_count) - expert = flat_expert[valid] - destination = torch.div( - expert, - config.experts_per_rank, - rounding_mode="floor", - ) - order = torch.argsort(destination, stable=True) - destination = destination.index_select(0, order) - send_counts_tensor = torch.bincount( - destination, - minlength=config.ep_size, - ).to(torch.int64) - recv_counts_tensor = self._exchange_counts(send_counts_tensor) - return _DispatchPlan( - send_token=token[valid].index_select(0, order), - send_slot=slot[valid].index_select(0, order), - send_local_expert=expert.index_select(0, order).remainder( - config.experts_per_rank - ), - send_counts=tuple( - int(value) for value in send_counts_tensor.cpu().tolist() - ), - recv_counts=tuple( - int(value) for value in recv_counts_tensor.cpu().tolist() - ), - ) - - def run(self) -> RedispatchedGradOutput: - config = self.config - plan = self._plan() - - recv_token = self._all_to_all( - plan.send_token, - plan.send_counts, - plan.recv_counts, - ) - recv_slot = self._all_to_all( - plan.send_slot, - plan.send_counts, - plan.recv_counts, - ) - recv_expert = self._all_to_all( - plan.send_local_expert, - plan.send_counts, - plan.recv_counts, - ) - recv_grad_output = self._all_to_all( - self.request.grad_output.index_select(0, plan.send_token).float(), - plan.send_counts, - plan.recv_counts, - ) - recv_rank = torch.repeat_interleave( - torch.arange( - config.ep_size, - dtype=torch.int64, - device=self.request.device, - ), - torch.tensor( - plan.recv_counts, - dtype=torch.int64, - device=self.request.device, - ), - output_size=self.request.local_routes, - ) - if recv_grad_output.shape[0] != self.request.local_routes: - raise ValueError( - "route_metadata row count does not match the routes received " - "from the re-supplied topk_idx" - ) - - key = ( - ( - ( - recv_expert * config.ep_size - + recv_rank - ) - * int(config.max_tokens_per_rank) - + recv_token - ) - * config.top_k - + recv_slot - ) - compact_order = torch.argsort(key, stable=True) - actual_metadata = torch.stack( - ( - recv_expert.index_select(0, compact_order), - recv_rank.index_select(0, compact_order), - recv_token.index_select(0, compact_order), - recv_slot.index_select(0, compact_order), - ), - dim=1, - ).to(torch.int32) - if not torch.equal(actual_metadata, self.request.route_metadata): - raise ValueError( - "route_metadata does not match the re-supplied forward routes" - ) - - return RedispatchedGradOutput( - grad_output=recv_grad_output.index_select(0, compact_order), - ) - - -__all__ = [ - "Mxfp8BackwardRedispatch", - "RedispatchedGradOutput", -] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dprob.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dprob.py deleted file mode 100644 index 1cf9bf44a..000000000 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_dprob.py +++ /dev/null @@ -1,87 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Return pool-ordered dprob to the source ``(token, top-k)`` plane.""" - -from __future__ import annotations - -import torch -import torch.distributed as dist - -from ..._contracts import ValidatedBackwardRequest -from ._adapter import _decode_moe_tensor - - -def return_grad_topk_weights( - request: ValidatedBackwardRequest, - redispatched_grad_output: torch.Tensor, -) -> torch.Tensor: - """Compute semantic route gradients and return them to every source. - - The fused dGLU kernel consumes an MXFP8 materialization of ``grad_output``. - Using its in-kernel dprob would therefore expose quantization error through - an operation whose public contract specifies straight-through semantics. - Recompute only this scalar gradient from the original FP32 dY, the BF16 - forward stash, and the decoded FC2 weight. - """ - - config = request.config - metadata = request.route_metadata.to(torch.int64) - local_dprob = torch.zeros( - (request.local_routes,), - dtype=torch.float32, - device=request.device, - ) - fc2_weight = _decode_moe_tensor(request.fc2_weight) - gate, up = request.fc1_c.float().split( - config.intermediate_size, - dim=-1, - ) - if config.gate_up_clamp is not None: - gate = gate.clamp(max=config.gate_up_clamp) - up = up.clamp( - min=-config.gate_up_clamp, - max=config.gate_up_clamp, - ) - hidden = (gate * torch.sigmoid(gate)) * up - local_expert = metadata[:, 0] - for expert in range(config.experts_per_rank): - positions = torch.nonzero( - local_expert == expert, - as_tuple=False, - ).flatten() - if positions.numel() == 0: - continue - grad_output = redispatched_grad_output.index_select(0, positions) - expert_hidden = hidden.index_select(0, positions) - grad_hidden = grad_output @ fc2_weight[expert].transpose(0, 1) - local_dprob.index_copy_( - 0, - positions, - (grad_hidden * expert_hidden).sum(dim=-1), - ) - - global_dprob = torch.zeros( - ( - config.ep_size, - int(config.max_tokens_per_rank), - config.top_k, - ), - dtype=torch.float32, - device=request.device, - ) - if request.local_routes: - global_dprob[ - metadata[:, 1], - metadata[:, 2], - metadata[:, 3], - ] = local_dprob - if config.ep_size > 1: - dist.all_reduce(global_dprob, group=config.ep_group) - return global_dprob[ - config.ep_rank, - : request.token_count, - ] - - -__all__ = ["return_grad_topk_weights"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py index 9cf069be5..8680e37b0 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py @@ -21,6 +21,7 @@ @dataclass(frozen=True) class Mxfp8DgluResult: grad_activation: torch.Tensor + grad_topk_weights: torch.Tensor fc1_recompute: torch.Tensor fc1_recompute_sf: torch.Tensor fc1_col_output: torch.Tensor @@ -42,6 +43,10 @@ def launch_backward_dglu( grad_activation=inputs.output_activation[ : inputs.token_count ].float(), + # The dGLU epilogue has already returned source-order dprob through + # the symmetric token-communication plane. Own the public result so a + # later launch cannot overwrite it. + grad_topk_weights=inputs.dprob[: inputs.token_count].clone(), fc1_recompute=inputs.fc1_recompute, fc1_recompute_sf=inputs.fc1_recompute_sf, fc1_col_output=inputs.fc1_col_output, diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_layout.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_layout.py deleted file mode 100644 index e1ae287d5..000000000 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_layout.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Stateless lowering from the public backward stash to Rubin pool rows.""" - -from __future__ import annotations - -from dataclasses import dataclass - -import torch - -from ..._contracts import ForwardConfig, ValidatedBackwardRequest - -@dataclass(frozen=True) -class Mxfp8BackwardLayout: - """Stateless lowering from compact route metadata to the dGLU pool LUT.""" - - preact_row_lut: torch.Tensor - - @classmethod - def from_request( - cls, - request: ValidatedBackwardRequest, - ) -> "Mxfp8BackwardLayout": - config = request.config - metadata = request.route_metadata - if config.max_tokens_per_rank is None: - raise ValueError("MXFP8 backward requires max_tokens_per_rank") - - bounds = ( - (metadata[:, 0], 0, config.experts_per_rank, "local expert"), - (metadata[:, 1], 0, config.ep_size, "source rank"), - ( - metadata[:, 2], - 0, - config.max_tokens_per_rank, - "source token", - ), - (metadata[:, 3], 0, config.top_k, "source top-k slot"), - ) - for values, lower, upper, name in bounds: - if values.numel() and bool( - ((values < lower) | (values >= upper)).any().item() - ): - raise ValueError( - f"route_metadata contains an out-of-range {name}" - ) - - preact_row_lut = cls._build_preact_row_lut(config, metadata) - return cls( - preact_row_lut=preact_row_lut, - ) - - @staticmethod - def _build_preact_row_lut( - config: ForwardConfig, - metadata: torch.Tensor, - ) -> torch.Tensor: - lut = torch.full( - ( - config.ep_size, - int(config.max_tokens_per_rank), - config.top_k, - ), - -1, - dtype=torch.int32, - device=metadata.device, - ) - if metadata.shape[0]: - compact_rows = torch.arange( - metadata.shape[0], - dtype=torch.int32, - device=metadata.device, - ) - lut[ - metadata[:, 1].to(torch.int64), - metadata[:, 2].to(torch.int64), - metadata[:, 3].to(torch.int64), - ] = compact_rows - return lut - -__all__ = ["Mxfp8BackwardLayout"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_staging.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_staging.py deleted file mode 100644 index 4225aeda7..000000000 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_staging.py +++ /dev/null @@ -1,516 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Tensor staging for the explicit Rubin MXFP8 backward invocation.""" - -from __future__ import annotations - -import math - -import torch - -from ..._contracts import ValidatedBackwardRequest -from .._plan import PreparedResources -from .._workspace import padded_mxfp8_scale_columns -from ._adapter import ( - _GATE_UP_INTERLEAVE, - _decode_moe_tensor, - _interleave_gate_up_rows, - _quantize_plain_mxfp8, - _stack_blocked_scales, - _typed_view, - _zero_workspace_prefix, - _zero_workspace_range, -) -from ._backward_compile import ( - Mxfp8BackwardLaunchInputs, - PreparedMxfp8BackwardKernel, -) -from ._backward_layout import Mxfp8BackwardLayout - -_DATA_DTYPE = torch.float8_e4m3fn -_SCALE_DTYPE = torch.float8_e8m0fnu - - -def _typed_prefix_view( - byte_tensor: torch.Tensor, - dtype: torch.dtype, - shape: tuple[int, ...], -) -> torch.Tensor: - """Return a compact typed view of a prefix of a reusable byte region.""" - - nbytes = math.prod(shape) * dtype.itemsize - if nbytes > byte_tensor.numel(): - raise ValueError( - f"byte region has {byte_tensor.numel()} bytes, " - f"cannot provide {nbytes} bytes for shape={shape}, dtype={dtype}" - ) - return _typed_view(byte_tensor.narrow(0, 0, nbytes), dtype, shape) - - -def _k_major(tensor: torch.Tensor) -> torch.Tensor: - """Return the same logical ``(E,K,N)`` tensor with K stride one.""" - - return tensor.permute(0, 2, 1).contiguous().permute(0, 2, 1) - - -def _prepare_backward_weights( - request: ValidatedBackwardRequest, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Quantize W2^T and W1^T along their backward reduction axes.""" - - config = request.config - intermediate = config.intermediate_size - - w2_t = _decode_moe_tensor(request.fc2_weight).transpose(1, 2) - q_w2_t = _quantize_plain_mxfp8(w2_t, axis=1) - fc1_weight = _k_major(q_w2_t.data) - fc1_weight_sf = _stack_blocked_scales( - q_w2_t.scale.permute(0, 2, 1).contiguous() - ) - - w1_t = _decode_moe_tensor(request.fc1_weight).transpose(1, 2) - q_w1_t = _quantize_plain_mxfp8(w1_t, axis=1) - interleaved_w1 = _interleave_gate_up_rows( - q_w1_t.data, - intermediate, - ) - fc2_weight = _k_major(interleaved_w1) - - scale_blocks = intermediate // 32 - gate_sf = q_w1_t.scale[:, :scale_blocks, :] - up_sf = q_w1_t.scale[:, scale_blocks:, :] - interleaved_sf = ( - torch.stack( - ( - gate_sf.view(torch.uint8), - up_sf.view(torch.uint8), - ), - dim=2, - ) - .reshape( - q_w1_t.scale.shape[0], - 2 * scale_blocks, - q_w1_t.scale.shape[2], - ) - .view(_SCALE_DTYPE) - ) - fc2_weight_sf = _stack_blocked_scales( - interleaved_sf.permute(0, 2, 1).contiguous() - ) - return ( - fc1_weight, - fc1_weight_sf, - fc2_weight, - fc2_weight_sf, - ) - - -def _router_order_key( - source_token: torch.Tensor, - source_slot: torch.Tensor, - request: ValidatedBackwardRequest, - prepared: PreparedMxfp8BackwardKernel, -) -> tuple[torch.Tensor, int]: - """Return each route's position in the deterministic router's source run.""" - - elements_per_vector = 4 # The staged top-k index tensor is Int32. - token_comm = prepared.kernel.token_comm - router_ctas = int(token_comm.router_data_cta_count) - router_warps = int(token_comm.router_warps_per_cta) - threads_per_cta = router_warps * 32 - grid_threads = router_ctas * threads_per_cta - tile_span = elements_per_vector * grid_threads - maximum_elements = ( - int(request.config.max_tokens_per_rank) * request.config.top_k - ) - load_rounds = ( - maximum_elements + tile_span - 1 - ) // tile_span - elements_per_thread = load_rounds * elements_per_vector - - flat_index = source_token * request.config.top_k + source_slot - load_round = torch.div( - flat_index, - tile_span, - rounding_mode="floor", - ) - in_tile = flat_index.remainder(tile_span) - grid_thread = torch.div( - in_tile, - elements_per_vector, - rounding_mode="floor", - ) - register_index = ( - load_round * elements_per_vector - + in_tile.remainder(elements_per_vector) - ) - cta = torch.div( - grid_thread, - threads_per_cta, - rounding_mode="floor", - ) - thread_in_cta = grid_thread.remainder(threads_per_cta) - warp = torch.div(thread_in_cta, 32, rounding_mode="floor") - lane = thread_in_cta.remainder(32) - order_key = ( - ( - (cta * router_warps + warp) * elements_per_thread - + register_index - ) - * 32 - + lane - ) - order_span = ( - router_ctas * router_warps * elements_per_thread * 32 - ) - return order_key, order_span - - -def _stage_fc1_preact( - request: ValidatedBackwardRequest, - layout: Mxfp8BackwardLayout, - prepared: PreparedMxfp8BackwardKernel, - fc1_preact: torch.Tensor, -) -> None: - """Lower compact public stash rows into the upstream dGLU pool layout.""" - - config = prepared.config - expected_shape = ( - prepared.pool_token_capacity, - 2 * config.intermediate, - ) - if ( - fc1_preact.dtype is not torch.bfloat16 - or tuple(fc1_preact.shape) != expected_shape - or not fc1_preact.is_contiguous() - ): - raise ValueError( - "backward fc1_preact must be contiguous BF16 with shape " - f"{expected_shape}, got shape={tuple(fc1_preact.shape)}, " - f"dtype={fc1_preact.dtype}" - ) - fc1_preact.zero_() - - metadata = request.route_metadata.to(torch.int64) - compact_rows = layout.preact_row_lut[ - metadata[:, 1], - metadata[:, 2], - metadata[:, 3], - ].to(torch.int64) - if compact_rows.numel() and bool((compact_rows < 0).any().item()): - raise RuntimeError("backward preactivation LUT is incomplete") - - if config.intermediate % _GATE_UP_INTERLEAVE: - raise RuntimeError( - "backward preactivation requires intermediate_size divisible by " - f"{_GATE_UP_INTERLEAVE}" - ) - gate, up = request.fc1_c.split(config.intermediate, dim=1) - pairs = config.intermediate // _GATE_UP_INTERLEAVE - interleaved_preact = torch.stack( - ( - gate.reshape(-1, pairs, _GATE_UP_INTERLEAVE), - up.reshape(-1, pairs, _GATE_UP_INTERLEAVE), - ), - dim=2, - ).reshape(-1, 2 * config.intermediate) - - physical_offset = 0 - for expert in range(config.num_experts): - positions = torch.nonzero( - metadata[:, 0] == expert, - as_tuple=False, - ).flatten() - count = int(positions.numel()) - if count: - # Receiver pools concatenate source ranks in a destination-relative - # ring: local rank first, then increasing ranks with wraparound. - # The public stash is source-rank sorted, so restore pool order. - source_rank = metadata.index_select(0, positions)[:, 1] - source_token = metadata.index_select(0, positions)[:, 2] - source_slot = metadata.index_select(0, positions)[:, 3] - source_order, source_order_span = _router_order_key( - source_token, - source_slot, - request, - prepared, - ) - ring_position = ( - source_rank - request.config.ep_rank - ) % request.config.ep_size - route_key = ring_position * source_order_span + source_order - positions = positions.index_select( - 0, - torch.argsort(route_key, stable=True), - ) - if physical_offset + count > prepared.pool_token_capacity: - raise RuntimeError( - "backward preactivation rows exceed Rubin pool capacity" - ) - destination = fc1_preact.narrow(0, physical_offset, count) - destination.copy_( - interleaved_preact.index_select( - 0, - compact_rows.index_select(0, positions), - ) - ) - physical_offset += ( - count + config.token_padding_block - 1 - ) // config.token_padding_block * config.token_padding_block - - if physical_offset > prepared.pool_token_capacity: - raise RuntimeError( - "backward preactivation rows exceed Rubin pool capacity" - ) - - -def stage_backward( - request: ValidatedBackwardRequest, - layout: Mxfp8BackwardLayout, - prepared: PreparedMxfp8BackwardKernel, - resources: PreparedResources, -) -> Mxfp8BackwardLaunchInputs: - config = prepared.config - capacity = config.max_tokens_per_rank - token_count = request.token_count - hidden = config.hidden - top_k = config.top_k - - symmetric = resources.workspace.symmetric - local = resources.workspace.local - grad_out = _typed_view( - symmetric["activation_data"], - _DATA_DTYPE, - (capacity, hidden), - ) - grad_out_sf = _typed_view( - symmetric["activation_scale"], - _SCALE_DTYPE, - (capacity, padded_mxfp8_scale_columns(hidden)), - ) - topk_weights = _typed_view( - symmetric["topk_weights"], - torch.float32, - (capacity, top_k), - ) - topk_idx = _typed_view( - local["topk_idx"], - torch.int32, - (capacity, top_k), - ) - output_activation = _typed_view( - symmetric["output_data"], - torch.bfloat16, - (capacity, hidden), - ) - overflow_flag = _typed_view( - local["overflow_flag"], - torch.int32, - (1,), - ) - fc1_preact_shape = tuple( - int(extent) for extent in prepared.kernel.get_fc1_preact_shape() - ) - fc1_preact = _typed_view( - local["backward_fc1_preact"], - torch.bfloat16, - fc1_preact_shape, - ) - local_workspace = local["kernel_local_workspace"] - shared_workspace = symmetric["kernel_shared_workspace"] - - _zero_workspace_prefix( - local_workspace, - prepared.local_workspace_zero_bytes, - name="MXFP8 backward local workspace", - ) - _zero_workspace_prefix( - shared_workspace, - prepared.shared_workspace_zero_bytes, - name="MXFP8 backward shared workspace", - ) - _stage_fc1_preact( - request, - layout, - prepared, - fc1_preact, - ) - _zero_workspace_range( - shared_workspace, - prepared.pre_reduced_activation_offset, - token_count * prepared.pre_reduced_activation_bytes_per_token, - name="MXFP8 backward pre-reduced activation workspace", - ) - quantized_combine = config.combine_format != "bf16" - if quantized_combine: - if ( - prepared.pre_reduced_activation_sf_offset is None - or prepared.pre_reduced_activation_sf_bytes_per_token <= 0 - ): - raise RuntimeError( - "MXFP8 backward quantized combine requires scale workspace" - ) - _zero_workspace_range( - shared_workspace, - prepared.pre_reduced_activation_sf_offset, - token_count - * prepared.pre_reduced_activation_sf_bytes_per_token, - name="MXFP8 backward pre-reduced activation scale workspace", - ) - elif ( - prepared.pre_reduced_activation_sf_offset is not None - or prepared.pre_reduced_activation_sf_bytes_per_token != 0 - ): - raise RuntimeError( - "MXFP8 backward BF16 combine must not expose scale workspace" - ) - quantized_grad = _quantize_plain_mxfp8( - request.grad_output, - axis=1, - ) - grad_out.zero_() - grad_out[:token_count].copy_(quantized_grad.data) - grad_out_sf.zero_() - grad_out_sf[ - :token_count, - : quantized_grad.scale.shape[1], - ].copy_(quantized_grad.scale) - topk_idx.fill_(-1) - topk_idx[:token_count].copy_( - request.topk_idx.to(torch.int32) - ) - topk_weights.zero_() - topk_weights[:token_count].copy_( - request.topk_weights.float() - ) - output_activation.zero_() - overflow_flag.zero_() - - aux_shapes = { - name: tuple(int(extent) for extent in shape) - for name, shape in prepared.kernel.get_aux_output_shapes().items() - } - dprob = _typed_view( - symmetric["backward_dprob"], - torch.float32, - aux_shapes["dprob"], - ) - dprob.zero_() - operands_mode = request.config.backward_wgrad_mode == "operands" - expected_flags = ( - prepared.dfc2_recompute, - prepared.dfc2_col_output, - prepared.enable_grad_y2_col_quant, - ) - if operands_mode: - if expected_flags != (True, True, True): - raise RuntimeError( - "wgrad operands require every backward auxiliary output" - ) - # These allocations are intentionally not execution-plan workspace: - # the returned operand bundle must remain valid after later calls. - fc1_recompute = torch.zeros( - aux_shapes["fc1_recompute"], - dtype=_DATA_DTYPE, - device=request.device, - ) - fc1_recompute_sf = torch.full( - aux_shapes["fc1_recompute_sf"], - 127, - dtype=torch.uint8, - device=request.device, - ).view(_SCALE_DTYPE) - fc1_col_output = torch.zeros( - aux_shapes["fc1_col_output"], - dtype=_DATA_DTYPE, - device=request.device, - ) - fc1_col_output_sf = torch.full( - aux_shapes["fc1_col_output_sf"], - 127, - dtype=torch.uint8, - device=request.device, - ).view(_SCALE_DTYPE) - grad_y2 = torch.zeros( - aux_shapes["grad_y2"], - dtype=_DATA_DTYPE, - device=request.device, - ) - grad_y2_sf = torch.full( - aux_shapes["grad_y2_sf"], - 127, - dtype=torch.uint8, - device=request.device, - ) - else: - if expected_flags != (False, False, False): - raise RuntimeError( - "default backward must disable wgrad auxiliary outputs" - ) - # Preserve the existing mode-none allocation/performance behavior: - # disabled fixed-ABI arguments alias reusable plan scratch. - fc1_recompute = _typed_prefix_view( - local["backward_aux_data"], - _DATA_DTYPE, - aux_shapes["fc1_recompute"], - ) - fc1_col_output = _typed_prefix_view( - local["backward_aux_data"], - _DATA_DTYPE, - aux_shapes["fc1_col_output"], - ) - fc1_recompute_sf = _typed_prefix_view( - local["backward_aux_scale"], - _SCALE_DTYPE, - aux_shapes["fc1_recompute_sf"], - ) - fc1_col_output_sf = _typed_prefix_view( - local["backward_aux_scale"], - _SCALE_DTYPE, - aux_shapes["fc1_col_output_sf"], - ) - grad_y2 = _typed_prefix_view( - local["backward_aux_data"], - _DATA_DTYPE, - aux_shapes["grad_y2"], - ) - grad_y2_sf = _typed_prefix_view( - local["backward_aux_scale"], - torch.uint8, - aux_shapes["grad_y2_sf"], - ) - fc1_weight, fc1_weight_sf, fc2_weight, fc2_weight_sf = ( - _prepare_backward_weights(request) - ) - return Mxfp8BackwardLaunchInputs( - grad_out=grad_out, - grad_out_sf=grad_out_sf, - topk_idx=topk_idx, - topk_weights=topk_weights, - fc1_weight=fc1_weight, - fc1_weight_sf=fc1_weight_sf, - fc2_weight=fc2_weight, - fc2_weight_sf=fc2_weight_sf, - beta=torch.ones( - (config.num_experts,), - dtype=torch.float32, - device=request.device, - ), - fc1_preact=fc1_preact, - output_activation=output_activation, - overflow_flag=overflow_flag, - dprob=dprob, - fc1_recompute=fc1_recompute, - fc1_recompute_sf=fc1_recompute_sf, - fc1_col_output=fc1_col_output, - fc1_col_output_sf=fc1_col_output_sf, - grad_y2=grad_y2, - grad_y2_sf=grad_y2_sf, - local_workspace=local_workspace, - shared_workspace=shared_workspace, - token_count=token_count, - ) - - -__all__ = ["stage_backward"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_wgrad_export.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_wgrad_export.py deleted file mode 100644 index 6146cf296..000000000 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_wgrad_export.py +++ /dev/null @@ -1,226 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Materialize caller-owned grouped-wgrad operands from backward auxiliaries.""" - -from __future__ import annotations - -import torch - -from ..._contracts import ValidatedBackwardRequest -from ..._types import MoeEpWgradOperands -from ._backward_launch import Mxfp8DgluResult -from ._wgrad_layout import ( - assemble_dfc2_atom_scales, - assemble_discrete_col_requant_scales, - deinterleave_gate_up_columns, - pool_data_as_wgrad_a, - pool_data_as_wgrad_b, -) - - -def export_wgrad_operands( - request: ValidatedBackwardRequest, - dglu: Mxfp8DgluResult, -) -> MoeEpWgradOperands: - """Convert physical Rubin pools to the grouped-wgrad Tensor2D ABI.""" - - if request.config.backward_wgrad_mode != "operands": - raise ValueError("wgrad operands can only be exported in operands mode") - stash = request.wgrad_forward_stash - if stash is None: - raise ValueError("wgrad operand export requires a forward stash") - - padded_ends = tuple( - int(value) for value in stash.expert_offsets.detach().cpu().tolist() - ) - valid_counts = tuple( - int(value) - for value in stash.valid_route_counts.detach().cpu().tolist() - ) - padded_routes = padded_ends[-1] if padded_ends else 0 - config = request.config - - _validate_aux_shapes(dglu, padded_routes, config) - - # dC is emitted in the kernel's 32-column gate/up strip order. Upstream - # exports route-weighted h and unweighted dY, whose product is the same - # dW2 contract previously represented as unweighted h and weighted dY. - # Every conversion allocates fresh storage, so no returned tensor aliases - # the reusable execution plan. - dc_pool = deinterleave_gate_up_columns( - dglu.fc1_col_output[:padded_routes], - config.intermediate_size, - ) - fc1_b = pool_data_as_wgrad_b(dc_pool, padded_routes) - fc1_sfb = assemble_dfc2_atom_scales( - dglu.fc1_col_output_sf, - valid_counts, - padded_ends, - 2 * config.intermediate_size, - config.sf_padding_size, - deinterleave_gate_up=config.intermediate_size, - ) - - fc2_a = pool_data_as_wgrad_a( - dglu.fc1_recompute, - padded_routes, - ) - fc2_sfa = assemble_dfc2_atom_scales( - dglu.fc1_recompute_sf, - valid_counts, - padded_ends, - config.intermediate_size, - config.sf_padding_size, - ) - fc2_b = pool_data_as_wgrad_b( - dglu.grad_y2, - padded_routes, - ) - fc2_sfb = assemble_discrete_col_requant_scales( - dglu.grad_y2_sf, - valid_counts, - padded_ends, - config.hidden_size, - config.sf_padding_size, - ) - - operands = MoeEpWgradOperands( - fc1_a=stash.fc1_a, - fc1_sfa=stash.fc1_sfa, - fc1_b=fc1_b, - fc1_sfb=fc1_sfb, - fc2_a=fc2_a, - fc2_sfa=fc2_sfa, - fc2_b=fc2_b, - fc2_sfb=fc2_sfb, - expert_offsets=stash.expert_offsets, - valid_route_counts=stash.valid_route_counts, - route_metadata=stash.route_metadata, - ) - _validate_grouped_wgrad_abi(operands, config, padded_routes) - return operands - - -def _validate_aux_shapes( - dglu: Mxfp8DgluResult, - padded_routes: int, - config, -) -> None: - expected_columns = ( - ("fc1_recompute", dglu.fc1_recompute, config.intermediate_size), - ( - "fc1_col_output", - dglu.fc1_col_output, - 2 * config.intermediate_size, - ), - ( - "grad_y2", - dglu.grad_y2, - config.hidden_size, - ), - ) - for name, tensor, columns in expected_columns: - if tensor.ndim != 2 or tensor.shape[1] != columns: - raise RuntimeError( - f"{name} must have shape (pool_capacity, {columns}), " - f"got {tuple(tensor.shape)}" - ) - if tensor.shape[0] < padded_routes or not tensor.is_contiguous(): - raise RuntimeError( - f"{name} does not contain a contiguous {padded_routes}-row " - "pool prefix" - ) - if tensor.dtype is not torch.float8_e4m3fn: - raise TypeError(f"{name} must have dtype torch.float8_e4m3fn") - _require_alignment(name, tensor, 16) - - for name, tensor in ( - ("fc1_recompute_sf", dglu.fc1_recompute_sf), - ("fc1_col_output_sf", dglu.fc1_col_output_sf), - ): - if tensor.ndim != 2 or tensor.dtype is not torch.float8_e8m0fnu: - raise TypeError(f"{name} must be a rank-2 E8M0 tensor") - _require_alignment(name, tensor, 16) - if dglu.grad_y2_sf.ndim != 1 or dglu.grad_y2_sf.dtype is not torch.uint8: - raise TypeError("grad_y2_sf must be a rank-1 uint8 tensor") - _require_alignment("grad_y2_sf", dglu.grad_y2_sf, 16) - - -def _validate_grouped_wgrad_abi( - operands: MoeEpWgradOperands, - config, - padded_routes: int, -) -> None: - rounded_hidden = _round_up(config.hidden_size, 128) - rounded_intermediate = _round_up(config.intermediate_size, 128) - rounded_gate_up = _round_up(2 * config.intermediate_size, 128) - sf_columns = _round_up(padded_routes // 32, 4) - expected = ( - ("fc1_a", operands.fc1_a, (config.hidden_size, padded_routes)), - ("fc1_sfa", operands.fc1_sfa, (rounded_hidden, sf_columns)), - ( - "fc1_b", - operands.fc1_b, - (padded_routes, 2 * config.intermediate_size), - ), - ("fc1_sfb", operands.fc1_sfb, (rounded_gate_up, sf_columns)), - ( - "fc2_a", - operands.fc2_a, - (config.intermediate_size, padded_routes), - ), - ( - "fc2_sfa", - operands.fc2_sfa, - (rounded_intermediate, sf_columns), - ), - ("fc2_b", operands.fc2_b, (padded_routes, config.hidden_size)), - ("fc2_sfb", operands.fc2_sfb, (rounded_hidden, sf_columns)), - ) - for name, tensor, shape in expected: - if tuple(tensor.shape) != shape: - raise RuntimeError( - f"grouped-wgrad {name} shape must be {shape}, " - f"got {tuple(tensor.shape)}" - ) - _require_alignment(name, tensor, 16) - - for name in ("fc1_a", "fc2_a"): - tensor = getattr(operands, name) - expected_stride = (padded_routes, 1) - if padded_routes and tensor.stride() != expected_stride: - raise RuntimeError( - f"grouped-wgrad {name} strides must be " - f"{expected_stride}, got {tensor.stride()}" - ) - for name in ("fc1_b", "fc2_b"): - tensor = getattr(operands, name) - expected_stride = (1, padded_routes) - if padded_routes and tensor.stride() != expected_stride: - raise RuntimeError( - f"grouped-wgrad {name} strides must be " - f"{expected_stride}, got {tensor.stride()}" - ) - for name in ("fc1_sfa", "fc1_sfb", "fc2_sfa", "fc2_sfb"): - if not getattr(operands, name).is_contiguous(): - raise RuntimeError(f"grouped-wgrad {name} must be contiguous") - _require_alignment("expert_offsets", operands.expert_offsets, 4) - - -def _require_alignment( - name: str, - tensor: torch.Tensor, - alignment: int, -) -> None: - if tensor.data_ptr() % alignment: - raise RuntimeError( - f"{name} address is not {alignment}-byte aligned" - ) - - -def _round_up(value: int, multiple: int) -> int: - return (value + multiple - 1) // multiple * multiple - - -__all__ = ["export_wgrad_operands"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py index 9cbd18fc3..6b4a3bd3e 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py @@ -34,6 +34,8 @@ class PreparedMxfp8Kernel: col_quant_sf_elements: int token_src_metadata_offset: int token_src_metadata_bytes: int + col_quant_sizes_offset: int | None + col_quant_sizes_bytes: int pre_reduced_activation_offset: int | None pre_reduced_activation_bytes_per_token: int pre_reduced_activation_sf_offset: int | None @@ -52,6 +54,9 @@ class CompiledMxfp8Kernel: _COMPILE_LOCK = threading.RLock() _COMPILE_CACHE: dict[tuple, CompiledMxfp8Kernel] = {} _TOKEN_SRC_METADATA_REGION = "nvlink.token_comm.token_src_metadata" +_COL_QUANT_SIZES_REGION = ( + "rubin.glu_mxfp8.mega.col_quant_expert_token_sizes" +) _PRE_REDUCED_ACTIVATION_REGION = ( "nvlink.token_comm.pre_reduced_activation" ) @@ -272,6 +277,35 @@ def prepare_kernel( if config.enable_col_quant else 0 ) + if config.enable_col_quant: + col_quant_sizes_region = device_workspace.region( + _COL_QUANT_SIZES_REGION + ) + if col_quant_sizes_region.buffer_space != "local": + raise RuntimeError( + "Rubin col-quant expert-size snapshot must reside in " + "local workspace" + ) + col_quant_sizes_offset = int( + device_workspace.offset(_COL_QUANT_SIZES_REGION) + ) + col_quant_sizes_bytes = int( + device_workspace.nbytes(_COL_QUANT_SIZES_REGION) + ) + expected_sizes_bytes = config.num_experts * torch.int32.itemsize + if col_quant_sizes_bytes != expected_sizes_bytes: + raise RuntimeError( + "Rubin col-quant expert-size snapshot has " + f"{col_quant_sizes_bytes} bytes, expected " + f"{expected_sizes_bytes}" + ) + if col_quant_sizes_offset + col_quant_sizes_bytes > local_bytes: + raise RuntimeError( + "Rubin col-quant expert-size snapshot exceeds local workspace" + ) + else: + col_quant_sizes_offset = None + col_quant_sizes_bytes = 0 requirements = WorkspaceRequirements.for_mxfp8( forward_config, kernel_local_workspace_bytes=local_bytes, @@ -307,6 +341,8 @@ def prepare_kernel( col_quant_sf_elements=col_quant_sf_elements, token_src_metadata_offset=token_src_metadata_offset, token_src_metadata_bytes=token_src_metadata_bytes, + col_quant_sizes_offset=col_quant_sizes_offset, + col_quant_sizes_bytes=col_quant_sizes_bytes, pre_reduced_activation_offset=pre_reduced_activation_offset, pre_reduced_activation_bytes_per_token=( pre_reduced_activation_bytes_per_token diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_stash.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_stash.py deleted file mode 100644 index 27409e837..000000000 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_stash.py +++ /dev/null @@ -1,302 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Eager-only ownership and materialization for the MXFP8 FC1 stash.""" - -from __future__ import annotations - -from dataclasses import dataclass - -import torch -import torch.distributed as dist - -from ..._contracts import ForwardConfig, ValidatedForwardRequest -from ..._types import MoeEpWgradForwardStash -from .._workspace import _align_up -from ._adapter import _GATE_UP_INTERLEAVE, Mxfp8LaunchInputs -from ._compile import PreparedMxfp8Kernel -from ._wgrad_layout import ( - assemble_discrete_col_requant_scales, - cumulative_padded_offsets, - pool_data_as_wgrad_a, -) - -_TOKEN_SRC_METADATA_BYTES = 8 - - -@dataclass(frozen=True) -class Mxfp8StashPlan: - """One launch's logical counts over an instance-owned raw C buffer.""" - - buffer: torch.Tensor - expert_counts: tuple[int, ...] - expert_offsets: tuple[int, ...] - local_routes: int - - -class Mxfp8ForwardStash: - """Own a padded high-watermark C buffer and compact public results.""" - - def __init__(self, config: ForwardConfig, device: torch.device) -> None: - if not config.generate_c: - raise ValueError("Mxfp8ForwardStash requires generate_c=True") - self.config = config - self.device = torch.device(device) - self._buffer: torch.Tensor | None = None - self._allocation_count = 0 - self._closed = False - - @property - def capacity(self) -> int: - return 0 if self._buffer is None else int(self._buffer.shape[0]) - - @property - def allocation_count(self) -> int: - return self._allocation_count - - def _local_expert_counts( - self, - request: ValidatedForwardRequest, - ) -> tuple[int, ...]: - flat_experts = request.topk_idx.reshape(-1).to(torch.int64) - valid_experts = flat_experts[flat_experts >= 0] - counts = torch.bincount( - valid_experts, - minlength=self.config.num_experts, - ).to(torch.int64) - if self.config.ep_size > 1: - if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError( - "distributed generate_c route counting requires an " - "initialized torch.distributed process group" - ) - dist.all_reduce( - counts, - op=dist.ReduceOp.SUM, - group=self.config.ep_group, - ) - begin = self.config.ep_rank * self.config.experts_per_rank - end = begin + self.config.experts_per_rank - return tuple(int(value) for value in counts[begin:end].cpu().tolist()) - - def prepare( - self, - request: ValidatedForwardRequest, - *, - pool_token_capacity: int, - ) -> Mxfp8StashPlan: - if self._closed: - raise RuntimeError("MXFP8 forward stash is closed") - if request.device != self.device: - raise ValueError( - f"MXFP8 forward stash is bound to {self.device}, " - f"got {request.device}" - ) - if torch.cuda.is_current_stream_capturing(): - raise NotImplementedError( - "MoeEp generate_c=True is eager-only and does not support " - "CUDA graph capture" - ) - if pool_token_capacity <= 0: - raise ValueError("pool_token_capacity must be positive") - - expert_counts = self._local_expert_counts(request) - expert_offsets = [] - padded_routes = 0 - token_padding = ( - self.config.token_padding_size - if self.config.backward_wgrad_mode == "operands" - else 128 - ) - for count in expert_counts: - expert_offsets.append(padded_routes) - padded_routes += _align_up( - count, - token_padding, - ) - - if padded_routes > pool_token_capacity: - raise RuntimeError( - "forward stash route layout exceeds Rubin pool capacity: " - f"{padded_routes} > {pool_token_capacity}" - ) - # The upstream training kernel validates the receiver-domain C tensor - # against its complete pool shape, even though only active expert rows - # are materialized for the public stash. - required_capacity = pool_token_capacity - if self._buffer is None or self.capacity < required_capacity: - self._buffer = torch.empty( - required_capacity, - 2 * self.config.intermediate_size, - dtype=torch.bfloat16, - device=self.device, - ) - self._allocation_count += 1 - - return Mxfp8StashPlan( - buffer=self._buffer, - expert_counts=expert_counts, - expert_offsets=tuple(expert_offsets), - local_routes=sum(expert_counts), - ) - - def materialize( - self, - plan: Mxfp8StashPlan, - inputs: Mxfp8LaunchInputs, - prepared: PreparedMxfp8Kernel, - ) -> tuple[ - torch.Tensor, - torch.Tensor, - MoeEpWgradForwardStash | None, - ]: - """Compact padded kernel rows into the documented public stash.""" - - if self._closed: - raise RuntimeError("MXFP8 forward stash is closed") - if inputs.fc1_c is not plan.buffer: - raise ValueError("launch inputs do not use this stash plan's buffer") - if prepared.token_src_metadata_bytes != ( - prepared.pool_token_capacity * _TOKEN_SRC_METADATA_BYTES - ): - raise RuntimeError("unexpected token_src_metadata byte size") - - local_routes = plan.local_routes - if local_routes == 0: - fc1_c = torch.empty( - (0, 2 * self.config.intermediate_size), - dtype=torch.bfloat16, - device=self.device, - ) - route_metadata = torch.empty( - (0, 4), - dtype=torch.int32, - device=self.device, - ) - else: - physical_rows = torch.cat( - tuple( - torch.arange( - count, - dtype=torch.int64, - device=self.device, - ) - + offset - for count, offset in zip( - plan.expert_counts, - plan.expert_offsets, - ) - if count - ) - ) - local_experts = torch.repeat_interleave( - torch.arange( - self.config.experts_per_rank, - dtype=torch.int64, - device=self.device, - ), - torch.tensor( - plan.expert_counts, - dtype=torch.int64, - device=self.device, - ), - output_size=local_routes, - ) - - metadata_region = inputs.shared_workspace.narrow( - 0, - prepared.token_src_metadata_offset, - prepared.token_src_metadata_bytes, - ) - packed_metadata = metadata_region.view(torch.int64).index_select( - 0, - physical_rows, - ) - src_tokens = packed_metadata & 0xFFFFFFFF - high = packed_metadata >> 32 - src_ranks = (high >> 16) & 0xFFFF - src_slots = high & 0xFFFF - - order_key = ( - ( - local_experts * self.config.ep_size - + src_ranks - ) - * int(self.config.max_tokens_per_rank) - + src_tokens - ) * self.config.top_k + src_slots - order = torch.argsort(order_key, stable=True) - physical_rows = physical_rows.index_select(0, order) - local_experts = local_experts.index_select(0, order) - src_ranks = src_ranks.index_select(0, order) - src_tokens = src_tokens.index_select(0, order) - src_slots = src_slots.index_select(0, order) - - raw_fc1_c = plan.buffer.index_select(0, physical_rows) - pairs = self.config.intermediate_size // _GATE_UP_INTERLEAVE - gate_up_blocks = raw_fc1_c.reshape( - local_routes, - pairs, - 2, - _GATE_UP_INTERLEAVE, - ) - gate = gate_up_blocks[:, :, 0, :].reshape( - local_routes, - self.config.intermediate_size, - ) - up = gate_up_blocks[:, :, 1, :].reshape( - local_routes, - self.config.intermediate_size, - ) - fc1_c = torch.cat((gate, up), dim=1) - route_metadata = torch.stack( - (local_experts, src_ranks, src_tokens, src_slots), - dim=1, - ).to(torch.int32) - - wgrad_stash = None - if self.config.backward_wgrad_mode == "operands": - if inputs.col_quant_data is None or inputs.col_quant_sf is None: - raise RuntimeError( - "wgrad operand mode requires forward column-requant outputs" - ) - padded_ends, expert_offsets = cumulative_padded_offsets( - plan.expert_counts, - self.config.token_padding_size, - self.device, - ) - padded_routes = padded_ends[-1] if padded_ends else 0 - fc1_a = pool_data_as_wgrad_a( - inputs.col_quant_data, - padded_routes, - ) - fc1_sfa = assemble_discrete_col_requant_scales( - inputs.col_quant_sf, - plan.expert_counts, - padded_ends, - self.config.hidden_size, - self.config.sf_padding_size, - ) - valid_route_counts = torch.tensor( - plan.expert_counts, - dtype=torch.int32, - device=self.device, - ) - wgrad_stash = MoeEpWgradForwardStash( - fc1_a=fc1_a, - fc1_sfa=fc1_sfa, - expert_offsets=expert_offsets, - valid_route_counts=valid_route_counts, - route_metadata=route_metadata, - ) - return fc1_c, route_metadata, wgrad_stash - - def close(self) -> None: - self._buffer = None - self._closed = True - - -__all__ = [ - "Mxfp8ForwardStash", - "Mxfp8StashPlan", -] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py new file mode 100644 index 000000000..f90890b3d --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py @@ -0,0 +1,278 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Ordinary/capturable launch path over fixed MXFP8 training resources.""" + +from __future__ import annotations + +import torch + +from ..._types import MoeEpTrainingWgradOperands +from .._runtime import _runtime_debug +from .._workspace import padded_mxfp8_scale_columns +from ._adapter import ( + Mxfp8LaunchInputs, + _typed_view, +) +from ._backward_compile import ( + Mxfp8BackwardLaunchInputs, + build_backward_runtime_kwargs, + compile_backward_or_get, +) +from ._compile import compile_or_get +from ._launch import build_runtime_kwargs +from ._training_resources import ( + Mxfp8TrainingExecutionViews, + Mxfp8TrainingResourceOwner, +) + + +def _zero_pre_reduced(inputs, prepared) -> None: + capacity = prepared.config.max_tokens_per_rank + offset = prepared.pre_reduced_activation_offset + bytes_per_token = prepared.pre_reduced_activation_bytes_per_token + if offset is not None and bytes_per_token: + inputs.shared_workspace.narrow( + 0, + offset, + capacity * bytes_per_token, + ).zero_() + sf_offset = prepared.pre_reduced_activation_sf_offset + sf_bytes_per_token = prepared.pre_reduced_activation_sf_bytes_per_token + if sf_offset is not None and sf_bytes_per_token: + inputs.shared_workspace.narrow( + 0, + sf_offset, + capacity * sf_bytes_per_token, + ).zero_() + + +def _activation_views( + execution: Mxfp8TrainingExecutionViews, + *, + backward: bool, + capacity: int, + hidden: int, +) -> tuple[torch.Tensor, torch.Tensor]: + workspace = ( + execution.backward.workspace + if backward + else execution.forward.workspace + ) + return ( + _typed_view( + workspace.symmetric["activation_data"], + torch.float8_e4m3fn, + (capacity, hidden), + ), + _typed_view( + workspace.symmetric["activation_scale"], + torch.float8_e8m0fnu, + (capacity, padded_mxfp8_scale_columns(hidden)), + ), + ) + + +def _write_expert_offsets( + execution: Mxfp8TrainingExecutionViews, + padding: int, +) -> None: + snapshot = execution.forward_expert_size_snapshot + if snapshot is None: + raise RuntimeError( + "training forward requires the persistent expert-size snapshot" + ) + counts = execution.slot.valid_route_counts + offsets = execution.slot.expert_offsets + counts.copy_(snapshot) + torch.add(counts, padding - 1, out=offsets) + torch.div(offsets, padding, rounding_mode="floor", out=offsets) + offsets.mul_(padding) + torch.cumsum(offsets, dim=0, out=offsets) + + +def launch_training_forward( + owner: Mxfp8TrainingResourceOwner, + execution: Mxfp8TrainingExecutionViews, + activation: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, +) -> torch.Tensor: + """Stage and launch one fixed-slot forward without host-visible routing.""" + + prepared = owner.forward_prepared + config = prepared.config + capacity = config.max_tokens_per_rank + slot = execution.slot + _runtime_debug( + "training-forward.begin", + slot=execution.slot.index, + token_count=int(activation.shape[0]), + ) + activation_data, activation_sf = _activation_views( + execution, + backward=False, + capacity=capacity, + hidden=config.hidden, + ) + _runtime_debug("training-forward.stage.begin", slot=execution.slot.index) + owner.stager.stage( + activation, + topk_idx, + topk_weights, + activation_data, + activation_sf, + slot.routing_topk_idx, + slot.routing_topk_weights, + ) + _runtime_debug("training-forward.stage.end", slot=execution.slot.index) + slot.forward_output.zero_() + slot.forward_overflow.zero_() + if slot.col_quant_data is not None: + slot.col_quant_data.zero_() + if slot.col_quant_sf is not None: + slot.col_quant_sf.zero_() + _runtime_debug("training-forward.reset.end", slot=execution.slot.index) + + workspace = execution.forward.workspace + inputs = Mxfp8LaunchInputs( + activation=activation_data, + activation_sf=activation_sf, + topk_indices=slot.routing_topk_idx, + topk_scores=slot.routing_topk_weights, + weights=owner.weight_bindings.forward, + fc1_c=slot.fc1_preact, + output_data=slot.forward_output, + col_quant_data=slot.col_quant_data, + col_quant_sf=slot.col_quant_sf, + overflow_flag=slot.forward_overflow, + local_workspace=workspace.local["kernel_local_workspace"], + shared_workspace=workspace.symmetric["kernel_shared_workspace"], + token_count=int(activation.shape[0]), + ) + _zero_pre_reduced(inputs, prepared) + _runtime_debug("training-forward.compile.begin", slot=execution.slot.index) + compiled = compile_or_get( + prepared, + inputs, + execution.forward, + ) + _runtime_debug("training-forward.compile.end", slot=execution.slot.index) + _runtime_debug("training-forward.launch.begin", slot=execution.slot.index) + compiled.callable( + **build_runtime_kwargs(inputs, execution.forward) + ) + _runtime_debug("training-forward.launch.end", slot=execution.slot.index) + _runtime_debug("training-forward.offsets.begin", slot=execution.slot.index) + _write_expert_offsets(execution, config.token_padding_block) + _runtime_debug("training-forward.offsets.end", slot=execution.slot.index) + _runtime_debug("training-forward.end", slot=execution.slot.index) + return slot.forward_output[: inputs.token_count] + + +def launch_training_backward( + owner: Mxfp8TrainingResourceOwner, + execution: Mxfp8TrainingExecutionViews, + grad_output: torch.Tensor, +) -> tuple[ + torch.Tensor, + torch.Tensor, + MoeEpTrainingWgradOperands, +]: + """Stage and launch one fixed-slot backward using forward's raw pool.""" + + prepared = owner.backward_prepared + config = prepared.config + capacity = config.max_tokens_per_rank + slot = execution.slot + token_count = int(grad_output.shape[0]) + _runtime_debug( + "training-backward.begin", + slot=execution.slot.index, + token_count=token_count, + ) + activation_data, activation_sf = _activation_views( + execution, + backward=True, + capacity=capacity, + hidden=config.hidden, + ) + _runtime_debug("training-backward.stage.begin", slot=execution.slot.index) + owner.stager.stage( + grad_output, + slot.routing_topk_idx[:token_count], + slot.routing_topk_weights[:token_count], + activation_data, + activation_sf, + slot.routing_topk_idx, + slot.routing_topk_weights, + ) + _runtime_debug("training-backward.stage.end", slot=execution.slot.index) + + slot.backward_output.zero_() + slot.grad_activation.zero_() + slot.backward_overflow.zero_() + slot.dprob.zero_() + slot.fc1_recompute.zero_() + slot.fc1_recompute_sf.view(torch.uint8).fill_(127) + slot.fc1_col_output.zero_() + slot.fc1_col_output_sf.view(torch.uint8).fill_(127) + slot.grad_y2.zero_() + slot.grad_y2_sf.fill_(127) + _runtime_debug("training-backward.reset.end", slot=execution.slot.index) + + workspace = execution.backward.workspace + weights = owner.weight_bindings.backward + inputs = Mxfp8BackwardLaunchInputs( + grad_out=activation_data, + grad_out_sf=activation_sf, + topk_idx=slot.routing_topk_idx, + topk_weights=slot.routing_topk_weights, + fc1_weight=weights.fc1_weight, + fc1_weight_sf=weights.fc1_weight_sf, + fc2_weight=weights.fc2_weight, + fc2_weight_sf=weights.fc2_weight_sf, + beta=owner.beta, + fc1_preact=slot.fc1_preact, + output_activation=slot.backward_output, + overflow_flag=slot.backward_overflow, + dprob=slot.dprob, + fc1_recompute=slot.fc1_recompute, + fc1_recompute_sf=slot.fc1_recompute_sf, + fc1_col_output=slot.fc1_col_output, + fc1_col_output_sf=slot.fc1_col_output_sf, + grad_y2=slot.grad_y2, + grad_y2_sf=slot.grad_y2_sf, + local_workspace=workspace.local["kernel_local_workspace"], + shared_workspace=workspace.symmetric["kernel_shared_workspace"], + token_count=token_count, + ) + _zero_pre_reduced(inputs, prepared) + _runtime_debug("training-backward.compile.begin", slot=execution.slot.index) + compiled = compile_backward_or_get( + prepared, + inputs, + execution.backward, + ) + _runtime_debug("training-backward.compile.end", slot=execution.slot.index) + _runtime_debug("training-backward.launch.begin", slot=execution.slot.index) + compiled.callable( + **build_backward_runtime_kwargs(inputs, execution.backward) + ) + _runtime_debug("training-backward.launch.end", slot=execution.slot.index) + slot.grad_activation.copy_(slot.backward_output) + _runtime_debug("training-backward.wgrad-export.begin", slot=execution.slot.index) + operands = owner.wgrad_exporter.export(slot) + _runtime_debug("training-backward.wgrad-export.end", slot=execution.slot.index) + _runtime_debug("training-backward.end", slot=execution.slot.index) + return ( + slot.grad_activation[:token_count], + slot.dprob[:token_count], + operands, + ) + + +__all__ = [ + "launch_training_backward", + "launch_training_forward", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py new file mode 100644 index 000000000..873e08e68 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py @@ -0,0 +1,1336 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Fixed-capacity slot/lane resources for graph-capable MXFP8 training.""" + +from __future__ import annotations + +import hashlib +import math +import threading +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Any, Mapping, Optional + +import torch +import torch.distributed as dist + +from ..._contracts import ForwardConfig +from ..._types import MoeEpTrainingWeights +from .._comm import SymmetricMemoryProvider +from .._plan import PreparedResources +from .._runtime import ( + RuntimeHandle, + RuntimeManager, + _RuntimeWatchdog, + _runtime_debug, + get_runtime_manager, +) +from .._workspace import ( + BufferRegion, + LocalMemoryProvider, + WorkspaceOwner, + WorkspaceRequirements, + WorkspaceViews, +) +from ._adapter import _typed_k_major_view, _typed_view +from ._backward_compile import PreparedMxfp8BackwardKernel +from ._compile import PreparedMxfp8Kernel +from ._fingerprint import canonical_json_sha256, source_tree_sha256 +from ._training_stage import Mxfp8TrainingStager +from ._training_weights import Mxfp8TrainingWeightBindings +from ._training_wgrad import Mxfp8TrainingWgradExporter + + +_DATA_DTYPE = torch.float8_e4m3fn +_SCALE_DTYPE = torch.float8_e8m0fnu + +_ROUTING_SYMMETRIC = frozenset({"topk_weights"}) +_ROUTING_LOCAL = frozenset({"topk_idx"}) +_FORWARD_SLOT_SYMMETRIC = frozenset({"output_data", *_ROUTING_SYMMETRIC}) +_FORWARD_SLOT_LOCAL = frozenset( + {"overflow_flag", "col_quant_data", "col_quant_sf", *_ROUTING_LOCAL} +) +_BACKWARD_SLOT_SYMMETRIC = frozenset( + {"output_data", "backward_dprob", *_ROUTING_SYMMETRIC} +) +_BACKWARD_SLOT_LOCAL = frozenset( + {"overflow_flag", "backward_aux_data", "backward_aux_scale", *_ROUTING_LOCAL} +) + + +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +def _align_scale_columns(token_capacity: int) -> int: + return _round_up((token_capacity + 31) // 32, 4) + + +def _lane_name( + lane: int, + phase: str, + space: str, + name: str, +) -> str: + return f"lane.{lane}.{phase}.{space}.{name}" + + +def _slot_name( + slot: int, + phase: str, + space: str, + name: str, +) -> str: + return f"slot.{slot}.{phase}.{space}.{name}" + + +def _custom_slot_name(slot: int, name: str) -> str: + return f"slot.{slot}.persistent.local.{name}" + + +def _custom_slot_symmetric_name(slot: int, name: str) -> str: + return f"slot.{slot}.persistent.symmetric.{name}" + + +def _clone_region(name: str, region: BufferRegion) -> BufferRegion: + return BufferRegion( + name=name, + nbytes=region.nbytes, + alignment=region.alignment, + ) + + +def _region_map( + requirements: WorkspaceRequirements, + space: str, +) -> dict[str, BufferRegion]: + regions = ( + requirements.symmetric_regions + if space == "symmetric" + else requirements.local_regions + ) + return {region.name: region for region in regions} + + +def _required_region( + requirements: WorkspaceRequirements, + space: str, + name: str, +) -> BufferRegion: + try: + return _region_map(requirements, space)[name] + except KeyError as exc: + raise ValueError( + f"{space} workspace requirements do not contain {name!r}" + ) from exc + + +def _add_lane_regions( + output: list[BufferRegion], + requirements: WorkspaceRequirements, + *, + lane: int, + phase: str, + space: str, + slot_names: frozenset[str], +) -> None: + regions = ( + requirements.symmetric_regions + if space == "symmetric" + else requirements.local_regions + ) + for region in regions: + if region.name in slot_names: + continue + if phase == "backward" and space == "local" and region.name == ( + "backward_fc1_preact" + ): + # The graph path aliases forward's raw receiver pool directly. + continue + output.append( + _clone_region( + _lane_name(lane, phase, space, region.name), + region, + ) + ) + + +def build_training_workspace_requirements( + config: ForwardConfig, + forward: PreparedMxfp8Kernel, + backward: PreparedMxfp8BackwardKernel, + *, + slot_count: int, + lane_count: int, +) -> WorkspaceRequirements: + """Build one deterministic root layout for N slots and M lanes.""" + + for name, value in (("slot_count", slot_count), ("lane_count", lane_count)): + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + if not config.generate_c: + raise ValueError("training resources require generate_c=True") + if forward.pool_token_capacity != backward.pool_token_capacity: + raise ValueError( + "forward/backward pool capacities must match, got " + f"{forward.pool_token_capacity} and " + f"{backward.pool_token_capacity}" + ) + + forward_requirements = forward.workspace_requirements + backward_requirements = backward.workspace_requirements + symmetric_regions: list[BufferRegion] = [] + local_regions: list[BufferRegion] = [] + + for lane in range(lane_count): + local_regions.extend( + ( + BufferRegion( + _lane_name( + lane, + "finalizer", + "local", + "global_overflow", + ), + torch.int32.itemsize, + 16, + ), + BufferRegion( + _lane_name( + lane, + "finalizer", + "local", + "overflow_ok", + ), + torch.bool.itemsize, + 16, + ), + ) + ) + _add_lane_regions( + symmetric_regions, + forward_requirements, + lane=lane, + phase="forward", + space="symmetric", + slot_names=_FORWARD_SLOT_SYMMETRIC, + ) + _add_lane_regions( + local_regions, + forward_requirements, + lane=lane, + phase="forward", + space="local", + slot_names=_FORWARD_SLOT_LOCAL, + ) + _add_lane_regions( + symmetric_regions, + backward_requirements, + lane=lane, + phase="backward", + space="symmetric", + slot_names=_BACKWARD_SLOT_SYMMETRIC, + ) + _add_lane_regions( + local_regions, + backward_requirements, + lane=lane, + phase="backward", + space="local", + slot_names=_BACKWARD_SLOT_LOCAL, + ) + + forward_symmetric = _region_map(forward_requirements, "symmetric") + forward_local = _region_map(forward_requirements, "local") + backward_symmetric = _region_map(backward_requirements, "symmetric") + backward_local = _region_map(backward_requirements, "local") + fc1_c_shape = tuple( + int(extent) + for extent in forward.kernel.get_aux_output_shapes()["fc1_c"] + ) + fc1_c_bytes = math.prod(fc1_c_shape) * torch.bfloat16.itemsize + backward_preact = _required_region( + backward_requirements, + "local", + "backward_fc1_preact", + ) + if fc1_c_bytes != backward_preact.nbytes: + raise ValueError( + "forward fc1_c and backward preactivation byte sizes differ: " + f"{fc1_c_bytes} != {backward_preact.nbytes}" + ) + aux_shapes = { + name: tuple(int(extent) for extent in shape) + for name, shape in backward.kernel.get_aux_output_shapes().items() + } + aux_dtypes = { + "fc1_recompute": _DATA_DTYPE, + "fc1_recompute_sf": _SCALE_DTYPE, + "fc1_col_output": _DATA_DTYPE, + "fc1_col_output_sf": _SCALE_DTYPE, + "grad_y2": _DATA_DTYPE, + "grad_y2_sf": torch.uint8, + } + scale_columns = _align_scale_columns(forward.pool_token_capacity) + wgrad_shapes = { + "wgrad_fc1_b": ( + forward.pool_token_capacity, + 2 * config.intermediate_size, + ), + "wgrad_fc1_sfa": ( + _round_up(config.hidden_size, 128), + scale_columns, + ), + "wgrad_fc1_sfb": ( + _round_up(2 * config.intermediate_size, 128), + scale_columns, + ), + "wgrad_fc2_a": ( + config.intermediate_size, + forward.pool_token_capacity, + ), + "wgrad_fc2_sfa": ( + _round_up(config.intermediate_size, 128), + scale_columns, + ), + "wgrad_fc2_sfb": ( + _round_up(config.hidden_size, 128), + scale_columns, + ), + } + + for slot in range(slot_count): + for name in sorted(_FORWARD_SLOT_SYMMETRIC): + if name in _ROUTING_SYMMETRIC: + continue + symmetric_regions.append( + _clone_region( + _slot_name(slot, "forward", "symmetric", name), + forward_symmetric[name], + ) + ) + for name in sorted(_BACKWARD_SLOT_SYMMETRIC): + if name in _ROUTING_SYMMETRIC: + continue + symmetric_regions.append( + _clone_region( + _slot_name(slot, "backward", "symmetric", name), + backward_symmetric[name], + ) + ) + for name in sorted(_FORWARD_SLOT_LOCAL): + if name in _ROUTING_LOCAL: + continue + region = forward_local.get(name) + if region is not None: + local_regions.append( + _clone_region( + _slot_name(slot, "forward", "local", name), + region, + ) + ) + for name in sorted(_BACKWARD_SLOT_LOCAL): + if name in _ROUTING_LOCAL: + continue + local_regions.append( + _clone_region( + _slot_name(slot, "backward", "local", name), + backward_local[name], + ) + ) + local_regions.extend( + ( + BufferRegion( + _custom_slot_name(slot, "fc1_preact"), + fc1_c_bytes, + alignment=128, + ), + BufferRegion( + _custom_slot_name(slot, "routing_topk_idx"), + int(config.max_tokens_per_rank) + * config.top_k + * torch.int32.itemsize, + alignment=16, + ), + BufferRegion( + _custom_slot_name(slot, "valid_route_counts"), + config.experts_per_rank * torch.int32.itemsize, + alignment=16, + ), + BufferRegion( + _custom_slot_name(slot, "expert_offsets"), + config.experts_per_rank * torch.int32.itemsize, + alignment=16, + ), + BufferRegion( + _custom_slot_name(slot, "grad_activation"), + int(config.max_tokens_per_rank) + * config.hidden_size + * torch.float32.itemsize, + alignment=16, + ), + ) + ) + symmetric_regions.append( + BufferRegion( + _custom_slot_symmetric_name(slot, "routing_topk_weights"), + int(config.max_tokens_per_rank) + * config.top_k + * torch.float32.itemsize, + alignment=16, + ) + ) + for name, dtype in aux_dtypes.items(): + local_regions.append( + BufferRegion( + _custom_slot_name(slot, name), + math.prod(aux_shapes[name]) * dtype.itemsize, + alignment=128 if name != "grad_y2_sf" else 16, + ) + ) + for name, shape in wgrad_shapes.items(): + local_regions.append( + BufferRegion( + _custom_slot_name(slot, name), + math.prod(shape), + alignment=128, + ) + ) + + return WorkspaceRequirements( + max_tokens_per_rank=int(config.max_tokens_per_rank), + symmetric_regions=tuple(symmetric_regions), + local_regions=tuple(local_regions), + ) + + +def _harmonize_symmetric_regions( + requirements: WorkspaceRequirements, + runtime: RuntimeHandle, + device: torch.device, +) -> WorkspaceRequirements: + """Make every peer-visible region size and offset identical on all ranks.""" + + if runtime.world_size <= 1: + return requirements + + regions = requirements.symmetric_regions + count = torch.tensor([len(regions)], dtype=torch.int64, device=device) + minimum_count = count.clone() + maximum_count = count.clone() + dist.all_reduce(minimum_count, op=dist.ReduceOp.MIN, group=runtime.group) + dist.all_reduce(maximum_count, op=dist.ReduceOp.MAX, group=runtime.group) + if int(minimum_count.item()) != int(maximum_count.item()): + raise RuntimeError( + "symmetric workspace region counts differ across EP ranks: " + f"min={int(minimum_count.item())}, max={int(maximum_count.item())}" + ) + + metadata = "\0".join( + f"{region.name}:{region.alignment}" for region in regions + ).encode() + signature_value = int.from_bytes( + hashlib.blake2b(metadata, digest_size=8).digest(), + "little", + ) & ((1 << 63) - 1) + signature = torch.tensor( + [signature_value], + dtype=torch.int64, + device=device, + ) + minimum_signature = signature.clone() + maximum_signature = signature.clone() + dist.all_reduce( + minimum_signature, + op=dist.ReduceOp.MIN, + group=runtime.group, + ) + dist.all_reduce( + maximum_signature, + op=dist.ReduceOp.MAX, + group=runtime.group, + ) + if int(minimum_signature.item()) != int(maximum_signature.item()): + raise RuntimeError( + "symmetric workspace region names, order, or alignments differ " + "across EP ranks: " + f"local_signature={signature_value}, " + "local_regions=" + f"{tuple((region.name, region.alignment) for region in regions)}" + ) + + local_sizes = torch.tensor( + [region.nbytes for region in regions], + dtype=torch.int64, + device=device, + ) + maximum_sizes = local_sizes.clone() + dist.all_reduce(maximum_sizes, op=dist.ReduceOp.MAX, group=runtime.group) + harmonized_sizes = tuple(int(value) for value in maximum_sizes.cpu().tolist()) + changes = tuple( + f"{region.name}:{region.nbytes}->{harmonized_size}" + for region, harmonized_size in zip(regions, harmonized_sizes) + if region.nbytes != harmonized_size + ) + _runtime_debug( + "training-resources.symmetric-layout-harmonized", + region_count=len(regions), + changed_regions=changes, + ) + if not changes: + return requirements + + return WorkspaceRequirements( + max_tokens_per_rank=requirements.max_tokens_per_rank, + symmetric_regions=tuple( + BufferRegion( + region.name, + harmonized_size, + alignment=region.alignment, + ) + for region, harmonized_size in zip(regions, harmonized_sizes) + ), + local_regions=requirements.local_regions, + ) + + +def _block_scaled_tensor_abi(tensor) -> dict[str, object]: + return { + "format": tensor.format.value, + "axis": int(tensor.axis), + "logical_shape": list(tensor.logical_shape), + "data": { + "shape": list(tensor.data.shape), + "stride": list(tensor.data.stride()), + "dtype": str(tensor.data.dtype), + }, + "scale": { + "shape": list(tensor.scale.shape), + "stride": list(tensor.scale.stride()), + "dtype": str(tensor.scale.dtype), + }, + } + + +def _workspace_abi(requirements: WorkspaceRequirements) -> dict[str, object]: + def regions(values) -> list[dict[str, object]]: + return [ + { + "name": region.name, + "nbytes": int(region.nbytes), + "alignment": int(region.alignment), + } + for region in values + ] + + return { + "max_tokens_per_rank": requirements.max_tokens_per_rank, + "symmetric_regions": regions(requirements.symmetric_regions), + "local_regions": regions(requirements.local_regions), + } + + +def _prepared_kernel_abi(prepared) -> dict[str, object]: + kernel = prepared.kernel + return { + "name": str(kernel.name()), + "architecture": list(prepared.architecture), + "effective_config": prepared.config.effective_config( + prepared.launch_cluster_count + ), + "launch": { + "cluster_count": int(prepared.launch_cluster_count), + "threads_per_cta": int(kernel.threads_per_cta), + "occupancy": int(getattr(kernel, "occupancy", 1)), + "smem_capacity": int(getattr(kernel, "smem_capacity", 0)), + }, + "workspace": _workspace_abi(prepared.workspace_requirements), + "pool_token_capacity": int(prepared.pool_token_capacity), + } + + +def _build_training_abi_facts( + config: ForwardConfig, + forward: PreparedMxfp8Kernel, + backward: PreparedMxfp8BackwardKernel, + weights: MoeEpTrainingWeights, + requirements: WorkspaceRequirements, + *, + slot_count: int, + lane_count: int, + source_tree_digest: str | None = None, +) -> dict[str, object]: + """Return rank-independent JSON-safe facts for one training resource ABI.""" + + if source_tree_digest is None: + source_root = Path(__file__).resolve().parents[1] / "cutedsl_src" + source_tree_digest = source_tree_sha256(source_root) + weight_facts = { + name: _block_scaled_tensor_abi(getattr(weights, name)) + for name in ( + "forward_fc1", + "forward_fc2", + "backward_w2_transpose", + "backward_w1_transpose", + ) + } + return { + "schema_version": 1, + "source_tree_sha256": source_tree_digest, + "ep": { + "size": int(config.ep_size), + "global_ranks": list(config.ep_global_ranks), + }, + "geometry": { + "num_experts": int(config.num_experts), + "experts_per_rank": int(config.experts_per_rank), + "hidden": int(config.hidden_size), + "intermediate": int(config.intermediate_size), + "top_k": int(config.top_k), + "max_tokens_per_rank": int(config.max_tokens_per_rank), + "max_recv_size_per_rank": int( + forward.config.max_recv_size_per_rank + ), + }, + "policy": { + "drop_on_overflow": bool(config.drop_on_overflow), + "combine_format": config.combine_format, + "output_format": config.output_format, + "apply_topk_in_fc1": bool(config.apply_topk_in_fc1), + "gate_up_clamp": config.gate_up_clamp, + }, + "resources": { + "slot_count": int(slot_count), + "lane_count": int(lane_count), + "workspace": _workspace_abi(requirements), + }, + "weights": weight_facts, + "forward_kernel": _prepared_kernel_abi(forward), + "backward_kernel": _prepared_kernel_abi(backward), + } + + +def _verify_training_abi_across_ranks( + facts: dict[str, object], + runtime: RuntimeHandle, + device: torch.device, +) -> str: + """Collectively reject rank-divergent training ABI before allocation.""" + + digest = canonical_json_sha256(facts) + if runtime.world_size <= 1: + return digest + digest_value = int(digest[:16], 16) & ((1 << 63) - 1) + minimum = torch.tensor([digest_value], dtype=torch.int64, device=device) + maximum = minimum.clone() + dist.all_reduce(minimum, op=dist.ReduceOp.MIN, group=runtime.group) + dist.all_reduce(maximum, op=dist.ReduceOp.MAX, group=runtime.group) + if int(minimum.item()) == int(maximum.item()): + return digest + + rank_digests: list[Any] = [None] * runtime.world_size + dist.all_gather_object(rank_digests, digest, group=runtime.group) + raise RuntimeError( + "MoeEp training ABI differs across expert-parallel ranks before " + "workspace allocation: " + f"digests={rank_digests}, local_facts={facts}" + ) + + +@dataclass(frozen=True) +class Mxfp8TrainingSlotViews: + """Persistent tensors that survive from forward through wgrad consumption.""" + + index: int + routing_topk_idx: torch.Tensor + routing_topk_weights: torch.Tensor + fc1_preact: torch.Tensor + col_quant_data: torch.Tensor | None + col_quant_sf: torch.Tensor | None + valid_route_counts: torch.Tensor + expert_offsets: torch.Tensor + forward_output: torch.Tensor + backward_output: torch.Tensor + grad_activation: torch.Tensor + dprob: torch.Tensor + forward_overflow: torch.Tensor + backward_overflow: torch.Tensor + fc1_recompute: torch.Tensor + fc1_recompute_sf: torch.Tensor + fc1_col_output: torch.Tensor + fc1_col_output_sf: torch.Tensor + grad_y2: torch.Tensor + grad_y2_sf: torch.Tensor + wgrad_fc1_b: torch.Tensor + wgrad_fc1_sfa: torch.Tensor + wgrad_fc1_sfb: torch.Tensor + wgrad_fc2_a: torch.Tensor + wgrad_fc2_sfa: torch.Tensor + wgrad_fc2_sfb: torch.Tensor + + +@dataclass(frozen=True) +class Mxfp8TrainingExecutionViews: + """One slot bound to one mutable execution lane.""" + + slot: Mxfp8TrainingSlotViews + forward: PreparedResources + backward: PreparedResources + forward_expert_size_snapshot: torch.Tensor | None + + +class Mxfp8TrainingResourceOwner: + """Own one combined symmetric/local root for N slots and M lanes.""" + + def __init__( + self, + config: ForwardConfig, + device: torch.device, + forward: PreparedMxfp8Kernel, + backward: PreparedMxfp8BackwardKernel, + weights: MoeEpTrainingWeights, + *, + slot_count: int, + lane_count: int, + runtime_manager: Optional[RuntimeManager] = None, + symmetric_provider: Optional[SymmetricMemoryProvider] = None, + local_provider: Optional[LocalMemoryProvider] = None, + ) -> None: + self.config = config + self.device = torch.device(device) + self.forward_prepared = forward + self.backward_prepared = backward + self.weight_bindings = Mxfp8TrainingWeightBindings(weights) + self.stager = Mxfp8TrainingStager(config.hidden_size, config.top_k) + self.wgrad_exporter = Mxfp8TrainingWgradExporter( + experts=config.experts_per_rank, + hidden=config.hidden_size, + intermediate=config.intermediate_size, + sf_padding=backward.config.sf_padding_block, + ) + self.beta = torch.ones( + (config.experts_per_rank,), + dtype=torch.float32, + device=self.device, + ) + self.slot_count = slot_count + self.lane_count = lane_count + self.requirements = build_training_workspace_requirements( + config, + forward, + backward, + slot_count=slot_count, + lane_count=lane_count, + ) + self._runtime_manager = runtime_manager or get_runtime_manager() + self._symmetric_provider = symmetric_provider + self._local_provider = local_provider + self._runtime: RuntimeHandle | None = None + self._workspace: WorkspaceOwner | None = None + self._abi_fingerprint: str | None = None + self._closed = False + self._lock = threading.RLock() + + @property + def prepared(self) -> bool: + return ( + not self._closed + and self._runtime is not None + and self._workspace is not None + and self._workspace.allocated + ) + + def prepare(self) -> None: + with self._lock: + if self._closed: + raise RuntimeError("training resources are closed") + if self.prepared: + return + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "training resources must be prepared before CUDA graph capture" + ) + _runtime_debug( + "training-resources.prepare.begin", + slot_count=self.slot_count, + lane_count=self.lane_count, + local_bytes=self.requirements.local_layout.total_bytes + if hasattr(self.requirements, "local_layout") + else sum( + region.nbytes for region in self.requirements.local_regions + ), + symmetric_bytes=sum( + region.nbytes for region in self.requirements.symmetric_regions + ), + ) + _runtime_debug("training-resources.runtime-acquire.begin") + runtime = self._runtime_manager.acquire(self.config, self.device) + _runtime_debug( + "training-resources.runtime-acquire.end", + runtime_ref_count=getattr( + self._runtime_manager, + "ref_count", + "?", + ), + ) + self._runtime = runtime + try: + layout_watchdog = _RuntimeWatchdog( + "training-resources.symmetric-layout-harmonize" + ) + layout_watchdog.start() + _runtime_debug( + "training-resources.symmetric-layout-harmonize.begin" + ) + try: + self.requirements = _harmonize_symmetric_regions( + self.requirements, + runtime, + self.device, + ) + finally: + layout_watchdog.close() + _runtime_debug( + "training-resources.symmetric-layout-harmonize.end" + ) + if runtime.world_size > 1: + abi_watchdog = _RuntimeWatchdog( + "training-resources.abi-handshake" + ) + abi_watchdog.start() + _runtime_debug("training-resources.abi-handshake.begin") + try: + abi_facts = _build_training_abi_facts( + self.config, + self.forward_prepared, + self.backward_prepared, + self.weight_bindings.weights, + self.requirements, + slot_count=self.slot_count, + lane_count=self.lane_count, + ) + self._abi_fingerprint = ( + _verify_training_abi_across_ranks( + abi_facts, + runtime, + self.device, + ) + ) + finally: + abi_watchdog.close() + _runtime_debug( + "training-resources.abi-handshake.end", + fingerprint=self._abi_fingerprint, + ) + _runtime_debug("training-resources.workspace-create.begin") + workspace = WorkspaceOwner( + self.requirements, + runtime, + symmetric_provider=self._symmetric_provider, + local_provider=self._local_provider, + ) + _runtime_debug( + "training-resources.workspace-create.end", + local_bytes=workspace.local_layout.total_bytes, + symmetric_bytes=workspace.symmetric_layout.total_bytes, + ) + self._workspace = workspace + allocation_watchdog = _RuntimeWatchdog( + "training-resources.workspace-allocate" + ) + allocation_watchdog.start() + try: + workspace.ensure_allocated() + finally: + allocation_watchdog.close() + _runtime_debug("training-resources.workspace-allocate.end") + if runtime.world_size > 1: + # Symmetric-root zeroing is asynchronous. No rank may + # enter the first device barrier until every peer has + # completed allocation and root initialization. + stream_watchdog = _RuntimeWatchdog( + "training-resources.stream-synchronize" + ) + stream_watchdog.start() + _runtime_debug( + "training-resources.stream-synchronize.begin" + ) + try: + torch.cuda.current_stream(self.device).synchronize() + finally: + stream_watchdog.close() + _runtime_debug("training-resources.stream-synchronize.end") + + barrier_watchdog = _RuntimeWatchdog( + "training-resources.rank-barrier" + ) + barrier_watchdog.start() + _runtime_debug("training-resources.rank-barrier.begin") + try: + dist.barrier(group=runtime.group) + finally: + barrier_watchdog.close() + _runtime_debug("training-resources.rank-barrier.end") + except Exception: + if self._workspace is not None: + self._workspace.close() + self._workspace = None + runtime.close() + self._runtime = None + raise + _runtime_debug("training-resources.prepare.end") + + def _flat_views(self, token_count: int) -> WorkspaceViews: + self.prepare() + assert self._workspace is not None + return self._workspace.views(token_count) + + @staticmethod + def _phase_workspace( + flat: WorkspaceViews, + requirements: WorkspaceRequirements, + *, + slot: int, + lane: int, + phase: str, + ) -> WorkspaceViews: + symmetric = {} + local = {} + slot_symmetric = ( + _FORWARD_SLOT_SYMMETRIC + if phase == "forward" + else _BACKWARD_SLOT_SYMMETRIC + ) + slot_local = ( + _FORWARD_SLOT_LOCAL + if phase == "forward" + else _BACKWARD_SLOT_LOCAL + ) + for region in requirements.symmetric_regions: + if region.name in _ROUTING_SYMMETRIC: + symmetric[region.name] = flat.symmetric[ + _custom_slot_symmetric_name(slot, "routing_topk_weights") + ] + continue + scope_name = ( + _slot_name(slot, phase, "symmetric", region.name) + if region.name in slot_symmetric + else _lane_name(lane, phase, "symmetric", region.name) + ) + symmetric[region.name] = flat.symmetric[scope_name] + for region in requirements.local_regions: + if region.name in _ROUTING_LOCAL: + local[region.name] = flat.local[ + _custom_slot_name(slot, "routing_topk_idx") + ] + continue + if phase == "backward" and region.name == "backward_fc1_preact": + local[region.name] = flat.local[ + _custom_slot_name(slot, "fc1_preact") + ] + continue + scope_name = ( + _slot_name(slot, phase, "local", region.name) + if region.name in slot_local + else _lane_name(lane, phase, "local", region.name) + ) + local[region.name] = flat.local[scope_name] + return WorkspaceViews( + token_count=flat.token_count, + symmetric=MappingProxyType(symmetric), + local=MappingProxyType(local), + peer_mapping=flat.peer_mapping, + ) + + def _slot_views( + self, + flat: WorkspaceViews, + slot: int, + ) -> Mxfp8TrainingSlotViews: + config = self.config + capacity = int(config.max_tokens_per_rank) + fwd_shapes = { + name: tuple(int(extent) for extent in shape) + for name, shape in self.forward_prepared.kernel.get_aux_output_shapes().items() + } + bwd_shapes = { + name: tuple(int(extent) for extent in shape) + for name, shape in self.backward_prepared.kernel.get_aux_output_shapes().items() + } + scale_columns = _align_scale_columns( + self.forward_prepared.pool_token_capacity + ) + + def local_bytes(name: str) -> torch.Tensor: + return flat.local[_custom_slot_name(slot, name)] + + col_quant_data = None + col_quant_sf = None + col_data_name = _slot_name( + slot, + "forward", + "local", + "col_quant_data", + ) + if col_data_name in flat.local: + col_quant_data = _typed_k_major_view( + flat.local[col_data_name], + _DATA_DTYPE, + fwd_shapes["col_quant_data"], + ) + col_quant_sf = _typed_view( + flat.local[ + _slot_name( + slot, + "forward", + "local", + "col_quant_sf", + ) + ], + torch.uint8, + fwd_shapes["col_quant_sf"], + ) + + return Mxfp8TrainingSlotViews( + index=slot, + routing_topk_idx=_typed_view( + local_bytes("routing_topk_idx"), + torch.int32, + (capacity, config.top_k), + ), + routing_topk_weights=_typed_view( + flat.symmetric[ + _custom_slot_symmetric_name( + slot, + "routing_topk_weights", + ) + ], + torch.float32, + (capacity, config.top_k), + ), + fc1_preact=_typed_view( + local_bytes("fc1_preact"), + torch.bfloat16, + fwd_shapes["fc1_c"], + ), + col_quant_data=col_quant_data, + col_quant_sf=col_quant_sf, + valid_route_counts=_typed_view( + local_bytes("valid_route_counts"), + torch.int32, + (config.experts_per_rank,), + ), + expert_offsets=_typed_view( + local_bytes("expert_offsets"), + torch.int32, + (config.experts_per_rank,), + ), + forward_output=_typed_view( + flat.symmetric[ + _slot_name( + slot, + "forward", + "symmetric", + "output_data", + ) + ], + torch.bfloat16, + (capacity, config.hidden_size), + ), + backward_output=_typed_view( + flat.symmetric[ + _slot_name( + slot, + "backward", + "symmetric", + "output_data", + ) + ], + torch.bfloat16, + (capacity, config.hidden_size), + ), + grad_activation=_typed_view( + local_bytes("grad_activation"), + torch.float32, + (capacity, config.hidden_size), + ), + dprob=_typed_view( + flat.symmetric[ + _slot_name( + slot, + "backward", + "symmetric", + "backward_dprob", + ) + ], + torch.float32, + bwd_shapes["dprob"], + ), + forward_overflow=_typed_view( + flat.local[ + _slot_name( + slot, + "forward", + "local", + "overflow_flag", + ) + ], + torch.int32, + (1,), + ), + backward_overflow=_typed_view( + flat.local[ + _slot_name( + slot, + "backward", + "local", + "overflow_flag", + ) + ], + torch.int32, + (1,), + ), + fc1_recompute=_typed_view( + local_bytes("fc1_recompute"), + _DATA_DTYPE, + bwd_shapes["fc1_recompute"], + ), + fc1_recompute_sf=_typed_view( + local_bytes("fc1_recompute_sf"), + _SCALE_DTYPE, + bwd_shapes["fc1_recompute_sf"], + ), + fc1_col_output=_typed_view( + local_bytes("fc1_col_output"), + _DATA_DTYPE, + bwd_shapes["fc1_col_output"], + ), + fc1_col_output_sf=_typed_view( + local_bytes("fc1_col_output_sf"), + _SCALE_DTYPE, + bwd_shapes["fc1_col_output_sf"], + ), + grad_y2=_typed_k_major_view( + local_bytes("grad_y2"), + _DATA_DTYPE, + bwd_shapes["grad_y2"], + ), + grad_y2_sf=_typed_view( + local_bytes("grad_y2_sf"), + torch.uint8, + bwd_shapes["grad_y2_sf"], + ), + wgrad_fc1_b=_typed_k_major_view( + local_bytes("wgrad_fc1_b"), + _DATA_DTYPE, + ( + self.forward_prepared.pool_token_capacity, + 2 * config.intermediate_size, + ), + ), + wgrad_fc1_sfa=_typed_view( + local_bytes("wgrad_fc1_sfa"), + _SCALE_DTYPE, + (_round_up(config.hidden_size, 128), scale_columns), + ), + wgrad_fc1_sfb=_typed_view( + local_bytes("wgrad_fc1_sfb"), + _SCALE_DTYPE, + ( + _round_up(2 * config.intermediate_size, 128), + scale_columns, + ), + ), + wgrad_fc2_a=_typed_view( + local_bytes("wgrad_fc2_a"), + _DATA_DTYPE, + ( + config.intermediate_size, + self.forward_prepared.pool_token_capacity, + ), + ), + wgrad_fc2_sfa=_typed_view( + local_bytes("wgrad_fc2_sfa"), + _SCALE_DTYPE, + ( + _round_up(config.intermediate_size, 128), + scale_columns, + ), + ), + wgrad_fc2_sfb=_typed_view( + local_bytes("wgrad_fc2_sfb"), + _SCALE_DTYPE, + (_round_up(config.hidden_size, 128), scale_columns), + ), + ) + + def views( + self, + *, + slot: int, + lane: int, + token_count: int, + ) -> Mxfp8TrainingExecutionViews: + with self._lock: + if slot < 0 or slot >= self.slot_count: + raise ValueError( + f"slot {slot} is outside [0, {self.slot_count})" + ) + if lane < 0 or lane >= self.lane_count: + raise ValueError( + f"lane {lane} is outside [0, {self.lane_count})" + ) + flat = self._flat_views(token_count) + forward_workspace = self._phase_workspace( + flat, + self.forward_prepared.workspace_requirements, + slot=slot, + lane=lane, + phase="forward", + ) + backward_workspace = self._phase_workspace( + flat, + self.backward_prepared.workspace_requirements, + slot=slot, + lane=lane, + phase="backward", + ) + snapshot = None + if self.forward_prepared.col_quant_sizes_offset is not None: + snapshot_bytes = forward_workspace.local[ + "kernel_local_workspace" + ].narrow( + 0, + self.forward_prepared.col_quant_sizes_offset, + self.forward_prepared.col_quant_sizes_bytes, + ) + snapshot = _typed_view( + snapshot_bytes, + torch.int32, + (self.config.experts_per_rank,), + ) + assert self._runtime is not None + return Mxfp8TrainingExecutionViews( + slot=self._slot_views(flat, slot), + forward=PreparedResources( + runtime=self._runtime, + workspace=forward_workspace, + ), + backward=PreparedResources( + runtime=self._runtime, + workspace=backward_workspace, + ), + forward_expert_size_snapshot=snapshot, + ) + + def refresh_weights(self) -> None: + """Enqueue fixed-address layout refreshes from the bound MXFP8 pack.""" + + with self._lock: + if self._closed: + raise RuntimeError("training resources are closed") + self.weight_bindings.refresh() + + def close(self) -> None: + with self._lock: + if self._closed: + return + if self._workspace is not None: + self._workspace.close() + self._workspace = None + if self._runtime is not None: + self._runtime.close() + self._runtime = None + self._closed = True + + def finalize_overflow( + self, + slots: tuple[int, ...], + *, + lane: int, + ) -> torch.Tensor: + """Aggregate slot flags and apply the public error/drop policy.""" + + if not slots: + raise ValueError("finalize_overflow requires at least one slot") + if len(set(slots)) != len(slots): + raise ValueError("finalize_overflow slots must be unique") + for slot in slots: + if slot < 0 or slot >= self.slot_count: + raise ValueError( + f"slot {slot} is outside [0, {self.slot_count})" + ) + if lane < 0 or lane >= self.lane_count: + raise ValueError(f"lane {lane} is outside [0, {self.lane_count})") + flat = self._flat_views(0) + global_overflow = _typed_view( + flat.local[ + _lane_name( + lane, + "finalizer", + "local", + "global_overflow", + ) + ], + torch.int32, + (1,), + ) + global_overflow.zero_() + for slot in slots: + for phase in ("forward", "backward"): + flag = _typed_view( + flat.local[ + _slot_name( + slot, + phase, + "local", + "overflow_flag", + ) + ], + torch.int32, + (1,), + ) + torch.maximum(global_overflow, flag, out=global_overflow) + assert self._runtime is not None + if self._runtime.world_size > 1: + dist.all_reduce( + global_overflow, + op=dist.ReduceOp.MAX, + group=self._runtime.group, + ) + if not self.config.drop_on_overflow: + assert_async = getattr(torch, "_assert_async", None) + if assert_async is None: + raise RuntimeError( + "drop_on_overflow=False training resources require " + "torch._assert_async" + ) + overflow_ok = _typed_view( + flat.local[ + _lane_name( + lane, + "finalizer", + "local", + "overflow_ok", + ) + ], + torch.bool, + (1,), + ) + torch.eq(global_overflow, 0, out=overflow_ok) + assert_async( + overflow_ok, + "Rubin MegaMoE receive route-pool overflow; " + "the fixed-slot outputs are invalid", + ) + return global_overflow + + +__all__ = [ + "Mxfp8TrainingExecutionViews", + "Mxfp8TrainingResourceOwner", + "Mxfp8TrainingSlotViews", + "build_training_workspace_requirements", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py new file mode 100644 index 000000000..16a05501d --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py @@ -0,0 +1,199 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Cached fused MXFP8 staging for fixed-address training resources.""" + +from __future__ import annotations + +import threading + +import torch + +from ._launch import _to_cute + + +class Mxfp8TrainingStager: + """Own one compile cache; steady-state staging is allocation-free.""" + + def __init__(self, hidden: int, top_k: int) -> None: + self.hidden = int(hidden) + self.top_k = int(top_k) + self._compiled: dict[tuple, object] = {} + self._lock = threading.RLock() + + def _validate( + self, + source: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + output: torch.Tensor, + output_sf: torch.Tensor, + output_topk_idx: torch.Tensor, + output_topk_weights: torch.Tensor, + ) -> int: + if source.dtype not in (torch.bfloat16, torch.float32): + raise TypeError( + "training staging source must be BF16 or FP32, " + f"got {source.dtype}" + ) + if source.ndim != 2 or source.shape[1] != self.hidden: + raise ValueError( + f"training staging source must have shape (T, {self.hidden})" + ) + if not source.is_contiguous(): + raise ValueError("training staging source must be contiguous") + token_count = int(source.shape[0]) + if topk_idx.shape != (token_count, self.top_k): + raise ValueError("training staging topk_idx shape mismatch") + if topk_idx.dtype is not torch.int32 or not topk_idx.is_contiguous(): + raise TypeError("training staging topk_idx must be contiguous Int32") + if topk_weights.shape != topk_idx.shape: + raise ValueError("training staging topk_weights shape mismatch") + if ( + topk_weights.dtype is not torch.float32 + or not topk_weights.is_contiguous() + ): + raise TypeError( + "training staging topk_weights must be contiguous FP32" + ) + if ( + output.dtype is not torch.float8_e4m3fn + or output.ndim != 2 + or output.shape[1] != self.hidden + or not output.is_contiguous() + ): + raise ValueError( + "training staging output must be contiguous E4M3 " + f"(capacity, {self.hidden})" + ) + if token_count > output.shape[0]: + raise ValueError( + f"token count {token_count} exceeds capacity {output.shape[0]}" + ) + logical_sf_columns = self.hidden // 32 + if ( + output_sf.dtype is not torch.float8_e8m0fnu + or output_sf.ndim != 2 + or output_sf.shape[0] != output.shape[0] + or output_sf.shape[1] < logical_sf_columns + or not output_sf.is_contiguous() + ): + raise ValueError("training staging output_sf has an invalid ABI") + for name, tensor, dtype in ( + ("output_topk_idx", output_topk_idx, torch.int32), + ("output_topk_weights", output_topk_weights, torch.float32), + ): + if ( + tensor.shape != (output.shape[0], self.top_k) + or tensor.dtype is not dtype + or not tensor.is_contiguous() + ): + raise ValueError(f"training staging {name} has an invalid ABI") + devices = { + source.device, + topk_idx.device, + topk_weights.device, + output.device, + output_sf.device, + output_topk_idx.device, + output_topk_weights.device, + } + if len(devices) != 1: + raise ValueError("all training staging tensors must share one device") + return token_count + + def stage( + self, + source: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + output: torch.Tensor, + output_sf: torch.Tensor, + output_topk_idx: torch.Tensor, + output_topk_weights: torch.Tensor, + ) -> None: + """Enqueue tail reset plus one fused quant-and-routing launch.""" + + token_count = self._validate( + source, + topk_idx, + topk_weights, + output, + output_sf, + output_topk_idx, + output_topk_weights, + ) + output_sf.zero_() + routing_in_place = ( + topk_idx.data_ptr() == output_topk_idx.data_ptr() + and topk_weights.data_ptr() == output_topk_weights.data_ptr() + ) + routing_partially_aliased = ( + topk_idx.data_ptr() == output_topk_idx.data_ptr() + ) != ( + topk_weights.data_ptr() == output_topk_weights.data_ptr() + ) + if routing_partially_aliased: + raise ValueError( + "training staging routing inputs must either both alias " + "their outputs or neither alias" + ) + if not routing_in_place: + output_topk_idx.fill_(-1) + output_topk_weights.zero_() + if token_count == 0: + return + + logical_sf_columns = self.hidden // 32 + import cuda.bindings.driver as cuda + + stream = torch.cuda.current_stream(source.device) + args = ( + _to_cute(source, dynamic_layout=False), + _to_cute(topk_idx, assumed_align=4, dynamic_layout=False), + _to_cute(topk_weights, assumed_align=4, dynamic_layout=False), + _to_cute(output[:token_count], dynamic_layout=False), + _to_cute( + output_sf[:token_count, :logical_sf_columns], + assumed_align=4, + dynamic_layout=False, + ), + _to_cute( + output_topk_idx[:token_count], + assumed_align=4, + dynamic_layout=False, + ), + _to_cute( + output_topk_weights[:token_count], + assumed_align=4, + dynamic_layout=False, + ), + cuda.CUstream(stream.cuda_stream), + ) + key = ( + source.device.index, + source.dtype, + token_count, + self.hidden, + self.top_k, + tuple(output_sf.stride()), + ) + with self._lock: + compiled = self._compiled.get(key) + if compiled is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "MXFP8 training stager must be compiled before " + "CUDA graph capture" + ) + import cutlass.cute as cute + + from ._training_stage_kernel import Mxfp8TrainingStageKernel + + kernel = Mxfp8TrainingStageKernel(self.hidden, self.top_k) + compiled = cute.compile(kernel, *args) + self._compiled[key] = compiled + compiled(*args) + + +__all__ = ["Mxfp8TrainingStager"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py new file mode 100644 index 000000000..c92cf6035 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py @@ -0,0 +1,134 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""One-launch BF16/FP32 to MXFP8 staging for fixed training resources.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Float32, Int32 + +from ..cutedsl_src.helpers.constants import Fp32Max, Fp8E4M3RcpLimit +from ..cutedsl_src.helpers.ptx_helpers import cvt_f32_to_fp8_to_f32 + + +class Mxfp8TrainingStageKernel: + """Quantize one token row per CTA and repack its routing metadata.""" + + _threads_per_cta = 128 + _sf_vec = 32 + + def __init__(self, hidden: int, top_k: int) -> None: + self.hidden = int(hidden) + self.top_k = int(top_k) + if self.hidden <= 0 or self.hidden % self._sf_vec: + raise ValueError( + "MXFP8 training stage requires hidden divisible by 32" + ) + if self.top_k <= 0 or self.top_k > self._threads_per_cta: + raise ValueError( + "MXFP8 training stage requires " + f"1 <= top_k <= {self._threads_per_cta}" + ) + + @cute.jit + def __call__( + self, + source: cute.Tensor, + topk_idx: cute.Tensor, + topk_weights: cute.Tensor, + output: cute.Tensor, + output_sf: cute.Tensor, + output_topk_idx: cute.Tensor, + output_topk_weights: cute.Tensor, + stream: cuda.CUstream, + ) -> None: + self._kernel( + source, + topk_idx, + topk_weights, + output, + output_sf, + output_topk_idx, + output_topk_weights, + ).launch( + grid=[source.shape[0], 1, 1], + block=[self._threads_per_cta, 1, 1], + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.kernel + def _kernel( + self, + source: cute.Tensor, + topk_idx: cute.Tensor, + topk_weights: cute.Tensor, + output: cute.Tensor, + output_sf: cute.Tensor, + output_topk_idx: cute.Tensor, + output_topk_weights: cute.Tensor, + ) -> None: + token = cute.arch.block_idx()[0] + tid = cute.arch.thread_idx()[0] + hidden: cutlass.Constexpr[int] = self.hidden + sf_vec: cutlass.Constexpr[int] = self._sf_vec + threads: cutlass.Constexpr[int] = self._threads_per_cta + block_count: cutlass.Constexpr[int] = hidden // sf_vec + rounds: cutlass.Constexpr[int] = ( + block_count + threads - 1 + ) // threads + + for block_round in cutlass.range_constexpr(rounds): + block = tid + Int32(block_round * threads) + if block < Int32(block_count): + values = cute.make_rmem_tensor((sf_vec,), Float32) + absmax = Float32(0.0) + for element in cutlass.range_constexpr(sf_vec): + value = Float32( + source[ + token, + block * Int32(sf_vec) + Int32(element), + ] + ) + values[element] = value + absmax = cute.arch.fmax( + absmax, + cute.arch.fmax(value, -value), + ) + + scale_f32 = Float32( + cvt_f32_to_fp8_to_f32( + absmax * Float32(Fp8E4M3RcpLimit), + cutlass.Float8E8M0FNU, + ) + ) + scale = scale_f32.to(cutlass.Float8E8M0FNU) + reciprocal = cute.arch.fmin( + cute.arch.rcp_approx(scale_f32), + Float32(Fp32Max), + ) + reciprocal = reciprocal * cute.arch.fmin( + scale_f32 * Float32(1.0e30), + Float32(1.0), + ) + for element in cutlass.range_constexpr(sf_vec): + output[ + token, + block * Int32(sf_vec) + Int32(element), + ] = ( + values[element] * reciprocal + ).to(cutlass.Float8E4M3FN) + output_sf[token, block] = scale + + if tid < Int32(self.top_k): + output_topk_idx[token, tid] = Int32(topk_idx[token, tid]) + output_topk_weights[token, tid] = Float32( + topk_weights[token, tid] + ) + + +__all__ = ["Mxfp8TrainingStageKernel"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py new file mode 100644 index 000000000..afb693534 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py @@ -0,0 +1,356 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Stable, allocation-free layout staging for pre-quantized training weights.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from ..._types import BlockScaledTensor, MoeEpTrainingWeights +from ._adapter import Mxfp8Weights + + +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +def _empty_k_major_like(tensor: torch.Tensor) -> torch.Tensor: + if tensor.ndim != 3: + raise ValueError( + f"K-major training weight must be rank 3, got {tensor.ndim}" + ) + experts, reduction, output = tensor.shape + return torch.empty_strided( + tensor.shape, + (reduction * output, 1, reduction), + dtype=tensor.dtype, + device=tensor.device, + ) + + +def _empty_blocked_scales( + source: BlockScaledTensor, + *, + raw_rows: int, + raw_columns: int, + dtype: torch.dtype, +) -> torch.Tensor: + experts = source.data.shape[0] + packed_bytes = _round_up(raw_rows, 128) * _round_up(raw_columns, 4) + return torch.empty( + (experts, packed_bytes), + dtype=torch.uint8, + device=source.device, + ).view(dtype) + + +def _copy_k_major( + target: torch.Tensor, + source: torch.Tensor, +) -> None: + if target.shape != source.shape: + raise ValueError( + f"K-major copy shape mismatch: {target.shape} != {source.shape}" + ) + target.copy_(source) + + +def _copy_gate_up_interleaved_last( + target: torch.Tensor, + source: torch.Tensor, + intermediate: int, +) -> None: + """Copy ``(E,K,gate||up)`` into 32-column gate/up strip order.""" + + experts, reduction, gate_up = source.shape + if target.shape != source.shape or gate_up != 2 * intermediate: + raise ValueError("forward FC1 training weight shape mismatch") + pairs = intermediate // 32 + source_view = source.view( + experts, + reduction, + 2, + pairs, + 32, + ).permute(0, 1, 3, 2, 4) + target_view = target.as_strided( + (experts, reduction, pairs, 2, 32), + ( + target.stride(0), + target.stride(1), + 64 * target.stride(2), + 32 * target.stride(2), + target.stride(2), + ), + ) + target_view.copy_(source_view) + + +def _copy_gate_up_interleaved_reduction( + target: torch.Tensor, + source: torch.Tensor, + intermediate: int, +) -> None: + """Copy ``(E,gate||up,N)`` into 32-row gate/up strip order.""" + + experts, gate_up, output = source.shape + if target.shape != source.shape or gate_up != 2 * intermediate: + raise ValueError("backward W1-transpose training weight shape mismatch") + pairs = intermediate // 32 + source_view = source.view( + experts, + 2, + pairs, + 32, + output, + ).permute(0, 2, 1, 3, 4) + target_view = target.as_strided( + (experts, pairs, 2, 32, output), + ( + target.stride(0), + 64 * target.stride(1), + 32 * target.stride(1), + target.stride(1), + target.stride(2), + ), + ) + target_view.copy_(source_view) + + +def _copy_blocked_scales_plain( + target: torch.Tensor, + source: torch.Tensor, + *, + raw_rows: int, + raw_columns: int, +) -> None: + """Pack public ``(E,Kblocks,N)`` scales for a non-interleaved weight.""" + + experts = source.shape[0] + if tuple(source.shape) != (experts, raw_columns, raw_rows): + raise ValueError( + "plain training scale shape mismatch: " + f"{tuple(source.shape)} != " + f"{(experts, raw_columns, raw_rows)}" + ) + if raw_rows % 128 or raw_columns % 4: + raise ValueError( + "training scale pack requires rows divisible by 128 and " + "columns divisible by 4" + ) + row_blocks = raw_rows // 128 + column_blocks = raw_columns // 4 + source_view = source.view( + torch.uint8, + ).view( + experts, + column_blocks, + 4, + row_blocks, + 4, + 32, + ).permute(0, 3, 1, 5, 4, 2) + target.view(torch.uint8).view( + experts, + row_blocks, + column_blocks, + 32, + 4, + 4, + ).copy_(source_view) + + +def _copy_blocked_scales_gate_up_rows( + target: torch.Tensor, + source: torch.Tensor, + *, + intermediate: int, + reduction_blocks: int, +) -> None: + """Pack forward FC1 scales after 32-row gate/up interleave.""" + + experts = source.shape[0] + raw_rows = 2 * intermediate + raw_columns = reduction_blocks + if tuple(source.shape) != (experts, raw_columns, raw_rows): + raise ValueError("forward FC1 training scale shape mismatch") + if intermediate % 64 or raw_columns % 4: + raise ValueError( + "forward FC1 scale pack requires intermediate divisible by 64 " + "and reduction blocks divisible by 4" + ) + row_blocks = raw_rows // 128 + column_blocks = raw_columns // 4 + source_view = source.view(torch.uint8).view( + experts, + column_blocks, + 4, + 2, + row_blocks, + 2, + 32, + ).permute(0, 4, 1, 6, 5, 3, 2) + target.view(torch.uint8).view( + experts, + row_blocks, + column_blocks, + 32, + 2, + 2, + 4, + ).copy_(source_view) + + +def _copy_blocked_scales_gate_up_columns( + target: torch.Tensor, + source: torch.Tensor, + *, + intermediate: int, + output: int, +) -> None: + """Pack backward W1-transpose scales with interleaved K blocks.""" + + experts = source.shape[0] + reduction_blocks = intermediate // 32 + if tuple(source.shape) != ( + experts, + 2 * reduction_blocks, + output, + ): + raise ValueError("backward W1-transpose training scale shape mismatch") + if output % 128 or reduction_blocks % 2: + raise ValueError( + "backward W1-transpose scale pack requires output divisible by " + "128 and intermediate divisible by 64" + ) + row_blocks = output // 128 + column_blocks = reduction_blocks // 2 + source_view = source.view(torch.uint8).view( + experts, + 2, + column_blocks, + 2, + row_blocks, + 4, + 32, + ).permute(0, 4, 2, 6, 5, 3, 1) + target.view(torch.uint8).view( + experts, + row_blocks, + column_blocks, + 32, + 4, + 2, + 2, + ).copy_(source_view) + + +@dataclass(frozen=True) +class Mxfp8BackwardWeights: + """Kernel names follow the two backward FC stages.""" + + fc1_weight: torch.Tensor + fc1_weight_sf: torch.Tensor + fc2_weight: torch.Tensor + fc2_weight_sf: torch.Tensor + + +class Mxfp8TrainingWeightBindings: + """Stable staging tensors refreshed from four pre-quantized sources.""" + + def __init__(self, weights: MoeEpTrainingWeights) -> None: + self.weights = weights + fwd_fc1 = weights.forward_fc1 + fwd_fc2 = weights.forward_fc2 + bwd_w2t = weights.backward_w2_transpose + bwd_w1t = weights.backward_w1_transpose + + self.forward = Mxfp8Weights( + fc1_weight=_empty_k_major_like(fwd_fc1.data), + fc1_weight_sf=_empty_blocked_scales( + fwd_fc1, + raw_rows=fwd_fc1.data.shape[2], + raw_columns=fwd_fc1.data.shape[1] // 32, + dtype=torch.uint8, + ), + fc2_weight=_empty_k_major_like(fwd_fc2.data), + fc2_weight_sf=_empty_blocked_scales( + fwd_fc2, + raw_rows=fwd_fc2.data.shape[2], + raw_columns=fwd_fc2.data.shape[1] // 32, + dtype=torch.uint8, + ), + ) + self.backward = Mxfp8BackwardWeights( + fc1_weight=_empty_k_major_like(bwd_w2t.data), + fc1_weight_sf=_empty_blocked_scales( + bwd_w2t, + raw_rows=bwd_w2t.data.shape[2], + raw_columns=bwd_w2t.data.shape[1] // 32, + dtype=torch.float8_e8m0fnu, + ), + fc2_weight=_empty_k_major_like(bwd_w1t.data), + fc2_weight_sf=_empty_blocked_scales( + bwd_w1t, + raw_rows=bwd_w1t.data.shape[2], + raw_columns=bwd_w1t.data.shape[1] // 32, + dtype=torch.float8_e8m0fnu, + ), + ) + self.refresh() + + def refresh(self) -> None: + """Enqueue fixed-address layout copies; safe to record in a graph.""" + + fwd_fc1 = self.weights.forward_fc1 + fwd_fc2 = self.weights.forward_fc2 + bwd_w2t = self.weights.backward_w2_transpose + bwd_w1t = self.weights.backward_w1_transpose + intermediate = fwd_fc2.data.shape[1] + + _copy_gate_up_interleaved_last( + self.forward.fc1_weight, + fwd_fc1.data, + intermediate, + ) + _copy_blocked_scales_gate_up_rows( + self.forward.fc1_weight_sf, + fwd_fc1.scale, + intermediate=intermediate, + reduction_blocks=fwd_fc1.data.shape[1] // 32, + ) + _copy_k_major(self.forward.fc2_weight, fwd_fc2.data) + _copy_blocked_scales_plain( + self.forward.fc2_weight_sf, + fwd_fc2.scale, + raw_rows=fwd_fc2.data.shape[2], + raw_columns=fwd_fc2.data.shape[1] // 32, + ) + + _copy_k_major(self.backward.fc1_weight, bwd_w2t.data) + _copy_blocked_scales_plain( + self.backward.fc1_weight_sf, + bwd_w2t.scale, + raw_rows=bwd_w2t.data.shape[2], + raw_columns=bwd_w2t.data.shape[1] // 32, + ) + _copy_gate_up_interleaved_reduction( + self.backward.fc2_weight, + bwd_w1t.data, + intermediate, + ) + _copy_blocked_scales_gate_up_columns( + self.backward.fc2_weight_sf, + bwd_w1t.scale, + intermediate=intermediate, + output=bwd_w1t.data.shape[2], + ) + + +__all__ = [ + "Mxfp8BackwardWeights", + "Mxfp8TrainingWeightBindings", +] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py new file mode 100644 index 000000000..0a1fd851e --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py @@ -0,0 +1,176 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Fixed-address WGrad operand materialization for the training graph path.""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING + +import torch + +from ..._types import MoeEpTrainingWgradOperands +from ._launch import _to_cute + +if TYPE_CHECKING: + from ._training_resources import Mxfp8TrainingSlotViews + + +class Mxfp8TrainingWgradExporter: + """Own scale-expansion compiles; every export writes existing buffers.""" + + def __init__( + self, + *, + experts: int, + hidden: int, + intermediate: int, + sf_padding: int = 128, + ) -> None: + self.experts = int(experts) + self.hidden = int(hidden) + self.intermediate = int(intermediate) + self.sf_padding = int(sf_padding) + self._compiled: dict[tuple[int, int, int | None], object] = {} + self._lock = threading.RLock() + + @staticmethod + def _copy_gate_up_data( + target: torch.Tensor, + source: torch.Tensor, + intermediate: int, + ) -> None: + pool_rows = source.shape[0] + pairs = intermediate // 32 + source_view = source.view(pool_rows, pairs, 2, 32).permute( + 0, + 2, + 1, + 3, + ) + target_view = target.as_strided( + (pool_rows, 2, pairs, 32), + ( + target.stride(0), + intermediate * target.stride(1), + 32 * target.stride(1), + target.stride(1), + ), + ) + target_view.copy_(source_view) + + def _expand_scales( + self, + source: torch.Tensor, + counts: torch.Tensor, + offsets: torch.Tensor, + output: torch.Tensor, + *, + non_k_size: int, + deinterleave_gate_up: int | None = None, + ) -> None: + if source.dtype not in (torch.uint8, torch.float8_e8m0fnu): + raise TypeError("WGrad source scales must use Uint8 or E8M0") + if output.dtype is not torch.float8_e8m0fnu: + raise TypeError("WGrad output scales must use E8M0") + source_bytes = source.view(torch.uint8).reshape(-1) + output_bytes = output.view(torch.uint8).reshape(-1) + key = (int(non_k_size), self.sf_padding, deinterleave_gate_up) + import cuda.bindings.driver as cuda + + stream = torch.cuda.current_stream(output.device) + args = ( + _to_cute(source_bytes, dynamic_layout=False), + _to_cute(counts, assumed_align=4, dynamic_layout=False), + _to_cute(offsets, assumed_align=4, dynamic_layout=False), + _to_cute(output_bytes, dynamic_layout=False), + cuda.CUstream(stream.cuda_stream), + ) + with self._lock: + compiled = self._compiled.get(key) + if compiled is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "WGrad scale expansion must be compiled before " + "CUDA graph capture" + ) + import cutlass.cute as cute + + from ._training_wgrad_kernel import ( + Mxfp8TrainingScaleExpandKernel, + ) + + kernel = Mxfp8TrainingScaleExpandKernel( + non_k_size=non_k_size, + expert_count=self.experts, + source_sf_padding=self.sf_padding, + deinterleave_gate_up=deinterleave_gate_up, + ) + compiled = cute.compile(kernel, *args) + self._compiled[key] = compiled + compiled(*args) + + def export( + self, + slot: "Mxfp8TrainingSlotViews", + ) -> MoeEpTrainingWgradOperands: + """Write and return the fixed-capacity grouped-WGrad operands.""" + + if slot.col_quant_data is None or slot.col_quant_sf is None: + raise RuntimeError("training WGrad export requires forward col-quant") + pool_rows = slot.fc1_recompute.shape[0] + if slot.col_quant_data.shape[0] != pool_rows: + raise RuntimeError("forward/backward WGrad pool capacities differ") + + self._copy_gate_up_data( + slot.wgrad_fc1_b, + slot.fc1_col_output, + self.intermediate, + ) + slot.wgrad_fc2_a.copy_(slot.fc1_recompute.transpose(0, 1)) + self._expand_scales( + slot.col_quant_sf, + slot.valid_route_counts, + slot.expert_offsets, + slot.wgrad_fc1_sfa, + non_k_size=self.hidden, + ) + self._expand_scales( + slot.fc1_col_output_sf, + slot.valid_route_counts, + slot.expert_offsets, + slot.wgrad_fc1_sfb, + non_k_size=2 * self.intermediate, + deinterleave_gate_up=self.intermediate, + ) + self._expand_scales( + slot.fc1_recompute_sf, + slot.valid_route_counts, + slot.expert_offsets, + slot.wgrad_fc2_sfa, + non_k_size=self.intermediate, + ) + self._expand_scales( + slot.grad_y2_sf, + slot.valid_route_counts, + slot.expert_offsets, + slot.wgrad_fc2_sfb, + non_k_size=self.hidden, + ) + + return MoeEpTrainingWgradOperands( + fc1_a=slot.col_quant_data.transpose(0, 1), + fc1_sfa=slot.wgrad_fc1_sfa, + fc1_b=slot.wgrad_fc1_b, + fc1_sfb=slot.wgrad_fc1_sfb, + fc2_a=slot.wgrad_fc2_a, + fc2_sfa=slot.wgrad_fc2_sfa, + fc2_b=slot.grad_y2, + fc2_sfb=slot.wgrad_fc2_sfb, + expert_offsets=slot.expert_offsets, + valid_route_counts=slot.valid_route_counts, + ) + + +__all__ = ["Mxfp8TrainingWgradExporter"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py new file mode 100644 index 000000000..c30a4f4f3 --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py @@ -0,0 +1,177 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Device expansion of per-expert SF atoms into fixed-capacity WGrad ABI.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Int32 + + +class Mxfp8TrainingScaleExpandKernel: + """Expand 128-row SF extents into 256-row fixed-capacity segments.""" + + _threads = 256 + _atom_bytes = 512 + _neutral_e8m0 = 127 + + def __init__( + self, + *, + non_k_size: int, + expert_count: int, + source_sf_padding: int, + deinterleave_gate_up: int | None = None, + ) -> None: + self.non_k_size = int(non_k_size) + self.expert_count = int(expert_count) + self.source_sf_padding = int(source_sf_padding) + self.deinterleave_gate_up = ( + None + if deinterleave_gate_up is None + else int(deinterleave_gate_up) + ) + if self.non_k_size <= 0 or self.non_k_size % 128: + raise ValueError("WGrad scale non-K size must be divisible by 128") + if self.expert_count <= 0: + raise ValueError("WGrad scale expansion requires experts") + if ( + self.source_sf_padding <= 0 + or self.source_sf_padding % 128 + ): + raise ValueError( + "WGrad source SF padding must be a positive multiple of 128" + ) + if ( + self.deinterleave_gate_up is not None + and self.non_k_size != 2 * self.deinterleave_gate_up + ): + raise ValueError("gate/up scale deinterleave size mismatch") + + @cute.jit + def __call__( + self, + source: cute.Tensor, + valid_counts: cute.Tensor, + expert_offsets: cute.Tensor, + output: cute.Tensor, + stream: cuda.CUstream, + ) -> None: + output_bytes = cute.size(output) + self._kernel( + source, + valid_counts, + expert_offsets, + output, + ).launch( + grid=[ + output_bytes // self._threads, + 1, + 1, + ], + block=[self._threads, 1, 1], + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.kernel + def _kernel( + self, + source: cute.Tensor, + valid_counts: cute.Tensor, + expert_offsets: cute.Tensor, + output: cute.Tensor, + ) -> None: + linear = ( + cute.arch.block_idx()[0] * Int32(self._threads) + + cute.arch.thread_idx()[0] + ) + + atom_bytes: cutlass.Constexpr[int] = self._atom_bytes + non_k_atoms: cutlass.Constexpr[int] = self.non_k_size // 128 + atom = linear // Int32(atom_bytes) + byte_in_atom = linear % Int32(atom_bytes) + value = cutlass.Uint8(self._neutral_e8m0) + target_atom_base = Int32(0) + source_atom_base = Int32(0) + previous_end = Int32(0) + + for expert in cutlass.range_constexpr(self.expert_count): + end = Int32(expert_offsets[expert]) + target_token_atoms = (end - previous_end) // Int32(128) + source_token_atoms = ( + ( + Int32(valid_counts[expert]) + + Int32(self.source_sf_padding - 1) + ) + // Int32(self.source_sf_padding) + ) * Int32(self.source_sf_padding // 128) + target_atom_count = ( + Int32(non_k_atoms) * target_token_atoms + ) + in_expert = ( + (atom >= target_atom_base) + & (atom < target_atom_base + target_atom_count) + ) + if in_expert & (target_token_atoms > Int32(0)): + relative_atom = atom - target_atom_base + hidden_atom = relative_atom // target_token_atoms + token_atom = relative_atom % target_token_atoms + if token_atom < source_token_atoms: + source_hidden_atom = hidden_atom + source_byte = byte_in_atom + if cutlass.const_expr( + self.deinterleave_gate_up is not None + ): + lane = byte_in_atom // Int32(16) + byte_tail = byte_in_atom % Int32(16) + group = byte_tail // Int32(4) + column_lane = byte_tail % Int32(4) + feature = ( + hidden_atom * Int32(128) + + group * Int32(32) + + lane + ) + intermediate = Int32(self.deinterleave_gate_up) + source_feature = Int32(0) + if feature < intermediate: + source_feature = ( + (feature // Int32(32)) * Int32(64) + + feature % Int32(32) + ) + else: + up_feature = feature - intermediate + source_feature = ( + (up_feature // Int32(32)) * Int32(64) + + Int32(32) + + up_feature % Int32(32) + ) + source_hidden_atom = source_feature // Int32(128) + source_feature_in_atom = source_feature % Int32(128) + source_byte = ( + (source_feature_in_atom % Int32(32)) + * Int32(16) + + (source_feature_in_atom // Int32(32)) + * Int32(4) + + column_lane + ) + source_atom = ( + source_atom_base + + source_hidden_atom * source_token_atoms + + token_atom + ) + value = source[ + source_atom * Int32(atom_bytes) + source_byte + ] + target_atom_base += target_atom_count + source_atom_base += Int32(non_k_atoms) * source_token_atoms + previous_end = end + + output[linear] = value + + +__all__ = ["Mxfp8TrainingScaleExpandKernel"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_wgrad_layout.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_wgrad_layout.py deleted file mode 100644 index cb2a2fa03..000000000 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_wgrad_layout.py +++ /dev/null @@ -1,410 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""MXFP8 pool-to-grouped-wgrad data and scale layout conversion.""" - -from __future__ import annotations - -import torch - -from .._workspace import _align_up - -_SF_VEC_SIZE = 32 -_SF_ATOM_ROWS = 128 -_SF_ATOM_COLUMNS = 4 -_SF_ATOM_BYTES = _SF_ATOM_ROWS * _SF_ATOM_COLUMNS -_GATE_UP_INTERLEAVE = 32 -_NEUTRAL_E8M0 = 127 - - -def cumulative_padded_offsets( - valid_counts: tuple[int, ...], - padding: int, - device: torch.device, -) -> tuple[tuple[int, ...], torch.Tensor]: - """Return cumulative padded expert ends as host values and device Int32.""" - - ends = [] - total = 0 - for count in valid_counts: - if count < 0: - raise ValueError("expert route counts must be non-negative") - total += _align_up(count, padding) - ends.append(total) - return tuple(ends), torch.tensor( - ends, - dtype=torch.int32, - device=device, - ) - - -def pool_data_as_wgrad_a( - pool_data: torch.Tensor, - padded_routes: int, -) -> torch.Tensor: - """Copy pool ``(K,M)`` bytes into contiguous logical wgrad A ``(M,K)``.""" - - _validate_pool_prefix(pool_data, padded_routes) - return pool_data[:padded_routes].transpose(0, 1).contiguous() - - -def pool_data_as_wgrad_b( - pool_data: torch.Tensor, - padded_routes: int, -) -> torch.Tensor: - """Copy pool ``(K,N)`` bytes into a K-major logical wgrad B ``(K,N)``.""" - - _validate_pool_prefix(pool_data, padded_routes) - return ( - pool_data[:padded_routes] - .transpose(0, 1) - .contiguous() - .transpose(0, 1) - ) - - -def deinterleave_gate_up_columns( - tensor: torch.Tensor, - intermediate: int, -) -> torch.Tensor: - """Convert 32-column ``gate,up`` strips to logical ``gate || up``.""" - - expected = 2 * intermediate - if tensor.ndim != 2 or tensor.shape[1] != expected: - raise ValueError( - f"gate/up tensor must have shape (rows, {expected}), " - f"got {tuple(tensor.shape)}" - ) - if intermediate % _GATE_UP_INTERLEAVE: - raise ValueError( - "intermediate size must be divisible by " - f"{_GATE_UP_INTERLEAVE}" - ) - pairs = intermediate // _GATE_UP_INTERLEAVE - blocks = tensor.reshape( - tensor.shape[0], - pairs, - 2, - _GATE_UP_INTERLEAVE, - ) - gate = blocks[:, :, 0, :].reshape(tensor.shape[0], intermediate) - up = blocks[:, :, 1, :].reshape(tensor.shape[0], intermediate) - return torch.cat((gate, up), dim=1) - - -def assemble_discrete_col_requant_scales( - packed_scales: torch.Tensor, - valid_counts: tuple[int, ...], - padded_ends: tuple[int, ...], - non_k_size: int, - sf_padding: int, -) -> torch.Tensor: - """Assemble upstream col-requant SF atoms for grouped wgrad. - - Since upstream revision 71d5fc1, each expert already emits - ``(non-K/128, K/128, 512)`` atoms, which is grouped-wgrad order. - """ - - return _assemble_atom_scales( - packed_scales, - valid_counts, - padded_ends, - non_k_size, - sf_padding, - source_hidden_major=True, - source_name="col-requant", - ) - - -def assemble_dfc2_atom_scales( - packed_scales: torch.Tensor, - valid_counts: tuple[int, ...], - padded_ends: tuple[int, ...], - non_k_size: int, - sf_padding: int, - *, - deinterleave_gate_up: int | None = None, -) -> torch.Tensor: - """Reorder dFC2 epilogue atoms from token-major to grouped-wgrad order.""" - - return _assemble_atom_scales( - packed_scales, - valid_counts, - padded_ends, - non_k_size, - sf_padding, - source_hidden_major=False, - source_name="dFC2", - deinterleave_gate_up=deinterleave_gate_up, - ) - - -def _assemble_atom_scales( - packed_scales: torch.Tensor, - valid_counts: tuple[int, ...], - padded_ends: tuple[int, ...], - non_k_size: int, - sf_padding: int, - *, - source_hidden_major: bool, - source_name: str, - deinterleave_gate_up: int | None = None, -) -> torch.Tensor: - """Expand compact per-expert SF atoms to the data-padded K extent.""" - - if len(valid_counts) != len(padded_ends): - raise ValueError("expert count and padded offset lengths must match") - _validate_padded_ends(padded_ends) - if sf_padding % _SF_ATOM_ROWS: - raise ValueError("scale padding must be divisible by 128") - padded_non_k = _align_up(non_k_size, _SF_ATOM_ROWS) - non_k_atoms = padded_non_k // _SF_ATOM_ROWS - flat_u8 = packed_scales.view(torch.uint8).reshape(-1) - expert_parts = [] - previous_end = 0 - source_byte_offset = 0 - for count, end in zip(valid_counts, padded_ends): - target_extent = end - previous_end - if count < 0 or count > target_extent: - raise ValueError("valid expert routes exceed their padded extent") - if target_extent % _SF_ATOM_ROWS: - raise ValueError( - "data-padded expert extents must be multiples of 128" - ) - source_extent = _align_up(count, sf_padding) - source_token_atoms = source_extent // _SF_ATOM_ROWS - target_token_atoms = target_extent // _SF_ATOM_ROWS - if source_token_atoms > target_token_atoms: - raise ValueError("scale-padded extent exceeds data-padded extent") - source_byte_count = ( - source_token_atoms * non_k_atoms * _SF_ATOM_BYTES - ) - if source_byte_offset + source_byte_count > flat_u8.numel(): - raise ValueError( - f"{source_name} scale output is smaller than its layout" - ) - target_raw = torch.full( - (padded_non_k, target_token_atoms * _SF_ATOM_COLUMNS), - _NEUTRAL_E8M0, - dtype=torch.uint8, - device=flat_u8.device, - ) - if source_byte_count: - source = flat_u8.narrow( - 0, - source_byte_offset, - source_byte_count, - ) - if source_hidden_major: - source = source.reshape( - non_k_atoms, - source_token_atoms, - _SF_ATOM_BYTES, - ) - else: - source = ( - source.reshape( - source_token_atoms, - non_k_atoms, - _SF_ATOM_BYTES, - ) - .permute(1, 0, 2) - .contiguous() - ) - source_raw = _from_blocked_bytes( - source.reshape(-1), - padded_non_k, - source_token_atoms * _SF_ATOM_COLUMNS, - ) - target_raw[:, : source_raw.shape[1]].copy_(source_raw) - if deinterleave_gate_up is not None: - target_raw = ( - deinterleave_gate_up_columns( - target_raw.transpose(0, 1), - deinterleave_gate_up, - ) - .transpose(0, 1) - .contiguous() - ) - expert_parts.append(_to_blocked_bytes(target_raw)) - source_byte_offset += source_byte_count - previous_end = end - - total_routes = padded_ends[-1] if padded_ends else 0 - scale_columns = _align_up(total_routes // _SF_VEC_SIZE, 4) - if expert_parts: - assembled_u8 = torch.cat(expert_parts) - else: - assembled_u8 = flat_u8.new_empty((0,)) - expected = padded_non_k * scale_columns - if assembled_u8.numel() != expected: - raise RuntimeError( - f"assembled {source_name} scale size mismatch: " - f"{assembled_u8.numel()} != {expected}" - ) - return assembled_u8.reshape(padded_non_k, scale_columns).view( - torch.float8_e8m0fnu - ) - - -def assemble_plain_col_scales( - col_scales: torch.Tensor, - valid_counts: tuple[int, ...], - padded_ends: tuple[int, ...], - non_k_size: int, - sf_padding: int, - *, - deinterleave_gate_up: int | None = None, -) -> torch.Tensor: - """Assemble plain ``(K/32,N)`` col scales into grouped-wgrad SF atoms.""" - - if len(valid_counts) != len(padded_ends): - raise ValueError("expert count and padded offset lengths must match") - _validate_padded_ends(padded_ends) - if col_scales.ndim != 2 or col_scales.shape[1] != non_k_size: - raise ValueError( - "plain column scales must have shape " - f"(rows, {non_k_size}), got {tuple(col_scales.shape)}" - ) - if sf_padding % _SF_VEC_SIZE: - raise ValueError("scale padding must be divisible by 32") - - source_u8 = col_scales.view(torch.uint8) - expert_parts = [] - previous_end = 0 - sf_row = 0 - for count, end in zip(valid_counts, padded_ends): - padded_extent = end - previous_end - if padded_extent % _SF_VEC_SIZE: - raise ValueError("padded expert extents must be divisible by 32") - valid_sf_rows = (count + _SF_VEC_SIZE - 1) // _SF_VEC_SIZE - padded_sf_rows = padded_extent // _SF_VEC_SIZE - if count < 0 or count > padded_extent: - raise ValueError("valid expert routes exceed their padded extent") - if sf_row + valid_sf_rows > source_u8.shape[0]: - raise ValueError("plain column scale output is too short") - - raw = torch.full( - (non_k_size, padded_sf_rows), - _NEUTRAL_E8M0, - dtype=torch.uint8, - device=col_scales.device, - ) - if valid_sf_rows: - source = source_u8[ - sf_row : sf_row + valid_sf_rows, - :, - ] - if deinterleave_gate_up is not None: - source = deinterleave_gate_up_columns( - source, - deinterleave_gate_up, - ) - raw[:, :valid_sf_rows].copy_(source.transpose(0, 1)) - expert_parts.append(_to_blocked_bytes(raw)) - sf_row += _align_up(count, sf_padding) // _SF_VEC_SIZE - previous_end = end - - padded_non_k = _align_up(non_k_size, _SF_ATOM_ROWS) - total_routes = padded_ends[-1] if padded_ends else 0 - scale_columns = _align_up(total_routes // _SF_VEC_SIZE, 4) - if expert_parts: - assembled_u8 = torch.cat(expert_parts) - else: - assembled_u8 = source_u8.new_empty((0,)) - expected = padded_non_k * scale_columns - if assembled_u8.numel() != expected: - raise RuntimeError( - "assembled plain column scale size mismatch: " - f"{assembled_u8.numel()} != {expected}" - ) - return assembled_u8.reshape(padded_non_k, scale_columns).view( - torch.float8_e8m0fnu - ) - - -def _to_blocked_bytes(raw_scale: torch.Tensor) -> torch.Tensor: - rows, columns = raw_scale.shape - if rows == 0 or columns == 0: - return raw_scale.new_empty((0,), dtype=torch.uint8) - padded_rows = _align_up(rows, _SF_ATOM_ROWS) - padded_columns = _align_up(columns, _SF_ATOM_COLUMNS) - padded = torch.full( - (padded_rows, padded_columns), - _NEUTRAL_E8M0, - dtype=torch.uint8, - device=raw_scale.device, - ) - padded[:rows, :columns].copy_(raw_scale) - blocks = padded.view( - padded_rows // _SF_ATOM_ROWS, - _SF_ATOM_ROWS, - padded_columns // _SF_ATOM_COLUMNS, - _SF_ATOM_COLUMNS, - ).permute(0, 2, 1, 3) - return ( - blocks.reshape(-1, 4, 32, 4) - .transpose(1, 2) - .reshape(-1) - ) - - -def _from_blocked_bytes( - packed_scale: torch.Tensor, - rows: int, - columns: int, -) -> torch.Tensor: - """Invert the grouped-wgrad 128x4 scale-atom swizzle.""" - - padded_rows = _align_up(rows, _SF_ATOM_ROWS) - padded_columns = _align_up(columns, _SF_ATOM_COLUMNS) - expected = padded_rows * padded_columns - flat = packed_scale.view(torch.uint8).reshape(-1) - if flat.numel() != expected: - raise ValueError( - f"blocked scale has {flat.numel()} bytes, expected {expected}" - ) - if expected == 0: - return flat.new_empty((rows, columns)) - row_atoms = padded_rows // _SF_ATOM_ROWS - column_atoms = padded_columns // _SF_ATOM_COLUMNS - raw = ( - flat.reshape(row_atoms * column_atoms, 32, 4, 4) - .transpose(1, 2) - .reshape(row_atoms, column_atoms, _SF_ATOM_ROWS, _SF_ATOM_COLUMNS) - .permute(0, 2, 1, 3) - .reshape(padded_rows, padded_columns) - ) - return raw[:rows, :columns] - - -def _validate_pool_prefix( - pool_data: torch.Tensor, - padded_routes: int, -) -> None: - if pool_data.ndim != 2 or not pool_data.is_contiguous(): - raise ValueError("pool data must be a contiguous rank-2 tensor") - if padded_routes < 0 or padded_routes > pool_data.shape[0]: - raise ValueError( - f"padded route count {padded_routes} exceeds pool capacity " - f"{pool_data.shape[0]}" - ) - - -def _validate_padded_ends(padded_ends: tuple[int, ...]) -> None: - previous = 0 - for end in padded_ends: - if end < previous: - raise ValueError("expert offsets must be non-decreasing") - previous = end - - -__all__ = [ - "assemble_dfc2_atom_scales", - "assemble_discrete_col_requant_scales", - "assemble_plain_col_scales", - "cumulative_padded_offsets", - "deinterleave_gate_up_columns", - "pool_data_as_wgrad_a", - "pool_data_as_wgrad_b", -] diff --git a/python/cudnn/moe_ep/_types.py b/python/cudnn/moe_ep/_types.py index 1525060ab..e1bfb48de 100644 --- a/python/cudnn/moe_ep/_types.py +++ b/python/cudnn/moe_ep/_types.py @@ -8,7 +8,7 @@ import operator from dataclasses import dataclass from enum import Enum -from typing import Tuple, Union +from typing import Any, Tuple, Union import torch @@ -210,71 +210,26 @@ def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: @dataclass(frozen=True) -class MoeEpWgradForwardStash: - """Caller-owned forward state required to form expert-local wgrads. - - ``fc1_a`` and ``fc1_sfa`` represent the MXFP8 ``x.T`` operand. Valid - routes for each local expert occupy the beginning of its padded range; - ``expert_offsets`` contains cumulative padded end offsets and - ``valid_route_counts`` contains the corresponding unpadded row counts. - Scale factors use the blocked layout consumed by grouped wgrad, with - logical 1x32 scaling and physical 128x4 scale tiles. - ``route_metadata`` is the compact identity table returned by forward, - using ``(local_expert, src_rank, src_token, src_slot)`` rows. It validates - that the stash belongs to the matching routed call; it is not padded or - row-aligned with the operands' K dimension. +class MoeEpTrainingWeights: + """Stable MXFP8 bindings for forward and dgrad GEMMs. + + Forward consumes ``forward_fc1`` with logical shape ``(E,H,2I)`` and + ``forward_fc2`` with ``(E,I,H)``. Backward consumes independently + quantized transposes: ``backward_w2_transpose=(E,H,I)`` for + ``dH=dY@W2.T`` and ``backward_w1_transpose=(E,2I,H)`` for + ``dX=dC@W1.T``. Every tensor is block-scaled along logical axis 1, the + reduction axis of its corresponding GEMM. """ - fc1_a: torch.Tensor - fc1_sfa: torch.Tensor - expert_offsets: torch.Tensor - valid_route_counts: torch.Tensor - route_metadata: torch.Tensor + forward_fc1: BlockScaledTensor + forward_fc2: BlockScaledTensor + backward_w2_transpose: BlockScaledTensor + backward_w1_transpose: BlockScaledTensor @dataclass(frozen=True) -class MoeEpWgradOperands: - """Caller-owned MXFP8 operands for expert-local grouped wgrad GEMMs. - - The represented operations are ``dW1 = fc1_a @ fc1_b`` and - ``dW2 = fc2_a @ fc2_b``. For total padded route extent ``K``, their - logical shapes are ``fc1_a=(H,K)``, ``fc1_b=(K,2I)``, - ``fc2_a=(I,K)``, and ``fc2_b=(K,H)``. Each scale tensor uses grouped - wgrad's blocked 1x32 layout: ``(round_up(non-K,128), round_up(K/32,4))``. - The shared expert metadata has the same meaning as in - :class:`MoeEpWgradForwardStash`. - - Attributes: - fc1_a: E4M3 data for the FC1 A operand, logically ``x.T`` with shape - ``(H, K)``. ``x`` is the activation dispatched to each local - expert. The K dimension concatenates the experts' independently - padded route ranges. - fc1_sfa: E8M0 scales for ``fc1_a``. Each logical scale covers 32 - consecutive K elements of one hidden-feature row. - fc1_b: E4M3 data for the FC1 B operand, logically - ``dC=[d_gate | d_up]`` with shape ``(K, 2I)``. ``dC`` is the - gradient of the pre-SwiGLU FC1 accumulator; columns use the public - gate-then-up order rather than the kernel's internal strip - interleave. - fc1_sfb: E8M0 scales for ``fc1_b``. Each logical scale covers 32 - consecutive K rows for one gate/up feature column. - fc2_a: E4M3 data for the FC2 A operand, logically ``(p*h).T`` with - shape ``(I, K)``. ``h=SwiGLU(C)`` and ``p`` is the route's FP32 - router score, applied exactly once before column quantization. - fc2_sfa: E8M0 scales for ``fc2_a``. Each logical scale covers 32 - consecutive K elements of one intermediate-feature row. - fc2_b: E4M3 data for the FC2 B operand, logically unweighted ``dY`` - with shape ``(K, H)``. ``dY`` is the routed FC2 output gradient. - fc2_sfb: E8M0 scales for ``fc2_b``. Each logical scale covers 32 - consecutive K rows for one hidden-feature column. - expert_offsets: Int32 cumulative padded K-end offset for every local - expert. Adjacent equal offsets represent an empty expert. - valid_route_counts: Int32 unpadded route count for every local expert; - valid rows occupy the beginning of each padded expert range. - route_metadata: Compact Int32 route identity table with columns - ``(local_expert, src_rank, src_token, src_slot)``. It identifies - the routed call but is not padded or row-aligned with K. - """ +class MoeEpTrainingWgradOperands: + """Fixed-capacity MXFP8 operands produced by the training resource path.""" fc1_a: torch.Tensor fc1_sfa: torch.Tensor @@ -286,7 +241,193 @@ class MoeEpWgradOperands: fc2_sfb: torch.Tensor expert_offsets: torch.Tensor valid_route_counts: torch.Tensor - route_metadata: torch.Tensor + + +@dataclass(frozen=True) +class MoeEpTrainingSlot: + """Opaque index of one persistent forward/backward training slot.""" + + index: int + _resource_token: object + + +@dataclass(frozen=True) +class MoeEpExecutionLane: + """Opaque index of one mutable per-stream execution lane.""" + + index: int + _resource_token: object + + +class MoeEpTrainingResources: + """TE-owned lease on fixed-capacity training slots and execution lanes.""" + + def __init__( + self, + *, + owner: Any, + operator_token: object, + weights: MoeEpTrainingWeights, + slot_count: int, + lane_count: int, + device: torch.device, + ) -> None: + self._owner = owner + self._operator_token = operator_token + self._resource_token = object() + self.weights = weights + self.device = torch.device(device) + self.slots = tuple( + MoeEpTrainingSlot(index, self._resource_token) + for index in range(slot_count) + ) + self.lanes = tuple( + MoeEpExecutionLane(index, self._resource_token) + for index in range(lane_count) + ) + self._closed = False + + @property + def closed(self) -> bool: + return self._closed + + def _check_binding( + self, + operator_token: object, + slot: MoeEpTrainingSlot, + lane: MoeEpExecutionLane, + ) -> None: + if self._closed: + raise RuntimeError("MoeEp training resources are closed") + if self._operator_token is not operator_token: + raise ValueError("training resources belong to another MoeEp instance") + if ( + slot._resource_token is not self._resource_token + or slot not in self.slots + ): + raise ValueError("training slot does not belong to these resources") + if ( + lane._resource_token is not self._resource_token + or lane not in self.lanes + ): + raise ValueError("execution lane does not belong to these resources") + + def refresh_weights(self) -> None: + """Enqueue fixed-address weight-layout refreshes on the current stream. + + Call after every in-place data+scale update and before the first + forward/backward that consumes that version. The caller must establish + stream/event ordering, must not refresh between a matching forward and + backward, and must not overlap refresh with any consumer of these + resources. Replacing source storage requires closing the old operator, + creating a new ``MoeEp`` instance and resources, and capturing a new + graph. This method may itself be captured, in which case replay executes + only the recorded device transforms. + """ + + if self._closed: + raise RuntimeError("MoeEp training resources are closed") + self._owner.refresh_weights() + + def forward( + self, + slot: MoeEpTrainingSlot, + lane: MoeEpExecutionLane, + activation: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ) -> torch.Tensor: + """Run the fixed-slot forward in ordinary or capture mode.""" + + self._check_binding(self._operator_token, slot, lane) + execution = self._owner.views( + slot=slot.index, + lane=lane.index, + token_count=int(activation.shape[0]), + ) + from ._megamoe_backend.mxfp8._training_execute import ( + launch_training_forward, + ) + + return launch_training_forward( + self._owner, + execution, + activation, + topk_idx, + topk_weights, + ) + + def backward( + self, + slot: MoeEpTrainingSlot, + lane: MoeEpExecutionLane, + grad_output: torch.Tensor, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + MoeEpTrainingWgradOperands, + ]: + """Run fixed-slot dgrad/dprob in ordinary or capture mode.""" + + self._check_binding(self._operator_token, slot, lane) + execution = self._owner.views( + slot=slot.index, + lane=lane.index, + token_count=int(grad_output.shape[0]), + ) + from ._megamoe_backend.mxfp8._training_execute import ( + launch_training_backward, + ) + + grad_activation, grad_topk_weights, operands = ( + launch_training_backward( + self._owner, + execution, + grad_output, + ) + ) + return grad_activation, grad_topk_weights, operands + + def finalize_overflow( + self, + slots: Tuple[MoeEpTrainingSlot, ...], + lane: MoeEpExecutionLane | None = None, + ) -> torch.Tensor: + """Aggregate one computation group's flags and apply its policy.""" + + if self._closed: + raise RuntimeError("MoeEp training resources are closed") + if lane is None: + lane = self.lanes[0] + if ( + not isinstance(lane, MoeEpExecutionLane) + or lane._resource_token is not self._resource_token + or lane not in self.lanes + ): + raise ValueError( + "overflow execution lane does not belong to these resources" + ) + slot_indices = [] + for slot in slots: + if ( + not isinstance(slot, MoeEpTrainingSlot) + or slot._resource_token is not self._resource_token + or slot not in self.slots + ): + raise ValueError( + "overflow slot does not belong to these resources" + ) + slot_indices.append(slot.index) + return self._owner.finalize_overflow( + tuple(slot_indices), + lane=lane.index, + ) + + def close(self) -> None: + if self._closed: + return + self._owner.close() + self._closed = True MoeTensor = Union[torch.Tensor, BlockScaledTensor] @@ -294,8 +435,11 @@ class MoeEpWgradOperands: __all__ = [ "BlockScaledTensor", - "MoeEpWgradForwardStash", - "MoeEpWgradOperands", + "MoeEpExecutionLane", + "MoeEpTrainingResources", + "MoeEpTrainingSlot", + "MoeEpTrainingWeights", + "MoeEpTrainingWgradOperands", "MoeFormat", "MoeTensor", "parse_format", diff --git a/python/cudnn/moe_ep/_validation.py b/python/cudnn/moe_ep/_validation.py index 9ee259fa3..f2b725001 100644 --- a/python/cudnn/moe_ep/_validation.py +++ b/python/cudnn/moe_ep/_validation.py @@ -1,12 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Pure public-contract validation for :mod:`cudnn.moe_ep`. - -This module intentionally depends only on PyTorch and the lightweight public -API types. It must not import CuTeDSL, CUDA Python, NVSHMEM, or the private -MegaMoE runtime. -""" +"""Pure public-contract validation for :mod:`cudnn.moe_ep`.""" from __future__ import annotations @@ -14,20 +9,20 @@ import torch -from ._contracts import ( - ForwardConfig, - ValidatedBackwardRequest, - ValidatedForwardRequest, -) +from ._contracts import ForwardConfig, ValidatedForwardRequest from ._types import ( BlockScaledTensor, - MoeEpWgradForwardStash, + MoeEpTrainingWeights, MoeFormat, MoeTensor, ) -def _replace_axis(shape: Tuple[int, ...], axis: int, extent: int) -> Tuple[int, ...]: +def _replace_axis( + shape: Tuple[int, ...], + axis: int, + extent: int, +) -> Tuple[int, ...]: result = list(shape) result[axis] = extent return tuple(result) @@ -37,10 +32,6 @@ def _ceil_div(value: int, divisor: int) -> int: return (value + divisor - 1) // divisor -def _round_up(value: int, multiple: int) -> int: - return _ceil_div(value, multiple) * multiple - - def _require_torch_dtype(name: str) -> torch.dtype: dtype = getattr(torch, name, None) if dtype is None: @@ -81,21 +72,19 @@ def _validate_tensor_representation( logical_shape = _logical_shape(tensor) if logical_shape != expected_logical_shape: raise ValueError( - f"{name} logical shape must be {expected_logical_shape}, got {logical_shape}" + f"{name} logical shape must be {expected_logical_shape}, " + f"got {logical_shape}" ) - if isinstance(tensor, torch.Tensor): _validate_strided(name, tensor) if not tensor.is_floating_point(): raise ValueError(f"{name} must be floating point, got {tensor.dtype}") return - if not isinstance(tensor, BlockScaledTensor): raise ValueError( f"{name} must be a torch.Tensor or BlockScaledTensor, " f"got {type(tensor).__name__}" ) - if tensor.axis != 1: raise ValueError(f"{name} block-scaled axis must be 1, got {tensor.axis}") _validate_strided(f"{name}.data", tensor.data) @@ -112,7 +101,6 @@ def _validate_tensor_representation( block_size = 16 expected_data_dtype = torch.uint8 expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") - expected_data_shape = _replace_axis( expected_logical_shape, tensor.axis, @@ -168,8 +156,6 @@ def _validate_routes( *, validate_expert_ids: bool, ) -> None: - """Validate the public routing plane shared by forward and backward.""" - if not isinstance(topk_idx, torch.Tensor): raise ValueError( f"topk_idx must be a torch.Tensor, got {type(topk_idx).__name__}" @@ -181,7 +167,6 @@ def _validate_routes( ) _validate_strided("topk_idx", topk_idx) _validate_strided("topk_weights", topk_weights) - route_shape = (token_count, config.top_k) if tuple(topk_idx.shape) != route_shape: raise ValueError( @@ -223,7 +208,7 @@ def validate_forward( *, validate_expert_ids: bool = True, ) -> ValidatedForwardRequest: - """Validate public forward semantics without importing a device backend.""" + """Validate inference-forward semantics without importing a device backend.""" activation_shape = _logical_shape(activation) if len(activation_shape) != 2 or activation_shape[1] != config.hidden_size: @@ -232,20 +217,25 @@ def validate_forward( f"got {activation_shape}" ) token_count = activation_shape[0] - expected_fc1 = ( - config.experts_per_rank, - config.hidden_size, - 2 * config.intermediate_size, + _validate_tensor_representation("activation", activation, activation_shape) + _validate_tensor_representation( + "fc1_weight", + fc1_weight, + ( + config.experts_per_rank, + config.hidden_size, + 2 * config.intermediate_size, + ), ) - expected_fc2 = ( - config.experts_per_rank, - config.intermediate_size, - config.hidden_size, + _validate_tensor_representation( + "fc2_weight", + fc2_weight, + ( + config.experts_per_rank, + config.intermediate_size, + config.hidden_size, + ), ) - _validate_tensor_representation("activation", activation, activation_shape) - _validate_tensor_representation("fc1_weight", fc1_weight, expected_fc1) - _validate_tensor_representation("fc2_weight", fc2_weight, expected_fc2) - _validate_routes( config, token_count, @@ -253,7 +243,6 @@ def validate_forward( topk_weights, validate_expert_ids=False, ) - device = _tensor_device(activation) for name, tensor in ( ("fc1_weight", fc1_weight), @@ -264,12 +253,6 @@ def validate_forward( tensor_device = _tensor_device(tensor) if tensor_device != device: raise ValueError(f"{name} must be on {device}, got {tensor_device}") - - # Boolean compaction plus the host-visible ``item()`` below is not CUDA - # graph capturable. Eager calls (including the mandatory pre-capture - # warmup) retain strict validation. During capture/replay, callers must - # preserve that validated invariant: every route is -1 or a valid global - # expert ID. if device.type == "cuda": with torch.cuda.device(device): capturing = torch.cuda.is_current_stream_capturing() @@ -277,7 +260,6 @@ def validate_forward( capturing = False if validate_expert_ids and not capturing: _validate_expert_ids(config, topk_idx) - return ValidatedForwardRequest( config=config, activation=activation, @@ -290,334 +272,76 @@ def validate_forward( ) -def _validate_wgrad_forward_stash( +def validate_training_weights( config: ForwardConfig, - stash: MoeEpWgradForwardStash, - route_metadata: torch.Tensor, - device: torch.device, -) -> None: - """Validate caller-owned forward operands and their route identity.""" + weights: MoeEpTrainingWeights, +) -> torch.device: + """Validate fixed MXFP8 weight bindings used by training resources.""" - if not isinstance(stash, MoeEpWgradForwardStash): + if not isinstance(weights, MoeEpTrainingWeights): raise TypeError( - "wgrad_forward_stash must be a MoeEpWgradForwardStash, " - f"got {type(stash).__name__}" + "weights must be a MoeEpTrainingWeights, " + f"got {type(weights).__name__}" ) - - tensors = ( - ("fc1_a", stash.fc1_a), - ("fc1_sfa", stash.fc1_sfa), - ("expert_offsets", stash.expert_offsets), - ("valid_route_counts", stash.valid_route_counts), - ("route_metadata", stash.route_metadata), - ) - for name, tensor in tensors: - if not isinstance(tensor, torch.Tensor): - raise TypeError( - f"wgrad_forward_stash.{name} must be a torch.Tensor, " - f"got {type(tensor).__name__}" - ) - _validate_strided(f"wgrad_forward_stash.{name}", tensor) - if tensor.device != device: - raise ValueError( - f"wgrad_forward_stash.{name} must be on {device}, " - f"got {tensor.device}" - ) - - e4m3_dtype = _require_torch_dtype("float8_e4m3fn") - e8m0_dtype = _require_torch_dtype("float8_e8m0fnu") - if stash.fc1_a.dtype is not e4m3_dtype: - raise TypeError( - "wgrad_forward_stash.fc1_a must have dtype " - f"{e4m3_dtype}, got {stash.fc1_a.dtype}" - ) - if stash.fc1_sfa.dtype is not e8m0_dtype: - raise TypeError( - "wgrad_forward_stash.fc1_sfa must have dtype " - f"{e8m0_dtype}, got {stash.fc1_sfa.dtype}" - ) - - expert_shape = (config.experts_per_rank,) - for name, tensor in ( - ("expert_offsets", stash.expert_offsets), - ("valid_route_counts", stash.valid_route_counts), - ): - if tuple(tensor.shape) != expert_shape: - raise ValueError( - f"wgrad_forward_stash.{name} shape must be {expert_shape}, " - f"got {tuple(tensor.shape)}" - ) - if tensor.dtype is not torch.int32: - raise TypeError( - f"wgrad_forward_stash.{name} must have dtype torch.int32, " - f"got {tensor.dtype}" - ) - - if tuple(stash.route_metadata.shape) != tuple(route_metadata.shape): - raise ValueError( - "wgrad_forward_stash.route_metadata shape must match " - "route_metadata" - ) - if stash.route_metadata.dtype is not torch.int32: - raise TypeError( - "wgrad_forward_stash.route_metadata must have dtype torch.int32, " - f"got {stash.route_metadata.dtype}" - ) - if not torch.equal(stash.route_metadata, route_metadata): - raise ValueError( - "wgrad_forward_stash route identity does not match route_metadata" - ) - - offsets = [int(value) for value in stash.expert_offsets.cpu().tolist()] - counts = [int(value) for value in stash.valid_route_counts.cpu().tolist()] - previous = 0 - for expert, (offset, count) in enumerate(zip(offsets, counts)): - padded_routes = offset - previous - if offset < previous: - raise ValueError( - "wgrad_forward_stash.expert_offsets must be non-decreasing" - ) - if count < 0 or count > padded_routes: - raise ValueError( - "wgrad_forward_stash.valid_route_counts must fit each " - f"expert's padded range; expert {expert} has count={count} " - f"and capacity={padded_routes}" - ) - expected_padded_routes = _round_up( - count, - config.token_padding_size, - ) - if padded_routes != expected_padded_routes: - raise ValueError( - "wgrad_forward_stash expert ranges must use the canonical " - f"{config.token_padding_size}-row padding; expert {expert} " - f"has capacity={padded_routes}, expected=" - f"{expected_padded_routes}" - ) - previous = offset - - padded_route_count = offsets[-1] if offsets else 0 - expected_fc1_a = (config.hidden_size, padded_route_count) - if tuple(stash.fc1_a.shape) != expected_fc1_a: - raise ValueError( - "wgrad_forward_stash.fc1_a shape must be " - f"{expected_fc1_a}, got {tuple(stash.fc1_a.shape)}" - ) - if padded_route_count % 32: - raise ValueError( - "wgrad_forward_stash padded route count must be divisible by 32" - ) - expected_fc1_sfa = ( - _round_up(config.hidden_size, 128), - _round_up(padded_route_count // 32, 4), - ) - if tuple(stash.fc1_sfa.shape) != expected_fc1_sfa: - raise ValueError( - "wgrad_forward_stash.fc1_sfa shape must be " - f"{expected_fc1_sfa}, got {tuple(stash.fc1_sfa.shape)}" - ) - if padded_route_count and not stash.fc1_a.is_contiguous(): - raise ValueError( - "wgrad_forward_stash.fc1_a must use compact (K, 1) strides" - ) - if not stash.fc1_sfa.is_contiguous(): - raise ValueError( - "wgrad_forward_stash.fc1_sfa must be contiguous" - ) - for name, tensor, alignment in ( - ("fc1_a", stash.fc1_a, 16), - ("fc1_sfa", stash.fc1_sfa, 16), - ("expert_offsets", stash.expert_offsets, 4), - ("valid_route_counts", stash.valid_route_counts, 4), - ): - if tensor.data_ptr() % alignment: - raise ValueError( - f"wgrad_forward_stash.{name} must be " - f"{alignment}-byte aligned" - ) - - local_routes = int(route_metadata.shape[0]) - if sum(counts) != local_routes: - raise ValueError( - "wgrad_forward_stash.valid_route_counts must sum to the " - "route_metadata row count" - ) - if local_routes: - local_experts = route_metadata[:, 0].to(torch.int64) - if bool( + expected = ( + ( + "weights.forward_fc1", + weights.forward_fc1, ( - (local_experts < 0) - | (local_experts >= config.experts_per_rank) - ).any().item() - ): - raise ValueError( - "route_metadata contains out-of-range local expert ids" - ) - metadata_counts = torch.bincount( - local_experts, - minlength=config.experts_per_rank, - ) - expected_counts = stash.valid_route_counts.to(torch.int64) - if not torch.equal(metadata_counts, expected_counts): - raise ValueError( - "wgrad_forward_stash.valid_route_counts do not match " - "route_metadata" - ) - expected_experts = torch.repeat_interleave( - torch.arange( config.experts_per_rank, - dtype=torch.int64, - device=device, + config.hidden_size, + 2 * config.intermediate_size, ), - expected_counts, - output_size=local_routes, - ) - if not torch.equal(local_experts, expected_experts): - raise ValueError( - "route_metadata rows must be grouped by local expert" - ) - - src_ranks = route_metadata[:, 1] - src_tokens = route_metadata[:, 2] - src_slots = route_metadata[:, 3] - if bool(((src_ranks < 0) | (src_ranks >= config.ep_size)).any().item()): - raise ValueError("route_metadata contains out-of-range source ranks") - if bool((src_tokens < 0).any().item()): - raise ValueError("route_metadata contains negative source tokens") - if config.max_tokens_per_rank is not None and bool( - (src_tokens >= config.max_tokens_per_rank).any().item() - ): - raise ValueError("route_metadata contains out-of-range source tokens") - if bool(((src_slots < 0) | (src_slots >= config.top_k)).any().item()): - raise ValueError("route_metadata contains out-of-range source slots") - - -def validate_backward( - config: ForwardConfig, - grad_output: torch.Tensor, - fc1_weight: MoeTensor, - fc2_weight: MoeTensor, - topk_idx: torch.Tensor, - topk_weights: torch.Tensor, - fc1_c: torch.Tensor, - route_metadata: torch.Tensor, - *, - wgrad_forward_stash: MoeEpWgradForwardStash | None = None, -) -> ValidatedBackwardRequest: - """Validate public backward semantics without importing a device backend.""" - - if not isinstance(grad_output, torch.Tensor): - raise TypeError( - "grad_output must be a torch.Tensor, " - f"got {type(grad_output).__name__}" - ) - _validate_strided("grad_output", grad_output) - if grad_output.ndim != 2 or grad_output.shape[1] != config.hidden_size: - raise ValueError( - f"grad_output shape must be (T, {config.hidden_size}), " - f"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}" - ) - token_count = int(grad_output.shape[0]) - expected_fc1 = ( - config.experts_per_rank, - config.hidden_size, - 2 * config.intermediate_size, - ) - expected_fc2 = ( - config.experts_per_rank, - config.intermediate_size, - config.hidden_size, - ) - _validate_tensor_representation("fc1_weight", fc1_weight, expected_fc1) - _validate_tensor_representation("fc2_weight", fc2_weight, expected_fc2) - - _validate_routes( - config, - token_count, - topk_idx, - topk_weights, - validate_expert_ids=True, + ), + ( + "weights.forward_fc2", + weights.forward_fc2, + ( + config.experts_per_rank, + config.intermediate_size, + config.hidden_size, + ), + ), + ( + "weights.backward_w2_transpose", + weights.backward_w2_transpose, + ( + config.experts_per_rank, + config.hidden_size, + config.intermediate_size, + ), + ), + ( + "weights.backward_w1_transpose", + weights.backward_w1_transpose, + ( + config.experts_per_rank, + 2 * config.intermediate_size, + config.hidden_size, + ), + ), ) - - if not isinstance(route_metadata, torch.Tensor): - raise TypeError( - "route_metadata must be a torch.Tensor, " - f"got {type(route_metadata).__name__}" - ) - _validate_strided("route_metadata", route_metadata) - if route_metadata.ndim != 2 or route_metadata.shape[1] != 4: - raise ValueError( - "route_metadata shape must be (local_routes, 4), " - f"got {tuple(route_metadata.shape)}" - ) - if route_metadata.dtype is not torch.int32: - raise TypeError( - "route_metadata must have dtype torch.int32, " - f"got {route_metadata.dtype}" - ) - local_routes = int(route_metadata.shape[0]) - expected_fc1_c = (local_routes, 2 * config.intermediate_size) - if not isinstance(fc1_c, torch.Tensor): - raise TypeError( - f"fc1_c must be a torch.Tensor, got {type(fc1_c).__name__}" - ) - _validate_strided("fc1_c", fc1_c) - if tuple(fc1_c.shape) != expected_fc1_c: - raise ValueError( - f"fc1_c shape must be {expected_fc1_c}, got {tuple(fc1_c.shape)}" - ) - if fc1_c.dtype is not torch.bfloat16: - raise TypeError( - f"fc1_c must have dtype torch.bfloat16, got {fc1_c.dtype}" - ) - - device = grad_output.device - for name, tensor in ( - ("fc1_weight", fc1_weight), - ("fc2_weight", fc2_weight), - ("topk_idx", topk_idx), - ("topk_weights", topk_weights), - ("grad_output", grad_output), - ("fc1_c", fc1_c), - ("route_metadata", route_metadata), - ): - tensor_device = _tensor_device(tensor) - if tensor_device != device: + for name, tensor, shape in expected: + _validate_tensor_representation(name, tensor, shape) + if not isinstance(tensor, BlockScaledTensor): + raise TypeError( + f"{name} must be an MXFP8 BlockScaledTensor for " + "fixed training resources" + ) + if tensor.format is not MoeFormat.MXFP8: + raise NotImplementedError( + f"{name} must use format='mxfp8', got {tensor.format.value!r}" + ) + if not tensor.data.is_contiguous() or not tensor.scale.is_contiguous(): raise ValueError( - f"{name} must be on {device}, got {tensor_device}" + f"{name} data and scale must be contiguous for fixed " + "training weight binding" ) - - if config.backward_wgrad_mode == "operands": - _validate_wgrad_forward_stash( - config, - wgrad_forward_stash, - route_metadata, - device, - ) - elif wgrad_forward_stash is not None: - raise ValueError( - "wgrad_forward_stash is only accepted when " - "backward_wgrad_mode='operands'" - ) - - return ValidatedBackwardRequest( - config=config, - grad_output=grad_output, - fc1_weight=fc1_weight, - fc2_weight=fc2_weight, - topk_idx=topk_idx, - topk_weights=topk_weights, - fc1_c=fc1_c, - route_metadata=route_metadata, - token_count=token_count, - local_routes=local_routes, - device=device, - wgrad_forward_stash=wgrad_forward_stash, - ) + device = weights.forward_fc1.device + for name, tensor, _shape in expected[1:]: + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}") + return device -__all__ = ["validate_backward", "validate_forward"] +__all__ = ["validate_forward", "validate_training_weights"] diff --git a/python/cudnn/moe_ep/api.py b/python/cudnn/moe_ep/api.py index cb6f5674d..03273859a 100644 --- a/python/cudnn/moe_ep/api.py +++ b/python/cudnn/moe_ep/api.py @@ -15,22 +15,24 @@ import threading import warnings from numbers import Real -from typing import Literal, Optional, Tuple, Union +from typing import Optional, Union import torch import torch.distributed as dist -from ._contracts import ForwardConfig, ValidatedBackwardRequest +from ._contracts import ForwardConfig from ._tuning import MoeEpTuningConfig from ._types import ( BlockScaledTensor, - MoeEpWgradForwardStash, - MoeEpWgradOperands, + MoeEpExecutionLane, + MoeEpTrainingResources, + MoeEpTrainingSlot, + MoeEpTrainingWeights, MoeFormat, MoeTensor, parse_format as _parse_format, ) -from ._validation import validate_backward, validate_forward +from ._validation import validate_forward, validate_training_weights def _resolve_ep_topology( ep_group: Optional[dist.ProcessGroup], @@ -62,6 +64,26 @@ def _resolve_ep_topology( return ep_size, ep_rank, ep_global_ranks +def _validate_training_assert_capability(config: ForwardConfig) -> None: + """Fail before allocation when graph error-mode primitives are unavailable.""" + + if config.drop_on_overflow: + return + if not callable(getattr(torch, "_assert_async", None)): + raise RuntimeError( + "drop_on_overflow=False training resources require callable " + "torch._assert_async before CUDA Graph capture" + ) + if config.ep_size <= 1: + return + backend = dist.get_backend(config.ep_group) + if backend != dist.Backend.NCCL and str(backend).lower() != "nccl": + raise NotImplementedError( + "drop_on_overflow=False EP2+ training resources require an NCCL " + "process group for the captured scalar global overflow OR" + ) + + class MoeEp: """Fused SwiGLU MoE operator with contiguous expert parallel sharding. @@ -77,24 +99,10 @@ class MoeEp: ``apply_topk_in_fc1=True``. Native NVFP4 operands and NVFP4 combine/output are not executable. - 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. - - With ``backward_wgrad_mode="operands"``, ``generate_c=True``, - ``token_padding_size=256``, and ``sf_padding_size=128`` are required. - Forward additionally returns a caller-owned - :class:`MoeEpWgradForwardStash`; backward accepts that exact routed-call - stash by keyword and additionally returns :class:`MoeEpWgradOperands`. - This opt-in path is available under the Rubin MXFP8 backward capability - gates documented below. + ``__call__`` is the inference-only forward surface. Training uses + :meth:`prepare_training_resources`; the returned fixed-slot resource handle + provides ordinary/capturable ``forward`` and ``backward`` methods without + compact host-visible stashes. The backend is created lazily on the first supported forward call. Valid combinations outside the current backend capability matrix fail explicitly @@ -118,8 +126,6 @@ def __init__( combine_format: Union[MoeFormat, str] = MoeFormat.BF16, apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, - generate_c: bool = False, - backward_wgrad_mode: Literal["none", "operands"] = "none", token_padding_size: int = 128, sf_padding_size: int = 128, tuning: Optional[MoeEpTuningConfig] = None, @@ -155,17 +161,6 @@ def __init__( raise ValueError("drop_on_overflow must be a bool") if not isinstance(apply_topk_in_fc1, bool): raise ValueError("apply_topk_in_fc1 must be a bool") - if not isinstance(generate_c, bool): - raise ValueError("generate_c must be a bool") - if backward_wgrad_mode not in ("none", "operands"): - raise ValueError( - "backward_wgrad_mode must be 'none' or 'operands', " - f"got {backward_wgrad_mode!r}" - ) - if backward_wgrad_mode == "operands" and not generate_c: - raise ValueError( - "backward_wgrad_mode='operands' requires generate_c=True" - ) for name, value in ( ("token_padding_size", token_padding_size), ("sf_padding_size", sf_padding_size), @@ -174,16 +169,6 @@ def __init__( raise ValueError( f"{name} must be a positive integer, got {value!r}" ) - if backward_wgrad_mode == "operands" and token_padding_size != 256: - raise ValueError( - "backward_wgrad_mode='operands' requires " - "token_padding_size=256" - ) - if backward_wgrad_mode == "operands" and sf_padding_size != 128: - raise ValueError( - "backward_wgrad_mode='operands' requires " - "sf_padding_size=128" - ) if sf_padding_size % 128: raise ValueError( "sf_padding_size must be a positive multiple of 128, " @@ -226,8 +211,6 @@ def __init__( self.combine_format = _parse_format(combine_format) self.apply_topk_in_fc1 = apply_topk_in_fc1 self.gate_up_clamp = None if gate_up_clamp is None else abs(gate_up_clamp) - self.generate_c = generate_c - self.backward_wgrad_mode = backward_wgrad_mode self.token_padding_size = token_padding_size self.sf_padding_size = sf_padding_size self.tuning = MoeEpTuningConfig() if tuning is None else tuning @@ -266,16 +249,18 @@ def __init__( combine_format=self.combine_format.value, apply_topk_in_fc1=self.apply_topk_in_fc1, gate_up_clamp=self.gate_up_clamp, - generate_c=self.generate_c, + generate_c=False, token_padding_size=self.token_padding_size, sf_padding_size=self.sf_padding_size, tuning=self.tuning, - backward_wgrad_mode=self.backward_wgrad_mode, + backward_wgrad_mode="none", ) self._forward_backend = None self._forward_backend_device = None self._validated_topk_idx = None self._validated_topk_version = None + self._operator_token = object() + self._training_resources: MoeEpTrainingResources | None = None self._closed = False @staticmethod @@ -287,7 +272,7 @@ def _tensor_version(tensor: torch.Tensor) -> int | None: except RuntimeError: return None - def _get_backend(self, request, *, backward: bool): + def _get_backend(self, request): """Create and cache the private backend on first supported use.""" with self._lifecycle_lock: @@ -305,10 +290,7 @@ def _get_backend(self, request, *, backward: bool): ) _backend.validate_config(self._forward_config) - if backward: - _backend.validate_backward_request(request) - else: - _backend.validate_request(request) + _backend.validate_request(request) if self._forward_backend is None: self._forward_backend = _backend.create_backend( @@ -318,26 +300,6 @@ def _get_backend(self, request, *, backward: bool): self._forward_backend_device = request.device return self._forward_backend - def _count_local_routes(self, request: ValidatedBackwardRequest) -> int: - """Number of valid routes this rank's experts receive. - - The request already passed expert-id validation. Data-dependent: - single-rank counts locally, while EP exchanges per-rank route counts - (the same exchange the device dispatch performs). - """ - - flat = request.topk_idx.reshape(-1).to(torch.int64) - expert = flat[flat != -1] - 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, @@ -345,26 +307,15 @@ def __call__( fc2_weight: MoeTensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor, - ) -> Union[ - MoeTensor, - Tuple[MoeTensor, torch.Tensor, torch.Tensor], - Tuple[ - MoeTensor, - torch.Tensor, - torch.Tensor, - MoeEpWgradForwardStash, - ], - ]: + ) -> MoeTensor: """Validate and dispatch one fused MoE+EP forward call. 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 constructed with ``generate_c=True``. In - ``backward_wgrad_mode="operands"``, the latter tuple has a fourth - ``MoeEpWgradForwardStash`` item. + Training callers must use :meth:`prepare_training_resources` and the + returned fixed-slot resource handle. """ with self._lifecycle_lock: @@ -395,7 +346,7 @@ def __call__( else: self._validated_topk_idx = None self._validated_topk_version = None - return self._get_backend(request, backward=False).forward(request) + return self._get_backend(request).forward(request) def warmup( self, @@ -432,65 +383,82 @@ def warmup( if device.type == "cuda": torch.cuda.synchronize(device) - def backward( + def prepare_training_resources( self, - grad_output: torch.Tensor, - fc1_weight: MoeTensor, - fc2_weight: MoeTensor, - topk_idx: torch.Tensor, - topk_weights: torch.Tensor, - fc1_c: torch.Tensor, - route_metadata: torch.Tensor, + weights: MoeEpTrainingWeights, *, - wgrad_forward_stash: Optional[MoeEpWgradForwardStash] = None, - ) -> Union[ - Tuple[torch.Tensor, torch.Tensor], - Tuple[torch.Tensor, torch.Tensor, MoeEpWgradOperands], - ]: - """Validate and dispatch one device MoE backward call. - - Requires ``generate_c=True``; consumes the forward stash - (``fc1_c``, ``route_metadata``) plus the re-supplied weights and - routing inputs. Returns float32 - ``(grad_activation, grad_topk_weights)``. In - ``backward_wgrad_mode="operands"``, ``wgrad_forward_stash`` is - required and the return tuple has a third ``MoeEpWgradOperands`` item. - The Rubin MXFP8 device path supports BF16/MXFP8 combine and BF16 output - for any positive EP size under its documented capability gates. - Backward has hardware acceptance coverage at EP1/EP2/EP4. Forward and - backward both quantize each FP32 route accumulator directly to MXFP8 - before top-k reduction. + slot_count: int = 2, + lane_count: int = 1, + ) -> MoeEpTrainingResources: + """Bind MXFP8 weights and allocate fixed-capacity training resources. + + This collective preparation must run on every EP rank before CUDA + Graph capture. The returned handle owns persistent microbatch slots + and mutable per-stream execution lanes; closing the operator also + closes the handle. A closed handle cannot be replaced on this + operator; create a new ``MoeEp`` instance to bind new weight storage. """ with self._lifecycle_lock: if self._closed: raise RuntimeError("MoeEp is closed") - if not self.generate_c: + for name, value in ( + ("slot_count", slot_count), + ("lane_count", lane_count), + ): + if ( + isinstance(value, bool) + or not isinstance(value, int) + or value <= 0 + ): + raise ValueError( + f"{name} must be a positive integer, got {value!r}" + ) + if self._training_resources is not None: + if not self._training_resources.closed: + raise RuntimeError("MoeEp training resources already exist") raise RuntimeError( - "backward requires the operator to be constructed with " - "generate_c=True" + "MoeEp training resources were closed; create a new " + "MoeEp instance before preparing replacement weights" ) - backward_request = validate_backward( + device = validate_training_weights( self._forward_config, - grad_output, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights, - fc1_c, - route_metadata, - wgrad_forward_stash=wgrad_forward_stash, + weights, ) - local_routes = self._count_local_routes(backward_request) - if backward_request.local_routes != local_routes: + _validate_training_assert_capability(self._forward_config) + from . import _backend + + _backend.validate_config(self._forward_config) + if ( + self._forward_backend is not None + and device != self._forward_backend_device + ): raise ValueError( - "route_metadata row count must match the routes received " - "from the re-supplied topk_idx" + f"MoeEp backend is bound to " + f"{self._forward_backend_device}; got {device}" ) - return self._get_backend(backward_request, backward=True).backward( - backward_request + if self._forward_backend is None: + self._forward_backend = _backend.create_backend( + self._forward_config, + device, + ) + self._forward_backend_device = device + owner = self._forward_backend.prepare_training_resources( + weights, + slot_count=slot_count, + lane_count=lane_count, + ) + resources = MoeEpTrainingResources( + owner=owner, + operator_token=self._operator_token, + weights=weights, + slot_count=slot_count, + lane_count=lane_count, + device=device, ) + self._training_resources = resources + return resources def close(self) -> None: """Release compiled-backend instance resources; idempotent.""" @@ -506,6 +474,9 @@ def close(self) -> None: self._forward_backend_device = None self._validated_topk_idx = None self._validated_topk_version = None + if self._training_resources is not None: + self._training_resources.close() + self._training_resources = None self._closed = True def __enter__(self) -> "MoeEp": @@ -539,8 +510,10 @@ def __del__(self) -> None: __all__ = [ "BlockScaledTensor", "MoeEp", - "MoeEpWgradForwardStash", - "MoeEpWgradOperands", + "MoeEpExecutionLane", + "MoeEpTrainingResources", + "MoeEpTrainingSlot", + "MoeEpTrainingWeights", "MoeFormat", "MoeTensor", ] diff --git a/test/python/moe_ep/moe_ep_backward_support.py b/test/python/moe_ep/moe_ep_backward_support.py deleted file mode 100644 index 390daa42d..000000000 --- a/test/python/moe_ep/moe_ep_backward_support.py +++ /dev/null @@ -1,212 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Shared reference helpers for MoE EP backward tests.""" - -from __future__ import annotations - -import torch - -from moe_ep.moe_ep_forward_support import _reference_args -from moe_ep.moe_ep_reference import MoeEpReference - -__all__ = [ - "_assert_backward_matches", - "_dense_wgrads_from_operands", - "_expected_backward", - "_grad_output", - "_reference_backward", -] - - -_BACKWARD_CLOSE_KWARGS = ( - {"rtol": 0.15, "atol": 0.125}, # grad_activation is BF16-rounded. - {"rtol": 0.15, "atol": 0.125}, # router-weight gradient. -) - - -def _round_up(value: int, multiple: int) -> int: - return (value + multiple - 1) // multiple * multiple - - -def _unpack_wgrad_scale_part( - packed: torch.Tensor, - rows: int, - columns: int, -) -> torch.Tensor: - """Invert grouped-wgrad's 128x4 scale-atom swizzle.""" - - padded_rows = _round_up(rows, 128) - padded_columns = _round_up(columns, 4) - row_atoms = padded_rows // 128 - column_atoms = padded_columns // 4 - atom_count = row_atoms * column_atoms - expected = padded_rows * padded_columns - if packed.numel() != expected: - raise ValueError( - f"packed scale part has {packed.numel()} bytes, expected {expected}" - ) - blocked = ( - packed.reshape(atom_count, 32, 4, 4) - .transpose(1, 2) - .reshape(row_atoms, column_atoms, 128, 4) - .permute(0, 2, 1, 3) - .reshape(padded_rows, padded_columns) - ) - return blocked[:rows, :columns].view(torch.float8_e8m0fnu).float() - - -def _dequantize_wgrad_operand( - data: torch.Tensor, - scales: torch.Tensor, - expert_offsets: torch.Tensor, - *, - k_dim: int, -) -> torch.Tensor: - """Decode one public grouped-wgrad operand without launching a GEMM.""" - - if data.ndim != 2 or k_dim not in (0, 1): - raise ValueError("wgrad operand must be rank 2 with k_dim 0 or 1") - non_k = int(data.shape[1 - k_dim]) - padded_non_k = _round_up(non_k, 128) - flat_scales = scales.view(torch.uint8).reshape(-1) - output = torch.empty(data.shape, dtype=torch.float32, device=data.device) - ends = [int(value) for value in expert_offsets.detach().cpu().tolist()] - previous = 0 - scale_byte_offset = 0 - for end in ends: - extent = end - previous - scale_columns = _round_up(extent // 32, 4) - scale_byte_count = padded_non_k * scale_columns - part = flat_scales.narrow( - 0, - scale_byte_offset, - scale_byte_count, - ) - logical_scale = _unpack_wgrad_scale_part( - part, - non_k, - extent // 32, - ) - if k_dim == 1: - expanded_scale = logical_scale.repeat_interleave(32, dim=1) - output[:, previous:end] = ( - data[:, previous:end].float() * expanded_scale - ) - else: - expanded_scale = logical_scale.repeat_interleave( - 32, - dim=1, - ).transpose(0, 1) - output[previous:end, :] = ( - data[previous:end, :].float() * expanded_scale - ) - previous = end - scale_byte_offset += scale_byte_count - if previous != data.shape[k_dim]: - raise ValueError("expert offsets do not cover the operand K dimension") - if scale_byte_offset != flat_scales.numel(): - raise ValueError("expert offsets do not cover the scale tensor") - return output - - -def _dense_wgrads_from_operands(operands): - """Reference grouped matmuls over the exported operand ABI.""" - - fc1_a = _dequantize_wgrad_operand( - operands.fc1_a, - operands.fc1_sfa, - operands.expert_offsets, - k_dim=1, - ) - fc1_b = _dequantize_wgrad_operand( - operands.fc1_b, - operands.fc1_sfb, - operands.expert_offsets, - k_dim=0, - ) - fc2_a = _dequantize_wgrad_operand( - operands.fc2_a, - operands.fc2_sfa, - operands.expert_offsets, - k_dim=1, - ) - fc2_b = _dequantize_wgrad_operand( - operands.fc2_b, - operands.fc2_sfb, - operands.expert_offsets, - k_dim=0, - ) - fc1_parts = [] - fc2_parts = [] - previous = 0 - for end_value in operands.expert_offsets.detach().cpu().tolist(): - end = int(end_value) - fc1_parts.append( - fc1_a[:, previous:end] @ fc1_b[previous:end, :] - ) - fc2_parts.append( - fc2_a[:, previous:end] @ fc2_b[previous:end, :] - ) - previous = end - return torch.stack(fc1_parts), torch.stack(fc2_parts) - - -def _reference_backward(config) -> MoeEpReference: - options = dict(config) - options.pop("tuning", None) - options.pop("sf_padding_size", None) - options["intermediate_format"] = "mxfp8" - options["backward_operand_format"] = "mxfp8" - return MoeEpReference(**options) - - -def _grad_output( - device: torch.device, - token_count: int, - *, - seed: int, -) -> torch.Tensor: - generator = torch.Generator(device=device).manual_seed(seed) - return ( - torch.randn( - token_count, - 128, - generator=generator, - dtype=torch.float32, - device=device, - ) - / 8 - ) - - -def _expected_backward(reference, grad_output, args, stash): - return reference.backward( - grad_output, - *_reference_args(args)[1:], - *stash, - ) - - -def _assert_backward_matches(actual, expected, topk_idx) -> None: - assert len(actual) == len(expected) == 2 - for name, gradient, reference, close_kwargs in zip( - ("grad_activation", "grad_topk_weights"), - actual, - expected, - _BACKWARD_CLOSE_KWARGS, - ): - assert gradient.shape == reference.shape - assert gradient.dtype == torch.float32 - assert torch.isfinite(gradient).all() - torch.testing.assert_close( - gradient, - reference, - msg=lambda default, name=name: ( - f"{name} does not match the backward reference\n{default}" - ), - **close_kwargs, - ) - - dropped = topk_idx == -1 - assert actual[1][dropped].eq(0).all() diff --git a/test/python/moe_ep/moe_ep_distributed_workers.py b/test/python/moe_ep/moe_ep_distributed_workers.py index 9c8e7b5a0..0b0442b73 100644 --- a/test/python/moe_ep/moe_ep_distributed_workers.py +++ b/test/python/moe_ep/moe_ep_distributed_workers.py @@ -10,29 +10,27 @@ import torch import torch.distributed as dist -from moe_ep.moe_ep_backward_support import ( +from moe_ep.moe_ep_test_support import ( _assert_backward_matches, - _dense_wgrads_from_operands, - _expected_backward, - _grad_output, - _reference_backward, -) -from moe_ep.moe_ep_forward_support import ( _assert_matches_reference, + _assert_wgrads_match_reference, + _fixed_training_reference, + _fixed_training_weights, _forward_config, + _grad_output, _output_as_float, - _reference_args, _reference_forward, + make_distributed_forward_inputs, + quantize_mxfp8, ) -from moe_ep.moe_ep_test_data import make_distributed_forward_inputs __all__ = [ - "_distributed_backward_worker", + "_distributed_backward_reference_worker", "_distributed_output_worker", + "_distributed_subgroup_backward_reference_worker", "_distributed_subgroup_output_worker", - "_distributed_wgrad_worker", + "_run_backward_reference_case", "_run_forward_output_case", - "_run_wgrad_operand_case", ] @@ -45,7 +43,7 @@ def _run_forward_output_case( combine_format: str = "bf16", expected_global_ranks: tuple[int, ...] | None = None, ) -> None: - """Run forward parity and dropped-route checks on one initialized EP group.""" + """Run inference-forward parity and dropped-route checks.""" from cudnn import MoeEp @@ -59,20 +57,29 @@ def _run_forward_output_case( expected = _reference_forward(args, **config) op = MoeEp(**config) try: - assert op.ep_rank == ep_rank - if expected_global_ranks is not None: - assert op.ep_global_ranks == expected_global_ranks - actual = op(*args) + actual_snapshot = _output_as_float(actual).clone() torch.cuda.synchronize(device) - _assert_matches_reference(actual, expected) args[3].fill_(-1) dropped = op(*args) + dropped_snapshot = _output_as_float(dropped).clone() torch.cuda.synchronize(device) - assert _output_as_float(dropped).eq(0).all() dist.barrier(group=ep_group) + assertion_error = None + try: + assert op.ep_rank == ep_rank + if expected_global_ranks is not None: + assert op.ep_global_ranks == expected_global_ranks + _assert_matches_reference(actual_snapshot, expected) + assert dropped_snapshot.eq(0).all() + except BaseException as error: + assertion_error = error + dist.barrier(group=ep_group) + if assertion_error is not None: + raise assertion_error + op.close() op = None dist.barrier(group=ep_group) @@ -138,10 +145,10 @@ def _distributed_subgroup_output_worker( ep_group = subgroups[subgroup_index] ep_rank = dist.get_rank(ep_group) ep_size = dist.get_world_size(ep_group) - assert tuple( + actual_global_ranks = tuple( dist.get_global_rank(ep_group, group_rank) for group_rank in range(ep_size) - ) == subgroup_memberships[subgroup_index] + ) _run_forward_output_case( device=device, @@ -151,338 +158,264 @@ def _distributed_subgroup_output_worker( expected_global_ranks=subgroup_memberships[subgroup_index], ) dist.barrier() + assert actual_global_ranks == subgroup_memberships[subgroup_index] finally: if dist.is_initialized(): dist.destroy_process_group() -def _distributed_backward_worker( - rank: int, - world_size: int, - init_file: str, - combine_format: str, - gate_up_clamp: float | None = None, -) -> None: - """Run distributed forward stashing and backward reference parity.""" - - from cudnn import MoeEp - - device = torch.device("cuda", rank) - torch.cuda.set_device(device) - dist.init_process_group( - backend="nccl", - init_method=f"file://{init_file}", - rank=rank, - world_size=world_size, - device_id=device, - timeout=timedelta(seconds=300), - ) - op = None - try: - args = make_distributed_forward_inputs(rank, world_size, device) - config = _forward_config( - num_experts=2 * world_size, - ep_group=dist.group.WORLD, - max_tokens_per_rank=8, - generate_c=True, - combine_format=combine_format, - gate_up_clamp=gate_up_clamp, - ) - reference = _reference_backward(config) - grad_output = _grad_output( - device, - args[3].shape[0], - seed=20260820 + rank, - ) - op = MoeEp(**config) - _, fc1_c, route_metadata = op(*args) - stash = (fc1_c, route_metadata) - expected = _expected_backward(reference, grad_output, args, stash) - - first = op.backward(grad_output, *args[1:], *stash) - second = op.backward(grad_output, *args[1:], *stash) - torch.cuda.synchronize(device) - - # Complete collective work before local assertions. A failure before - # this barrier would leave peer ranks waiting for process-group timeout. - dist.barrier() - _assert_backward_matches(first, expected, args[3]) - _assert_backward_matches(second, expected, args[3]) - - op.close() - op = None - finally: - if op is not None: - op.close() - if dist.is_initialized(): - dist.destroy_process_group() - - -def _make_wgrad_inputs( - rank: int, - world_size: int, +def _make_distributed_backward_inputs( + ep_rank: int, + ep_size: int, device: torch.device, ): - """Build routes with negative weights, drops, and one empty local expert.""" - - args = list(make_distributed_forward_inputs(rank, world_size, device)) - token_count = args[3].shape[0] - local_expert = 2 * rank - remote_expert = 2 * ((rank + 1) % world_size) - topk_idx = torch.full( - (token_count, 2), - -1, - dtype=torch.int32, - device=device, - ) - topk_weights = torch.empty( - (token_count, 2), - dtype=torch.bfloat16, - device=device, - ) - weight_rows = ( - (0.5, -0.25), - (7.0, -1.25), - (1.5, -9.0), - ) - for token in range(token_count): - pattern = token % 3 - if pattern == 0: - topk_idx[token] = torch.tensor( - (local_expert, remote_expert), - dtype=torch.int32, - device=device, - ) - elif pattern == 1: - topk_idx[token, 1] = local_expert - else: - topk_idx[token, 0] = remote_expert - topk_weights[token] = torch.tensor( - weight_rows[pattern], - dtype=torch.bfloat16, + """Build a minimal local/remote/drop case with one empty local expert.""" + + generator = torch.Generator(device=device).manual_seed(20260828 + ep_rank) + local_experts, token_count, hidden, intermediate = 2, 2, 128, 256 + activation = ( + torch.randn( + token_count, + hidden, + generator=generator, device=device, ) - args[3] = topk_idx - args[4] = topk_weights - return tuple(args) - - -def _source_route_expert( - source_rank: int, - token: int, - slot: int, - world_size: int, -) -> int: - pattern = token % 3 - if pattern == 0: - return ( - 2 * source_rank - if slot == 0 - else 2 * ((source_rank + 1) % world_size) + / 4 + ).to(torch.bfloat16) + fc1_weight = quantize_mxfp8( + torch.randn( + local_experts, + hidden, + 2 * intermediate, + generator=generator, + device=device, ) - if pattern == 1: - return 2 * source_rank if slot == 1 else -1 - return 2 * ((source_rank + 1) % world_size) if slot == 0 else -1 - - -def _assert_local_operand_metadata( - operands, - reference_operands, - *, - rank: int, - world_size: int, -) -> None: - local_experts = 2 - assert operands.expert_offsets.shape == (local_experts,) - assert operands.valid_route_counts.shape == (local_experts,) - assert torch.equal( - operands.route_metadata, - reference_operands.route_metadata, + / 8, + axis=1, ) - assert torch.equal( - operands.valid_route_counts, - reference_operands.valid_route_counts, + fc2_weight = quantize_mxfp8( + torch.randn( + local_experts, + intermediate, + hidden, + generator=generator, + device=device, + ) + / 8, + axis=1, ) - assert torch.equal( - operands.expert_offsets, - reference_operands.expert_offsets, + local_expert = ep_rank * local_experts + remote_expert = ((ep_rank + 1) % ep_size) * local_experts + topk_idx = torch.tensor( + [[local_expert, remote_expert], [-1, local_expert]], + dtype=torch.int32, + device=device, ) - - metadata = operands.route_metadata - if metadata.numel(): - assert metadata[:, 0].ge(0).all() - assert metadata[:, 0].lt(local_experts).all() - counts = torch.bincount( - metadata[:, 0].to(torch.int64), - minlength=local_experts, - ).to(torch.int32) - assert torch.equal(operands.valid_route_counts, counts) - assert counts[0] > 0 - assert counts[1] == 0 - - expected_offsets = [] - padded_end = 0 - for count in counts.tolist(): - padded_end += ((count + 255) // 256) * 256 - expected_offsets.append(padded_end) - assert operands.expert_offsets.tolist() == expected_offsets - - for local_expert, source_rank, token, slot in metadata.tolist(): - global_expert = _source_route_expert( - source_rank, - token, - slot, - world_size, - ) - assert global_expert != -1 - assert global_expert // local_experts == rank - assert global_expert % local_experts == local_expert + topk_weights = torch.tensor( + [[0.625, 0.375], [0.0, 1.0]], + dtype=torch.float32, + device=device, + ) + grad_output = _grad_output( + device, + token_count, + seed=20260901 + ep_rank, + ) + return ( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ), grad_output -def _run_grouped_wgrad( - operands, - prefix: str, - *, - wgrad_tensor=None, - accumulate_on_output: bool = False, -): - import cudnn - - return cudnn.grouped_gemm_wgrad_wrapper_sm100( - a_tensor=getattr(operands, f"{prefix}_a"), - b_tensor=getattr(operands, f"{prefix}_b"), - sfa_tensor=getattr(operands, f"{prefix}_sfa"), - sfb_tensor=getattr(operands, f"{prefix}_sfb"), - offsets_tensor=operands.expert_offsets, - output_mode="dense", - wgrad_tensor=wgrad_tensor, - wgrad_dtype=torch.bfloat16, - acc_dtype=torch.float32, - mma_tiler_mn=(128, 128), - cluster_shape_mn=(1, 1), - sf_vec_size=32, - accumulate_on_output=accumulate_on_output, - )["wgrad_tensor"] - - -def _run_wgrad_operand_case( +def _run_backward_reference_case( *, device: torch.device, ep_group, - rank: int, - world_size: int, + ep_rank: int, + ep_size: int, + combine_format: str = "bf16", + gate_up_clamp: float | None = None, + expected_global_ranks: tuple[int, ...] | None = None, ) -> None: - """Exercise production FC1/FC2 operands through grouped wgrad.""" + """Run fixed-resource training after the independent distributed oracle.""" - from cudnn import MoeEp, MoeEpWgradOperands + from cudnn import MoeEp - args = _make_wgrad_inputs(rank, world_size, device) - config = _forward_config( - num_experts=2 * world_size, + args, grad_output = _make_distributed_backward_inputs( + ep_rank, + ep_size, + device, + ) + num_experts = 2 * ep_size + max_recv_size_per_rank = 3 + + # Finish all collective reference work, including dense local dW, before + # constructing or launching the production operator. + expected = _fixed_training_reference( + args, + grad_output, + combine_format=combine_format, + gate_up_clamp=gate_up_clamp, ep_group=ep_group, - max_tokens_per_rank=8, - generate_c=True, - backward_wgrad_mode="operands", - token_padding_size=256, - sf_padding_size=128, + num_experts=num_experts, + max_recv_size_per_rank=max_recv_size_per_rank, + drop_on_overflow=True, ) - reference = _reference_backward(config) - grad_output = _grad_output( - device, - args[3].shape[0], - seed=20260821 + rank, + expected_y, expected_dx, expected_dprob, expected_wgrads = expected + expected_dense_wgrads = expected_wgrads.dense_wgrads() + weights = _fixed_training_weights(args) + + op = MoeEp( + num_experts=num_experts, + hidden_size=128, + intermediate_size=256, + top_k=2, + ep_group=ep_group, + max_tokens_per_rank=args[0].shape[0], + max_recv_size_per_rank=max_recv_size_per_rank, + drop_on_overflow=True, + combine_format=combine_format, + gate_up_clamp=gate_up_clamp, ) - - with MoeEp(**config) as op: - output, fc1_c, route_metadata, forward_stash = op(*args) - ( - reference_output, - reference_fc1_c, - reference_metadata, - reference_stash, - ) = reference(*_reference_args(args)) - _assert_matches_reference(output, reference_output) - assert torch.equal(route_metadata, reference_metadata) - assert forward_stash.route_metadata is route_metadata - - backward = op.backward( - grad_output, - *args[1:], - fc1_c, - route_metadata, - wgrad_forward_stash=forward_stash, + try: + resources = op.prepare_training_resources( + weights, + slot_count=1, + lane_count=1, + ) + slot = resources.slots[0] + lane = resources.lanes[0] + resources.refresh_weights() + actual_y = resources.forward( + slot, + lane, + args[0], + args[3], + args[4], ) - reference_backward = reference.backward( + actual_dx, actual_dprob, actual_wgrads = resources.backward( + slot, + lane, grad_output, - *_reference_args(args)[1:], - reference_fc1_c, - reference_metadata, - wgrad_forward_stash=reference_stash, ) + overflow = resources.finalize_overflow((slot,), lane) + torch.cuda.synchronize(device) - _assert_backward_matches(backward[:2], reference_backward[:2], args[3]) - operands = backward[2] - reference_operands = reference_backward[2] - assert isinstance(operands, MoeEpWgradOperands) - assert operands.route_metadata is forward_stash.route_metadata - assert operands.expert_offsets is forward_stash.expert_offsets - assert operands.valid_route_counts is forward_stash.valid_route_counts - _assert_local_operand_metadata( - operands, - reference_operands, - rank=rank, - world_size=world_size, - ) + # No rank may enter a local assertion while a peer is still inside a + # collective kernel. A second barrier keeps cleanup aligned on failure. + dist.barrier(group=ep_group) + assertion_error = None + try: + assert op.ep_rank == ep_rank + assert op.ep_size == ep_size + if expected_global_ranks is not None: + assert op.ep_global_ranks == expected_global_ranks + assert overflow.eq(0).all() + assert args[3][0, 0] // 2 == ep_rank + assert args[3][0, 1] // 2 == (ep_rank + 1) % ep_size + assert args[3].eq(-1).any() + assert expected_wgrads.valid_route_counts[1].eq(0) + assert actual_wgrads.valid_route_counts[1].eq(0) + _assert_matches_reference(actual_y, expected_y) + _assert_backward_matches( + (actual_dx, actual_dprob), + (expected_dx, expected_dprob), + args[3], + ) + _assert_wgrads_match_reference( + actual_wgrads, + expected_wgrads, + expected_dense=expected_dense_wgrads, + ) + except BaseException as error: + assertion_error = error + dist.barrier(group=ep_group) + if assertion_error is not None: + raise assertion_error + finally: + op.close() - expected_wgrads = reference_operands.dense_wgrads() - decoded_wgrads = _dense_wgrads_from_operands(operands) - for actual, expected in zip(decoded_wgrads, expected_wgrads): - torch.testing.assert_close( - actual, - expected, - rtol=0.15, - atol=0.125, - ) - assert actual[1].eq(0).all() - - # The in-tree grouped-wgrad implementation is an SM100 kernel. Rubin - # validates the producer numerics above; the direct SM100 ABI test covers - # execution and accumulate_on_output with the same public bundle layout. - if torch.cuda.get_device_capability(device) != (10, 0): - return - - for prefix, expected in zip(("fc1", "fc2"), expected_wgrads): - actual = _run_grouped_wgrad(operands, prefix) - initial = torch.full_like(actual, 0.25) - accumulated = _run_grouped_wgrad( - operands, - prefix, - wgrad_tensor=initial, - accumulate_on_output=True, - ) - torch.cuda.synchronize(device) - assert accumulated is initial - torch.testing.assert_close( - actual.float(), - expected, - rtol=0.15, - atol=0.125, - ) - torch.testing.assert_close( - accumulated.float(), - expected + 0.25, - rtol=0.15, - atol=0.125, - ) - assert actual[1].eq(0).all() - assert accumulated[1].eq(0.25).all() +def _distributed_subgroup_backward_reference_worker( + global_rank: int, + global_world_size: int, + init_file: str, +) -> None: + """Run fixed-resource backward in two non-contiguous EP2 groups.""" + + device = torch.device("cuda", global_rank) + torch.cuda.set_device(device) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=global_rank, + world_size=global_world_size, + device_id=device, + timeout=timedelta(minutes=10), + ) + ep_group = None + try: + subgroup_memberships = ((0, 2), (1, 3)) + # Every WORLD rank must create every subgroup in the same order. + subgroups = [ + dist.new_group( + list(members), + backend="nccl", + timeout=timedelta(minutes=10), + ) + for members in subgroup_memberships + ] + subgroup_index = global_rank % 2 + expected_global_ranks = subgroup_memberships[subgroup_index] + ep_group = subgroups[subgroup_index] + ep_rank = dist.get_rank(ep_group) + ep_size = dist.get_world_size(ep_group) + actual_global_ranks = tuple( + dist.get_global_rank(ep_group, group_rank) + for group_rank in range(ep_size) + ) + assert ep_size == len(expected_global_ranks) + assert ep_rank == expected_global_ranks.index(global_rank) + assert actual_global_ranks == expected_global_ranks -def _distributed_wgrad_worker( + _run_backward_reference_case( + device=device, + ep_group=ep_group, + ep_rank=ep_rank, + ep_size=ep_size, + combine_format="bf16", + expected_global_ranks=expected_global_ranks, + ) + finally: + if dist.is_initialized(): + try: + # Keep both independent groups alive until all work is done, + # then collectively finalize the process-local runtime. + dist.barrier() + from cudnn.moe_ep._megamoe_backend._runtime import ( + get_runtime_manager, + ) + + get_runtime_manager().shutdown() + dist.barrier() + finally: + if ep_group is not None: + dist.destroy_process_group(ep_group) + dist.destroy_process_group() + + +def _distributed_backward_reference_worker( rank: int, world_size: int, init_file: str, + combine_format: str, + gate_up_clamp: float | None = None, ) -> None: + """Initialize one local rank and run distributed training parity.""" + device = torch.device("cuda", rank) torch.cuda.set_device(device) dist.init_process_group( @@ -491,16 +424,26 @@ def _distributed_wgrad_worker( rank=rank, world_size=world_size, device_id=device, - timeout=timedelta(seconds=600), + timeout=timedelta(minutes=10), ) try: - _run_wgrad_operand_case( + _run_backward_reference_case( device=device, ep_group=dist.group.WORLD, - rank=rank, - world_size=world_size, + ep_rank=rank, + ep_size=world_size, + combine_format=combine_format, + gate_up_clamp=gate_up_clamp, ) - dist.barrier() finally: if dist.is_initialized(): - dist.destroy_process_group() + try: + dist.barrier() + from cudnn.moe_ep._megamoe_backend._runtime import ( + get_runtime_manager, + ) + + get_runtime_manager().shutdown() + dist.barrier() + finally: + dist.destroy_process_group() diff --git a/test/python/moe_ep/moe_ep_forward_support.py b/test/python/moe_ep/moe_ep_forward_support.py deleted file mode 100644 index cd70583bf..000000000 --- a/test/python/moe_ep/moe_ep_forward_support.py +++ /dev/null @@ -1,319 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Shared support for MoE EP forward tests.""" - -from __future__ import annotations - -import pytest -import torch -import torch.distributed as dist -import torch.nn.functional as F - -from moe_ep.moe_ep_reference import ( - BlockScaledTensor as ReferenceBlockScaledTensor, - MoeEpReference, - MoeFormat, - forward_combine_round_trip, - quantize_blockwise, -) -from moe_ep.moe_ep_test_data import quantize_mxfp8 - -_DEFAULT_FORWARD_CONFIG = { - "num_experts": 2, - "hidden_size": 128, - "intermediate_size": 256, - "top_k": 2, - "max_tokens_per_rank": 5, - "apply_topk_in_fc1": True, - "combine_format": "bf16", - "output_format": "bf16", -} -_REFERENCE_CLOSE_KWARGS = {"rtol": 0.05, "atol": 0.0625} - -__all__ = [ - "_assert_matches_reference", - "_forward_config", - "_make_forward_case", - "_naive_reference", - "_output_as_float", - "_reference_forward", - "_replay_cuda_graph", - "_require_distributed_sm107", - "_sm107_device", - "_stress_backend_reuse", -] - - -def _forward_config(**overrides): - return {**_DEFAULT_FORWARD_CONFIG, **overrides} - - -def _output_as_float(output): - if isinstance(output, torch.Tensor): - return output.float() - return output.dequantize() - - -def _assert_matches_reference(actual, expected): - torch.testing.assert_close( - _output_as_float(actual), - _output_as_float(expected), - **_REFERENCE_CLOSE_KWARGS, - ) - - -def _naive_reference( - activation, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights, - *, - apply_topk_in_fc1, - clamp=None, - combine_format=MoeFormat.BF16, - intermediate_format=None, - apply_topk_after_combine=False, -): - 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, - device=activation.device, - ) - 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 - if intermediate_format is not None: - intermediate = quantize_blockwise( - intermediate, - intermediate_format, - ).dequantize() - result = intermediate @ fc2_weight[expert].float() - if not apply_topk_in_fc1 and not apply_topk_after_combine: - result = result * route_weight - result = forward_combine_round_trip(result, combine_format) - if not apply_topk_in_fc1 and apply_topk_after_combine: - result = result * route_weight - combine[token, slot] = result - return combine.sum(dim=1).to(torch.bfloat16) - - -def _as_reference_tensor(tensor): - if isinstance(tensor, torch.Tensor): - return tensor - return ReferenceBlockScaledTensor( - data=tensor.data, - scale=tensor.scale, - format=tensor.format.value, - logical_shape=tensor.logical_shape, - axis=tensor.axis, - ) - - -def _reference_args(args): - return ( - _as_reference_tensor(args[0]), - _as_reference_tensor(args[1]), - _as_reference_tensor(args[2]), - args[3], - args[4], - ) - - -def _reference_forward(args, **overrides): - # Rubin's fused FC1 epilogue stores the post-SwiGLU intermediate as MXFP8 - # before FC2 consumes it. Keep MoeEpReference's default raw semantics for - # its standalone tests, but model the device precision for API comparisons. - config = _forward_config(**overrides) - config.pop("tuning", None) - config.setdefault("intermediate_format", "mxfp8") - return MoeEpReference(**config)(*_reference_args(args)) - - -def _sm107_device() -> torch.device: - if not torch.cuda.is_available(): - pytest.skip("Rubin MXFP8 forward requires CUDA") - device = torch.device("cuda", 0) - if torch.cuda.get_device_capability(device) != (10, 7): - pytest.skip("Rubin MXFP8 forward requires exactly SM107 (compute capability 10.7)") - return device - - -def _require_distributed_sm107(world_size: int) -> None: - if not dist.is_available() or not dist.is_nccl_available(): - pytest.skip("multi-GPU Rubin MXFP8 forward requires NCCL") - if torch.cuda.device_count() < world_size: - pytest.skip(f"multi-GPU Rubin MXFP8 forward requires {world_size} GPUs") - if any( - torch.cuda.get_device_capability(index) != (10, 7) - for index in range(world_size) - ): - pytest.skip( - "multi-GPU Rubin MXFP8 forward requires exactly SM107 " - "(compute capability 10.7) on every rank" - ) - try: - import nvshmem.core # noqa: F401 - except (ImportError, OSError): - pytest.skip("multi-GPU Rubin MXFP8 forward requires NVSHMEM") - - -def _make_forward_case( - device: torch.device, - *, - experts: int, - tokens: int, - hidden: int, - intermediate: int, - top_k: int, - index_dtype: torch.dtype, - weight_dtype: torch.dtype, -): - """Build a deterministic supported case for the shape/format matrix.""" - - seed = ( - 20260811 - + experts * 1009 - + tokens * 101 - + hidden * 11 - + intermediate - + top_k - ) - generator = torch.Generator(device=device).manual_seed(seed) - activation = quantize_mxfp8( - torch.randn(tokens, hidden, generator=generator, device=device), - axis=1, - ) - fc1_weight = quantize_mxfp8( - torch.randn( - experts, - hidden, - 2 * intermediate, - generator=generator, - device=device, - ) - / 8, - axis=1, - ) - fc2_weight = quantize_mxfp8( - torch.randn( - experts, - intermediate, - hidden, - generator=generator, - device=device, - ) - / 8, - axis=1, - ) - topk_idx = ( - torch.arange(tokens * top_k, device=device) - .reshape(tokens, top_k) - .remainder(experts) - .to(index_dtype) - ) - topk_weights = torch.arange( - 1, - tokens * top_k + 1, - dtype=torch.float32, - device=device, - ).reshape(tokens, top_k) - topk_weights /= topk_weights.sum(dim=1, keepdim=True) - return ( - activation, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights.to(weight_dtype), - ) - - -def _stress_backend_reuse( - op, - args, - original_topk_idx, - original_topk_weights, - device, - *, - check_weight_refresh, -): - backend = op._forward_backend - assert backend is not None - compiled = backend._compiled - plan_workspace = backend._plan._workspace - weight_refresh_count = ( - backend._adapter.weight_refresh_count if check_weight_refresh else None - ) - alternate_stream = torch.cuda.Stream(device=device) - - for iteration in range(100): - args[3].copy_(original_topk_idx) - args[4].copy_( - original_topk_weights * float((iteration % 7) + 1) / 7.0 - ) - if iteration % 10 == 0: - args[3].fill_(-1) - stream = ( - torch.cuda.current_stream(device) - if iteration % 2 == 0 - else alternate_stream - ) - with torch.cuda.stream(stream): - stressed = op(*args) - stream.synchronize() - if iteration % 10 == 0: - assert _output_as_float(stressed).eq(0).all() - else: - assert torch.isfinite(_output_as_float(stressed)).all() - assert backend._compiled is compiled - assert backend._plan._workspace is plan_workspace - if weight_refresh_count is not None: - assert backend._adapter.weight_refresh_count == weight_refresh_count - - -def _replay_cuda_graph( - op, - args, - original_topk_idx, - expected, - device, - *, - synchronize_ranks=None, -): - synchronize_ranks = synchronize_ranks or (lambda: None) - op.warmup(*args) - synchronize_ranks() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_output = op(*args) - synchronize_ranks() - - for replay in range(20): - if replay % 2: - args[3].fill_(-1) - else: - args[3].copy_(original_topk_idx) - synchronize_ranks() - graph.replay() - torch.cuda.synchronize(device) - if replay % 2: - assert _output_as_float(graph_output).eq(0).all() - else: - _assert_matches_reference(graph_output, expected) diff --git a/test/python/moe_ep/moe_ep_test_data.py b/test/python/moe_ep/moe_ep_test_data.py deleted file mode 100644 index bcd6c6000..000000000 --- a/test/python/moe_ep/moe_ep_test_data.py +++ /dev/null @@ -1,168 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Deterministic input data and quantization helpers for MoE EP tests.""" - -from __future__ import annotations - -import torch -import torch.nn.functional as F - - -def make_forward_inputs(device: torch.device): - """Build one deterministic MXFP8 forward case.""" - - generator = torch.Generator(device=device).manual_seed(20260811) - experts, tokens, hidden, intermediate = 2, 5, 128, 256 - activation = quantize_mxfp8( - torch.randn(tokens, hidden, generator=generator, device=device), - axis=1, - ) - fc1_weight = quantize_mxfp8( - torch.randn( - experts, - hidden, - 2 * intermediate, - generator=generator, - device=device, - ) - / 8, - axis=1, - ) - fc2_weight = quantize_mxfp8( - torch.randn( - experts, - intermediate, - hidden, - generator=generator, - device=device, - ) - / 8, - axis=1, - ) - topk_idx = torch.tensor( - [[0, 1], [1, 0], [0, -1], [1, 0], [0, 1]], - dtype=torch.int32, - device=device, - ) - topk_weights = torch.tensor( - [ - [0.75, 0.25], - [0.625, 0.375], - [1.0, 0.0], - [0.5, 0.5], - [0.875, 0.125], - ], - dtype=torch.bfloat16, - device=device, - ) - return activation, fc1_weight, fc2_weight, topk_idx, topk_weights - - -def make_distributed_forward_inputs( - rank: int, - world_size: int, - device: torch.device, -): - """Build rank-local inputs with one local and one remote route per token.""" - - generator = torch.Generator(device=device).manual_seed(20260811 + rank) - # Vary local shapes without exceeding the distributed tests' - # max_tokens_per_rank=8 contract at EP sizes above seven. - local_experts, tokens, hidden, intermediate = ( - 2, - rank % 7 + 2, - 128, - 256, - ) - activation = quantize_mxfp8( - torch.randn(tokens, hidden, generator=generator, device=device), - axis=1, - ) - fc1_weight = quantize_mxfp8( - torch.randn( - local_experts, - hidden, - 2 * intermediate, - generator=generator, - device=device, - ) - / 8, - axis=1, - ) - fc2_weight = quantize_mxfp8( - torch.randn( - local_experts, - intermediate, - hidden, - generator=generator, - device=device, - ) - / 8, - axis=1, - ) - remote_rank = (rank + 1) % world_size - topk_idx = torch.tensor( - [ - [ - rank * local_experts + token % local_experts, - remote_rank * local_experts + (token + 1) % local_experts, - ] - for token in range(tokens) - ], - dtype=torch.int32, - device=device, - ) - topk_weights = torch.tensor( - [[0.625, 0.375]], - dtype=torch.bfloat16, - device=device, - ).expand(tokens, -1).contiguous() - return activation, fc1_weight, fc2_weight, topk_idx, topk_weights - - -def quantize_mxfp8(tensor: torch.Tensor, *, axis: int = -1): - """Return a public logical MXFP8 tensor (E4M3 payload + E8M0 scales).""" - - from cudnn import BlockScaledTensor - - axis = axis % tensor.ndim - logical_shape = tuple(tensor.shape) - logical_extent = logical_shape[axis] - moved = tensor.float().movedim(axis, -1) - block_count = (logical_extent + 31) // 32 - padded_extent = block_count * 32 - if padded_extent != logical_extent: - moved = F.pad(moved, (0, padded_extent - logical_extent)) - - blocks = moved.reshape(*moved.shape[:-1], block_count, 32) - raw_scale = blocks.abs().amax(dim=-1) / 448.0 - safe_scale = torch.where(raw_scale > 0, raw_scale, 1.0) - power_of_two_scale = torch.where( - raw_scale > 0, - torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), - torch.zeros_like(raw_scale), - ) - scale = power_of_two_scale.to(torch.float8_e8m0fnu) - reciprocal = torch.where(scale.float() > 0, scale.float().reciprocal(), 0.0) - payload = ( - (blocks * reciprocal.unsqueeze(-1)) - .clamp(-448.0, 448.0) - .to(torch.float8_e4m3fn) - .reshape(*moved.shape)[..., :logical_extent] - ) - - return BlockScaledTensor( - data=payload.movedim(-1, axis).contiguous(), - scale=scale.movedim(-1, axis).contiguous(), - format="mxfp8", - logical_shape=logical_shape, - axis=axis, - ) - - -__all__ = [ - "make_distributed_forward_inputs", - "make_forward_inputs", - "quantize_mxfp8", -] diff --git a/test/python/moe_ep/moe_ep_test_support.py b/test/python/moe_ep/moe_ep_test_support.py new file mode 100644 index 000000000..52f6d47cb --- /dev/null +++ b/test/python/moe_ep/moe_ep_test_support.py @@ -0,0 +1,1402 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared data, forward, and backward support for MoE EP tests.""" + +from __future__ import annotations + +# Common + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from moe_ep.moe_ep_reference import ( + BlockScaledTensor as ReferenceBlockScaledTensor, + MoeEpReference, + MoeFormat, + forward_combine_round_trip, + quantize_blockwise, +) + +__all__ = [ + "_assert_backward_matches", + "_assert_fixed_training_drop_overflow_result", + "_assert_fixed_training_matches_reference", + "_assert_matches_reference", + "_assert_training_graph_tails_are_reset", + "_assert_training_weight_sources_changed", + "_assert_wgrads_match_reference", + "_capture_fixed_training_batch", + "_copy_training_weight_sources_", + "_dense_wgrads_from_operands", + "_expected_backward", + "_fixed_training_case", + "_fixed_training_drop_overflow_case", + "_fixed_training_drop_overflow_reference", + "_fixed_training_reference", + "_fixed_training_weights", + "_forward_config", + "_grad_output", + "_make_forward_case", + "_naive_reference", + "_output_as_float", + "_prefill_training_graph_sentinels", + "_reference_backward", + "_reference_forward", + "_replay_cuda_graph", + "_require_distributed_sm107", + "_run_fixed_training_batch", + "_sm107_device", + "_stress_backend_reuse", + "_training_public_pointers", + "_training_source_pointers", + "_training_weight_source_pointers", + "_training_weight_source_values", + "_TrainingResourceContractOwner", + "_training_abi_prepared", + "_training_config", + "_training_contract_resources", + "_training_inputs", + "_training_prepared_pair", + "_training_staging_tensors", + "_training_weight_defect", + "_training_weights", + "make_distributed_forward_inputs", + "make_forward_inputs", + "quantize_mxfp8", +] + + +# Data + + +def make_forward_inputs(device: torch.device): + """Build one deterministic MXFP8 forward case.""" + + generator = torch.Generator(device=device).manual_seed(20260811) + experts, tokens, hidden, intermediate = 2, 5, 128, 256 + activation = quantize_mxfp8( + torch.randn(tokens, hidden, generator=generator, device=device), + axis=1, + ) + fc1_weight = quantize_mxfp8( + torch.randn( + experts, + hidden, + 2 * intermediate, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + fc2_weight = quantize_mxfp8( + torch.randn( + experts, + intermediate, + hidden, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + topk_idx = torch.tensor( + [[0, 1], [1, 0], [0, -1], [1, 0], [0, 1]], + dtype=torch.int32, + device=device, + ) + topk_weights = torch.tensor( + [ + [0.75, 0.25], + [0.625, 0.375], + [1.0, 0.0], + [0.5, 0.5], + [0.875, 0.125], + ], + dtype=torch.bfloat16, + device=device, + ) + return activation, fc1_weight, fc2_weight, topk_idx, topk_weights + + +def make_distributed_forward_inputs( + rank: int, + world_size: int, + device: torch.device, +): + """Build rank-local inputs with one local and one remote route per token.""" + + generator = torch.Generator(device=device).manual_seed(20260811 + rank) + # Vary local shapes without exceeding the distributed tests' + # max_tokens_per_rank=8 contract at EP sizes above seven. + local_experts, tokens, hidden, intermediate = ( + 2, + rank % 7 + 2, + 128, + 256, + ) + activation = quantize_mxfp8( + torch.randn(tokens, hidden, generator=generator, device=device), + axis=1, + ) + fc1_weight = quantize_mxfp8( + torch.randn( + local_experts, + hidden, + 2 * intermediate, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + fc2_weight = quantize_mxfp8( + torch.randn( + local_experts, + intermediate, + hidden, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + remote_rank = (rank + 1) % world_size + topk_idx = torch.tensor( + [ + [ + rank * local_experts + token % local_experts, + remote_rank * local_experts + (token + 1) % local_experts, + ] + for token in range(tokens) + ], + dtype=torch.int32, + device=device, + ) + topk_weights = ( + torch.tensor( + [[0.625, 0.375]], + dtype=torch.bfloat16, + device=device, + ) + .expand(tokens, -1) + .contiguous() + ) + return activation, fc1_weight, fc2_weight, topk_idx, topk_weights + + +def quantize_mxfp8(tensor: torch.Tensor, *, axis: int = -1): + """Return a public logical MXFP8 tensor (E4M3 payload + E8M0 scales).""" + + from cudnn import BlockScaledTensor + + axis = axis % tensor.ndim + logical_shape = tuple(tensor.shape) + logical_extent = logical_shape[axis] + moved = tensor.float().movedim(axis, -1) + block_count = (logical_extent + 31) // 32 + padded_extent = block_count * 32 + if padded_extent != logical_extent: + moved = F.pad(moved, (0, padded_extent - logical_extent)) + + blocks = moved.reshape(*moved.shape[:-1], block_count, 32) + raw_scale = blocks.abs().amax(dim=-1) / 448.0 + safe_scale = torch.where(raw_scale > 0, raw_scale, 1.0) + power_of_two_scale = torch.where( + raw_scale > 0, + torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), + torch.zeros_like(raw_scale), + ) + scale = power_of_two_scale.to(torch.float8_e8m0fnu) + reciprocal = torch.where(scale.float() > 0, scale.float().reciprocal(), 0.0) + payload = (blocks * reciprocal.unsqueeze(-1)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn).reshape(*moved.shape)[..., :logical_extent] + + return BlockScaledTensor( + data=payload.movedim(-1, axis).contiguous(), + scale=scale.movedim(-1, axis).contiguous(), + format="mxfp8", + logical_shape=logical_shape, + axis=axis, + ) + + +# Training setup + + +def _training_config(**overrides): + from cudnn.moe_ep._contracts import ForwardConfig + from cudnn.moe_ep._tuning import MoeEpTuningConfig + + values = { + "num_experts": 2, + "hidden_size": 128, + "intermediate_size": 256, + "top_k": 2, + "experts_per_rank": 2, + "ep_size": 1, + "ep_rank": 0, + "ep_group": None, + "ep_global_ranks": (), + "max_tokens_per_rank": 4, + "max_recv_size_per_rank": 4, + "drop_on_overflow": True, + "output_format": "bf16", + "combine_format": "bf16", + "apply_topk_in_fc1": True, + "gate_up_clamp": None, + "generate_c": True, + "token_padding_size": 128, + "sf_padding_size": 128, + "tuning": MoeEpTuningConfig(), + "backward_wgrad_mode": "operands", + } + values.update(overrides) + return ForwardConfig(**values) + + +def _training_inputs(): + return ( + torch.randn(2, 128, dtype=torch.bfloat16), + torch.randn(2, 128, 512, dtype=torch.bfloat16), + torch.randn(2, 256, 128, dtype=torch.bfloat16), + torch.tensor([[0, -1], [1, 0]], dtype=torch.int32), + torch.randn(2, 2, dtype=torch.float32), + ) + + +def _training_prepared_pair(config, pool_rows: int = 512): + from cudnn.moe_ep._megamoe_backend._workspace import WorkspaceRequirements + + forward_shapes = { + "fc1_c": (pool_rows, 512), + "col_quant_data": (pool_rows, 128), + "col_quant_sf": (2048,), + } + backward_shapes = { + "dprob": (4, 2), + "fc1_recompute": (pool_rows, 256), + "fc1_recompute_sf": (256, 8), + "fc1_col_output": (pool_rows, 512), + "fc1_col_output_sf": (512, 8), + "grad_y2": (pool_rows, 128), + "grad_y2_sf": (2048,), + } + forward = SimpleNamespace( + pool_token_capacity=pool_rows, + workspace_requirements=WorkspaceRequirements.for_mxfp8( + config, + kernel_local_workspace_bytes=1024, + kernel_shared_workspace_bytes=2048, + col_quant_data_bytes=pool_rows * 128, + col_quant_sf_bytes=2048, + ), + kernel=SimpleNamespace(get_aux_output_shapes=lambda: forward_shapes), + col_quant_sizes_offset=0, + col_quant_sizes_bytes=8, + ) + backward = SimpleNamespace( + pool_token_capacity=pool_rows, + config=SimpleNamespace(sf_padding_block=128), + workspace_requirements=WorkspaceRequirements.for_mxfp8( + config, + kernel_local_workspace_bytes=3072, + kernel_shared_workspace_bytes=4096, + backward_fc1_preact_bytes=pool_rows * 512 * 2, + backward_dprob_bytes=4 * 2 * 4, + backward_aux_data_bytes=pool_rows * 512, + backward_aux_scale_bytes=512 * 8, + ), + kernel=SimpleNamespace(get_aux_output_shapes=lambda: backward_shapes), + ) + return forward, backward + + +def _training_abi_prepared(name: str, max_recv_size: int = 4): + from cudnn.moe_ep._megamoe_backend._workspace import ( + BufferRegion, + WorkspaceRequirements, + ) + + workspace = WorkspaceRequirements( + max_tokens_per_rank=4, + symmetric_regions=(BufferRegion("symmetric", 256),), + local_regions=(BufferRegion("local", 128),), + ) + kernel_config = SimpleNamespace( + max_recv_size_per_rank=max_recv_size, + effective_config=lambda cluster_count: { + "name": name, + "max_recv_size_per_rank": max_recv_size, + "launch_cluster_count": cluster_count, + }, + ) + return SimpleNamespace( + kernel=SimpleNamespace( + name=lambda: name, + threads_per_cta=128, + occupancy=1, + smem_capacity=1024, + ), + architecture=(10, 7), + config=kernel_config, + launch_cluster_count=16, + workspace_requirements=workspace, + pool_token_capacity=512, + ) + + +def _training_weights(args=None): + from cudnn.moe_ep import MoeEpTrainingWeights + from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( + _quantize_plain_mxfp8, + ) + + if args is None: + args = _training_inputs() + return MoeEpTrainingWeights( + forward_fc1=_quantize_plain_mxfp8(args[1], axis=1), + forward_fc2=_quantize_plain_mxfp8(args[2], axis=1), + backward_w2_transpose=_quantize_plain_mxfp8( + args[2].transpose(1, 2).contiguous(), + axis=1, + ), + backward_w1_transpose=_quantize_plain_mxfp8( + args[1].transpose(1, 2).contiguous(), + axis=1, + ), + ) + + +def _training_empty_block_scaled_like(tensor, *, axis: int, format: str): + import cudnn + + logical_shape = tensor.logical_shape + data_shape = list(logical_shape) + scale_shape = list(logical_shape) + if format == "mxfp8": + data_dtype = tensor.data.dtype + scale_dtype = tensor.scale.dtype + scale_shape[axis] = (logical_shape[axis] + 31) // 32 + else: + data_dtype = torch.uint8 + scale_dtype = tensor.data.dtype + data_shape[axis] = (logical_shape[axis] + 1) // 2 + scale_shape[axis] = (logical_shape[axis] + 15) // 16 + return cudnn.BlockScaledTensor( + data=torch.empty(tuple(data_shape), dtype=data_dtype, device=tensor.device), + scale=torch.empty( + tuple(scale_shape), + dtype=scale_dtype, + device=tensor.device, + ), + format=format, + logical_shape=logical_shape, + axis=axis, + ) + + +def _training_same_shape_noncontiguous(tensor: torch.Tensor) -> torch.Tensor: + result = tensor.transpose(-2, -1).contiguous().transpose(-2, -1) + assert tuple(result.shape) == tuple(tensor.shape) + assert not result.is_contiguous() + return result + + +def _training_weight_defect(weights, field: str, defect: str): + tensor = getattr(weights, field) + expected_shape = tensor.logical_shape + if defect == "plain_tensor": + invalid = torch.empty( + expected_shape, + dtype=torch.bfloat16, + device=tensor.device, + ) + error_type = TypeError + message = f"weights.{field} must be an MXFP8 BlockScaledTensor for " "fixed training resources" + elif defect == "logical_shape": + wrong_shape = (expected_shape[0] - 1, *expected_shape[1:]) + invalid = replace( + tensor, + data=tensor.data[: wrong_shape[0]].contiguous(), + scale=tensor.scale[: wrong_shape[0]].contiguous(), + logical_shape=wrong_shape, + ) + error_type = ValueError + message = f"weights.{field} logical shape must be {expected_shape}, " f"got {wrong_shape}" + elif defect == "axis": + invalid = _training_empty_block_scaled_like( + tensor, + axis=2, + format="mxfp8", + ) + error_type = ValueError + message = f"weights.{field} block-scaled axis must be 1, got 2" + elif defect == "format": + invalid = _training_empty_block_scaled_like( + tensor, + axis=1, + format="nvfp4", + ) + error_type = NotImplementedError + message = f"weights.{field} must use format='mxfp8', got 'nvfp4'" + elif defect == "device": + invalid = replace( + tensor, + data=torch.empty_like(tensor.data, device="meta"), + scale=torch.empty_like(tensor.scale, device="meta"), + ) + error_type = ValueError + message = f"weights.{field} must be on cpu, got meta" + else: + part = "data" if defect == "data_noncontiguous" else "scale" + invalid = replace( + tensor, + **{part: _training_same_shape_noncontiguous(getattr(tensor, part))}, + ) + error_type = ValueError + message = f"weights.{field} data and scale must be contiguous for fixed " "training weight binding" + return replace(weights, **{field: invalid}), error_type, message + + +class _TrainingResourceContractOwner: + def __init__(self, *, slot_count: int = 2, lane_count: int = 1) -> None: + self.slot_count = slot_count + self.lane_count = lane_count + self.close_calls = 0 + self.refresh_calls = 0 + self.views_calls = 0 + + def refresh_weights(self) -> None: + self.refresh_calls += 1 + + def views(self, **kwargs): + del kwargs + self.views_calls += 1 + raise AssertionError("binding rejection must happen before owner views") + + def _flat_views(self, token_count: int): + del token_count + raise AssertionError("invalid finalization must fail before workspace access") + + def finalize_overflow(self, slots, *, lane): + from cudnn.moe_ep._megamoe_backend.mxfp8._training_resources import ( + Mxfp8TrainingResourceOwner, + ) + + return Mxfp8TrainingResourceOwner.finalize_overflow( + self, + slots, + lane=lane, + ) + + def close(self) -> None: + self.close_calls += 1 + + +def _training_contract_resources( + *, + owner=None, + slot_count: int = 2, + lane_count: int = 1, +): + from cudnn.moe_ep import MoeEpTrainingResources + + if owner is None: + owner = _TrainingResourceContractOwner( + slot_count=slot_count, + lane_count=lane_count, + ) + resources = MoeEpTrainingResources( + owner=owner, + operator_token=object(), + weights=SimpleNamespace(mock_training_weights=True), + slot_count=slot_count, + lane_count=lane_count, + device=torch.device("cpu"), + ) + return resources, owner + + +def _training_staging_tensors(*, capacity: int | None = None): + activation, _, _, topk_idx, topk_weights = make_forward_inputs(torch.device("cpu")) + source = activation.dequantize(torch.bfloat16).contiguous() + token_count, hidden = source.shape + top_k = topk_idx.shape[1] + if capacity is None: + capacity = token_count + return { + "source": source, + "topk_idx": topk_idx, + "topk_weights": topk_weights.float().contiguous(), + "output": torch.empty( + (capacity, hidden), + dtype=torch.float8_e4m3fn, + ), + "output_sf": torch.empty( + (capacity, hidden // 32), + dtype=torch.float8_e8m0fnu, + ), + "output_topk_idx": torch.empty( + (capacity, top_k), + dtype=torch.int32, + ), + "output_topk_weights": torch.empty( + (capacity, top_k), + dtype=torch.float32, + ), + } + + +# Forward + + +_DEFAULT_FORWARD_CONFIG = { + "num_experts": 2, + "hidden_size": 128, + "intermediate_size": 256, + "top_k": 2, + "max_tokens_per_rank": 5, + "apply_topk_in_fc1": True, + "combine_format": "bf16", + "output_format": "bf16", +} +_REFERENCE_CLOSE_KWARGS = {"rtol": 0.05, "atol": 0.0625} + + +def _forward_config(**overrides): + return {**_DEFAULT_FORWARD_CONFIG, **overrides} + + +def _output_as_float(output): + if isinstance(output, torch.Tensor): + return output.float() + return output.dequantize() + + +def _assert_matches_reference(actual, expected): + torch.testing.assert_close( + _output_as_float(actual), + _output_as_float(expected), + **_REFERENCE_CLOSE_KWARGS, + ) + + +def _naive_reference( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + *, + apply_topk_in_fc1, + clamp=None, + combine_format=MoeFormat.BF16, + intermediate_format=None, + apply_topk_after_combine=False, +): + 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, + device=activation.device, + ) + 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 + if intermediate_format is not None: + intermediate = quantize_blockwise( + intermediate, + intermediate_format, + ).dequantize() + result = intermediate @ fc2_weight[expert].float() + if not apply_topk_in_fc1 and not apply_topk_after_combine: + result = result * route_weight + result = forward_combine_round_trip(result, combine_format) + if not apply_topk_in_fc1 and apply_topk_after_combine: + result = result * route_weight + combine[token, slot] = result + return combine.sum(dim=1).to(torch.bfloat16) + + +def _as_reference_tensor(tensor): + if isinstance(tensor, torch.Tensor): + return tensor + return ReferenceBlockScaledTensor( + data=tensor.data, + scale=tensor.scale, + format=tensor.format.value, + logical_shape=tensor.logical_shape, + axis=tensor.axis, + ) + + +def _reference_args(args): + return ( + _as_reference_tensor(args[0]), + _as_reference_tensor(args[1]), + _as_reference_tensor(args[2]), + args[3], + args[4], + ) + + +def _reference_forward(args, **overrides): + # Rubin's fused FC1 epilogue stores the post-SwiGLU intermediate as MXFP8 + # before FC2 consumes it. Keep MoeEpReference's default raw semantics for + # its standalone tests, but model the device precision for API comparisons. + config = _forward_config(**overrides) + config.pop("tuning", None) + config.setdefault("intermediate_format", "mxfp8") + return MoeEpReference(**config)(*_reference_args(args)) + + +def _sm107_device() -> torch.device: + if not torch.cuda.is_available(): + pytest.skip("Rubin MXFP8 forward requires CUDA") + device = torch.device("cuda", 0) + if torch.cuda.get_device_capability(device) != (10, 7): + pytest.skip("Rubin MXFP8 forward requires exactly SM107 (compute capability 10.7)") + return device + + +def _require_distributed_sm107(world_size: int) -> None: + if not dist.is_available() or not dist.is_nccl_available(): + pytest.skip("multi-GPU Rubin MXFP8 forward requires NCCL") + if torch.cuda.device_count() < world_size: + pytest.skip(f"multi-GPU Rubin MXFP8 forward requires {world_size} GPUs") + if any(torch.cuda.get_device_capability(index) != (10, 7) for index in range(world_size)): + pytest.skip("multi-GPU Rubin MXFP8 forward requires exactly SM107 " "(compute capability 10.7) on every rank") + try: + import nvshmem.core # noqa: F401 + except (ImportError, OSError): + pytest.skip("multi-GPU Rubin MXFP8 forward requires NVSHMEM") + + +def _make_forward_case( + device: torch.device, + *, + experts: int, + tokens: int, + hidden: int, + intermediate: int, + top_k: int, + index_dtype: torch.dtype, + weight_dtype: torch.dtype, +): + """Build a deterministic supported case for the shape/format matrix.""" + + seed = 20260811 + experts * 1009 + tokens * 101 + hidden * 11 + intermediate + top_k + generator = torch.Generator(device=device).manual_seed(seed) + activation = quantize_mxfp8( + torch.randn(tokens, hidden, generator=generator, device=device), + axis=1, + ) + fc1_weight = quantize_mxfp8( + torch.randn( + experts, + hidden, + 2 * intermediate, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + fc2_weight = quantize_mxfp8( + torch.randn( + experts, + intermediate, + hidden, + generator=generator, + device=device, + ) + / 8, + axis=1, + ) + topk_idx = torch.arange(tokens * top_k, device=device).reshape(tokens, top_k).remainder(experts).to(index_dtype) + topk_weights = torch.arange( + 1, + tokens * top_k + 1, + dtype=torch.float32, + device=device, + ).reshape(tokens, top_k) + topk_weights /= topk_weights.sum(dim=1, keepdim=True) + return ( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights.to(weight_dtype), + ) + + +def _stress_backend_reuse( + op, + args, + original_topk_idx, + original_topk_weights, + device, + *, + check_weight_refresh, +): + backend = op._forward_backend + assert backend is not None + compiled = backend._compiled + plan_workspace = backend._plan._workspace + weight_refresh_count = backend._adapter.weight_refresh_count if check_weight_refresh else None + alternate_stream = torch.cuda.Stream(device=device) + + for iteration in range(100): + args[3].copy_(original_topk_idx) + args[4].copy_(original_topk_weights * float((iteration % 7) + 1) / 7.0) + if iteration % 10 == 0: + args[3].fill_(-1) + stream = torch.cuda.current_stream(device) if iteration % 2 == 0 else alternate_stream + with torch.cuda.stream(stream): + stressed = op(*args) + stream.synchronize() + if iteration % 10 == 0: + assert _output_as_float(stressed).eq(0).all() + else: + assert torch.isfinite(_output_as_float(stressed)).all() + assert backend._compiled is compiled + assert backend._plan._workspace is plan_workspace + if weight_refresh_count is not None: + assert backend._adapter.weight_refresh_count == weight_refresh_count + + +def _replay_cuda_graph( + op, + args, + original_topk_idx, + expected, + device, + *, + synchronize_ranks=None, +): + synchronize_ranks = synchronize_ranks or (lambda: None) + op.warmup(*args) + synchronize_ranks() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = op(*args) + synchronize_ranks() + + for replay in range(20): + if replay % 2: + args[3].fill_(-1) + else: + args[3].copy_(original_topk_idx) + synchronize_ranks() + graph.replay() + torch.cuda.synchronize(device) + if replay % 2: + assert _output_as_float(graph_output).eq(0).all() + else: + _assert_matches_reference(graph_output, expected) + + +# Backward + + +_BACKWARD_CLOSE_KWARGS = ( + {"rtol": 0.15, "atol": 0.125}, # grad_activation is BF16-rounded. + {"rtol": 0.15, "atol": 0.125}, # router-weight gradient. +) +_WGRAD_CLOSE_KWARGS = {"rtol": 0.2, "atol": 0.25} + + +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +def _unpack_wgrad_scale_part( + packed: torch.Tensor, + rows: int, + columns: int, +) -> torch.Tensor: + """Invert grouped-wgrad's 128x4 scale-atom swizzle.""" + + padded_rows = _round_up(rows, 128) + padded_columns = _round_up(columns, 4) + row_atoms = padded_rows // 128 + column_atoms = padded_columns // 4 + atom_count = row_atoms * column_atoms + expected = padded_rows * padded_columns + if packed.numel() != expected: + raise ValueError(f"packed scale part has {packed.numel()} bytes, expected {expected}") + blocked = ( + packed.reshape(atom_count, 32, 4, 4).transpose(1, 2).reshape(row_atoms, column_atoms, 128, 4).permute(0, 2, 1, 3).reshape(padded_rows, padded_columns) + ) + return blocked[:rows, :columns].view(torch.float8_e8m0fnu).float() + + +def _dequantize_wgrad_operand( + data: torch.Tensor, + scales: torch.Tensor, + expert_offsets: torch.Tensor, + *, + k_dim: int, +) -> torch.Tensor: + """Decode one public grouped-wgrad operand without launching a GEMM.""" + + if data.ndim != 2 or k_dim not in (0, 1): + raise ValueError("wgrad operand must be rank 2 with k_dim 0 or 1") + non_k = int(data.shape[1 - k_dim]) + padded_non_k = _round_up(non_k, 128) + flat_scales = scales.view(torch.uint8).reshape(-1) + output = torch.zeros(data.shape, dtype=torch.float32, device=data.device) + ends = [int(value) for value in expert_offsets.detach().cpu().tolist()] + k_capacity = int(data.shape[k_dim]) + previous = 0 + scale_byte_offset = 0 + for end in ends: + if end < previous or end > k_capacity: + raise ValueError("expert offsets must be nondecreasing and fit the operand " f"K capacity ({k_capacity})") + extent = end - previous + if extent % 32: + raise ValueError("each padded expert K extent must be divisible by 32") + if extent == 0: + continue + scale_columns = _round_up(extent // 32, 4) + scale_byte_count = padded_non_k * scale_columns + if scale_byte_offset + scale_byte_count > flat_scales.numel(): + raise ValueError("expert offsets exceed the scale tensor") + part = flat_scales.narrow( + 0, + scale_byte_offset, + scale_byte_count, + ) + logical_scale = _unpack_wgrad_scale_part( + part, + non_k, + extent // 32, + ) + if k_dim == 1: + expanded_scale = logical_scale.repeat_interleave(32, dim=1) + output[:, previous:end] = data[:, previous:end].float() * expanded_scale + else: + expanded_scale = logical_scale.repeat_interleave( + 32, + dim=1, + ).transpose(0, 1) + output[previous:end, :] = data[previous:end, :].float() * expanded_scale + previous = end + scale_byte_offset += scale_byte_count + + if previous < k_capacity: + capacity_tail = data.narrow(k_dim, previous, k_capacity - previous) + if bool(capacity_tail.float().ne(0).any().item()): + raise ValueError("unused WGrad operand capacity tail must contain zero data") + scale_tail = flat_scales[scale_byte_offset:] + if scale_tail.numel() and bool(scale_tail.ne(127).any().item()): + raise ValueError("unused WGrad operand capacity tail must contain neutral E8M0 scales") + return output + + +def _dense_wgrads_from_operands(operands): + """Reference grouped matmuls over the exported operand ABI.""" + + fc1_a = _dequantize_wgrad_operand( + operands.fc1_a, + operands.fc1_sfa, + operands.expert_offsets, + k_dim=1, + ) + fc1_b = _dequantize_wgrad_operand( + operands.fc1_b, + operands.fc1_sfb, + operands.expert_offsets, + k_dim=0, + ) + fc2_a = _dequantize_wgrad_operand( + operands.fc2_a, + operands.fc2_sfa, + operands.expert_offsets, + k_dim=1, + ) + fc2_b = _dequantize_wgrad_operand( + operands.fc2_b, + operands.fc2_sfb, + operands.expert_offsets, + k_dim=0, + ) + fc1_parts = [] + fc2_parts = [] + ends = [int(value) for value in operands.expert_offsets.detach().cpu().tolist()] + valid_counts = [int(value) for value in operands.valid_route_counts.detach().cpu().tolist()] + if len(ends) != len(valid_counts): + raise ValueError("expert offsets and valid route counts must have equal size") + + previous = 0 + for expert, (end, valid_count) in enumerate(zip(ends, valid_counts)): + extent = end - previous + if valid_count < 0 or valid_count > extent: + raise ValueError(f"expert {expert} valid route count {valid_count} exceeds " f"its padded extent {extent}") + valid_end = previous + valid_count + for name, tensor, k_dim in ( + ("fc1_a", fc1_a, 1), + ("fc1_b", fc1_b, 0), + ("fc2_a", fc2_a, 1), + ("fc2_b", fc2_b, 0), + ): + padding = tensor.narrow(k_dim, valid_end, end - valid_end) + if bool(padding.ne(0).any().item()): + raise ValueError(f"{name} expert {expert} padded rows must decode to zero") + fc1_parts.append(fc1_a[:, previous:valid_end] @ fc1_b[previous:valid_end, :]) + fc2_parts.append(fc2_a[:, previous:valid_end] @ fc2_b[previous:valid_end, :]) + previous = end + return torch.stack(fc1_parts), torch.stack(fc2_parts) + + +def _reference_backward(config) -> MoeEpReference: + options = dict(config) + for production_only in ( + "drop_on_overflow", + "ep_global_ranks", + "ep_rank", + "ep_size", + "experts_per_rank", + "max_recv_size_per_rank", + "sf_padding_size", + "tuning", + ): + options.pop(production_only, None) + options["intermediate_format"] = "mxfp8" + options["backward_operand_format"] = "mxfp8" + return MoeEpReference(**options) + + +def _fixed_training_weights(args): + """Build the four stable MXFP8 source packs required by training.""" + + from cudnn.moe_ep import MoeEpTrainingWeights + from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( + _quantize_plain_mxfp8, + ) + + fc1_weight = args[1] + fc2_weight = args[2] + dense_fc1 = fc1_weight if isinstance(fc1_weight, torch.Tensor) else fc1_weight.dequantize() + dense_fc2 = fc2_weight if isinstance(fc2_weight, torch.Tensor) else fc2_weight.dequantize() + return MoeEpTrainingWeights( + forward_fc1=(_quantize_plain_mxfp8(dense_fc1, axis=1) if isinstance(fc1_weight, torch.Tensor) else fc1_weight), + forward_fc2=(_quantize_plain_mxfp8(dense_fc2, axis=1) if isinstance(fc2_weight, torch.Tensor) else fc2_weight), + backward_w2_transpose=_quantize_plain_mxfp8( + dense_fc2.transpose(1, 2).contiguous(), + axis=1, + ), + backward_w1_transpose=_quantize_plain_mxfp8( + dense_fc1.transpose(1, 2).contiguous(), + axis=1, + ), + ) + + +def _fixed_training_reference( + args, + grad_output, + *, + combine_format, + gate_up_clamp, + ep_group=None, + num_experts=None, + **config_overrides, +): + """Run the standalone oracle for EP1 or a distributed EP group.""" + + ep_size = 1 if ep_group is None else dist.get_world_size(ep_group) + local_experts = int(args[1].shape[0]) + if num_experts is None: + num_experts = local_experts * ep_size + reference_config = _forward_config(**config_overrides) + reference_config.update( + num_experts=num_experts, + hidden_size=int(args[0].shape[1]), + intermediate_size=int(args[2].shape[1]), + top_k=int(args[3].shape[1]), + max_tokens_per_rank=config_overrides.get( + "max_tokens_per_rank", + int(args[0].shape[0]), + ), + ep_group=ep_group, + combine_format=combine_format, + gate_up_clamp=gate_up_clamp, + generate_c=True, + backward_wgrad_mode="operands", + # The standalone operand oracle's legacy ABI uses 256-row + # segments. Production fixed resources use 128-row segments; + # their represented dense gradients are compared below. + token_padding_size=256, + ) + reference = _reference_backward(reference_config) + reference_args = _reference_args(args) + output, fc1_c, route_metadata, forward_stash = reference(*reference_args) + grad_activation, grad_topk_weights, wgrad_operands = reference.backward( + grad_output, + *reference_args[1:], + fc1_c, + route_metadata, + wgrad_forward_stash=forward_stash, + ) + return ( + output, + grad_activation, + grad_topk_weights, + wgrad_operands, + ) + + +def _grad_output( + device: torch.device, + token_count: int, + *, + seed: int, +) -> torch.Tensor: + generator = torch.Generator(device=device).manual_seed(seed) + return ( + torch.randn( + token_count, + 128, + generator=generator, + dtype=torch.float32, + device=device, + ) + / 8 + ) + + +def _expected_backward(reference, grad_output, args, stash): + return reference.backward( + grad_output, + *_reference_args(args)[1:], + *stash, + ) + + +def _assert_backward_matches(actual, expected, topk_idx) -> None: + assert len(actual) == len(expected) == 2 + for name, gradient, reference, close_kwargs in zip( + ("grad_activation", "grad_topk_weights"), + actual, + expected, + _BACKWARD_CLOSE_KWARGS, + ): + assert gradient.shape == reference.shape + assert gradient.dtype == torch.float32 + assert torch.isfinite(gradient).all() + torch.testing.assert_close( + gradient, + reference, + msg=lambda default, name=name: (f"{name} does not match the backward reference\n{default}"), + **close_kwargs, + ) + + dropped = topk_idx == -1 + assert actual[1][dropped].eq(0).all() + + +def _assert_wgrads_match_reference( + actual, + expected, + *, + expected_dense=None, +) -> None: + """Compare fixed-capacity production operands with standalone dense dW.""" + + torch.testing.assert_close( + actual.valid_route_counts, + expected.valid_route_counts, + rtol=0, + atol=0, + msg="valid route counts differ from the independent reference", + ) + actual_dense = _dense_wgrads_from_operands(actual) + if expected_dense is None: + expected_dense = expected.dense_wgrads() + for name, actual_dw, expected_dw in zip( + ("grad_fc1_weight", "grad_fc2_weight"), + actual_dense, + expected_dense, + ): + torch.testing.assert_close( + actual_dw, + expected_dw, + msg=lambda default, name=name: (f"{name} does not match the independent reference\n{default}"), + **_WGRAD_CLOSE_KWARGS, + ) + + +_TRAINING_WGRAD_DATA_FIELDS = ("fc1_a", "fc1_b", "fc2_a", "fc2_b") +_TRAINING_WGRAD_SF_FIELDS = ("fc1_sfa", "fc1_sfb", "fc2_sfa", "fc2_sfb") +_TRAINING_WEIGHT_FIELDS = ( + "forward_fc1", + "forward_fc2", + "backward_w2_transpose", + "backward_w1_transpose", +) + + +def _fixed_training_case(device): + args = list(make_forward_inputs(device)) + args[0] = args[0].dequantize(torch.bfloat16) + args[4] = args[4].float() + args[3].fill_(-1) + args[4].zero_() + args[3][0, 0] = 0 + args[4][0, 0] = 1 + grad_output = _grad_output( + device, + args[0].shape[0], + seed=20260828, + ) + return args, grad_output + + +def _assert_fixed_training_matches_reference( + actual, + expected, + topk_idx, +) -> None: + actual_y, actual_dx, actual_dprob, actual_wgrads = actual + expected_y, expected_dx, expected_dprob, expected_wgrads = expected + _assert_matches_reference(actual_y, expected_y) + _assert_backward_matches( + (actual_dx, actual_dprob), + (expected_dx, expected_dprob), + topk_idx, + ) + _assert_wgrads_match_reference(actual_wgrads, expected_wgrads) + + +def _run_fixed_training_batch(resources, lane, cases): + """Run refresh, ordered forwards/backwards, and one overflow finalization.""" + + resources.refresh_weights() + outputs = [resources.forward(slot, lane, args[0], args[3], args[4]) for slot, args, _ in cases] + backwards = [resources.backward(slot, lane, grad_output) for slot, _, grad_output in cases] + overflow = resources.finalize_overflow( + tuple(slot for slot, _, _ in cases), + lane, + ) + return tuple( + SimpleNamespace( + y=output, + dx=backward[0], + dprob=backward[1], + wgrads=backward[2], + overflow=overflow, + ) + for output, backward in zip(outputs, backwards) + ) + + +def _capture_fixed_training_batch(resources, lane, cases, capture_stream): + """Capture the shared fixed-training sequence for one or more slots.""" + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + actuals = _run_fixed_training_batch(resources, lane, cases) + capture_stream.synchronize() + return SimpleNamespace( + graph=graph, + actuals=actuals, + public_pointers=tuple(_training_public_pointers(actual) for actual in actuals), + ) + + +def _training_source_pointers(case) -> dict[str, int]: + return { + name: getattr(case, name).data_ptr() + for name in ( + "activation", + "topk_idx", + "topk_weights", + "grad_output", + ) + } + + +def _training_weight_source_pointers(weights) -> dict[str, int]: + return {f"{name}.{part}": getattr(getattr(weights, name), part).data_ptr() for name in _TRAINING_WEIGHT_FIELDS for part in ("data", "scale")} + + +def _training_weight_source_values(weights) -> dict[str, torch.Tensor]: + return {f"{name}.{part}": getattr(getattr(weights, name), part).clone() for name in _TRAINING_WEIGHT_FIELDS for part in ("data", "scale")} + + +def _assert_training_weight_sources_changed(weights, previous) -> None: + for name in _TRAINING_WEIGHT_FIELDS: + for part in ("data", "scale"): + assert not torch.equal( + getattr(getattr(weights, name), part), + previous[f"{name}.{part}"], + ) + + +def _training_public_pointers(actual) -> dict[str, int]: + pointers = { + "y": actual.y.data_ptr(), + "dx": actual.dx.data_ptr(), + "dprob": actual.dprob.data_ptr(), + "overflow": actual.overflow.data_ptr(), + } + pointers.update( + { + f"wgrads.{name}": getattr(actual.wgrads, name).data_ptr() + for name in ( + *_TRAINING_WGRAD_DATA_FIELDS, + *_TRAINING_WGRAD_SF_FIELDS, + "expert_offsets", + "valid_route_counts", + ) + } + ) + return pointers + + +def _prefill_training_graph_sentinels(slot_views, actual) -> None: + """Poison every history-sensitive full-capacity destination.""" + + slot_views.routing_topk_idx.fill_(0x1A2B3C) + slot_views.routing_topk_weights.fill_(31.25) + slot_views.forward_output.fill_(29.0) + slot_views.backward_output.fill_(-27.0) + slot_views.grad_activation.fill_(23.0) + slot_views.dprob.fill_(-19.0) + slot_views.expert_offsets.fill_(-17) + slot_views.valid_route_counts.fill_(-13) + for name in _TRAINING_WGRAD_DATA_FIELDS: + getattr(actual.wgrads, name).fill_(1.0) + for name in _TRAINING_WGRAD_SF_FIELDS: + getattr(actual.wgrads, name).view(torch.uint8).fill_(0) + + +def _assert_training_graph_tails_are_reset( + slot_views, + actual, + *, + token_count: int, + capacity: int, +) -> None: + if token_count < capacity: + assert slot_views.routing_topk_idx[token_count:].eq(-1).all() + assert slot_views.routing_topk_weights[token_count:].eq(0).all() + assert slot_views.forward_output[token_count:].eq(0).all() + assert slot_views.backward_output[token_count:].eq(0).all() + assert slot_views.grad_activation[token_count:].eq(0).all() + assert slot_views.dprob[token_count:].eq(0).all() + + counts = actual.wgrads.valid_route_counts.detach().cpu().tolist() + expected_offsets = [] + offset = 0 + for count in counts: + offset += (int(count) + 127) // 128 * 128 + expected_offsets.append(offset) + assert actual.wgrads.expert_offsets.detach().cpu().tolist() == expected_offsets + + +def _copy_training_weight_sources_(destination, source) -> None: + for name in _TRAINING_WEIGHT_FIELDS: + destination_pack = getattr(destination, name) + source_pack = getattr(source, name) + destination_pack.data.copy_(source_pack.data) + destination_pack.scale.copy_(source_pack.scale) + + +def _fixed_training_drop_overflow_case(device): + base_args, base_grad_output = _fixed_training_case(device) + topk_idx = torch.tensor( + [[0, 1]], + dtype=torch.int32, + device=device, + ) + topk_weights = torch.tensor( + [[0.75, 0.25]], + dtype=torch.float32, + device=device, + ) + args = ( + base_args[0][:1].clone(), + base_args[1], + base_args[2], + topk_idx, + topk_weights, + ) + return args, base_grad_output[:1].clone() + + +def _fixed_training_drop_overflow_reference( + args, + grad_output, + *, + drop_expert1, +): + reference_topk_idx = args[3].clone() + if drop_expert1: + assert reference_topk_idx.shape == (1, 2) + assert reference_topk_idx.detach().cpu().tolist() == [[0, 1]] + reference_topk_idx[0, 1] = -1 + reference_args = ( + args[0], + args[1], + args[2], + reference_topk_idx, + args[4].clone(), + ) + return ( + _fixed_training_reference( + reference_args, + grad_output, + combine_format="bf16", + gate_up_clamp=None, + ), + reference_topk_idx, + ) + + +def _assert_fixed_training_drop_overflow_result( + actual, + expected, + reference_topk_idx, + *, + expected_overflow, +): + assert actual.overflow.eq(expected_overflow).all() + _assert_fixed_training_matches_reference( + (actual.y, actual.dx, actual.dprob, actual.wgrads), + expected, + reference_topk_idx, + ) + + if expected_overflow: + assert reference_topk_idx[0, 1].eq(-1) + assert actual.dprob[0, 1].eq(0) + assert actual.wgrads.valid_route_counts.detach().cpu().tolist() == [1, 0] + # Expert 0 owns the first 128-row padded segment. Expert 1 starts at + # pool capacity and therefore has no retained segment or dense dW. + assert actual.wgrads.expert_offsets.detach().cpu().tolist() == [128, 128] + dense_dw1, dense_dw2 = _dense_wgrads_from_operands(actual.wgrads) + assert dense_dw1[1].eq(0).all() + assert dense_dw2[1].eq(0).all() diff --git a/test/python/moe_ep/probe_moe_ep_training_graph.py b/test/python/moe_ep/probe_moe_ep_training_graph.py new file mode 100644 index 000000000..e3eea1c4d --- /dev/null +++ b/test/python/moe_ep/probe_moe_ep_training_graph.py @@ -0,0 +1,994 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Fixed-resource SM107 multi-rank CUDA Graph communication probe. + +Run from the project root on one node with two Rubin GPUs. The container +launcher selects an architecture-native Python/PyTorch environment:: + + data/script/run_moe_ep_bf16_combine_container.sh \ + --my-version-root "$PWD/my-version" \ + training-graph-probe + +Set ``MOE_EP_GRAPH_PROBE_NPROC=4`` (or another local EP size) on the host to +reuse the same probe beyond EP2. + +The probe exercises only the public fixed-resource ordinary/capture path, +including fixed-address staging/reset operations, forward/backward CuTeDSL +callables, and a one-scalar NCCL overflow OR. +""" + +from __future__ import annotations + +import argparse +import gc +import os +import socket +import time +from contextlib import contextmanager +from datetime import timedelta + +import torch +import torch.distributed as dist + +from cudnn import MoeEp, MoeEpTrainingWeights +from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( + _quantize_plain_mxfp8, +) +from cudnn.moe_ep._megamoe_backend._runtime import ( + _RuntimeWatchdog, + get_runtime_manager, +) + + +def _debug_phase(rank: int, phase: str) -> None: + if os.environ.get("MOE_EP_DEBUG_RUNTIME", "0") != "1": + return + print( + "[moe-ep-probe] " + f"time={time.monotonic():.6f} host={socket.gethostname()} " + f"pid={os.getpid()} rank={rank} phase={phase}", + flush=True, + ) + + +@contextmanager +def _debug_phase_scope(rank: int, phase: str): + _debug_phase(rank, f"{phase}.begin") + try: + yield + finally: + _debug_phase(rank, f"{phase}.end") + + +def _synchronize_with_watchdog( + rank: int, + device: torch.device, + phase: str, +) -> None: + watchdog = _RuntimeWatchdog(phase) + watchdog.start() + _debug_phase(rank, f"{phase}.begin") + try: + torch.cuda.synchronize(device) + finally: + watchdog.close() + _debug_phase(rank, f"{phase}.end") + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--diagnostic-replays", type=int, default=2) + parser.add_argument("--burst-replays", type=int, default=100) + parser.add_argument( + "--multistream-replays", + type=int, + default=10, + help=( + "two-lane cross-stream graph replays; use a larger value such as " + "100 for dedicated stress runs" + ), + ) + parser.add_argument( + "--max-recv-size-per-rank", + type=int, + default=1, + help=( + "bounded receive capacity; must remain below the forced-overflow " + "route count so the probe retains overflow coverage" + ), + ) + parser.add_argument( + "--cycles", + type=int, + default=2, + help=( + "create/capture/destroy cycles; the first is exhaustive and later " + "cycles use a minimal replay to verify teardown/re-init" + ), + ) + parser.add_argument("--timeout-seconds", type=int, default=600) + parser.add_argument( + "--skip-multistream", + action="store_true", + help="skip the two-lane ordered cross-stream resource probe", + ) + parser.add_argument( + "--expect-overflow-assert", + action="store_true", + help=( + "run only the fatal drop_on_overflow=False graph assertion probe; " + "success requires every rank to observe the expected CUDA error" + ), + ) + return parser.parse_args() + + +def _require_positive(name: str, value: int) -> None: + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + + +def _assert_replay_tensor( + name: str, + actual: torch.Tensor, + expected: torch.Tensor, +) -> None: + """Compare graph replay outputs with dtype-appropriate semantics.""" + + low_precision = { + torch.float8_e4m3fn, + torch.float8_e5m2, + torch.float8_e8m0fnu, + torch.uint8, + torch.int32, + torch.int64, + } + if actual.dtype in low_precision: + if not torch.equal(actual, expected): + raise AssertionError( + f"{name} is not bitwise equal after graph replay" + ) + return + torch.testing.assert_close( + actual, + expected, + rtol=1e-5, + atol=1e-6, + msg=f"{name} differs after graph replay", + ) + + +def _make_inputs( + rank: int, + device: torch.device, +) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: + token_count = 8 + hidden = 128 + intermediate = 256 + experts_per_rank = 2 + top_k = 2 + generator = torch.Generator(device=device).manual_seed(20260828 + rank) + + activation = ( + torch.randn( + token_count, + hidden, + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + / 8 + ) + fc1_weight = ( + torch.randn( + experts_per_rank, + hidden, + 2 * intermediate, + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + / 16 + ) + fc2_weight = ( + torch.randn( + experts_per_rank, + intermediate, + hidden, + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + / 16 + ) + topk_idx = torch.full( + (token_count, top_k), + -1, + dtype=torch.int32, + device=device, + ) + topk_weights = torch.zeros( + (token_count, top_k), + dtype=torch.float32, + device=device, + ) + # Exactly one route is received by each rank during eager warmup, so every + # positive max_recv_size_per_rank remains within capacity. + topk_idx[0, 0] = rank * experts_per_rank + topk_weights[0, 0] = 1.0 + grad_output = ( + torch.randn( + token_count, + hidden, + dtype=torch.float32, + device=device, + generator=generator, + ) + / 8 + ) + return ( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ), grad_output + + +def _route_pattern( + kind: str, + rank: int, + world_size: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + token_count = 8 + top_k = 2 + experts_per_rank = 2 + indices = torch.full( + (token_count, top_k), + -1, + dtype=torch.int32, + device=device, + ) + weights = torch.zeros( + (token_count, top_k), + dtype=torch.float32, + device=device, + ) + if kind == "local": + indices[0, 0] = rank * experts_per_rank + weights[0, 0] = 1.0 + elif kind == "remote": + peer = (rank + 1) % world_size + indices[0, 0] = peer * experts_per_rank + weights[0, 0] = 1.0 + elif kind == "overflow": + # Every source sends all routes to rank 0. Its raw receive count is + # therefore much larger than max_recv_size_per_rank=1. + indices.fill_(0) + weights.fill_(0.5) + else: + raise ValueError(f"unknown route pattern {kind!r}") + return indices, weights + + +def _make_two_slot_inputs( + rank: int, + world_size: int, + device: torch.device, +) -> tuple[ + tuple[torch.Tensor, ...], + torch.Tensor, + tuple[torch.Tensor, ...], + torch.Tensor, + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], +]: + args0, grad0 = _make_inputs(rank, device) + local = _route_pattern("local", rank, world_size, device) + remote = _route_pattern("remote", rank, world_size, device) + args0 = (*args0[:3], local[0].clone(), local[1].clone()) + args1 = ( + args0[0].clone(), + args0[1], + args0[2], + remote[0].clone(), + remote[1].clone(), + ) + return args0, grad0, args1, grad0.clone(), local, remote + + +def _make_training_weights( + args: tuple[torch.Tensor, ...], +) -> MoeEpTrainingWeights: + return MoeEpTrainingWeights( + forward_fc1=_quantize_plain_mxfp8(args[1], axis=1), + forward_fc2=_quantize_plain_mxfp8(args[2], axis=1), + backward_w2_transpose=_quantize_plain_mxfp8( + args[2].transpose(1, 2).contiguous(), + axis=1, + ), + backward_w1_transpose=_quantize_plain_mxfp8( + args[1].transpose(1, 2).contiguous(), + axis=1, + ), + ) + + +def _make_operator( + *, + world_size: int, + group, + max_recv_size_per_rank: int, + drop_on_overflow: bool, +) -> MoeEp: + return MoeEp( + num_experts=2 * world_size, + hidden_size=128, + intermediate_size=256, + top_k=2, + ep_group=group, + max_tokens_per_rank=8, + max_recv_size_per_rank=max_recv_size_per_rank, + drop_on_overflow=drop_on_overflow, + combine_format="bf16", + ) + + +def _close_probe_operator( + *, + device: torch.device, + group, + op: MoeEp, +) -> None: + torch.cuda.synchronize(device) + dist.barrier(group=group) + op.close() + gc.collect() + torch.cuda.synchronize(device) + dist.barrier(group=group) + + +def _run_training_resource_probe( + *, + rank: int, + world_size: int, + device: torch.device, + group, + diagnostic_replays: int, + burst_replays: int, + max_recv_size_per_rank: int, + full_probe: bool, +) -> None: + """Exercise full graph behavior or a minimal teardown/re-init replay.""" + + args0, grad0, args1, grad1, local, remote = _make_two_slot_inputs( + rank, + world_size, + device, + ) + # Keep immutable baseline patterns separate from the graph-bound input + # tensors. Overflow injection mutates the latter in place. + weights = _make_training_weights(args0) + op = _make_operator( + world_size=world_size, + group=group, + max_recv_size_per_rank=max_recv_size_per_rank, + drop_on_overflow=True, + ) + graph = None + try: + resources = op.prepare_training_resources( + weights, + slot_count=2, + lane_count=1, + ) + slot0, slot1 = resources.slots + lane0 = resources.lanes[0] + + # Ordinary execution is the collective warmup for all fused staging, + # MegaMoE, and fixed-capacity WGrad export compile caches. + resources.refresh_weights() + y0 = resources.forward( + slot0, + lane0, + args0[0], + args0[3], + args0[4], + ) + y1 = resources.forward( + slot1, + lane0, + args1[0], + args1[3], + args1[4], + ) + dx0, dp0, operands0 = resources.backward(slot0, lane0, grad0) + dx1, dp1, operands1 = resources.backward(slot1, lane0, grad1) + overflow_status = resources.finalize_overflow((slot0, slot1)) + torch.cuda.synchronize(device) + dist.barrier(group=group) + if int(overflow_status.item()) != 0: + raise AssertionError("ordinary fixed-resource warmup overflowed") + + comparison_names = ( + "y0", + "y1", + "dx0", + "dx1", + "dprob0", + "dprob1", + "slot0.fc1_a", + "slot0.fc1_b", + "slot0.fc2_a", + "slot0.fc2_b", + "slot1.fc1_a", + "slot1.fc1_b", + "slot1.fc2_a", + "slot1.fc2_b", + ) + ordinary = { + name: tensor.clone() + for name, tensor in zip( + comparison_names, + ( + y0, + y1, + dx0, + dx1, + dp0, + dp1, + operands0.fc1_a, + operands0.fc1_b, + operands0.fc2_a, + operands0.fc2_b, + operands1.fc1_a, + operands1.fc1_b, + operands1.fc2_a, + operands1.fc2_b, + ), + ) + } + ordinary_offsets = ( + operands0.expert_offsets.clone(), + operands1.expert_offsets.clone(), + ) + + stream = torch.cuda.Stream(device=device) + stream.wait_stream(torch.cuda.current_stream(device)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + resources.refresh_weights() + graph_y0 = resources.forward( + slot0, + lane0, + args0[0], + args0[3], + args0[4], + ) + graph_y1 = resources.forward( + slot1, + lane0, + args1[0], + args1[3], + args1[4], + ) + graph_dx0, graph_dp0, graph_operands0 = resources.backward( + slot0, + lane0, + grad0, + ) + graph_dx1, graph_dp1, graph_operands1 = resources.backward( + slot1, + lane0, + grad1, + ) + graph_overflow = resources.finalize_overflow((slot0, slot1)) + dist.barrier(group=group) + + with torch.cuda.stream(stream): + graph.replay() + stream.synchronize() + dist.barrier(group=group) + if int(graph_overflow.item()) != 0: + raise AssertionError("captured fixed-resource graph overflowed") + + captured = { + name: tensor + for name, tensor in zip( + comparison_names, + ( + graph_y0, + graph_y1, + graph_dx0, + graph_dx1, + graph_dp0, + graph_dp1, + graph_operands0.fc1_a, + graph_operands0.fc1_b, + graph_operands0.fc2_a, + graph_operands0.fc2_b, + graph_operands1.fc1_a, + graph_operands1.fc1_b, + graph_operands1.fc2_a, + graph_operands1.fc2_b, + ), + ) + } + for name in comparison_names: + _assert_replay_tensor(name, captured[name], ordinary[name]) + torch.testing.assert_close( + graph_operands0.expert_offsets, + ordinary_offsets[0], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + graph_operands1.expert_offsets, + ordinary_offsets[1], + rtol=0, + atol=0, + ) + + if full_probe: + # Diagnostic mode aligns ranks after every replay and verifies that + # fixed-slot dprob reset prevents history accumulation. + dprob_reference = graph_dp0.clone() + for _ in range(diagnostic_replays): + with torch.cuda.stream(stream): + graph.replay() + stream.synchronize() + dist.barrier(group=group) + if int(graph_overflow.item()) != 0: + raise AssertionError( + "fixed-resource diagnostic replay overflowed" + ) + torch.testing.assert_close( + graph_dp0, + dprob_reference, + rtol=1e-5, + atol=1e-6, + ) + + # Production-like burst: no synchronization or host collective in + # the loop. The graph contains the captured scalar overflow OR. + with torch.cuda.stream(stream): + for _ in range(burst_replays): + graph.replay() + stream.synchronize() + dist.barrier(group=group) + if int(graph_overflow.item()) != 0: + raise AssertionError("fixed-resource replay burst overflowed") + torch.testing.assert_close( + graph_dp0, + ordinary["dprob0"], + rtol=1e-5, + atol=1e-6, + ) + + # Overflow both slots, then restore their distinct valid patterns. + overflow = _route_pattern("overflow", rank, world_size, device) + with torch.cuda.stream(stream): + args0[3].copy_(overflow[0]) + args0[4].copy_(overflow[1]) + args1[3].copy_(overflow[0]) + args1[4].copy_(overflow[1]) + graph.replay() + stream.synchronize() + dist.barrier(group=group) + if int(graph_overflow.item()) != 1: + raise AssertionError("fixed-resource overflow was not global") + + with torch.cuda.stream(stream): + args0[3].copy_(local[0]) + args0[4].copy_(local[1]) + args1[3].copy_(remote[0]) + args1[4].copy_(remote[1]) + graph.replay() + stream.synchronize() + dist.barrier(group=group) + recovered_overflow = int(graph_overflow.item()) + if recovered_overflow != 0: + raise AssertionError( + "fixed-resource graph did not recover: " + f"rank={rank}, global_overflow={recovered_overflow}, " + f"slot0_routing_restored=" + f"{torch.equal(args0[3], local[0])}, " + f"slot1_routing_restored=" + f"{torch.equal(args1[3], remote[0])}" + ) + + if rank == 0: + mode = "full" if full_probe else "reinit" + effective_burst = burst_replays if full_probe else 0 + print( + f"MOE_EP_EP{world_size}_TRAINING_RESOURCES_GRAPH_PASS " + f"mode={mode} burst={effective_burst}", + flush=True, + ) + finally: + if graph is not None: + del graph + _close_probe_operator(device=device, group=group, op=op) + + +def _run_multistream_resource_probe( + *, + rank: int, + world_size: int, + device: torch.device, + group, + replays: int, + max_recv_size_per_rank: int, +) -> None: + """Capture two independent lanes with deterministic cross-rank ordering.""" + + args0, grad0, args1, grad1, _, _ = _make_two_slot_inputs( + rank, + world_size, + device, + ) + op = _make_operator( + world_size=world_size, + group=group, + max_recv_size_per_rank=max_recv_size_per_rank, + drop_on_overflow=True, + ) + graph = None + try: + with _debug_phase_scope(rank, "multistream.prepare"): + resources = op.prepare_training_resources( + _make_training_weights(args0), + slot_count=2, + lane_count=2, + ) + slot0, slot1 = resources.slots + lane0, lane1 = resources.lanes + with _debug_phase_scope(rank, "multistream.refresh-weights"): + resources.refresh_weights() + + with _debug_phase_scope(rank, "multistream.lane0-forward"): + eager_y0 = resources.forward( + slot0, lane0, args0[0], args0[3], args0[4] + ) + with _debug_phase_scope(rank, "multistream.lane0-backward"): + eager_dx0, eager_dp0, _ = resources.backward( + slot0, + lane0, + grad0, + ) + with _debug_phase_scope(rank, "multistream.lane0-finalize"): + resources.finalize_overflow((slot0,), lane0) + _synchronize_with_watchdog( + rank, + device, + "multistream.lane0-synchronize", + ) + with _debug_phase_scope(rank, "multistream.lane0-barrier"): + dist.barrier(group=group) + + with _debug_phase_scope(rank, "multistream.lane1-forward"): + eager_y1 = resources.forward( + slot1, lane1, args1[0], args1[3], args1[4] + ) + with _debug_phase_scope(rank, "multistream.lane1-backward"): + eager_dx1, eager_dp1, _ = resources.backward( + slot1, + lane1, + grad1, + ) + with _debug_phase_scope(rank, "multistream.lane1-finalize"): + resources.finalize_overflow((slot1,), lane1) + _synchronize_with_watchdog( + rank, + device, + "multistream.lane1-synchronize", + ) + with _debug_phase_scope(rank, "multistream.lane1-barrier"): + dist.barrier(group=group) + expected = tuple( + tensor.clone() + for tensor in ( + eager_y0, + eager_dx0, + eager_dp0, + eager_y1, + eager_dx1, + eager_dp1, + ) + ) + + capture_stream = torch.cuda.Stream(device=device) + lane_stream0 = torch.cuda.Stream(device=device) + lane_stream1 = torch.cuda.Stream(device=device) + fork_event = torch.cuda.Event() + done_event0 = torch.cuda.Event() + done_event1 = torch.cuda.Event() + capture_stream.wait_stream(torch.cuda.current_stream(device)) + + # One outer graph visits two lane-bound streams, rejoins them, then + # emits exactly one NCCL overflow finalizer. The MegaMoE kernels use + # device-side cross-rank software synchronization and consume one CTA + # slot per SM. Launching both lanes concurrently can let different + # ranks schedule different lanes first, leaving each lane waiting for + # peers whose matching kernel cannot be scheduled. Chain lane 1 after + # lane 0 so every rank observes the same collective order while still + # validating independent per-stream lane storage and graph edges. + graph = torch.cuda.CUDAGraph() + capture_watchdog = _RuntimeWatchdog("multistream.capture") + capture_watchdog.start() + with _debug_phase_scope(rank, "multistream.capture"): + try: + with torch.cuda.graph(graph, stream=capture_stream): + fork_event.record(capture_stream) + lane_stream0.wait_event(fork_event) + with torch.cuda.stream(lane_stream0): + graph_y0 = resources.forward( + slot0, + lane0, + args0[0], + args0[3], + args0[4], + ) + graph_dx0, graph_dp0, _ = resources.backward( + slot0, + lane0, + grad0, + ) + done_event0.record(lane_stream0) + lane_stream1.wait_event(done_event0) + with torch.cuda.stream(lane_stream1): + graph_y1 = resources.forward( + slot1, + lane1, + args1[0], + args1[3], + args1[4], + ) + graph_dx1, graph_dp1, _ = resources.backward( + slot1, + lane1, + grad1, + ) + done_event1.record(lane_stream1) + capture_stream.wait_event(done_event1) + graph_overflow = resources.finalize_overflow( + (slot0, slot1), + lane0, + ) + finally: + capture_watchdog.close() + with _debug_phase_scope(rank, "multistream.capture-barrier"): + dist.barrier(group=group) + + with _debug_phase_scope(rank, "multistream.replay"): + with torch.cuda.stream(capture_stream): + for _ in range(replays): + graph.replay() + replay_watchdog = _RuntimeWatchdog( + "multistream.replay-synchronize" + ) + replay_watchdog.start() + with _debug_phase_scope( + rank, + "multistream.replay-synchronize", + ): + try: + capture_stream.synchronize() + finally: + replay_watchdog.close() + with _debug_phase_scope(rank, "multistream.replay-barrier"): + dist.barrier(group=group) + if int(graph_overflow.item()) != 0: + raise AssertionError("multi-stream fixed-resource graph overflowed") + + actual = ( + graph_y0, + graph_dx0, + graph_dp0, + graph_y1, + graph_dx1, + graph_dp1, + ) + for index, (value, reference) in enumerate(zip(actual, expected)): + _assert_replay_tensor( + f"multistream[{index}]", + value, + reference, + ) + if rank == 0: + print( + f"MOE_EP_EP{world_size}_MULTISTREAM_GRAPH_PASS " + f"replays={replays}", + flush=True, + ) + finally: + if graph is not None: + del graph + _close_probe_operator(device=device, group=group, op=op) + + +def _run_error_mode_assert_probe( + *, + rank: int, + world_size: int, + device: torch.device, + group, + max_recv_size_per_rank: int, +) -> None: + """Require a captured global overflow to assert on every rank.""" + + args, grad_output = _make_inputs(rank, device) + local = _route_pattern("local", rank, world_size, device) + overflow = _route_pattern("overflow", rank, world_size, device) + route_indices = local[0].clone() + route_weights = local[1].clone() + op = _make_operator( + world_size=world_size, + group=group, + max_recv_size_per_rank=max_recv_size_per_rank, + drop_on_overflow=False, + ) + resources = op.prepare_training_resources( + _make_training_weights(args), + slot_count=1, + lane_count=1, + ) + slot = resources.slots[0] + lane = resources.lanes[0] + + # Warm every kernel and prove the assertion accepts a valid execution. + resources.refresh_weights() + resources.forward( + slot, + lane, + args[0], + route_indices, + route_weights, + ) + resources.backward(slot, lane, grad_output) + resources.finalize_overflow((slot,), lane) + torch.cuda.synchronize(device) + dist.barrier(group=group) + + stream = torch.cuda.Stream(device=device) + stream.wait_stream(torch.cuda.current_stream(device)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + resources.refresh_weights() + resources.forward( + slot, + lane, + args[0], + route_indices, + route_weights, + ) + resources.backward(slot, lane, grad_output) + resources.finalize_overflow((slot,), lane) + dist.barrier(group=group) + + # A valid replay confirms capture before intentionally poisoning the + # context with the fatal error-mode assertion. + with torch.cuda.stream(stream): + graph.replay() + stream.synchronize() + dist.barrier(group=group) + + try: + with torch.cuda.stream(stream): + route_indices.copy_(overflow[0]) + route_weights.copy_(overflow[1]) + graph.replay() + stream.synchronize() + except BaseException as exc: + print( + f"MOE_EP_EP{world_size}_ERROR_MODE_ASSERT_PASS " + f"rank={rank} error={type(exc).__name__}", + flush=True, + ) + # CUDA device assertions poison the process context. Do not run Python + # destructors, NCCL collectives, or NVSHMEM finalization afterward. + os._exit(0) + + print( + f"MOE_EP_EP{world_size}_ERROR_MODE_ASSERT_MISSING rank={rank}", + flush=True, + ) + os._exit(1) + + +def main() -> None: + args = _parse_args() + _require_positive("diagnostic_replays", args.diagnostic_replays) + _require_positive("burst_replays", args.burst_replays) + _require_positive("multistream_replays", args.multistream_replays) + _require_positive("cycles", args.cycles) + _require_positive( + "max_recv_size_per_rank", + args.max_recv_size_per_rank, + ) + + world_size = int(os.environ.get("WORLD_SIZE", "1")) + rank = int(os.environ.get("RANK", "0")) + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + if world_size < 2: + raise RuntimeError( + f"this probe requires WORLD_SIZE >= 2, got {world_size}" + ) + forced_overflow_routes = world_size * 8 * 2 + if args.max_recv_size_per_rank >= forced_overflow_routes: + raise ValueError( + "max_recv_size_per_rank must remain below the probe's forced " + f"overflow route count {forced_overflow_routes}, got " + f"{args.max_recv_size_per_rank}" + ) + + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + capability = torch.cuda.get_device_capability(device) + if capability != (10, 7): + raise RuntimeError( + "this probe requires Rubin SM107; " + f"rank {rank} found compute capability {capability}" + ) + os.environ.setdefault("CUTE_DSL_ARCH", "sm_107a") + + dist.init_process_group( + backend="nccl", + init_method="env://", + device_id=device, + timeout=timedelta(seconds=args.timeout_seconds), + ) + try: + if args.expect_overflow_assert: + _run_error_mode_assert_probe( + rank=rank, + world_size=world_size, + device=device, + group=dist.group.WORLD, + max_recv_size_per_rank=args.max_recv_size_per_rank, + ) + raise AssertionError("fatal overflow assertion probe returned") + for cycle in range(args.cycles): + with _debug_phase_scope( + rank, + f"training-resources-cycle-{cycle}", + ): + _run_training_resource_probe( + rank=rank, + world_size=world_size, + device=device, + group=dist.group.WORLD, + diagnostic_replays=args.diagnostic_replays, + burst_replays=args.burst_replays, + max_recv_size_per_rank=args.max_recv_size_per_rank, + full_probe=cycle == 0, + ) + if not args.skip_multistream: + with _debug_phase_scope(rank, "multistream"): + _run_multistream_resource_probe( + rank=rank, + world_size=world_size, + device=device, + group=dist.group.WORLD, + replays=args.multistream_replays, + max_recv_size_per_rank=args.max_recv_size_per_rank, + ) + if rank == 0: + print( + f"MOE_EP_EP{world_size}_CUDA_GRAPH_PROBE_PASS", + flush=True, + ) + finally: + if dist.is_initialized(): + try: + with _debug_phase_scope(rank, "runtime-shutdown"): + get_runtime_manager().shutdown() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/test/python/moe_ep/test_moe_ep_backward.py b/test/python/moe_ep/test_moe_ep_backward.py index b58d3e9ae..a451db390 100644 --- a/test/python/moe_ep/test_moe_ep_backward.py +++ b/test/python/moe_ep/test_moe_ep_backward.py @@ -1,400 +1,312 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Core MoE EP backward contract, parity, and distributed tests.""" +"""Fixed-resource MoE EP backward and training-graph contracts.""" from __future__ import annotations +from contextlib import nullcontext import inspect import os -from dataclasses import replace +import threading +from dataclasses import fields +from pathlib import Path from types import SimpleNamespace +from unittest.mock import Mock +import cudnn import pytest import torch import torch.multiprocessing as mp -from cudnn.moe_ep import MoeEp -from cudnn.moe_ep._contracts import ForwardConfig -from cudnn.moe_ep._megamoe_backend import _capability -from cudnn.moe_ep._megamoe_backend.mxfp8._backend import Mxfp8Backend -from cudnn.moe_ep._megamoe_backend.mxfp8._backward_layout import ( - Mxfp8BackwardLayout, +from cudnn.moe_ep import ( + MoeEp, + MoeEpExecutionLane, + MoeEpTrainingResources, + MoeEpTrainingSlot, + MoeEpTrainingWgradOperands, ) -from cudnn.moe_ep._megamoe_backend.mxfp8._backward_staging import ( - _stage_fc1_preact, - stage_backward, +from cudnn.moe_ep._validation import validate_training_weights +from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( + _typed_k_major_view, ) -from cudnn.moe_ep._megamoe_backend._workspace import WorkspaceRequirements -from cudnn.moe_ep._megamoe_backend.mxfp8._backward_dispatch import ( - Mxfp8BackwardRedispatch, +from cudnn.moe_ep._megamoe_backend._workspace import ( + BufferRegion, + WorkspaceRequirements, ) -from cudnn.moe_ep._megamoe_backend.mxfp8._backward_dprob import ( - return_grad_topk_weights, +from cudnn.moe_ep._megamoe_backend.mxfp8._training_resources import ( + Mxfp8TrainingResourceOwner, + _build_training_abi_facts, + _harmonize_symmetric_regions, + _verify_training_abi_across_ranks, + build_training_workspace_requirements, +) +from cudnn.moe_ep._megamoe_backend.mxfp8._training_stage import ( + Mxfp8TrainingStager, +) +from cudnn.moe_ep._megamoe_backend.mxfp8._fingerprint import ( + canonical_json_sha256, +) +from cudnn.moe_ep._megamoe_backend.mxfp8._training_weights import ( + Mxfp8TrainingWeightBindings, ) from cudnn.moe_ep._tuning import MoeEpTuningConfig -from cudnn.moe_ep._validation import validate_backward -from moe_ep.moe_ep_backward_support import ( - _assert_backward_matches, - _expected_backward, - _grad_output, - _reference_backward, +from moe_ep.moe_ep_reference import ( + MoeEpReference, ) from moe_ep.moe_ep_distributed_workers import ( - _distributed_backward_worker, + _distributed_backward_reference_worker, + _distributed_subgroup_backward_reference_worker, ) -from moe_ep.moe_ep_forward_support import ( - _forward_config, +from moe_ep.moe_ep_test_support import ( + _assert_fixed_training_drop_overflow_result, + _assert_fixed_training_matches_reference, + _assert_training_graph_tails_are_reset, + _assert_training_weight_sources_changed, + _capture_fixed_training_batch, + _copy_training_weight_sources_, + _dense_wgrads_from_operands, + _fixed_training_case, + _fixed_training_drop_overflow_case, + _fixed_training_drop_overflow_reference, + _fixed_training_reference, + _fixed_training_weights, + _grad_output, + _prefill_training_graph_sentinels, _require_distributed_sm107, + _run_fixed_training_batch, _sm107_device, -) -from moe_ep.moe_ep_reference import ( - MoeFormat, - backward_combine_round_trip, - forward_combine_round_trip, -) -from moe_ep.moe_ep_test_data import ( + _training_public_pointers, + _training_source_pointers, + _training_weight_source_pointers, + _training_weight_source_values, + _TrainingResourceContractOwner, + _training_abi_prepared, + _training_config, + _training_contract_resources, + _training_inputs, + _training_prepared_pair, + _training_staging_tensors, + _training_weight_defect, + _training_weights, make_forward_inputs, - quantize_mxfp8, ) - -def _config(**overrides) -> ForwardConfig: - values = { - "num_experts": 2, - "hidden_size": 128, - "intermediate_size": 256, - "top_k": 2, - "experts_per_rank": 2, - "ep_size": 1, - "ep_rank": 0, - "ep_group": None, - "ep_global_ranks": (), - "max_tokens_per_rank": 4, - "output_format": "bf16", - "combine_format": "bf16", - "apply_topk_in_fc1": True, - "gate_up_clamp": None, - "generate_c": True, - "token_padding_size": 128, - "sf_padding_size": 128, - "tuning": MoeEpTuningConfig(), - } - values.update(overrides) - return ForwardConfig(**values) - - -def _inputs(): - activation = torch.randn(2, 128, dtype=torch.bfloat16) - fc1_weight = torch.randn(2, 128, 512, dtype=torch.bfloat16) - fc2_weight = torch.randn(2, 256, 128, dtype=torch.bfloat16) - topk_idx = torch.tensor([[0, -1], [1, 0]], dtype=torch.int32) - topk_weights = torch.randn(2, 2, dtype=torch.float32) - return activation, fc1_weight, fc2_weight, topk_idx, topk_weights - - -def _validate_backward( - config, - grad_output, - args, - fc1_c, - route_metadata, -): - return validate_backward( - config, - grad_output, - *args[1:], - fc1_c, - route_metadata, - ) - - -# Validation, layout, staging, and backend contracts. +# L0 contracts @pytest.mark.L0 -def test_validate_backward_builds_typed_request_and_checks_stash(): - config = _config() - args = _inputs() - grad_output = torch.randn(2, 128) - fc1_c = torch.randn(3, 512, dtype=torch.bfloat16) - route_metadata = torch.tensor( - [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], - dtype=torch.int32, - ) - - request = _validate_backward( - config, - grad_output, - args, - fc1_c, - route_metadata, - ) - - assert request.config is config - assert request.fc1_weight is args[1] - assert request.topk_idx is args[3] - assert request.local_routes == 3 - with pytest.raises(ValueError, match="fc1_c shape must be"): - _validate_backward( - config, - grad_output, - args, - fc1_c[:2], - route_metadata, +@pytest.mark.parametrize( + "case", + [ + pytest.param("backward-regions", id="backward-regions"), + pytest.param("slot-lane-layout", id="slot-lane-layout"), + ], +) +def test_training_workspace_layout_contract(case): + if case == "backward-regions": + requirements = WorkspaceRequirements.for_mxfp8( + _training_config(), + kernel_local_workspace_bytes=64, + kernel_shared_workspace_bytes=128, + backward_fc1_preact_bytes=1024, + backward_dprob_bytes=32, + backward_aux_data_bytes=512, + backward_aux_scale_bytes=256, ) - with pytest.raises(TypeError, match="route_metadata must have dtype"): - _validate_backward( + expected = ( + ("symmetric", "backward_dprob", 32, None), + ("local", "backward_fc1_preact", 1024, 128), + ("local", "backward_aux_data", 512, None), + ("local", "backward_aux_scale", 256, None), + ) + else: + config = _training_config() + forward, backward = _training_prepared_pair(config) + requirements = build_training_workspace_requirements( config, - grad_output, - args, - fc1_c, - route_metadata.to(torch.int64), + forward, + backward, + slot_count=2, + lane_count=1, + ) + expected = ( + ("symmetric", "lane.0.forward.symmetric.kernel_shared_workspace", None, None), + ("symmetric", "lane.0.backward.symmetric.kernel_shared_workspace", None, None), + ("symmetric", "slot.0.backward.symmetric.backward_dprob", None, None), + ("symmetric", "slot.1.backward.symmetric.backward_dprob", None, None), + ("local", "slot.0.persistent.local.fc1_preact", None, None), + ("local", "slot.1.persistent.local.fc1_preact", None, None), + ) + + regions = { + "symmetric": {region.name: region for region in requirements.symmetric_regions}, + "local": {region.name: region for region in requirements.local_regions}, + } + for storage, name, nbytes, alignment in expected: + region = regions[storage][name] + if nbytes is not None: + assert region.nbytes == nbytes + if alignment is not None: + assert region.alignment == alignment + + if case == "slot-lane-layout": + assert tuple(region.name for region in requirements.symmetric_regions if region.name.startswith("slot.0.")) == ( + "slot.0.forward.symmetric.output_data", + "slot.0.backward.symmetric.backward_dprob", + "slot.0.backward.symmetric.output_data", + "slot.0.persistent.symmetric.routing_topk_weights", ) @pytest.mark.L0 -def test_mxfp8_backward_layout_builds_public_preactivation_lut(): - config = _config() - args = _inputs() - route_metadata = torch.tensor( - [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], - dtype=torch.int32, - ) - request = _validate_backward( - config, - torch.randn(2, 128), - args, - torch.randn(3, 512, dtype=torch.bfloat16), - route_metadata, +def test_training_workspace_harmonizes_each_symmetric_region(monkeypatch): + requirements = WorkspaceRequirements( + max_tokens_per_rank=1, + symmetric_regions=( + BufferRegion("first", 1), + BufferRegion("second", 257), + ), + local_regions=(BufferRegion("local", 1),), ) + runtime = SimpleNamespace(world_size=2, group=object()) + + def all_reduce(tensor, *, op, group): + assert group is runtime.group + if tensor.numel() == 2 and op == torch.distributed.ReduceOp.MAX: + tensor.copy_(torch.tensor([257, 257], dtype=torch.int64)) - layout = Mxfp8BackwardLayout.from_request(request) + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + harmonized = _harmonize_symmetric_regions( + requirements, + runtime, + torch.device("cpu"), + ) - assert layout.preact_row_lut[0, 0, 0].item() == 0 - assert layout.preact_row_lut[0, 1, 1].item() == 1 - assert layout.preact_row_lut[0, 1, 0].item() == 2 + assert tuple(region.nbytes for region in harmonized.symmetric_regions) == (257, 257) + assert harmonized.local_regions == requirements.local_regions @pytest.mark.L0 -def test_mxfp8_backward_stages_compact_preactivation_into_pool_rows(): - config = _config() - args = _inputs() - route_metadata = torch.tensor( - [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], - dtype=torch.int32, +def test_training_abi_fingerprint_is_stable_and_structural(): + config = _training_config(ep_size=2, ep_global_ranks=(0, 1)) + forward = _training_abi_prepared("forward") + backward = _training_abi_prepared("backward") + weights = _training_weights() + requirements = WorkspaceRequirements( + max_tokens_per_rank=4, + symmetric_regions=(BufferRegion("symmetric", 256),), + local_regions=(BufferRegion("local", 128),), ) - fc1_c = torch.arange( - 3 * 512, - dtype=torch.float32, - ).reshape(3, 512).to(torch.bfloat16) - request = _validate_backward( + first = _build_training_abi_facts( config, - torch.randn(2, 128), - args, - fc1_c, - route_metadata, - ) - layout = Mxfp8BackwardLayout.from_request(request) - pool_capacity = 256 - fc1_preact = torch.empty( - pool_capacity, - 512, - dtype=torch.bfloat16, - ) - prepared = SimpleNamespace( - config=SimpleNamespace( - intermediate=256, - num_experts=2, - token_padding_block=128, - ), - kernel=SimpleNamespace( - token_comm=SimpleNamespace( - router_data_cta_count=1, - router_warps_per_cta=4, - ) - ), - pool_token_capacity=pool_capacity, + forward, + backward, + weights, + requirements, + slot_count=2, + lane_count=1, + source_tree_digest="source", + ) + second = _build_training_abi_facts( + config, + forward, + backward, + weights, + requirements, + slot_count=2, + lane_count=1, + source_tree_digest="source", + ) + changed = _build_training_abi_facts( + config, + forward, + backward, + weights, + requirements, + slot_count=2, + lane_count=2, + source_tree_digest="source", ) - _stage_fc1_preact(request, layout, prepared, fc1_preact) - - staged = fc1_preact - gate, up = fc1_c.split(256, dim=1) - expected = torch.stack( - (gate.reshape(3, 8, 32), up.reshape(3, 8, 32)), - dim=2, - ).reshape(3, 512) - torch.testing.assert_close(staged[0], expected[0], rtol=0, atol=0) - torch.testing.assert_close(staged[1], expected[1], rtol=0, atol=0) - torch.testing.assert_close(staged[128], expected[2], rtol=0, atol=0) - assert staged[2:128].eq(0).all() - assert staged[129:].eq(0).all() + assert canonical_json_sha256(first) == canonical_json_sha256(second) + assert canonical_json_sha256(first) != canonical_json_sha256(changed) @pytest.mark.L0 -def test_mxfp8_backward_stages_sources_in_destination_ring_order(): - metadata = torch.tensor( - [[0, 0, 0, 0], [0, 1, 0, 0]], - dtype=torch.int64, - ) - fc1_c = torch.stack( - ( - torch.cat( - ( - torch.full((32,), 10, dtype=torch.bfloat16), - torch.full((32,), 11, dtype=torch.bfloat16), - ) - ), - torch.cat( - ( - torch.full((32,), 20, dtype=torch.bfloat16), - torch.full((32,), 21, dtype=torch.bfloat16), - ) - ), - ) - ) - request = SimpleNamespace( - config=SimpleNamespace( - ep_rank=1, - ep_size=2, - max_tokens_per_rank=1, - top_k=1, - ), - route_metadata=metadata, - fc1_c=fc1_c, - ) - layout = SimpleNamespace( - preact_row_lut=torch.tensor([[[0]], [[1]]], dtype=torch.int32) - ) - pool_capacity = 128 - fc1_preact = torch.empty( - pool_capacity, - 64, - dtype=torch.bfloat16, - ) - prepared = SimpleNamespace( - config=SimpleNamespace( - intermediate=32, - num_experts=1, - token_padding_block=128, - ), - kernel=SimpleNamespace( - token_comm=SimpleNamespace( - router_data_cta_count=1, - router_warps_per_cta=4, - ) - ), - pool_token_capacity=pool_capacity, - ) - - _stage_fc1_preact(request, layout, prepared, fc1_preact) +def test_training_abi_handshake_rejects_rank_mismatch(monkeypatch): + runtime = SimpleNamespace(world_size=2, group=object()) - staged = fc1_preact - # Destination rank 1 receives source rank 1 before wrapped source rank 0. - expected = torch.stack( - (fc1_c[:, :32], fc1_c[:, 32:]), - dim=1, - ).reshape(2, 64) - torch.testing.assert_close(staged[0], expected[1], rtol=0, atol=0) - torch.testing.assert_close(staged[1], expected[0], rtol=0, atol=0) + def all_reduce(tensor, *, op, group): + assert group is runtime.group + if op == torch.distributed.ReduceOp.MAX: + tensor.add_(1) + def all_gather_object(output, value, *, group): + assert group is runtime.group + output[:] = [value, "different"] -@pytest.mark.L0 -def test_mxfp8_backward_stages_source_routes_in_router_vector_order(): - metadata = torch.tensor( - [ - [0, 0, 0, 0], - [0, 0, 1, 1], - [0, 0, 2, 0], - [0, 0, 3, 1], - [0, 0, 4, 0], - ], - dtype=torch.int64, - ) - fc1_c = torch.arange(5, dtype=torch.bfloat16).view(5, 1).expand( - 5, - 64, - ).contiguous() - request = SimpleNamespace( - config=SimpleNamespace( - ep_rank=0, - ep_size=1, - max_tokens_per_rank=5, - top_k=2, - ), - route_metadata=metadata, - fc1_c=fc1_c, - ) - preact_row_lut = torch.full((1, 5, 2), -1, dtype=torch.int32) - preact_row_lut[ - metadata[:, 1], - metadata[:, 2], - metadata[:, 3], - ] = torch.arange(5, dtype=torch.int32) - layout = SimpleNamespace(preact_row_lut=preact_row_lut) - pool_capacity = 128 - fc1_preact = torch.empty( - pool_capacity, - 64, - dtype=torch.bfloat16, - ) - prepared = SimpleNamespace( - config=SimpleNamespace( - intermediate=32, - num_experts=1, - token_padding_block=128, - ), - kernel=SimpleNamespace( - token_comm=SimpleNamespace( - router_data_cta_count=1, - router_warps_per_cta=4, - ) - ), - pool_token_capacity=pool_capacity, - ) + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + monkeypatch.setattr( + torch.distributed, + "all_gather_object", + all_gather_object, + ) + with pytest.raises(RuntimeError, match="ABI differs"): + _verify_training_abi_across_ranks( + {"schema_version": 1}, + runtime, + torch.device("cpu"), + ) - _stage_fc1_preact(request, layout, prepared, fc1_preact) - staged = fc1_preact - # Int32 router loads four adjacent routes per thread, then stable-sorts by - # register round and lane: flat routes 0,4,8 precede 3,7. - assert staged[:5, 0].tolist() == [0, 2, 4, 1, 3] +@pytest.mark.L0 +def test_training_resource_views_share_lane_scratch_but_not_slot_state(): + config = _training_config() + forward, backward = _training_prepared_pair(config) + class Runtime: + device = torch.device("cpu") + rank = 0 + world_size = 1 + nvshmem_enabled = False + closed = False -@pytest.mark.L0 -def test_mxfp8_backward_workspace_regions_are_explicit_and_symmetric(): - requirements = WorkspaceRequirements.for_mxfp8( - _config(), - kernel_local_workspace_bytes=64, - kernel_shared_workspace_bytes=128, - backward_fc1_preact_bytes=1024, - backward_dprob_bytes=32, - backward_aux_data_bytes=512, - backward_aux_scale_bytes=256, - ) - symmetric = { - region.name: region for region in requirements.symmetric_regions - } - local = {region.name: region for region in requirements.local_regions} + def ensure_open(self): + assert not self.closed - assert symmetric["backward_dprob"].nbytes == 32 - assert local["backward_fc1_preact"].nbytes == 1024 - assert local["backward_fc1_preact"].alignment == 128 - assert local["backward_aux_data"].nbytes == 512 - assert local["backward_aux_scale"].nbytes == 256 + def close(self): + self.closed = True - with pytest.raises(ValueError, match="must be enabled together"): - WorkspaceRequirements.for_mxfp8( - _config(), - kernel_local_workspace_bytes=64, - kernel_shared_workspace_bytes=128, - backward_dprob_bytes=32, - ) + runtime = Runtime() + runtime_manager = SimpleNamespace(acquire=lambda actual_config, actual_device: runtime) + weights = _training_weights() + owner = Mxfp8TrainingResourceOwner( + config, + torch.device("cpu"), + forward, + backward, + weights, + slot_count=2, + lane_count=1, + runtime_manager=runtime_manager, + ) + try: + first = owner.views(slot=0, lane=0, token_count=4) + second = owner.views(slot=1, lane=0, token_count=4) + assert first.forward.workspace.local["kernel_local_workspace"].data_ptr() == second.forward.workspace.local["kernel_local_workspace"].data_ptr() + assert first.slot.fc1_preact.data_ptr() != second.slot.fc1_preact.data_ptr() + assert first.slot.dprob.data_ptr() != second.slot.dprob.data_ptr() + assert first.forward_expert_size_snapshot is not None + assert first.forward_expert_size_snapshot.data_ptr() == (second.forward_expert_size_snapshot.data_ptr()) + finally: + owner.close() + + assert runtime.closed @pytest.mark.L0 -def test_rubin_adapter_source_tracks_current_kernel_signatures(): +def test_training_sources_track_adapter_grad_y2_and_dfc2_contracts(): from cudnn.moe_ep._megamoe_backend.mxfp8 import ( _backward_compile, _compile, @@ -402,753 +314,726 @@ def test_rubin_adapter_source_tracks_current_kernel_signatures(): forward_source = inspect.getsource(_compile.prepare_kernel) backward_source = inspect.getsource(_backward_compile.prepare_backward_kernel) - runtime_source = inspect.getsource( - _backward_compile.build_backward_runtime_kwargs - ) + runtime_source = inspect.getsource(_backward_compile.build_backward_runtime_kwargs) + dglu_source = _DGLU.read_text(encoding="utf-8") + dfc2_source = _DGLU_EPILOGUE.read_text(encoding="utf-8") - assert "apply_topk_in_fc1=config.apply_topk_in_fc1" not in forward_source assert "gate_up_clamp=config.gate_up_clamp" in backward_source assert "dfc2_recompute=dfc2_recompute" in backward_source - assert "dfc2_col_output=dfc2_col_output" in backward_source assert "enable_grad_y2_col_quant=enable_grad_y2_col_quant" in backward_source assert '"fc1_preact":' in runtime_source - overflow_runtime_source = runtime_source.split( - '"overflow_flag":', - 1, - )[1].split('"dprob":', 1)[0] - assert "dynamic_layout=False" in overflow_runtime_source - for output_name in ( - "dprob", - "fc1_recompute", - "fc1_recompute_sf", - "fc1_col_output", - "fc1_col_output_sf", - "grad_y2", - "grad_y2_sf", + assert '"dprob":' in runtime_source + assert "generate_c=config.generate_c" in forward_source + for contract in ( + "enable_grad_y2_col_quant", + "num_ctas_grad_y2_col_quant", + "grad_y2_sizes_region", + "_snapshot_grad_y2_expert_sizes", + "grad_y2_col_quant", + "grad_y2: cute.Tensor", + "grad_y2_sf: cute.Tensor", ): - assert f'"{output_name}":' in runtime_source + assert contract in dglu_source + assert dglu_source.index("self._snapshot_grad_y2_expert_sizes(tidx)") < dglu_source.index("self.token_comm.reset_tail()") + assert dglu_source.index("self._topk_reduce(") < dglu_source.index("self.grad_y2_col_quant(") + assert "def _stg_col_sf_atom_value(" in dfc2_source + assert "feature_atom = feature // cutlass.Int32(128)" in dfc2_source + assert "feature_lane * cutlass.Int32(16)" in dfc2_source + assert "feature_bank * cutlass.Int32(4)" in dfc2_source + assert "real_sf[feature_atom, token_atom, atom_byte]" in dfc2_source + assert "def tma_store_dfc2_outputs(" in dfc2_source + assert dfc2_source.count("self._stg_col_sf_atom_value(") >= 2 -@pytest.mark.L0 -def test_stage_backward_exposes_fixed_aux_shapes_and_resets_symmetric_dprob(): - config = _config() - args = _inputs() - route_metadata = torch.tensor( - [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], - dtype=torch.int32, - ) - request = _validate_backward( - config, - torch.randn(2, 128), - args, - torch.randn(3, 512, dtype=torch.bfloat16), - route_metadata, - ) - layout = Mxfp8BackwardLayout.from_request(request) - aux_shapes = { - "dprob": (4, 2), - "fc1_recompute": (8, 256), - "fc1_recompute_sf": (1, 256), - "fc1_col_output": (8, 512), - "fc1_col_output_sf": (1, 512), - "grad_y2": (8, 128), - "grad_y2_sf": (32,), - } - kernel = SimpleNamespace( - token_comm=SimpleNamespace( - router_data_cta_count=1, - router_warps_per_cta=4, - ), - get_fc1_preact_shape=lambda: (256, 512), - get_aux_output_shapes=lambda: aux_shapes, - ) - prepared = SimpleNamespace( - config=SimpleNamespace( - max_tokens_per_rank=4, - hidden=128, - top_k=2, - intermediate=256, - num_experts=2, - combine_format="bf16", - token_padding_block=128, - ), - kernel=kernel, - pool_token_capacity=256, - pre_reduced_activation_offset=0, - pre_reduced_activation_bytes_per_token=4, - pre_reduced_activation_sf_offset=None, - pre_reduced_activation_sf_bytes_per_token=0, - local_workspace_zero_bytes=0, - shared_workspace_zero_bytes=0, - dfc2_recompute=False, - dfc2_col_output=False, - enable_grad_y2_col_quant=False, - ) - symmetric_dprob = torch.full((32,), 0x7F, dtype=torch.uint8) - resources = SimpleNamespace( - workspace=SimpleNamespace( - symmetric={ - "activation_data": torch.empty(4 * 128, dtype=torch.uint8), - "activation_scale": torch.empty(4 * 16, dtype=torch.uint8), - "topk_weights": torch.empty(4 * 2 * 4, dtype=torch.uint8), - "output_data": torch.empty(4 * 128 * 2, dtype=torch.uint8), - "backward_dprob": symmetric_dprob, - "kernel_shared_workspace": torch.empty( - 64, - dtype=torch.uint8, - ), - }, - local={ - "topk_idx": torch.empty(4 * 2 * 4, dtype=torch.uint8), - "overflow_flag": torch.empty(4, dtype=torch.uint8), - "backward_fc1_preact": torch.empty( - 256 * 512 * 2, - dtype=torch.uint8, - ), - "backward_aux_data": torch.empty( - 8 * 512, - dtype=torch.uint8, - ), - "backward_aux_scale": torch.empty( - 512, - dtype=torch.uint8, - ), - "kernel_local_workspace": torch.empty( - 64, - dtype=torch.uint8, - ), - }, +# WGrad operand contracts + + +@pytest.mark.L1 +@pytest.mark.parametrize( + ("field", "defect"), + [ + pytest.param(field, "logical_shape", id=f"{field}-logical-shape") + for field in ( + "forward_fc1", + "forward_fc2", + "backward_w2_transpose", + "backward_w1_transpose", ) + ] + + [ + pytest.param("forward_fc1", defect, id=f"forward_fc1-{defect}") + for defect in ( + "plain_tensor", + "axis", + "format", + "data_noncontiguous", + "scale_noncontiguous", + ) + ], +) +def test_validate_training_weights_rejects_targeted_defects(field, defect): + invalid, error_type, message = _training_weight_defect( + _training_weights(), + field, + defect, ) + with pytest.raises(error_type) as exc_info: + validate_training_weights(_training_config(), invalid) + assert str(exc_info.value) == message - inputs = stage_backward(request, layout, prepared, resources) - - assert inputs.fc1_preact.shape == (256, 512) - assert inputs.fc1_preact.dtype is torch.bfloat16 - assert inputs.dprob.shape == (4, 2) - assert inputs.dprob.dtype is torch.float32 - assert inputs.dprob.eq(0).all() - assert inputs.fc1_recompute.shape == aux_shapes["fc1_recompute"] - assert inputs.fc1_recompute.dtype is torch.float8_e4m3fn - assert inputs.fc1_recompute_sf.shape == aux_shapes["fc1_recompute_sf"] - assert inputs.fc1_recompute_sf.dtype is torch.float8_e8m0fnu - assert inputs.fc1_col_output.shape == aux_shapes["fc1_col_output"] - assert inputs.fc1_col_output_sf.shape == aux_shapes["fc1_col_output_sf"] - assert ( - inputs.fc1_recompute.data_ptr() - == inputs.fc1_col_output.data_ptr() - ) - assert ( - inputs.fc1_recompute_sf.data_ptr() - == inputs.fc1_col_output_sf.data_ptr() - ) - assert inputs.grad_y2.shape == aux_shapes["grad_y2"] - assert inputs.grad_y2_sf.shape == aux_shapes["grad_y2_sf"] - assert inputs.grad_y2.data_ptr() == inputs.fc1_recompute.data_ptr() - assert inputs.grad_y2_sf.data_ptr() == inputs.fc1_recompute_sf.data_ptr() - operands_config = replace( - config, - backward_wgrad_mode="operands", - token_padding_size=256, +@pytest.mark.L1 +def test_validate_training_weights_rejects_cross_field_device_mismatch(): + invalid, error_type, message = _training_weight_defect( + _training_weights(), + "backward_w1_transpose", + "device", ) - operands_request = replace(request, config=operands_config) - operands_aux_shapes = { - "dprob": (4, 2), - "fc1_recompute": (512, 256), - "fc1_recompute_sf": (8, 256), - "fc1_col_output": (512, 512), - "fc1_col_output_sf": (8, 512), - "grad_y2": (512, 128), - "grad_y2_sf": (256 // 32 * 128,), + with pytest.raises(error_type) as exc_info: + validate_training_weights(_training_config(), invalid) + assert str(exc_info.value) == message + + +@pytest.mark.L1 +def test_validate_training_weights_accepts_complete_fixed_weight_set(): + assert validate_training_weights( + _training_config(), + _training_weights(), + ) == torch.device("cpu") + + +def _operator(**overrides) -> MoeEp: + values = { + "num_experts": 2, + "hidden_size": 128, + "intermediate_size": 256, + "top_k": 2, + "max_tokens_per_rank": 4, } - operands_kernel = SimpleNamespace( - token_comm=kernel.token_comm, - get_fc1_preact_shape=lambda: (512, 512), - get_aux_output_shapes=lambda: operands_aux_shapes, - ) - operands_prepared = SimpleNamespace( - **{ - **vars(prepared), - "config": SimpleNamespace( - **{ - **vars(prepared.config), - "token_padding_block": 256, - } - ), - "kernel": operands_kernel, - "pool_token_capacity": 512, - "dfc2_recompute": True, - "dfc2_col_output": True, - "enable_grad_y2_col_quant": True, - } - ) - operands_local = dict(resources.workspace.local) - operands_local["backward_fc1_preact"] = torch.empty( - 512 * 512 * 2, - dtype=torch.uint8, - ) - operands_resources = SimpleNamespace( - workspace=SimpleNamespace( - symmetric=resources.workspace.symmetric, - local=operands_local, - ) + values.update(overrides) + return MoeEp( + **values, ) - operand_inputs = stage_backward( - operands_request, - Mxfp8BackwardLayout.from_request(operands_request), - operands_prepared, - operands_resources, - ) - assert operand_inputs.fc1_recompute.shape == (512, 256) - assert operand_inputs.fc1_col_output.shape == (512, 512) - assert operand_inputs.grad_y2.shape == (512, 128) - assert operand_inputs.grad_y2_sf.shape == ( - 256 // 32 * 128, - ) - assert ( - operand_inputs.fc1_recompute.data_ptr() - != operands_local["backward_aux_data"].data_ptr() - ) - assert ( - operand_inputs.fc1_col_output.data_ptr() - != operands_local["backward_aux_data"].data_ptr() - ) - assert operand_inputs.fc1_recompute.eq(0).all() - assert operand_inputs.fc1_col_output.eq(0).all() - assert operand_inputs.grad_y2.eq(0).all() - assert operand_inputs.fc1_preact[128:256].eq(0).all() - assert torch.equal( - operand_inputs.fc1_preact[256], - torch.stack( - ( - request.fc1_c[2, :256].reshape(8, 32), - request.fc1_c[2, 256:].reshape(8, 32), - ), - dim=1, - ).reshape(512), +def _install_contract_backend( + monkeypatch, + *, + weights=None, + slot_count=1, + lane_count=1, +): + import cudnn.moe_ep._backend as backend_seam + import cudnn.moe_ep.api as api_module + + weights = weights or SimpleNamespace(mock_training_weights=True) + state = SimpleNamespace( + backends=[], + validate=Mock(return_value=torch.device("cpu")), ) + def create_backend(config, device): + del config, device + owner = _TrainingResourceContractOwner( + slot_count=slot_count, + lane_count=lane_count, + ) + backend = SimpleNamespace( + owner=owner, + prepare_training_resources=Mock(return_value=owner), + close=Mock(), + ) + state.backends.append(backend) + return backend + + monkeypatch.setattr(api_module, "validate_training_weights", state.validate) + monkeypatch.setattr(backend_seam, "validate_config", lambda config: None) + monkeypatch.setattr(backend_seam, "create_backend", create_backend) + return weights, state + @pytest.mark.L0 -@pytest.mark.parametrize("apply_topk_in_fc1", [True, False]) -def test_mxfp8_backward_recomputes_semantic_grad_topk_weights( - apply_topk_in_fc1, -): - args = list(_inputs()) - fc2_weight = torch.zeros_like(args[2]) - fc2_weight[0, 0, 0] = 2 - fc2_weight[1, 0, 0] = 4 - args[2] = fc2_weight - config = _config(apply_topk_in_fc1=apply_topk_in_fc1) - route_metadata = torch.tensor( - [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], - dtype=torch.int32, - ) - fc1_c = torch.zeros(3, 512, dtype=torch.bfloat16) - fc1_c[:, 0] = 1 - fc1_c[:, 256] = 1 - request = _validate_backward( - config, - torch.randn(2, 128), - args, - fc1_c, - route_metadata, - ) - redispatched_grad_output = torch.zeros(3, 128) - redispatched_grad_output[:, 0] = torch.tensor([1.0, 2.0, 3.0]) +def test_k_major_workspace_view_matches_upstream_token_major_abi(): + storage = torch.arange(12, dtype=torch.uint8) + view = _typed_k_major_view(storage, torch.uint8, (3, 4)) - grad_topk = return_grad_topk_weights( - request, - redispatched_grad_output, - ) - silu_one = torch.sigmoid(torch.tensor(1.0)) - torch.testing.assert_close( - grad_topk, - silu_one * torch.tensor([[2.0, 0.0], [12.0, 4.0]]), - rtol=1e-6, - atol=1e-6, - ) + assert view.shape == (3, 4) + assert view.stride() == (1, 3) + assert torch.equal(view, storage.reshape(4, 3).transpose(0, 1)) @pytest.mark.L0 -def test_mxfp8_backward_dprob_recompute_applies_gate_up_clamp(): - args = list(_inputs()) - fc2_weight = torch.zeros_like(args[2]) - fc2_weight[0, 0, 0] = 2 - fc2_weight[1, 0, 0] = 4 - args[2] = fc2_weight - config = _config(gate_up_clamp=0.5) - route_metadata = torch.tensor( - [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], - dtype=torch.int32, - ) - fc1_c = torch.zeros(3, 512, dtype=torch.bfloat16) - fc1_c[:, 0] = 2 - fc1_c[:, 256] = 2 - request = _validate_backward( - config, - torch.randn(2, 128), - args, - fc1_c, - route_metadata, - ) - redispatched_grad_output = torch.zeros(3, 128) - redispatched_grad_output[:, 0] = torch.tensor([1.0, 2.0, 3.0]) +def test_only_fixed_training_wgrad_types_are_public(): + expected = [f"fc{layer}_{part}" for layer in (1, 2) for part in ("a", "sfa", "b", "sfb")] + expected += ["expert_offsets", "valid_route_counts"] + assert [field.name for field in fields(MoeEpTrainingWgradOperands)] == expected + assert not hasattr(cudnn, "MoeEpWgradForwardStash") + assert not hasattr(cudnn, "MoeEpWgradOperands") - grad_topk = return_grad_topk_weights( - request, - redispatched_grad_output, - ) - clamped_hidden = 0.25 * torch.sigmoid(torch.tensor(0.5)) - torch.testing.assert_close( - grad_topk, - clamped_hidden * torch.tensor([[2.0, 0.0], [12.0, 4.0]]), - rtol=1e-6, - atol=1e-6, - ) +@pytest.mark.L0 +def test_prepare_training_resources_binds_weights_and_slot_lanes(monkeypatch): + weights = _training_weights() + _, state = _install_contract_backend( + monkeypatch, + weights=weights, + slot_count=2, + ) + operator = _operator() + resources = operator.prepare_training_resources( + weights, + slot_count=2, + lane_count=1, + ) + + assert isinstance(resources, MoeEpTrainingResources) + assert all(isinstance(slot, MoeEpTrainingSlot) for slot in resources.slots) + assert isinstance(resources.lanes[0], MoeEpExecutionLane) + resources.refresh_weights() + owner = state.backends[0].owner + assert owner.refresh_calls == 1 + operator.close() + assert resources.closed + assert owner.close_calls == 1 @pytest.mark.L0 -def test_mxfp8_grad_output_redispatch_uses_public_route_order(): - config = _config() - args = _inputs() - route_metadata = torch.tensor( - [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], - dtype=torch.int32, - ) - request = _validate_backward( - config, - torch.randn(2, 128), - args, - torch.randn(3, 512, dtype=torch.bfloat16), - route_metadata, - ) +def test_prepare_training_resources_rejects_plain_weights(): + with _operator() as operator: + with pytest.raises(TypeError, match="MoeEpTrainingWeights"): + operator.prepare_training_resources(_training_inputs()[1]) - actual = Mxfp8BackwardRedispatch(request).run() - expected_rows = torch.tensor([0, 1, 1], dtype=torch.int64) - torch.testing.assert_close( - actual.grad_output, - request.grad_output.index_select(0, expected_rows).float(), - rtol=0, - atol=0, - ) +@pytest.mark.L0 +def test_error_mode_requires_async_assert_before_prepare(monkeypatch): + from cudnn.moe_ep.api import _validate_training_assert_capability + + monkeypatch.setattr(torch, "_assert_async", None) + config = SimpleNamespace(drop_on_overflow=False, ep_size=1) + with pytest.raises(RuntimeError, match="callable torch._assert_async"): + _validate_training_assert_capability(config) + + _validate_training_assert_capability(SimpleNamespace(drop_on_overflow=True, ep_size=1)) @pytest.mark.L0 -def test_moe_ep_backward_delegates_validated_request(monkeypatch): - import cudnn.moe_ep._backend as backend_seam +def test_distributed_error_mode_requires_nccl(monkeypatch): + from cudnn.moe_ep.api import _validate_training_assert_capability - args = _inputs() - grad_output = torch.randn(2, 128) - fc1_c = torch.randn(3, 512, dtype=torch.bfloat16) - route_metadata = torch.tensor( - [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], - dtype=torch.int32, - ) - expected = ( - torch.empty(2, 128, dtype=torch.float32), - torch.empty(2, 2, dtype=torch.float32), + monkeypatch.setattr(torch, "_assert_async", lambda *args, **kwargs: None) + monkeypatch.setattr(torch.distributed, "get_backend", lambda group: "gloo") + config = SimpleNamespace( + drop_on_overflow=False, + ep_size=2, + ep_group=object(), ) + with pytest.raises(NotImplementedError, match="NCCL"): + _validate_training_assert_capability(config) - class Backend: - request = None - def backward(self, request): - self.request = request - return expected +@pytest.mark.L0 +def test_training_weight_refresh_keeps_destination_addresses_stable(): + weights = _training_weights() + bindings = Mxfp8TrainingWeightBindings(weights) + bindings.refresh() + tensors = ( + bindings.forward.fc1_weight, + bindings.forward.fc1_weight_sf, + bindings.forward.fc2_weight, + bindings.forward.fc2_weight_sf, + bindings.backward.fc1_weight, + bindings.backward.fc1_weight_sf, + bindings.backward.fc2_weight, + bindings.backward.fc2_weight_sf, + ) + pointers = tuple(tensor.data_ptr() for tensor in tensors) + snapshots = tuple(tensor.clone() for tensor in tensors) + + weights.forward_fc1.data.view(torch.uint8).bitwise_xor_(1) + bindings.refresh() + + assert tuple(tensor.data_ptr() for tensor in tensors) == pointers + assert not torch.equal(bindings.forward.fc1_weight, snapshots[0]) + for tensor in tensors: + assert tensor.is_contiguous() or tensor.stride(1) == 1 - def close(self): - pass - instance = Backend() - monkeypatch.setattr(backend_seam, "validate_config", lambda config: None) - monkeypatch.setattr( - backend_seam, - "validate_backward_request", - lambda request: None, - ) - monkeypatch.setattr( - backend_seam, - "create_backend", - lambda config, device: instance, +@pytest.mark.L0 +def test_reference_wgrad_math_remains_in_test_tree(): + torch.manual_seed(20260821) + tokens, hidden, intermediate = 3, 32, 32 + topk_idx = torch.tensor([[0, 2], [2, 0], [0, 2]], dtype=torch.int32) + topk_weights = torch.tensor([[0.5, 0.25], [0.0, 0.75], [1.0, 0.125]]) + activation, fc1_weight, fc2_weight, grad_output = ( + torch.randn(shape) / 8 + for shape in ( + (tokens, hidden), + (3, hidden, 2 * intermediate), + (3, intermediate, hidden), + (tokens, hidden), + ) ) - - operator = MoeEp( - num_experts=2, - hidden_size=128, - intermediate_size=256, + reference = MoeEpReference( + num_experts=3, + hidden_size=hidden, + intermediate_size=intermediate, top_k=2, - max_tokens_per_rank=4, + max_tokens_per_rank=tokens, generate_c=True, + backward_wgrad_mode="operands", + token_padding_size=256, ) - actual = operator.backward( + + _, fc1_c, metadata, stash = reference(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + _, _, operands = reference.backward( grad_output, - *args[1:], + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, fc1_c, - route_metadata, + metadata, + wgrad_forward_stash=stash, ) + dw1, dw2 = operands.dense_wgrads() - assert actual is expected - assert len(actual) == 2 - assert instance.request is not None - assert instance.request.local_routes == 3 - assert instance.request.fc1_c is fc1_c - assert instance.request.fc1_weight is args[1] + assert operands.valid_route_counts.tolist() == [3, 0, 3] + assert dw1.shape == (3, hidden, 2 * intermediate) + assert dw2.shape == (3, intermediate, hidden) + assert dw1[1].eq(0).all() + assert dw2[1].eq(0).all() -@pytest.mark.L0 -def test_moe_ep_backward_accepts_explicit_stashes_in_reordered_calls( +# Source contracts + + +_ROOT = Path(__file__).resolve().parents[3] +_CUTEDSL = _ROOT / "python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin" / "training/mega" +_DGLU = _CUTEDSL / "bwd_dglu/dglu_mxfp8_mega_moe_kernel.py" +_DGLU_EPILOGUE = _CUTEDSL / "bwd_dglu/dglu_mxfp8_fc12_epilogue.py" + + +# L1 fail-fast training-resource contracts + + +@pytest.mark.L1 +@pytest.mark.parametrize( + ("field", "value"), + [(field, value) for field in ("slot_count", "lane_count") for value in (0, True, 1.5)], +) +def test_prepare_training_resources_rejects_invalid_counts_before_backend( monkeypatch, + field, + value, ): import cudnn.moe_ep._backend as backend_seam + import cudnn.moe_ep.api as api_module - calls = [] - - class Backend: - def backward(self, request): - calls.append(request) - marker = request.fc1_c[0, 0].float().reshape(1) - return marker, marker + def unexpected_call(*args, **kwargs): + del args, kwargs + raise AssertionError("invalid counts must fail before weight/backend work") - def close(self): - pass - - backend = Backend() - monkeypatch.setattr(backend_seam, "validate_config", lambda config: None) monkeypatch.setattr( - backend_seam, - "validate_backward_request", - lambda request: None, - ) - monkeypatch.setattr( - backend_seam, - "create_backend", - lambda config, device: backend, - ) - operator = MoeEp( - num_experts=2, - hidden_size=128, - intermediate_size=256, - top_k=2, - max_tokens_per_rank=4, - generate_c=True, - ) - args = _inputs() - route_metadata = torch.tensor( - [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], - dtype=torch.int32, + api_module, + "validate_training_weights", + unexpected_call, ) - stash_a = torch.full((3, 512), 1.0, dtype=torch.bfloat16) - stash_b = torch.full((3, 512), 2.0, dtype=torch.bfloat16) - - markers = [] - for stash in (stash_a, stash_b, stash_b, stash_a): - result = operator.backward( - torch.randn(2, 128), - *args[1:], - stash, - route_metadata, + monkeypatch.setattr(backend_seam, "create_backend", unexpected_call) + counts = {"slot_count": 1, "lane_count": 1, field: value} + + with _operator() as operator, pytest.raises( + ValueError, + match=rf"{field} must be a positive integer", + ): + operator.prepare_training_resources( + SimpleNamespace(mock_training_weights=True), + **counts, ) - markers.append(result[0].item()) - assert markers == [1.0, 2.0, 2.0, 1.0] - assert calls[0].fc1_c is stash_a - assert calls[1].fc1_c is stash_b - assert calls[2].fc1_c is stash_b - assert calls[3].fc1_c is stash_a +@pytest.mark.L1 +def test_prepare_training_resources_rejects_duplicate_open_resources(monkeypatch): + weights, state = _install_contract_backend(monkeypatch) + + with _operator() as operator: + resources = operator.prepare_training_resources( + weights, + slot_count=1, + lane_count=1, + ) + with pytest.raises(RuntimeError, match="already exist"): + operator.prepare_training_resources( + weights, + slot_count=1, + lane_count=1, + ) -@pytest.mark.L0 -def test_moe_ep_backward_requires_generate_c(): - args = _inputs() - operator = MoeEp( - num_experts=2, - hidden_size=128, - intermediate_size=256, - top_k=2, - max_tokens_per_rank=4, - generate_c=False, + assert resources.closed + backend = state.backends[0] + state.validate.assert_called_once() + backend.prepare_training_resources.assert_called_once_with( + weights, + slot_count=1, + lane_count=1, ) + assert (backend.close.call_count, backend.owner.close_calls) == (1, 1) - with pytest.raises(RuntimeError, match="generate_c=True"): - operator.backward( - torch.randn(2, 128), - *args[1:], - torch.randn(3, 512, dtype=torch.bfloat16), - torch.zeros(3, 4, dtype=torch.int32), - ) +@pytest.mark.L1 +def test_closed_training_resources_require_a_new_operator(monkeypatch): + weights, state = _install_contract_backend(monkeypatch) -@pytest.mark.L0 -@pytest.mark.parametrize( - ("overrides", "message"), - [ - ({"output_format": "mxfp8"}, "output_format='bf16'"), - ({"apply_topk_in_fc1": False}, "apply_topk_in_fc1=True"), - ], -) -def test_mxfp8_backward_capability_rejects_unsupported_config( - monkeypatch, - overrides, - message, -): - config = _config(**overrides) - args = _inputs() - request = _validate_backward( - config, - torch.randn(2, 128), - args, - torch.randn(3, 512, dtype=torch.bfloat16), - torch.zeros(3, 4, dtype=torch.int32), - ) - monkeypatch.setattr(_capability, "_validate_device", lambda device: None) - monkeypatch.setattr( - _capability, - "_is_cuda_stream_capturing", - lambda device: False, + old_operator = _operator() + old_resources = old_operator.prepare_training_resources( + weights, + slot_count=1, + lane_count=1, ) - + old_resources.close() with pytest.raises( - NotImplementedError, - match=message, + RuntimeError, + match="create a new MoeEp instance", ): - _capability.validate_backward_request(request) + old_operator.prepare_training_resources( + weights, + slot_count=1, + lane_count=1, + ) + old_operator.close() + with _operator() as new_operator: + new_resources = new_operator.prepare_training_resources( + weights, + slot_count=1, + lane_count=1, + ) + assert not new_resources.closed -@pytest.mark.L0 -def test_mxfp8_backward_capability_accepts_gate_up_clamp(monkeypatch): - config = _config(gate_up_clamp=1.0) - args = _inputs() - request = _validate_backward( - config, - torch.randn(2, 128), - args, - torch.randn(3, 512, dtype=torch.bfloat16), - torch.zeros(3, 4, dtype=torch.int32), - ) - monkeypatch.setattr(_capability, "_validate_device", lambda device: None) - monkeypatch.setattr( - _capability, - "_is_cuda_stream_capturing", - lambda device: False, - ) + assert len(state.backends) == 2 + assert state.validate.call_count == 2 + assert all(backend.prepare_training_resources.call_count == 1 for backend in state.backends) - _capability.validate_backward_request(request) +@pytest.mark.L1 +def test_training_prepare_and_backend_close_reject_capture(monkeypatch): + from cudnn.moe_ep._megamoe_backend.mxfp8._backend import Mxfp8Backend -@pytest.mark.L0 -def test_mxfp8_backward_capability_accepts_ep_above_16(monkeypatch): - config = _config(ep_size=32, ep_rank=31) - args = _inputs() - request = _validate_backward( - config, - torch.randn(2, 128), - args, - torch.randn(3, 512, dtype=torch.bfloat16), - torch.zeros(3, 4, dtype=torch.int32), - ) - monkeypatch.setattr(_capability, "_validate_device", lambda device: None) monkeypatch.setattr( - _capability, - "_is_cuda_stream_capturing", - lambda device: False, + torch.cuda, + "is_current_stream_capturing", + lambda: True, ) - _capability.validate_backward_request(request) + owner = object.__new__(Mxfp8TrainingResourceOwner) + owner._lock = threading.RLock() + owner._closed = False + owner._runtime = None + owner._workspace = None + with pytest.raises( + RuntimeError, + match="must be prepared before CUDA graph capture", + ): + owner.prepare() + + backend = object.__new__(Mxfp8Backend) + backend._lock = threading.RLock() + backend._closed = False + backend.device = torch.device("cuda") + monkeypatch.setattr(torch.cuda, "device", lambda device: nullcontext()) + with pytest.raises(RuntimeError, match="cannot be closed during"): + backend.close() -@pytest.mark.L0 -def test_mxfp8_backward_capability_rejects_cuda_graph_capture(monkeypatch): - config = _config() - args = _inputs() - request = _validate_backward( - config, - torch.randn(2, 128), - args, - torch.randn(3, 512, dtype=torch.bfloat16), - torch.zeros(3, 4, dtype=torch.int32), +@pytest.mark.L1 +def test_training_resources_reject_foreign_and_forged_slot_lane_bindings(): + resources, owner = _training_contract_resources() + foreign, _ = _training_contract_resources() + slot = resources.slots[0] + lane = resources.lanes[0] + activation = torch.empty((0, 128), dtype=torch.bfloat16) + routing = ( + torch.empty((0, 2), dtype=torch.int32), + torch.empty((0, 2), dtype=torch.float32), ) - monkeypatch.setattr(_capability, "_validate_device", lambda device: None) - monkeypatch.setattr( - _capability, - "_is_cuda_stream_capturing", - lambda device: True, + checks = ( + ("training slot does not belong", resources.forward, (foreign.slots[0], lane, activation, *routing)), + ("training slot does not belong", resources.backward, (MoeEpTrainingSlot(99, slot._resource_token), lane, activation.float())), + ("execution lane does not belong", resources.forward, (slot, foreign.lanes[0], activation, *routing)), + ("execution lane does not belong", resources.backward, (slot, MoeEpExecutionLane(99, lane._resource_token), activation.float())), ) + for message, call, args in checks: + with pytest.raises(ValueError, match=message): + call(*args) - with pytest.raises(NotImplementedError, match="CUDA graph capture"): - _capability.validate_backward_request(request) + assert owner.views_calls == 0 -@pytest.mark.L0 -def test_mxfp8_backward_delegates_to_explicit_executor(monkeypatch): - import cudnn.moe_ep._megamoe_backend.mxfp8._backend as backend_module +@pytest.mark.L1 +def test_training_resources_reject_invalid_overflow_finalization(): + resources, _ = _training_contract_resources() + foreign, _ = _training_contract_resources() + slot = resources.slots[0] + lane = resources.lanes[0] + + with pytest.raises(ValueError, match="at least one slot"): + resources.finalize_overflow((), lane) + with pytest.raises(ValueError, match="slots must be unique"): + resources.finalize_overflow((slot, slot), lane) + with pytest.raises(ValueError, match="overflow slot does not belong"): + resources.finalize_overflow((foreign.slots[0],), lane) + with pytest.raises(ValueError, match="overflow execution lane does not belong"): + resources.finalize_overflow((slot,), foreign.lanes[0]) - config = _config() - args = _inputs() - request = _validate_backward( - config, - torch.randn(2, 128), - args, - torch.randn(3, 512, dtype=torch.bfloat16), - torch.tensor( - [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], - dtype=torch.int32, + +@pytest.mark.L1 +def test_training_resources_reject_calls_after_close_and_close_is_idempotent(): + resources, owner = _training_contract_resources() + slot = resources.slots[0] + lane = resources.lanes[0] + activation = torch.empty((0, 128), dtype=torch.bfloat16) + routing = ( + torch.empty((0, 2), dtype=torch.int32), + torch.empty((0, 2), dtype=torch.float32), + ) + + resources.close() + resources.close() + + assert resources.closed + assert owner.close_calls == 1 + calls = ( + resources.refresh_weights, + lambda: resources.forward(slot, lane, activation, *routing), + lambda: resources.backward(slot, lane, activation.float()), + lambda: resources.finalize_overflow((slot,), lane), + ) + for call in calls: + with pytest.raises(RuntimeError, match="resources are closed"): + call() + assert owner.refresh_calls == 0 + assert owner.views_calls == 0 + + +@pytest.mark.L1 +@pytest.mark.parametrize( + ("mismatch_reduce", "message"), + [ + (2, "region counts differ"), + (4, "names, order, or alignments differ"), + ], +) +def test_harmonize_symmetric_regions_rejects_collective_metadata_mismatch( + monkeypatch, + mismatch_reduce, + message, +): + requirements = WorkspaceRequirements( + max_tokens_per_rank=1, + symmetric_regions=( + BufferRegion("first", 64, alignment=128), + BufferRegion("second", 128, alignment=256), ), + local_regions=(), ) - expected = tuple(torch.empty(0) for _ in range(2)) + runtime = SimpleNamespace(world_size=2, group=object()) + reduce_calls = [] - class Executor: - def __init__(self, actual_config, actual_device): - assert actual_config is config - assert actual_device == torch.device("cpu") + def all_reduce(tensor, *, op, group): + assert group is runtime.group + reduce_calls.append(op) + if len(reduce_calls) == mismatch_reduce: + tensor.add_(1) - def run(self, actual_request): - assert actual_request is request - return expected + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) - def close(self): - pass + with pytest.raises(RuntimeError, match=message): + _harmonize_symmetric_regions( + requirements, + runtime, + torch.device("cpu"), + ) - monkeypatch.setattr( - backend_module, - "Mxfp8BackwardExecutor", - Executor, - ) - monkeypatch.setattr( - torch.cuda, - "is_current_stream_capturing", - lambda: False, - ) - class Stream: - def wait_event(self, event): - del event + assert len(reduce_calls) == mismatch_reduce - class Event: - def record(self, stream): - del stream - monkeypatch.setattr(torch.cuda, "current_stream", lambda device: Stream()) - monkeypatch.setattr(torch.cuda, "Event", Event) - backend = Mxfp8Backend(config, torch.device("cpu")) - assert backend.backward(request) is expected +_STAGER_FAILURES = { + "source-shape": (lambda t: t.update(source=t["source"][:, :-1].contiguous()), ValueError, r"source must have shape \(T, 128\)"), + "route-shape": (lambda t: t.update(topk_idx=t["topk_idx"][:, :-1].contiguous()), ValueError, "topk_idx shape mismatch"), + "weight-shape": (lambda t: t.update(topk_weights=t["topk_weights"][:, :-1].contiguous()), ValueError, "topk_weights shape mismatch"), + "route-dtype": (lambda t: t.update(topk_idx=t["topk_idx"].to(torch.int64)), TypeError, "contiguous Int32"), + "route-contiguity": (lambda t: t.update(topk_idx=t["topk_idx"].t().contiguous().t()), TypeError, "contiguous Int32"), + "weight-dtype": (lambda t: t.update(topk_weights=t["topk_weights"].to(torch.bfloat16)), TypeError, "contiguous FP32"), + "weight-contiguity": (lambda t: t.update(topk_weights=t["topk_weights"].t().contiguous().t()), TypeError, "contiguous FP32"), + "capacity": (lambda t: t.update(**{name: value[:4] for name, value in t.items() if name.startswith("output")}), ValueError, "token count 5 exceeds capacity 4"), + "device": (lambda t: t.update(source=torch.empty_like(t["source"], device="meta")), ValueError, "must share one device"), +} -# Single-rank and distributed backward numerical parity. +@pytest.mark.L1 +@pytest.mark.parametrize( + ("mutator", "error_type", "message"), + [ + pytest.param(*case, id=name) + for name, case in _STAGER_FAILURES.items() + ], +) +def test_training_stager_rejects_invalid_inputs(mutator, error_type, message): + tensors = _training_staging_tensors() + mutator(tensors) + with pytest.raises(error_type, match=message): + Mxfp8TrainingStager(hidden=128, top_k=2)._validate(**tensors) -def _make_reentrant_case_b(args, device): - activation = quantize_mxfp8( - args[0].dequantize(dtype=torch.float32) + 0.25, - axis=1, - ) - topk_idx = torch.tensor( - [[0, 0], [-1, -1], [0, -1], [0, 0], [0, -1]], - dtype=torch.int32, - device=device, - ) - topk_weights = torch.tensor( - [ - [0.625, 0.375], - [0.0, 0.0], - [1.0, 0.0], - [0.75, 0.25], - [1.0, 0.0], - ], - dtype=torch.bfloat16, - device=device, - ) - return activation, args[1], args[2], topk_idx, topk_weights +# L1 training graph @pytest.mark.L1 @pytest.mark.gpu_exclusive -@pytest.mark.parametrize("combine_format", ["bf16", "mxfp8"]) @pytest.mark.parametrize( - "gate_up_clamp", - [None, 0.5], - ids=["unclamped", "clamped"], + ( + "input_kind", + "combine_format", + "gate_up_clamp", + "top_k", + "tuning", + "all_dropped", + ), + [ + pytest.param("fixed", "bf16", None, 2, MoeEpTuningConfig(), False, id="bf16-unclamped"), + pytest.param("fixed", "mxfp8", 0.5, 2, MoeEpTuningConfig(), False, id="mxfp8-clamp-0.5"), + pytest.param("routed", "bf16", None, 1, MoeEpTuningConfig(), False, id="topk1-default-tuning"), + pytest.param( + "routed", + "bf16", + None, + 2, + MoeEpTuningConfig( + token_back_mode="reuse_dispatch_warps", + epi_flag_batch=(2, 2), + token_in_flag_batch=2, + group_hint=128, + ), + False, + id="topk2-nondefault-tuning", + ), + pytest.param("routed", "bf16", None, 2, MoeEpTuningConfig(), True, id="topk2-all-dropped"), + ], ) -def test_mxfp8_backward_ep1_matches_reference_and_resets_workspace( +def test_fixed_training_resources_ep1_matches_independent_reference( + input_kind, combine_format, gate_up_clamp, + top_k, + tuning, + all_dropped, ): device = _sm107_device() - args = make_forward_inputs(device) - config = _forward_config( - generate_c=True, + if input_kind == "fixed": + args, grad_output = _fixed_training_case(device) + max_recv_size = 1 + else: + base_args = make_forward_inputs(device) + args = ( + base_args[0].dequantize(torch.bfloat16), + base_args[1], + base_args[2], + base_args[3][:, :top_k].contiguous(), + base_args[4][:, :top_k].float().contiguous(), + ) + if all_dropped: + args[3].fill_(-1) + args[4].zero_() + grad_output = _grad_output(device, args[0].shape[0], seed=20260830) + max_recv_size = args[0].shape[0] * top_k + expected = _fixed_training_reference( + args, + grad_output, combine_format=combine_format, gate_up_clamp=gate_up_clamp, + tuning=tuning, ) - reference = _reference_backward(config) - grad_output = _grad_output(device, args[3].shape[0], seed=20260817) - with MoeEp(**config) as op: - _, fc1_c, route_metadata = op(*args) - stash = (fc1_c, route_metadata) - expected = _expected_backward(reference, grad_output, args, stash) - - first = op.backward(grad_output, *args[1:], *stash) - second = op.backward(grad_output, *args[1:], *stash) + with MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=top_k, + max_tokens_per_rank=args[0].shape[0], + max_recv_size_per_rank=max_recv_size, + drop_on_overflow=True, + combine_format=combine_format, + gate_up_clamp=gate_up_clamp, + tuning=tuning, + ) as op: + resources = op.prepare_training_resources( + _fixed_training_weights(args), + slot_count=1, + lane_count=1, + ) + slot = resources.slots[0] + lane = resources.lanes[0] + actual = _run_fixed_training_batch( + resources, + lane, + ((slot, args, grad_output),), + )[0] torch.cuda.synchronize(device) - _assert_backward_matches(first, expected, args[3]) - _assert_backward_matches(second, expected, args[3]) - - -@pytest.mark.L1 -@pytest.mark.gpu_exclusive -def test_mxfp8_backward_ep1_uses_explicit_stash_after_reordered_forwards(): - device = _sm107_device() - args_a = make_forward_inputs(device) - args_b = _make_reentrant_case_b(args_a, device) - config = _forward_config(generate_c=True) - reference = _reference_backward(config) - grad_a = _grad_output(device, args_a[3].shape[0], seed=20260818) - grad_b = _grad_output(device, args_b[3].shape[0], seed=20260819) - - with MoeEp(**config) as op: - _, fc1_c_a, metadata_a = op(*args_a) - _, fc1_c_b, metadata_b = op(*args_b) - cases = ( - (grad_b, args_b, (fc1_c_b, metadata_b)), - (grad_a, args_a, (fc1_c_a, metadata_a)), - (grad_b, args_b, (fc1_c_b, metadata_b)), - (grad_a, args_a, (fc1_c_a, metadata_a)), + assert actual.overflow.eq(0).all() + _assert_fixed_training_matches_reference( + (actual.y, actual.dx, actual.dprob, actual.wgrads), + expected, + args[3], ) - results = [] - for grad_output, args, stash in cases: - expected = _expected_backward(reference, grad_output, args, stash) - actual = op.backward(grad_output, *args[1:], *stash) - results.append((actual, expected, args[3])) - torch.cuda.synchronize(device) - for actual, expected, topk_idx in results: - _assert_backward_matches(actual, expected, topk_idx) + if all_dropped: + actual_dw1, actual_dw2 = _dense_wgrads_from_operands(actual.wgrads) + expected_dw1, expected_dw2 = expected[3].dense_wgrads() + zero_tensors = ( + actual.y, + expected[0], + actual.dx, + expected[1], + actual.dprob, + expected[2], + actual_dw1, + expected_dw1, + actual_dw2, + expected_dw2, + ) + assert all(tensor.eq(0).all() for tensor in zero_tensors) @pytest.mark.L1 @pytest.mark.gpu_exclusive -@pytest.mark.parametrize("world_size", [2, 4], ids=["ep2", "ep4"]) -@pytest.mark.parametrize("combine_format", ["bf16", "mxfp8"]) -def test_mxfp8_backward_multi_gpu_matches_reference( +@pytest.mark.parametrize( + ("world_size", "combine_format", "gate_up_clamp"), + [ + pytest.param(2, "bf16", None, id="ep2-bf16"), + pytest.param(2, "mxfp8", None, id="ep2-mxfp8"), + pytest.param(4, "bf16", None, id="ep4-bf16"), + pytest.param(4, "mxfp8", None, id="ep4-mxfp8"), + pytest.param(2, "bf16", 0.5, id="ep2-bf16-clamp-0.5"), + ], +) +def test_fixed_training_resources_multi_gpu_matches_independent_reference( world_size, combine_format, + gate_up_clamp, tmp_path, ): _require_distributed_sm107(world_size) os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") - init_file = ( - tmp_path - / f"{combine_format}_combine_mxfp8_backward_ep{world_size}.init" - ) + clamp_id = "none" if gate_up_clamp is None else str(gate_up_clamp) + init_file = tmp_path / f"backward_ep{world_size}_{combine_format}_clamp_{clamp_id}.init" mp.spawn( - _distributed_backward_worker, - args=(world_size, str(init_file), combine_format), + _distributed_backward_reference_worker, + args=( + world_size, + str(init_file), + combine_format, + gate_up_clamp, + ), nprocs=world_size, join=True, ) @@ -1156,31 +1041,448 @@ def test_mxfp8_backward_multi_gpu_matches_reference( @pytest.mark.L1 @pytest.mark.gpu_exclusive -def test_mxfp8_backward_ep2_gate_up_clamp_matches_reference(tmp_path): - world_size = 2 - _require_distributed_sm107(world_size) +def test_noncontiguous_ep2_fixed_training_matches_independent_reference( + tmp_path, +): + global_world_size = 4 + _require_distributed_sm107(global_world_size) os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") - init_file = tmp_path / "bf16_combine_clamped_mxfp8_backward_ep2.init" + init_file = tmp_path / "backward_two_noncontiguous_ep2.init" mp.spawn( - _distributed_backward_worker, - args=(world_size, str(init_file), "bf16", 0.5), - nprocs=world_size, + _distributed_subgroup_backward_reference_worker, + args=(global_world_size, str(init_file)), + nprocs=global_world_size, join=True, ) -@pytest.mark.L0 -def test_forward_and_backward_mxfp8_combine_are_direct_fp32(): - generator = torch.Generator().manual_seed(20260820) - accumulator = torch.randn(4, 128, generator=generator) * 3.25 - - backward = backward_combine_round_trip( - accumulator, - MoeFormat.MXFP8, - ) - forward = forward_combine_round_trip( - accumulator, - MoeFormat.MXFP8, +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +@pytest.mark.parametrize( + "case", + [ + pytest.param( + SimpleNamespace( + combine_format="bf16", + drop_on_overflow=True, + max_recv_size=1, + replay_count=20, + ), + id="bf16-drop", + ), + pytest.param( + SimpleNamespace( + combine_format="mxfp8", + drop_on_overflow=True, + max_recv_size=1, + replay_count=20, + ), + id="mxfp8-drop", + ), + pytest.param( + SimpleNamespace( + combine_format="bf16", + drop_on_overflow=False, + max_recv_size=2, + replay_count=2, + ), + id="bf16-error-no-overflow", + ), + ], +) +def test_fixed_training_resources_ep1_cuda_graph_replay(case): + device = _sm107_device() + if case.drop_on_overflow: + args0, grad0 = _fixed_training_case(device) + topk_idx1 = args0[3].clone() + topk_idx1[0, 0] = 1 + inputs = ( + (args0, grad0), + ( + ( + args0[0].clone(), + args0[1], + args0[2], + topk_idx1, + args0[4].clone(), + ), + grad0.clone(), + ), + ) + else: + inputs = (_fixed_training_drop_overflow_case(device),) + references = tuple( + _fixed_training_reference( + args, + grad_output, + combine_format=case.combine_format, + gate_up_clamp=None, + ) + for args, grad_output in inputs ) - torch.testing.assert_close(backward, forward, rtol=0, atol=0) + with MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=inputs[0][0][0].shape[0], + max_recv_size_per_rank=case.max_recv_size, + drop_on_overflow=case.drop_on_overflow, + combine_format=case.combine_format, + token_padding_size=128, + ) as op: + resources = op.prepare_training_resources( + _fixed_training_weights(inputs[0][0]), + slot_count=len(inputs), + lane_count=1, + ) + lane = resources.lanes[0] + batch = tuple((slot, args, grad_output) for slot, (args, grad_output) in zip(resources.slots, inputs)) + + def assert_batch(actuals): + for actual, (args, _), reference in zip( + actuals, + inputs, + references, + ): + assert actual.overflow.shape == (1,) + assert actual.overflow.dtype == torch.int32 + assert actual.overflow.eq(0).all() + _assert_fixed_training_matches_reference( + (actual.y, actual.dx, actual.dprob, actual.wgrads), + reference, + args[3], + ) + + eager_actuals = _run_fixed_training_batch(resources, lane, batch) + torch.cuda.synchronize(device) + assert_batch(eager_actuals) + + stream = torch.cuda.Stream(device=device) + stream.wait_stream(torch.cuda.current_stream(device)) + captured = _capture_fixed_training_batch( + resources, + lane, + batch, + stream, + ) + # In error mode, each replay executes the captured torch._assert_async + # with a false overflow condition; stable public nodes prove reuse. + for _ in range(case.replay_count): + captured.graph.replay() + torch.cuda.synchronize(device) + assert captured.public_pointers == tuple(_training_public_pointers(actual) for actual in captured.actuals) + assert_batch(captured.actuals) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_fixed_training_resources_ep1_two_shape_cuda_graph_contract(): + device = _sm107_device() + args, grad_large = _fixed_training_case(device) + max_tokens = int(args[0].shape[0]) + small_tokens = max_tokens - 2 + assert 0 < small_tokens < max_tokens + + large = SimpleNamespace( + name="large", + activation=args[0], + topk_idx=args[3], + topk_weights=args[4], + grad_output=grad_large, + ) + small = SimpleNamespace( + name="small", + activation=args[0][:small_tokens].clone(), + topk_idx=args[3][:small_tokens].clone(), + topk_weights=args[4][:small_tokens].clone(), + grad_output=grad_large[:small_tokens].clone(), + ) + assert all(getattr(large, name).data_ptr() != getattr(small, name).data_ptr() for name in ("activation", "topk_idx", "topk_weights", "grad_output")) + + weights = _fixed_training_weights(args) + weight_source_pointers = _training_weight_source_pointers(weights) + + with MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=max_tokens, + max_recv_size_per_rank=1, + drop_on_overflow=True, + ) as op: + resources = op.prepare_training_resources( + weights, + slot_count=1, + lane_count=1, + ) + slot = resources.slots[0] + lane = resources.lanes[0] + + def case_args(case): + return ( + case.activation, + weights.forward_fc1, + weights.forward_fc2, + case.topk_idx, + case.topk_weights, + ) + + def independent_reference(case): + return _fixed_training_reference( + case_args(case), + case.grad_output, + combine_format="bf16", + gate_up_clamp=None, + ) + + def warmup(case) -> None: + actual = _run_fixed_training_batch( + resources, + lane, + ((slot, case_args(case), case.grad_output),), + )[0] + torch.cuda.synchronize(device) + assert actual.overflow.eq(0).all(), f"{case.name} warmup overflowed" + _assert_fixed_training_matches_reference( + (actual.y, actual.dx, actual.dprob, actual.wgrads), + independent_reference(case), + case.topk_idx, + ) + + # Compile each static token-count specialization on the same resources, + # slot, and lane before either capture. + warmup(large) + warmup(small) + + capture_stream = torch.cuda.Stream(device=device) + capture_stream.wait_stream(torch.cuda.current_stream(device)) + + def capture(case): + # The shared sequence records refresh so replay observes in-place + # updates to all four bound source packs. + captured = _capture_fixed_training_batch( + resources, + lane, + ((slot, case_args(case), case.grad_output),), + capture_stream, + ) + return SimpleNamespace( + case=case, + graph=captured.graph, + actual=captured.actuals[0], + public_pointers=captured.public_pointers[0], + source_pointers=_training_source_pointers(case), + ) + + large_graph = capture(large) + small_graph = capture(small) + slot_views = resources._owner.views( + slot=slot.index, + lane=lane.index, + token_count=max_tokens, + ).slot + + def replay_and_check(captured): + _prefill_training_graph_sentinels(slot_views, captured.actual) + captured.graph.replay() + torch.cuda.synchronize(device) + + assert captured.actual.overflow.eq(0).all() + assert _training_public_pointers(captured.actual) == captured.public_pointers + assert _training_source_pointers(captured.case) == captured.source_pointers + assert _training_weight_source_pointers(weights) == weight_source_pointers + _assert_fixed_training_matches_reference( + ( + captured.actual.y, + captured.actual.dx, + captured.actual.dprob, + captured.actual.wgrads, + ), + independent_reference(captured.case), + captured.case.topk_idx, + ) + # The dense-dW check above decodes every expert segment and rejects + # nonzero expert padding, nonzero data capacity tails, or + # non-neutral scale tails left by the sentinels. + _assert_training_graph_tails_are_reset( + slot_views, + captured.actual, + token_count=int(captured.case.activation.shape[0]), + capacity=max_tokens, + ) + return captured.actual.y.clone() + + # The two graphs alias one persistent slot. Each replay must therefore + # fully replace the other shape's routing, gradients, and WGrad state. + for captured in (large_graph, small_graph, large_graph): + replay_and_check(captured) + + small_source_pointers = _training_source_pointers(small) + small.activation.mul_(-0.5) + small.topk_idx.fill_(-1) + small.topk_idx[0, 0] = 1 + small.topk_weights.zero_() + small.topk_weights[0, 0] = 0.625 + small.grad_output.mul_(-0.75) + assert _training_source_pointers(small) == small_source_pointers + + for captured in (small_graph, large_graph, small_graph): + replay_and_check(captured) + + old_large_y = replay_and_check(large_graph) + old_weight_values = _training_weight_source_values(weights) + generator = torch.Generator(device=device).manual_seed(20260829) + new_fc1 = ( + torch.randn( + weights.forward_fc1.logical_shape, + generator=generator, + device=device, + ) + / 16 + ) + new_fc2 = ( + torch.randn( + weights.forward_fc2.logical_shape, + generator=generator, + device=device, + ) + / 16 + ) + replacement = _fixed_training_weights( + ( + large.activation, + new_fc1, + new_fc2, + large.topk_idx, + large.topk_weights, + ) + ) + _copy_training_weight_sources_(weights, replacement) + + assert _training_weight_source_pointers(weights) == weight_source_pointers + _assert_training_weight_sources_changed(weights, old_weight_values) + + new_large_y = replay_and_check(large_graph) + assert not torch.equal(new_large_y, old_large_y) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_fixed_training_resources_ep1_drop_overflow_boundary_and_graph_transitions(): + device = _sm107_device() + args, grad_output = _fixed_training_drop_overflow_case(device) + assert args[0].shape[0] == 1 + assert args[3].detach().cpu().tolist() == [[0, 1]] + + references = { + expected_overflow: _fixed_training_drop_overflow_reference( + args, + grad_output, + drop_expert1=bool(expected_overflow), + ) + for expected_overflow in (0, 1) + } + + def assert_result(actual, expected_overflow): + expected, reference_topk_idx = references[expected_overflow] + _assert_fixed_training_drop_overflow_result( + actual, + expected, + reference_topk_idx, + expected_overflow=expected_overflow, + ) + + # The graph warmup below covers maxrecv=1 overflow; exercise the exact + # non-overflow boundary separately here. + with MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=1, + max_recv_size_per_rank=2, + drop_on_overflow=True, + token_padding_size=128, + ) as op: + resources = op.prepare_training_resources( + _fixed_training_weights(args), + slot_count=1, + lane_count=1, + ) + slot = resources.slots[0] + lane = resources.lanes[0] + actual = _run_fixed_training_batch( + resources, + lane, + ((slot, args, grad_output),), + )[0] + torch.cuda.synchronize(device) + + assert_result(actual, 0) + assert args[3][0, 1].eq(1) + assert actual.wgrads.valid_route_counts.detach().cpu().tolist() == [1, 1] + assert actual.wgrads.expert_offsets.detach().cpu().tolist() == [128, 256] + + overflow_routing = args[3].clone() + expert0_only_routing = overflow_routing.clone() + expert0_only_routing[0, 1] = -1 + routing_pointer = args[3].data_ptr() + + with MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=1, + max_recv_size_per_rank=1, + drop_on_overflow=True, + token_padding_size=128, + ) as op: + resources = op.prepare_training_resources( + _fixed_training_weights(args), + slot_count=1, + lane_count=1, + ) + slot = resources.slots[0] + lane = resources.lanes[0] + batch = ((slot, args, grad_output),) + + # Compile the fixed T=1 specialization and validate overflow eagerly + # before capturing the same forward/backward/finalize sequence. + warmup = _run_fixed_training_batch(resources, lane, batch)[0] + torch.cuda.synchronize(device) + assert_result(warmup, 1) + + capture_stream = torch.cuda.Stream(device=device) + capture_stream.wait_stream(torch.cuda.current_stream(device)) + captured = _capture_fixed_training_batch( + resources, + lane, + batch, + capture_stream, + ) + graph_actual = captured.actuals[0] + + for routing, expected_overflow in ( + (overflow_routing, 1), + (expert0_only_routing, 0), + (overflow_routing, 1), + ): + args[3].copy_(routing) + assert args[3].data_ptr() == routing_pointer + expected = _fixed_training_drop_overflow_reference( + args, + grad_output, + drop_expert1=bool(expected_overflow), + ) + captured.graph.replay() + torch.cuda.synchronize(device) + _assert_fixed_training_drop_overflow_result( + graph_actual, + *expected, + expected_overflow=expected_overflow, + ) + assert captured.public_pointers[0] == _training_public_pointers(graph_actual) diff --git a/test/python/moe_ep/test_moe_ep_cutedsl_grad_y2_source.py b/test/python/moe_ep/test_moe_ep_cutedsl_grad_y2_source.py deleted file mode 100644 index 221e0cf15..000000000 --- a/test/python/moe_ep/test_moe_ep_cutedsl_grad_y2_source.py +++ /dev/null @@ -1,63 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Source-only contracts for the upstream grad_y2 and dFC2 scale layout.""" - -from __future__ import annotations - -import ast -from pathlib import Path - -import pytest - - -_ROOT = Path(__file__).resolve().parents[3] -_CUTEDSL = ( - _ROOT - / "python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin" - / "training/mega" -) -_DGLU = _CUTEDSL / "bwd_dglu/dglu_mxfp8_mega_moe_kernel.py" -_DGLU_EPILOGUE = _CUTEDSL / "bwd_dglu/dglu_mxfp8_fc12_epilogue.py" - - -def _source_and_tree(path: Path) -> tuple[str, ast.Module]: - source = path.read_text(encoding="utf-8") - return source, ast.parse(source, filename=str(path)) - - -@pytest.mark.L0 -def test_dglu_source_exports_upstream_grad_y2_col_quant(): - source, tree = _source_and_tree(_DGLU) - - assert tree is not None - for contract in ( - "enable_grad_y2_col_quant", - "num_ctas_grad_y2_col_quant", - "grad_y2_sizes_region", - "_snapshot_grad_y2_expert_sizes", - "grad_y2_col_quant", - "grad_y2: cute.Tensor", - "grad_y2_sf: cute.Tensor", - ): - assert contract in source - assert source.index( - "self._snapshot_grad_y2_expert_sizes(tidx)" - ) < source.index("self.token_comm.reset_tail()") - assert source.index("self._topk_reduce(") < source.index( - "self.grad_y2_col_quant(" - ) - - -@pytest.mark.L0 -def test_dfc2_scale_source_uses_upstream_mn_major_atoms(): - source, tree = _source_and_tree(_DGLU_EPILOGUE) - - assert tree is not None - assert "def _stg_col_sf_atom_value(" in source - assert "MN-major 128-column × 4-token-block atom" in source - assert "atom_idx * Int64(512)" in source - assert "Int64(hidden_lane) * Int64(16)" in source - assert "Int64(hidden_bank) * Int64(4)" in source - assert "Int64(token_bank)" in source - assert source.count("self._stg_col_sf_atom_value(") >= 2 diff --git a/test/python/moe_ep/test_moe_ep_forward.py b/test/python/moe_ep/test_moe_ep_forward.py index ff5cc2092..1ab9cd23a 100644 --- a/test/python/moe_ep/test_moe_ep_forward.py +++ b/test/python/moe_ep/test_moe_ep_forward.py @@ -20,7 +20,7 @@ _distributed_output_worker, _distributed_subgroup_output_worker, ) -from moe_ep.moe_ep_forward_support import ( +from moe_ep.moe_ep_test_support import ( _assert_matches_reference, _forward_config, _make_forward_case, @@ -31,18 +31,15 @@ _require_distributed_sm107, _sm107_device, _stress_backend_reuse, + make_forward_inputs, + quantize_mxfp8, ) from moe_ep.moe_ep_reference import ( - BlockScaledTensor as ReferenceBlockScaledTensor, MoeEpReference, MoeFormat, forward_combine_round_trip, quantize_blockwise, ) -from moe_ep.moe_ep_test_data import ( - make_forward_inputs, - quantize_mxfp8, -) @@ -319,16 +316,6 @@ def test_combine_format_maps_to_contract_wire(public_format, wire_format): assert kernel_config.combine_format == wire_format -@pytest.mark.L0 -@pytest.mark.parametrize("combine_format", ["bf16", "mxfp8"]) -def test_megamoe_capability_enables_combine_formats(combine_format): - from cudnn import MoeEp - from cudnn.moe_ep._megamoe_backend._capability import validate_config - - with MoeEp(**_forward_config(combine_format=combine_format)) as op: - validate_config(op._forward_config) - - @pytest.mark.L0 def test_distributed_topk_can_exceed_local_expert_count(): from cudnn import MoeEp @@ -626,42 +613,12 @@ def unexpected_barrier(*, group): assert not backend._ep_launch_ready -@pytest.mark.L0 -def test_api_allocates_fresh_bf16_outputs_with_logical_shape(): - from cudnn import MoeEp - - device = _sm107_device() - args = make_forward_inputs(device) - activation, fc1_weight, fc2_weight = args[:3] - - assert activation.logical_shape == (5, 128) - assert fc1_weight.logical_shape == (2, 128, 512) - assert fc2_weight.logical_shape == (2, 256, 128) - - with MoeEp(**_forward_config()) as op: - first = op(*args) - snapshot = first.clone() - second = op(*args) - torch.cuda.synchronize(device) - - assert isinstance(first, torch.Tensor) - assert isinstance(second, torch.Tensor) - assert first.shape == second.shape == (5, 128) - assert first.dtype == second.dtype == torch.bfloat16 - assert first.device == second.device == device - assert first is not second - assert first.data_ptr() != second.data_ptr() - torch.testing.assert_close(first, snapshot, rtol=0, atol=0) - torch.testing.assert_close(first, second, rtol=0, atol=0) - - @pytest.mark.L0 @pytest.mark.parametrize( "kwargs", [ {"combine_format": "nvfp4"}, {"output_format": "mxfp8"}, - {"output_format": "nvfp4"}, {"apply_topk_in_fc1": False}, ], ) @@ -699,24 +656,36 @@ def test_training_megamoe_rejects_nvfp4_operand_before_cuda_query(monkeypatch): @pytest.mark.L0 -def test_fp8_activation_bf16_combine_forward_single_gpu(): +def test_bf16_forward_matches_reference_and_returns_fresh_outputs(): from cudnn import MoeEp device = _sm107_device() args = make_forward_inputs(device) + activation, fc1_weight, fc2_weight = args[:3] expected = _reference_forward(args) - with MoeEp(**_forward_config()) as op: - actual = op(*args) - torch.cuda.synchronize(device) + assert activation.logical_shape == (5, 128) + assert fc1_weight.logical_shape == (2, 128, 512) + assert fc2_weight.logical_shape == (2, 256, 128) + with MoeEp(**_forward_config()) as op: + first = op(*args) + snapshot = first.clone() + second = op(*args) args[3].fill_(-1) dropped = op(*args) torch.cuda.synchronize(device) - assert actual.shape == (5, 128) - assert actual.dtype == torch.bfloat16 - _assert_matches_reference(actual, expected) + assert isinstance(first, torch.Tensor) + assert isinstance(second, torch.Tensor) + assert first.shape == second.shape == (5, 128) + assert first.dtype == second.dtype == torch.bfloat16 + assert first.device == second.device == device + assert first is not second + assert first.data_ptr() != second.data_ptr() + torch.testing.assert_close(first, snapshot, rtol=0, atol=0) + torch.testing.assert_close(first, second, rtol=0, atol=0) + _assert_matches_reference(first, expected) assert dropped.eq(0).all() @@ -742,21 +711,25 @@ def test_mxfp8_combine_matches_direct_fp32_training_reference(): @pytest.mark.L1 @pytest.mark.gpu_exclusive @pytest.mark.parametrize( - "plain_mask", + ("plain_mask", "plain_dtype"), [ - (True, False, False), - (False, True, False), - (False, False, True), - (True, True, False), - (True, False, True), - (False, True, True), - (True, True, True), + pytest.param( + (True, False, False), + torch.bfloat16, + id="activation-bf16", + ), + pytest.param( + (False, True, False), + torch.float16, + id="fc1-fp16", + ), + pytest.param( + (False, False, True), + torch.float32, + id="fc2-fp32", + ), ], ) -@pytest.mark.parametrize( - "plain_dtype", - [torch.bfloat16, torch.float16, torch.float32], -) def test_plain_and_mixed_inputs_match_staged_reference( plain_mask, plain_dtype, @@ -859,65 +832,6 @@ def test_gate_up_clamp_matches_moe_ep_reference(): _assert_matches_reference(actual, expected) -@pytest.mark.L1 -@pytest.mark.gpu_exclusive -def test_generate_c_outputs_fc1_c_and_route_metadata(): - from cudnn import MoeEp - - device = _sm107_device() - args = make_forward_inputs(device) - config = _forward_config( - gate_up_clamp=1.25, - generate_c=True, - ) - expected_output, expected_fc1_c, expected_metadata = _reference_forward( - args, - **config, - ) - - with MoeEp(**config) as op: - first = op(*args) - output, fc1_c, route_metadata = first - fc1_c_snapshot = fc1_c.clone() - metadata_snapshot = route_metadata.clone() - - scaled_args = (*args[:4], args[4] * 0.25) - _, scaled_fc1_c, scaled_metadata = op(*scaled_args) - torch.cuda.synchronize(device) - - assert isinstance(first, tuple) - assert len(first) == 3 - assert output.shape == (5, 128) - assert output.dtype == torch.bfloat16 - assert fc1_c.shape == (9, 512) - assert fc1_c.dtype == torch.bfloat16 - assert route_metadata.shape == (9, 4) - assert route_metadata.dtype == torch.int32 - _assert_matches_reference(output, expected_output) - torch.testing.assert_close( - _output_as_float(fc1_c), - _output_as_float(expected_fc1_c), - rtol=0.01, - atol=0.01, - ) - torch.testing.assert_close( - route_metadata, - expected_metadata, - rtol=0, - atol=0, - ) - - # FC1 C is captured before clamp/SwiGLU and does not include router weights. - torch.testing.assert_close(scaled_fc1_c, fc1_c_snapshot, rtol=0, atol=0) - torch.testing.assert_close(scaled_metadata, metadata_snapshot, rtol=0, atol=0) - torch.testing.assert_close(fc1_c, fc1_c_snapshot, rtol=0, atol=0) - torch.testing.assert_close(route_metadata, metadata_snapshot, rtol=0, atol=0) - assert scaled_fc1_c is not fc1_c - assert scaled_metadata is not route_metadata - assert scaled_fc1_c.data_ptr() != fc1_c.data_ptr() - assert scaled_metadata.data_ptr() != route_metadata.data_ptr() - - @pytest.mark.L1 @pytest.mark.gpu_exclusive @pytest.mark.parametrize( @@ -1070,47 +984,6 @@ def test_nondefault_tuning_warmup_and_cuda_graph_replay(): ) -@pytest.mark.L0 -def test_reference_apply_topk_after_fc2_weights_after_combine_rounding(): - """Keep post-combine router weighting in reference-only semantics.""" - - device = torch.device("cpu") - args = make_forward_inputs(device) - # Duplicate routes make pre/post-combine weighting observably different. - args[3].copy_( - torch.tensor( - [[0, 0], [1, 1], [0, 0], [1, 1], [0, 0]], - dtype=torch.int32, - device=device, - ) - ) - args[4].copy_( - torch.tensor( - [[256.0, -255.0]] * 5, - dtype=torch.bfloat16, - device=device, - ) - ) - decoded_args = ( - args[0].dequantize(), - args[1].dequantize(), - args[2].dequantize(), - args[3], - args[4], - ) - expected = _naive_reference( - *decoded_args, - apply_topk_in_fc1=False, - intermediate_format=MoeFormat.MXFP8, - apply_topk_after_combine=True, - ) - pre_combine_weighting = _naive_reference( - *decoded_args, - apply_topk_in_fc1=False, - intermediate_format=MoeFormat.MXFP8, - ) - assert not torch.equal(expected, pre_combine_weighting) - @pytest.mark.L0 def test_forward_mxfp8_combine_is_direct_fp32(): generator = torch.Generator().manual_seed(20260819) @@ -1138,11 +1011,14 @@ def test_forward_mxfp8_combine_is_direct_fp32(): @pytest.mark.L1 @pytest.mark.gpu_exclusive @pytest.mark.parametrize( - "world_size", - [2, 3, 4], - ids=["ep2", "ep3", "ep4"], + ("world_size", "combine_format"), + [ + pytest.param(2, "bf16", id="ep2-bf16"), + pytest.param(2, "mxfp8", id="ep2-mxfp8"), + pytest.param(3, "mxfp8", id="ep3-mxfp8"), + pytest.param(4, "bf16", id="ep4-bf16"), + ], ) -@pytest.mark.parametrize("combine_format", ["bf16", "mxfp8"]) def test_mxfp8_forward_multi_gpu_matches_reference( world_size, combine_format, @@ -1295,35 +1171,6 @@ def test_column_requant_workspace_is_allocated_only_when_enabled(): # Reference and quantization self-checks. -@pytest.mark.L0 -def test_mxfp8_activation_representation(): - activation = quantize_mxfp8(torch.randn(3, 128), axis=1) - - assert activation.format.value == "mxfp8" - assert activation.logical_shape == (3, 128) - assert activation.axis == 1 - assert activation.data.shape == (3, 128) - assert activation.data.dtype == torch.float8_e4m3fn - assert activation.scale.shape == (3, 4) - assert activation.scale.dtype == torch.float8_e8m0fnu - assert torch.isfinite(activation.dequantize()).all() - - -@pytest.mark.L0 -def test_reference_mxfp8_block_scaled_round_trip(): - values = torch.linspace(-4.0, 4.0, 3 * 64).reshape(3, 64) - quantized = quantize_blockwise(values, MoeFormat.MXFP8) - - assert isinstance(quantized, ReferenceBlockScaledTensor) - assert quantized.format is MoeFormat.MXFP8 - assert quantized.logical_shape == (3, 64) - assert tuple(quantized.data.shape) == (3, 64) - assert tuple(quantized.scale.shape) == (3, 2) - assert quantized.scale.dtype == torch.float8_e8m0fnu - assert quantized.dequantize().shape == values.shape - assert torch.isfinite(quantized.dequantize()).all() - - @pytest.mark.L0 @pytest.mark.parametrize( "intermediate_format", @@ -1460,33 +1307,31 @@ def test_megamoe_capability_and_kernel_config_accept_ep_above_16(): @pytest.mark.L0 -def test_megamoe_capability_accepts_nonworld_subgroup_config(): - from cudnn.moe_ep._contracts import ForwardConfig - from cudnn.moe_ep._megamoe_backend._capability import validate_config - from cudnn.moe_ep._tuning import MoeEpTuningConfig +def test_ep32_peer_mapping_selects_vector_payload(): + from cutlass._mlir import ir - config = ForwardConfig( - num_experts=4, - hidden_size=128, - intermediate_size=256, - top_k=2, - experts_per_rank=2, - ep_size=2, - ep_rank=0, - ep_group=object(), - ep_global_ranks=(1, 3), - max_tokens_per_rank=8, - output_format="bf16", - combine_format="bf16", - apply_topk_in_fc1=True, - gate_up_clamp=None, - generate_c=False, - token_padding_size=128, - sf_padding_size=128, - tuning=MoeEpTuningConfig(), + from cudnn.moe_ep._megamoe_backend._comm import PeerMapping + from cudnn.moe_ep._megamoe_backend.cutedsl_src.communication.nvlink_domain.symmetric_buffer import ( + SymmetricBufferDevice, ) - validate_config(config) + offsets = tuple(index * 4096 for index in range(32)) + mapping = PeerMapping( + base_address=0x1000, + offsets=offsets, + rank=0, + ) + host = mapping.to_sym_buffer_host() + with ir.Context(): + device_type = SymmetricBufferDevice( + None, + host.max_ranks, + ).__get_mlir_types__()[0] + device_type_text = str(device_type) + + assert host.offsets == offsets + assert int(host.max_ranks) == 32 + assert device_type_text == "vector<32xi64>" @pytest.fixture @@ -1505,6 +1350,7 @@ def __init__(self, runtime_module, state=None): self._runtime_module = runtime_module self._state = state or runtime_module.RuntimeInitState.NOT_INITIALIZED self._world = None + self.initialize_count = 0 self.finalize_count = 0 def initialization_state(self): @@ -1512,6 +1358,7 @@ def initialization_state(self): def initialize(self, device, world): del device + self.initialize_count += 1 self._world = world self._state = self._runtime_module.RuntimeInitState.INITIALIZED @@ -1555,6 +1402,37 @@ def test_runtime_manager_shares_only_identical_subgroup(runtime_module): assert provider.finalize_count == 1 +@pytest.mark.L0 +def test_runtime_manager_keep_alive_reuses_until_explicit_shutdown(runtime_module): + world = runtime_module.RuntimeWorld( + rank=0, + size=2, + group=object(), + global_ranks=(0, 1), + ) + provider = _FakeRuntimeProvider(runtime_module) + manager = runtime_module.RuntimeManager( + provider_factory=lambda: provider, + world_resolver=lambda config: world, + keep_alive=True, + ) + + first = manager.acquire(object(), torch.device("cuda", 0)) + first.close() + assert manager.ref_count == 0 + assert provider.initialize_count == 1 + assert provider.finalize_count == 0 + + second = manager.acquire(object(), torch.device("cuda", 0)) + assert manager.ref_count == 1 + assert provider.initialize_count == 1 + second.close() + + manager.shutdown() + assert manager.ref_count == 0 + assert provider.finalize_count == 1 + + @pytest.mark.L0 def test_runtime_manager_rejects_different_same_geometry_subgroup( runtime_module, diff --git a/test/python/moe_ep/test_moe_ep_forward_multinode.py b/test/python/moe_ep/test_moe_ep_forward_multinode.py deleted file mode 100644 index 46efb85a4..000000000 --- a/test/python/moe_ep/test_moe_ep_forward_multinode.py +++ /dev/null @@ -1,209 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Torchrun-native multi-node MoE EP forward acceptance tests.""" - -from __future__ import annotations - -import os -from dataclasses import dataclass -from datetime import timedelta - -import pytest -import torch -import torch.distributed as dist - -from moe_ep.moe_ep_distributed_workers import ( - _run_forward_output_case, -) - - -pytestmark = [ - pytest.mark.L1, - pytest.mark.gpu_exclusive, - pytest.mark.moe_ep_multinode, -] - -_TORCHRUN_ENV = ("LOCAL_RANK", "LOCAL_WORLD_SIZE", "RANK", "WORLD_SIZE") -_PROCESS_GROUP_TIMEOUT = timedelta(minutes=10) - - -def _bind_torchrun_device_before_pytest_fixtures() -> None: - """Bind before the root conftest creates its session CUDA handle.""" - - value = os.environ.get("LOCAL_RANK") - if value is None or not torch.cuda.is_available(): - return - local_rank = int(value) - if 0 <= local_rank < torch.cuda.device_count(): - torch.cuda.set_device(local_rank) - - -_bind_torchrun_device_before_pytest_fixtures() - - -@dataclass(frozen=True) -class _TorchrunWorld: - rank: int - world_size: int - local_rank: int - local_world_size: int - device: torch.device - - -def _require_torchrun_environment() -> tuple[int, int, int, int]: - missing = [name for name in _TORCHRUN_ENV if name not in os.environ] - if missing: - pytest.skip( - "multi-node MoE EP forward requires torchrun environment variables: " - + ", ".join(missing) - ) - return ( - int(os.environ["RANK"]), - int(os.environ["WORLD_SIZE"]), - int(os.environ["LOCAL_RANK"]), - int(os.environ["LOCAL_WORLD_SIZE"]), - ) - - -@pytest.fixture(scope="session") -def torchrun_world(): - if not dist.is_available() or not dist.is_nccl_available(): - pytest.skip("multi-node Rubin MXFP8 forward requires NCCL") - - rank, world_size, local_rank, local_world_size = ( - _require_torchrun_environment() - ) - if local_rank < 0 or local_rank >= torch.cuda.device_count(): - pytest.skip( - f"torchrun LOCAL_RANK={local_rank} is not backed by a visible GPU" - ) - - device = torch.device("cuda", local_rank) - if torch.cuda.get_device_capability(device) != (10, 7): - pytest.skip( - "multi-node Rubin MXFP8 forward requires exactly SM107 " - "(compute capability 10.7) on every rank" - ) - try: - import nvshmem.core # noqa: F401 - except (ImportError, OSError): - pytest.skip("multi-node Rubin MXFP8 forward requires NVSHMEM") - - os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") - torch.cuda.set_device(device) - if dist.is_initialized(): - if dist.get_rank() != rank or dist.get_world_size() != world_size: - raise RuntimeError( - "existing process group does not match torchrun RANK/WORLD_SIZE" - ) - else: - dist.init_process_group( - backend="nccl", - init_method="env://", - device_id=device, - timeout=_PROCESS_GROUP_TIMEOUT, - ) - - context = _TorchrunWorld( - rank=rank, - world_size=world_size, - local_rank=local_rank, - local_world_size=local_world_size, - device=device, - ) - try: - yield context - finally: - if dist.is_initialized(): - dist.barrier() - dist.destroy_process_group() - - -@pytest.mark.parametrize( - ( - "ep_size", - "required_world_size", - "required_local_world_size", - "ep_global_ranks", - ), - [ - pytest.param( - 7, - 14, - 2, - tuple(range(0, 14, 2)), - id="ep7-world14", - ), - pytest.param( - 12, - 12, - 4, - tuple(range(12)), - id="ep12-world12", - ), - pytest.param( - 15, - 20, - 4, - tuple(rank for rank in range(20) if rank % 4 < 3), - id="ep15-world20", - ), - pytest.param( - 16, - 16, - 4, - tuple(range(16)), - id="ep16-world16", - ), - ], -) -@pytest.mark.parametrize("combine_format", ["bf16", "mxfp8"]) -def test_mxfp8_forward_multinode_matches_reference( - torchrun_world, - ep_size, - required_world_size, - required_local_world_size, - ep_global_ranks, - combine_format, -): - world = torchrun_world - if ( - world.world_size != required_world_size - or world.local_world_size != required_local_world_size - ): - pytest.skip( - f"EP{ep_size} requires torchrun WORLD_SIZE={required_world_size}, " - f"LOCAL_WORLD_SIZE={required_local_world_size}; got " - f"WORLD_SIZE={world.world_size}, " - f"LOCAL_WORLD_SIZE={world.local_world_size}" - ) - - if ep_size == world.world_size: - ep_group = dist.group.WORLD - else: - # All WORLD ranks must create subgroups in the same order, including - # idle ranks that are not members of this balanced EP group. - ep_group = dist.new_group( - list(ep_global_ranks), - backend="nccl", - timeout=_PROCESS_GROUP_TIMEOUT, - ) - - is_ep_member = world.rank in ep_global_ranks - try: - if is_ep_member: - ep_rank = dist.get_rank(ep_group) - _run_forward_output_case( - device=world.device, - ep_group=ep_group, - ep_rank=ep_rank, - ep_size=ep_size, - combine_format=combine_format, - expected_global_ranks=ep_global_ranks, - ) - dist.barrier() - finally: - if ep_group is not dist.group.WORLD and is_ep_member: - dist.destroy_process_group(ep_group) - dist.barrier() diff --git a/test/python/moe_ep/test_moe_ep_multinode.py b/test/python/moe_ep/test_moe_ep_multinode.py new file mode 100644 index 000000000..26b65636e --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_multinode.py @@ -0,0 +1,365 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Torchrun-native multi-node MoE EP forward/backward acceptance tests.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist + +from moe_ep.moe_ep_distributed_workers import ( + _run_backward_reference_case, + _run_forward_output_case, +) +from moe_ep.moe_ep_test_support import ( + _fixed_training_weights, + make_distributed_forward_inputs, +) + + +pytestmark = [ + pytest.mark.L1, + pytest.mark.gpu_exclusive, + pytest.mark.moe_ep_multinode, +] + +_TORCHRUN_ENV = ("LOCAL_RANK", "LOCAL_WORLD_SIZE", "RANK", "WORLD_SIZE") +_PROCESS_GROUP_TIMEOUT = timedelta(minutes=10) + + +def _bind_torchrun_device_before_pytest_fixtures() -> None: + """Bind before the root conftest creates its session CUDA handle.""" + + value = os.environ.get("LOCAL_RANK") + if value is None or not torch.cuda.is_available(): + return + local_rank = int(value) + if 0 <= local_rank < torch.cuda.device_count(): + torch.cuda.set_device(local_rank) + + +_bind_torchrun_device_before_pytest_fixtures() + + +@dataclass(frozen=True) +class _TorchrunWorld: + rank: int + world_size: int + local_rank: int + local_world_size: int + device: torch.device + + +def _require_torchrun_environment() -> tuple[int, int, int, int]: + missing = [name for name in _TORCHRUN_ENV if name not in os.environ] + if missing: + pytest.skip( + "multi-node MoE EP tests require torchrun environment variables: " + + ", ".join(missing) + ) + return ( + int(os.environ["RANK"]), + int(os.environ["WORLD_SIZE"]), + int(os.environ["LOCAL_RANK"]), + int(os.environ["LOCAL_WORLD_SIZE"]), + ) + + +@pytest.fixture(scope="session") +def torchrun_world(): + if not dist.is_available() or not dist.is_nccl_available(): + pytest.skip("multi-node Rubin MXFP8 tests require NCCL") + + rank, world_size, local_rank, local_world_size = ( + _require_torchrun_environment() + ) + if local_rank < 0 or local_rank >= torch.cuda.device_count(): + pytest.skip( + f"torchrun LOCAL_RANK={local_rank} is not backed by a visible GPU" + ) + + device = torch.device("cuda", local_rank) + if torch.cuda.get_device_capability(device) != (10, 7): + pytest.skip( + "multi-node Rubin MXFP8 tests require exactly SM107 " + "(compute capability 10.7) on every rank" + ) + try: + import nvshmem.core # noqa: F401 + except (ImportError, OSError): + pytest.skip("multi-node Rubin MXFP8 tests require NVSHMEM") + + os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") + torch.cuda.set_device(device) + if dist.is_initialized(): + if dist.get_rank() != rank or dist.get_world_size() != world_size: + raise RuntimeError( + "existing process group does not match torchrun RANK/WORLD_SIZE" + ) + else: + dist.init_process_group( + backend="nccl", + init_method="env://", + device_id=device, + timeout=_PROCESS_GROUP_TIMEOUT, + ) + + context = _TorchrunWorld( + rank=rank, + world_size=world_size, + local_rank=local_rank, + local_world_size=local_world_size, + device=device, + ) + try: + yield context + finally: + if dist.is_initialized(): + try: + dist.barrier() + from cudnn.moe_ep._megamoe_backend._runtime import ( + get_runtime_manager, + ) + + get_runtime_manager().shutdown() + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + ( + "ep_size", + "required_world_size", + "required_local_world_size", + "ep_global_ranks", + ), + [ + pytest.param( + 4, + 8, + 4, + (0, 1, 4, 5), + id="forward-ep4-world8", + ), + pytest.param( + 6, + 8, + 4, + (0, 1, 2, 4, 5, 6), + id="forward-ep6-world8", + ), + pytest.param( + 12, + 12, + 4, + tuple(range(12)), + id="forward-ep12-world12", + ), + pytest.param( + 16, + 16, + 4, + tuple(range(16)), + id="forward-ep16-world16", + ), + ], +) +@pytest.mark.parametrize("combine_format", ["bf16", "mxfp8"]) +def test_mxfp8_forward_multinode_matches_reference( + torchrun_world, + ep_size, + required_world_size, + required_local_world_size, + ep_global_ranks, + combine_format, +): + world = torchrun_world + if ( + world.world_size != required_world_size + or world.local_world_size != required_local_world_size + ): + pytest.skip( + f"EP{ep_size} requires torchrun WORLD_SIZE={required_world_size}, " + f"LOCAL_WORLD_SIZE={required_local_world_size}; got " + f"WORLD_SIZE={world.world_size}, " + f"LOCAL_WORLD_SIZE={world.local_world_size}" + ) + + if ep_size == world.world_size: + ep_group = dist.group.WORLD + else: + # All WORLD ranks must create subgroups in the same order, including + # idle ranks that are not members of this balanced EP group. + ep_group = dist.new_group( + list(ep_global_ranks), + backend="nccl", + timeout=_PROCESS_GROUP_TIMEOUT, + ) + + is_ep_member = world.rank in ep_global_ranks + try: + if is_ep_member: + ep_rank = dist.get_rank(ep_group) + _run_forward_output_case( + device=world.device, + ep_group=ep_group, + ep_rank=ep_rank, + ep_size=ep_size, + combine_format=combine_format, + expected_global_ranks=ep_global_ranks, + ) + dist.barrier() + finally: + if ep_group is not dist.group.WORLD and is_ep_member: + # The runtime manager intentionally remains alive after op.close(). + # Finalize it while this exact subgroup is still valid so the next + # parametrized case may create a fresh ProcessGroup object with the + # same membership. + from cudnn.moe_ep._megamoe_backend._runtime import ( + get_runtime_manager, + ) + + get_runtime_manager().shutdown() + dist.barrier(group=ep_group) + dist.destroy_process_group(ep_group) + dist.barrier() + + +@pytest.mark.parametrize( + ("ep_size", "required_world_size", "combine_format"), + [ + pytest.param( + 8, + 8, + "bf16", + id="backward-ep8-world8-bf16", + ), + pytest.param( + 8, + 8, + "mxfp8", + id="backward-ep8-world8-mxfp8", + ), + pytest.param( + 16, + 16, + "bf16", + id="backward-ep16-world16-bf16", + ), + pytest.param( + 16, + 16, + "mxfp8", + id="backward-ep16-world16-mxfp8", + ), + pytest.param( + 32, + 32, + "bf16", + id="backward-ep32-world32-bf16-minimal", + ), + ], +) +def test_fixed_training_resources_multinode_match_independent_reference( + torchrun_world, + ep_size, + required_world_size, + combine_format, +): + world = torchrun_world + if world.world_size != required_world_size: + pytest.skip( + f"EP{ep_size} requires torchrun WORLD_SIZE={required_world_size}; " + f"got WORLD_SIZE={world.world_size}" + ) + + _run_backward_reference_case( + device=world.device, + ep_group=dist.group.WORLD, + ep_rank=world.rank, + ep_size=ep_size, + combine_format=combine_format, + ) + + +@pytest.mark.parametrize( + ("rank_zero_lane_count", "other_lane_count"), + [ + pytest.param( + 2, + 1, + id="backward-ep8-world8-abi-mismatch", + ), + ], +) +def test_training_prepare_multinode_rejects_rank_abi_mismatch( + torchrun_world, + rank_zero_lane_count, + other_lane_count, +): + world = torchrun_world + if world.world_size != 8 or world.local_world_size != 4: + pytest.skip( + "EP8 ABI mismatch requires torchrun WORLD_SIZE=8, " + "LOCAL_WORLD_SIZE=4; got " + f"WORLD_SIZE={world.world_size}, " + f"LOCAL_WORLD_SIZE={world.local_world_size}" + ) + + from cudnn import MoeEp + + # A fixed helper rank gives every process byte-identical, locally valid + # weight packs. Only the locally valid lane count differs. + args = make_distributed_forward_inputs(0, 8, world.device) + weights = _fixed_training_weights(args) + op = MoeEp( + num_experts=16, + hidden_size=128, + intermediate_size=256, + top_k=2, + ep_group=dist.group.WORLD, + max_tokens_per_rank=8, + max_recv_size_per_rank=3, + drop_on_overflow=True, + combine_format="bf16", + ) + caught_error = None + try: + lane_count = ( + rank_zero_lane_count if world.rank == 0 else other_lane_count + ) + try: + op.prepare_training_resources( + weights, + slot_count=1, + lane_count=lane_count, + ) + except Exception as error: + caught_error = error + + # Reachable only after the collective prepare path returns or raises on + # every rank. The NCCL barrier blocks the host and launches no MoE + # kernel or device-side assertion. + dist.barrier( + group=dist.group.WORLD, + device_ids=[world.local_rank], + ) + + assert isinstance(caught_error, RuntimeError), ( + f"rank {world.rank} expected RuntimeError from collective prepare, " + f"got {caught_error!r}" + ) + message = str(caught_error) + assert ( + "symmetric workspace region counts differ" in message + or "ABI differs" in message + ), f"rank {world.rank} got unexpected prepare error: {message}" + finally: + op.close() diff --git a/test/python/moe_ep/test_moe_ep_wgrad_contract.py b/test/python/moe_ep/test_moe_ep_wgrad_contract.py deleted file mode 100644 index 4e307b5d6..000000000 --- a/test/python/moe_ep/test_moe_ep_wgrad_contract.py +++ /dev/null @@ -1,1170 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Focused L0 tests for the public MoeEP wgrad operand contract.""" - -from __future__ import annotations - -import os -from dataclasses import fields, replace -from types import SimpleNamespace - -import pytest -import torch -import torch.multiprocessing as mp - -from cudnn.moe_ep import ( - MoeEp, - MoeEpWgradForwardStash, - MoeEpWgradOperands, -) -from cudnn.moe_ep._megamoe_backend import _capability -from cudnn.moe_ep._megamoe_backend.mxfp8._backward_launch import ( - Mxfp8DgluResult, -) -from cudnn.moe_ep._megamoe_backend.mxfp8._backward_wgrad_export import ( - export_wgrad_operands, -) -from cudnn.moe_ep._megamoe_backend.mxfp8._config import Mxfp8KernelConfig -from cudnn.moe_ep._megamoe_backend.mxfp8._stash import ( - Mxfp8ForwardStash as Mxfp8ForwardStashOwner, -) -from cudnn.moe_ep._megamoe_backend.mxfp8._wgrad_layout import ( - assemble_dfc2_atom_scales, - assemble_discrete_col_requant_scales, - assemble_plain_col_scales, -) -from cudnn.moe_ep._validation import validate_backward -from moe_ep.moe_ep_backward_support import _dense_wgrads_from_operands -from moe_ep.moe_ep_distributed_workers import ( - _distributed_wgrad_worker, - _run_wgrad_operand_case, -) -from moe_ep.moe_ep_forward_support import ( - _require_distributed_sm107, - _sm107_device, -) -from moe_ep.moe_ep_reference import ( - MoeEpReference, - MoeFormat, - WgradOperandsReference, - quantize_blockwise, -) - - -def _inputs(): - activation = torch.randn(2, 128, dtype=torch.bfloat16) - fc1_weight = torch.randn(2, 128, 512, dtype=torch.bfloat16) - fc2_weight = torch.randn(2, 256, 128, dtype=torch.bfloat16) - topk_idx = torch.tensor([[0, -1], [1, 0]], dtype=torch.int32) - topk_weights = torch.randn(2, 2, dtype=torch.float32) - return activation, fc1_weight, fc2_weight, topk_idx, topk_weights - - -def _route_metadata(): - return torch.tensor( - [[0, 0, 0, 0], [0, 0, 1, 1], [1, 0, 1, 0]], - dtype=torch.int32, - ) - - -def _forward_stash(route_metadata=None): - if route_metadata is None: - route_metadata = _route_metadata() - return MoeEpWgradForwardStash( - fc1_a=torch.empty(128, 512, dtype=torch.float8_e4m3fn), - fc1_sfa=torch.empty(128, 16, dtype=torch.float8_e8m0fnu), - expert_offsets=torch.tensor([256, 512], dtype=torch.int32), - valid_route_counts=torch.tensor([2, 1], dtype=torch.int32), - route_metadata=route_metadata.clone(), - ) - - -def _operator(**overrides): - kwargs = { - "num_experts": 2, - "hidden_size": 128, - "intermediate_size": 256, - "top_k": 2, - "max_tokens_per_rank": 4, - "generate_c": True, - } - kwargs.update(overrides) - if ( - kwargs.get("backward_wgrad_mode") == "operands" - and "token_padding_size" not in overrides - ): - kwargs["token_padding_size"] = 256 - return MoeEp(**kwargs) - - -@pytest.mark.L0 -def test_wgrad_types_are_public_and_have_stable_fields(): - from cudnn import ( - MoeEpWgradForwardStash as TopLevelForwardStash, - MoeEpWgradOperands as TopLevelOperands, - ) - - assert TopLevelForwardStash is MoeEpWgradForwardStash - assert TopLevelOperands is MoeEpWgradOperands - assert [field.name for field in fields(MoeEpWgradForwardStash)] == [ - "fc1_a", - "fc1_sfa", - "expert_offsets", - "valid_route_counts", - "route_metadata", - ] - assert [field.name for field in fields(MoeEpWgradOperands)] == [ - "fc1_a", - "fc1_sfa", - "fc1_b", - "fc1_sfb", - "fc2_a", - "fc2_sfa", - "fc2_b", - "fc2_sfb", - "expert_offsets", - "valid_route_counts", - "route_metadata", - ] - - -@pytest.mark.L0 -def test_wgrad_mode_is_opt_in_and_requires_generate_c(): - with _operator(generate_c=False) as operator: - assert operator.backward_wgrad_mode == "none" - assert operator._forward_config.backward_wgrad_mode == "none" - - with pytest.raises(ValueError, match="must be 'none' or 'operands'"): - _operator(backward_wgrad_mode="weights") - with pytest.raises(ValueError, match="requires generate_c=True"): - _operator( - generate_c=False, - backward_wgrad_mode="operands", - ) - with pytest.raises(ValueError, match="requires token_padding_size=256"): - _operator( - backward_wgrad_mode="operands", - token_padding_size=128, - ) - with pytest.raises(ValueError, match="requires sf_padding_size=128"): - _operator( - backward_wgrad_mode="operands", - sf_padding_size=256, - ) - - with _operator(backward_wgrad_mode="operands") as operator: - assert operator.backward_wgrad_mode == "operands" - assert operator._forward_config.backward_wgrad_mode == "operands" - - -@pytest.mark.L0 -def test_validate_backward_checks_wgrad_stash_layout_and_route_identity(): - args = _inputs() - route_metadata = _route_metadata() - stash = _forward_stash(route_metadata) - with _operator(backward_wgrad_mode="operands") as operator: - with pytest.raises(TypeError, match="must be a MoeEpWgradForwardStash"): - validate_backward( - operator._forward_config, - torch.randn(2, 128), - *args[1:], - torch.randn(3, 512, dtype=torch.bfloat16), - route_metadata, - ) - - request = validate_backward( - operator._forward_config, - torch.randn(2, 128), - *args[1:], - torch.randn(3, 512, dtype=torch.bfloat16), - route_metadata, - wgrad_forward_stash=stash, - ) - assert request.wgrad_forward_stash is stash - - wrong_scale_shape = replace( - stash, - fc1_sfa=torch.empty(128, 15, dtype=torch.float8_e8m0fnu), - ) - with pytest.raises(ValueError, match="fc1_sfa shape must be"): - validate_backward( - operator._forward_config, - torch.randn(2, 128), - *args[1:], - torch.randn(3, 512, dtype=torch.bfloat16), - route_metadata, - wgrad_forward_stash=wrong_scale_shape, - ) - - wrong_a_stride = replace( - stash, - fc1_a=torch.empty( - 512, - 128, - dtype=torch.float8_e4m3fn, - ).transpose(0, 1), - ) - with pytest.raises(ValueError, match=r"compact \(K, 1\) strides"): - validate_backward( - operator._forward_config, - torch.randn(2, 128), - *args[1:], - torch.randn(3, 512, dtype=torch.bfloat16), - route_metadata, - wgrad_forward_stash=wrong_a_stride, - ) - - wrong_scale_stride = replace( - stash, - fc1_sfa=torch.empty( - 16, - 128, - dtype=torch.float8_e8m0fnu, - ).transpose(0, 1), - ) - with pytest.raises(ValueError, match="fc1_sfa must be contiguous"): - validate_backward( - operator._forward_config, - torch.randn(2, 128), - *args[1:], - torch.randn(3, 512, dtype=torch.bfloat16), - route_metadata, - wgrad_forward_stash=wrong_scale_stride, - ) - - noncanonical_offsets = replace( - stash, - expert_offsets=torch.tensor([512, 768], dtype=torch.int32), - ) - with pytest.raises(ValueError, match="canonical 256-row padding"): - validate_backward( - operator._forward_config, - torch.randn(2, 128), - *args[1:], - torch.randn(3, 512, dtype=torch.bfloat16), - route_metadata, - wgrad_forward_stash=noncanonical_offsets, - ) - - wrong_identity = _forward_stash(route_metadata.flip(0)) - with pytest.raises(ValueError, match="route identity"): - validate_backward( - operator._forward_config, - torch.randn(2, 128), - *args[1:], - torch.randn(3, 512, dtype=torch.bfloat16), - route_metadata, - wgrad_forward_stash=wrong_identity, - ) - - wrong_counts = MoeEpWgradForwardStash( - stash.fc1_a, - stash.fc1_sfa, - stash.expert_offsets, - torch.tensor([1, 2], dtype=torch.int32), - stash.route_metadata, - ) - with pytest.raises(ValueError, match="do not match route_metadata"): - validate_backward( - operator._forward_config, - torch.randn(2, 128), - *args[1:], - torch.randn(3, 512, dtype=torch.bfloat16), - route_metadata, - wgrad_forward_stash=wrong_counts, - ) - - -@pytest.mark.L0 -def test_default_mode_rejects_wgrad_stash_without_changing_default_contract(): - args = _inputs() - route_metadata = _route_metadata() - with _operator() as operator: - with pytest.raises(ValueError, match="only accepted"): - validate_backward( - operator._forward_config, - torch.randn(2, 128), - *args[1:], - torch.randn(3, 512, dtype=torch.bfloat16), - route_metadata, - wgrad_forward_stash=_forward_stash(route_metadata), - ) - - -@pytest.mark.L0 -def test_wgrad_mode_is_backend_capable_and_enables_forward_col_quant( - monkeypatch, -): - with _operator(backward_wgrad_mode="operands") as operator: - monkeypatch.setattr( - _capability, - "_validate_device", - lambda device: pytest.fail("device capability queried"), - ) - _capability.validate_config(operator._forward_config) - kernel_config = Mxfp8KernelConfig.from_forward_config( - operator._forward_config - ) - - assert kernel_config.enable_col_quant is True - assert kernel_config.token_padding_block == 256 - - -@pytest.mark.L0 -def test_forward_col_quant_runtime_uses_static_cute_layout(monkeypatch): - from cudnn.moe_ep._megamoe_backend.mxfp8 import _launch - - tensors = { - name: object() - for name in ( - "activation", - "activation_sf", - "topk_indices", - "topk_scores", - "fc1_weight", - "fc1_weight_sf", - "fc2_weight", - "fc2_weight_sf", - "fc1_c", - "output_data", - "col_quant_data", - "col_quant_sf", - "overflow_flag", - "local_workspace", - "shared_workspace", - ) - } - calls = {} - - def fake_to_cute( - tensor, - assumed_align=16, - *, - dynamic_layout=True, - ): - calls[id(tensor)] = (assumed_align, dynamic_layout) - return tensor - - monkeypatch.setattr(_launch, "_to_cute", fake_to_cute) - monkeypatch.setattr(_launch, "_to_cute_ptr", lambda tensor: tensor) - inputs = SimpleNamespace( - activation=tensors["activation"], - activation_sf=tensors["activation_sf"], - topk_indices=tensors["topk_indices"], - topk_scores=tensors["topk_scores"], - weights=SimpleNamespace( - fc1_weight=tensors["fc1_weight"], - fc1_weight_sf=tensors["fc1_weight_sf"], - fc2_weight=tensors["fc2_weight"], - fc2_weight_sf=tensors["fc2_weight_sf"], - ), - fc1_c=tensors["fc1_c"], - output_data=tensors["output_data"], - col_quant_data=tensors["col_quant_data"], - col_quant_sf=tensors["col_quant_sf"], - overflow_flag=tensors["overflow_flag"], - local_workspace=tensors["local_workspace"], - shared_workspace=tensors["shared_workspace"], - ) - resources = SimpleNamespace( - runtime=SimpleNamespace( - current_stream=lambda: SimpleNamespace(cuda_stream=0) - ), - workspace=SimpleNamespace( - peer_mapping=SimpleNamespace( - to_sym_buffer_host=lambda: object() - ) - ), - ) - - _launch.build_runtime_kwargs(inputs, resources) - - assert calls[id(tensors["col_quant_data"])] == (128, False) - assert calls[id(tensors["col_quant_sf"])] == (16, False) - assert calls[id(tensors["overflow_flag"])] == (4, False) - - -@pytest.mark.L0 -def test_opt_in_forward_and_backward_results_are_backend_representable( - monkeypatch, -): - import cudnn.moe_ep._backend as backend_seam - - args = _inputs() - route_metadata = _route_metadata() - fc1_c = torch.randn(3, 512, dtype=torch.bfloat16) - stash = _forward_stash(route_metadata) - forward_result = (torch.randn(2, 128), fc1_c, route_metadata, stash) - operands = MoeEpWgradOperands( - fc1_a=stash.fc1_a, - fc1_sfa=stash.fc1_sfa, - fc1_b=torch.empty(512, 512, dtype=torch.float8_e4m3fn), - fc1_sfb=torch.empty(512, 16, dtype=torch.float8_e8m0fnu), - fc2_a=torch.empty(256, 512, dtype=torch.float8_e4m3fn), - fc2_sfa=torch.empty(256, 16, dtype=torch.float8_e8m0fnu), - fc2_b=torch.empty(512, 128, dtype=torch.float8_e4m3fn), - fc2_sfb=torch.empty(128, 16, dtype=torch.float8_e8m0fnu), - expert_offsets=stash.expert_offsets, - valid_route_counts=stash.valid_route_counts, - route_metadata=stash.route_metadata, - ) - backward_result = ( - torch.randn(2, 128), - torch.randn(2, 2), - operands, - ) - - class Backend: - backward_request = None - - def forward(self, request): - return forward_result - - def backward(self, request): - self.backward_request = request - return backward_result - - def close(self): - pass - - backend = Backend() - monkeypatch.setattr(backend_seam, "validate_config", lambda config: None) - monkeypatch.setattr(backend_seam, "validate_request", lambda request: None) - monkeypatch.setattr( - backend_seam, - "validate_backward_request", - lambda request: None, - ) - monkeypatch.setattr( - backend_seam, - "create_backend", - lambda config, device: backend, - ) - - with _operator(backward_wgrad_mode="operands") as operator: - assert operator(*args) is forward_result - actual = operator.backward( - torch.randn(2, 128), - *args[1:], - fc1_c, - route_metadata, - wgrad_forward_stash=stash, - ) - - assert actual is backward_result - assert backend.backward_request.wgrad_forward_stash is stash - - -def _blocked_reference(raw: torch.Tensor) -> torch.Tensor: - rows, columns = raw.shape - if columns == 0: - return raw.new_empty((0,)) - padded_rows = (rows + 127) // 128 * 128 - padded_columns = (columns + 3) // 4 * 4 - padded = torch.full( - (padded_rows, padded_columns), - 127, - dtype=torch.uint8, - device=raw.device, - ) - padded[:rows, :columns] = raw - return ( - padded.view(padded_rows // 128, 128, padded_columns // 4, 4) - .permute(0, 2, 1, 3) - .reshape(-1, 4, 32, 4) - .transpose(1, 2) - .reshape(-1) - ) - - -@pytest.mark.L0 -def test_col_requant_scales_accept_upstream_hidden_atom_major_order(): - non_k = 256 - valid_counts = (33, 0, 129) - padded_ends = (256, 256, 512) - source0 = torch.arange(non_k * 4, dtype=torch.int32).to(torch.uint8).reshape( - non_k, - 4, - ) - source2 = ( - torch.arange(non_k * 8, dtype=torch.int32) + 37 - ).to(torch.uint8).reshape(non_k, 8) - target0 = torch.full((non_k, 8), 127, dtype=torch.uint8) - target0[:, :4] = source0 - blocked0 = _blocked_reference(target0) - blocked2 = _blocked_reference(source2) - - # Upstream col requant stores hidden atoms before token atoms inside each - # expert. Its 128-row SF padding is expanded to the 256-row data padding. - packed = torch.cat( - (_blocked_reference(source0), _blocked_reference(source2)) - ) - actual = assemble_discrete_col_requant_scales( - packed, - valid_counts, - padded_ends, - non_k, - 128, - ) - expected = torch.cat((blocked0, blocked2)).reshape(non_k, 16) - - assert actual.shape == (non_k, 16) - assert torch.equal(actual.view(torch.uint8), expected) - - -@pytest.mark.L0 -def test_dfc2_atom_scales_reorder_token_major_atoms_per_expert(): - non_k = 256 - valid_counts = (33, 0, 129) - padded_ends = (256, 256, 512) - source0 = torch.arange(non_k * 4, dtype=torch.int32).to(torch.uint8).reshape( - non_k, - 4, - ) - source2 = ( - torch.arange(non_k * 8, dtype=torch.int32) + 37 - ).to(torch.uint8).reshape(non_k, 8) - target0 = torch.full((non_k, 8), 127, dtype=torch.uint8) - target0[:, :4] = source0 - blocked0 = _blocked_reference(target0) - blocked2 = _blocked_reference(source2) - source_blocked0 = _blocked_reference(source0) - source_blocked2 = _blocked_reference(source2) - - # The dFC2 epilogue writes token atoms before hidden atoms. - physical0 = source_blocked0.reshape(2, 1, 512).permute(1, 0, 2).reshape(-1) - physical2 = source_blocked2.reshape(2, 2, 512).permute(1, 0, 2).reshape(-1) - actual = assemble_dfc2_atom_scales( - torch.cat((physical0, physical2)), - valid_counts, - padded_ends, - non_k, - 128, - ) - expected = torch.cat((blocked0, blocked2)).reshape(non_k, 16) - - assert actual.shape == (non_k, 16) - assert torch.equal(actual.view(torch.uint8), expected) - - -@pytest.mark.L0 -def test_dfc2_atom_scales_deinterleave_gate_up_before_repacking(): - intermediate = 256 - non_k = 2 * intermediate - logical_source = torch.arange( - non_k * 4, - dtype=torch.int32, - ).to(torch.uint8).reshape(non_k, 4) - gate = logical_source[:intermediate].reshape(8, 32, 4) - up = logical_source[intermediate:].reshape(8, 32, 4) - interleaved_source = torch.stack((gate, up), dim=1).reshape(non_k, 4) - target = torch.full((non_k, 8), 127, dtype=torch.uint8) - target[:, :4] = logical_source - - actual = assemble_dfc2_atom_scales( - _blocked_reference(interleaved_source), - (33,), - (256,), - non_k, - 128, - deinterleave_gate_up=intermediate, - ) - - assert torch.equal( - actual.view(torch.uint8), - _blocked_reference(target).reshape(non_k, 8), - ) - - -@pytest.mark.L0 -def test_plain_col_scales_assemble_256_data_padding_with_empty_expert(): - non_k = 256 - counts = (33, 0, 129) - padded_ends = (256, 256, 512) - source = torch.full((12, non_k), 127, dtype=torch.uint8) - source[:2] = torch.arange(2 * non_k, dtype=torch.int32).to( - torch.uint8 - ).reshape(2, non_k) - source[4:9] = ( - torch.arange(5 * non_k, dtype=torch.int32) + 19 - ).to(torch.uint8).reshape(5, non_k) - - raw0 = torch.full((non_k, 8), 127, dtype=torch.uint8) - raw0[:, :2] = source[:2].transpose(0, 1) - raw1 = torch.empty((non_k, 0), dtype=torch.uint8) - raw2 = torch.full((non_k, 8), 127, dtype=torch.uint8) - raw2[:, :5] = source[4:9].transpose(0, 1) - expected = torch.cat( - tuple( - _blocked_reference(raw) - for raw in (raw0, raw1, raw2) - ) - ).reshape(non_k, 16) - - actual = assemble_plain_col_scales( - source.view(torch.float8_e8m0fnu), - counts, - padded_ends, - non_k, - 128, - ) - - assert torch.equal(actual.view(torch.uint8), expected) - empty_plain = assemble_plain_col_scales( - torch.empty(0, non_k, dtype=torch.uint8).view( - torch.float8_e8m0fnu - ), - (0, 0), - (0, 0), - non_k, - 128, - ) - empty_discrete = assemble_discrete_col_requant_scales( - torch.empty(0, dtype=torch.uint8), - (0, 0), - (0, 0), - non_k, - 128, - ) - empty_dfc2 = assemble_dfc2_atom_scales( - torch.empty(0, dtype=torch.uint8), - (0, 0), - (0, 0), - non_k, - 128, - ) - assert empty_plain.shape == (non_k, 0) - assert empty_discrete.shape == (non_k, 0) - assert empty_dfc2.shape == (non_k, 0) - - -@pytest.mark.L0 -def test_forward_materializes_caller_owned_256_padded_operand_stash(): - operator = _operator(backward_wgrad_mode="operands") - config = operator._forward_config - owner = Mxfp8ForwardStashOwner(config, torch.device("cpu")) - request = SimpleNamespace( - topk_idx=torch.tensor([[0, -1], [1, 0]], dtype=torch.int32), - device=torch.device("cpu"), - ) - plan = owner.prepare(request, pool_token_capacity=512) - plan.buffer.zero_() - plan.buffer[0, 0] = 10 - plan.buffer[1, 0] = 20 - plan.buffer[256, 0] = 30 - - def pack(token, slot): - return token | (slot << 32) - - packed_metadata = torch.zeros(512, dtype=torch.int64) - packed_metadata[0] = pack(1, 1) - packed_metadata[1] = pack(0, 0) - packed_metadata[256] = pack(1, 0) - col_data = torch.zeros( - 512, - 128, - dtype=torch.float8_e4m3fn, - ) - col_sf = torch.full( - (512 // 32 * 128,), - 127, - dtype=torch.uint8, - ) - inputs = SimpleNamespace( - fc1_c=plan.buffer, - shared_workspace=packed_metadata.view(torch.uint8), - col_quant_data=col_data, - col_quant_sf=col_sf, - ) - prepared = SimpleNamespace( - token_src_metadata_offset=0, - token_src_metadata_bytes=512 * 8, - pool_token_capacity=512, - ) - - fc1_c, route_metadata, stash = owner.materialize( - plan, - inputs, - prepared, - ) - - assert stash is not None - assert stash.fc1_a.shape == (128, 512) - assert stash.fc1_a.stride(1) == 1 - assert stash.fc1_sfa.shape == (128, 16) - assert stash.expert_offsets.tolist() == [256, 512] - assert stash.valid_route_counts.tolist() == [2, 1] - assert stash.route_metadata is route_metadata - assert route_metadata.tolist() == [ - [0, 0, 0, 0], - [0, 0, 1, 1], - [1, 0, 1, 0], - ] - assert fc1_c[:, 0].tolist() == [20, 10, 30] - preserved = stash.fc1_a.clone() - col_data.fill_(1) - assert torch.equal(stash.fc1_a, preserved) - owner.close() - operator.close() - - -@pytest.mark.L0 -def test_backward_export_owns_outputs_and_uses_grouped_wgrad_strides(): - config = _operator( - backward_wgrad_mode="operands" - )._forward_config - stash = _forward_stash() - request = SimpleNamespace( - config=config, - wgrad_forward_stash=stash, - ) - pool_rows = 512 - sf_rows = 8 - aux = Mxfp8DgluResult( - grad_activation=torch.empty(2, 128), - fc1_recompute=torch.zeros( - pool_rows, - 256, - dtype=torch.float8_e4m3fn, - ), - fc1_recompute_sf=torch.full( - (sf_rows, 256), - 127, - dtype=torch.uint8, - ).view(torch.float8_e8m0fnu), - fc1_col_output=torch.zeros( - pool_rows, - 512, - dtype=torch.float8_e4m3fn, - ), - fc1_col_output_sf=torch.full( - (sf_rows, 512), - 127, - dtype=torch.uint8, - ).view(torch.float8_e8m0fnu), - grad_y2=torch.zeros( - pool_rows, - 128, - dtype=torch.float8_e4m3fn, - ), - grad_y2_sf=torch.full( - (pool_rows // 32 * 128,), - 127, - dtype=torch.uint8, - ), - ) - - operands = export_wgrad_operands(request, aux) - - assert operands.fc1_b.shape == (512, 512) - assert operands.fc1_b.stride(0) == 1 - assert operands.fc2_a.shape == (256, 512) - assert operands.fc2_a.stride(1) == 1 - assert operands.fc2_b.shape == (512, 128) - assert operands.fc2_b.stride(0) == 1 - assert operands.fc1_sfb.shape == (512, 16) - assert operands.fc2_sfa.shape == (256, 16) - assert operands.fc2_sfb.shape == (128, 16) - - preserved_fc1_b = operands.fc1_b.clone() - preserved_fc2_a = operands.fc2_a.clone() - preserved_fc2_b = operands.fc2_b.clone() - aux.fc1_col_output.fill_(1) - aux.fc1_recompute.fill_(1) - aux.grad_y2.fill_(1) - assert torch.equal(operands.fc1_b, preserved_fc1_b) - assert torch.equal(operands.fc2_a, preserved_fc2_a) - assert torch.equal(operands.fc2_b, preserved_fc2_b) - - -def _reference_wgrad_case(topk_idx, topk_weights): - torch.manual_seed(20260821) - token_count = topk_idx.shape[0] - hidden = intermediate = 32 - activation = torch.randn(token_count, hidden) / 4 - fc1_weight = torch.randn(3, hidden, 2 * intermediate) / 8 - fc2_weight = torch.randn(3, intermediate, hidden) / 8 - grad_output = torch.randn(token_count, hidden) / 8 - reference = MoeEpReference( - num_experts=3, - hidden_size=hidden, - intermediate_size=intermediate, - top_k=topk_idx.shape[1], - max_tokens_per_rank=token_count, - generate_c=True, - backward_wgrad_mode="operands", - token_padding_size=256, - ) - _, fc1_c, route_metadata, forward_stash = reference( - activation, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights, - ) - _, _, operands = reference.backward( - grad_output, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights, - fc1_c, - route_metadata, - wgrad_forward_stash=forward_stash, - ) - return ( - activation, - fc2_weight, - grad_output, - fc1_c, - route_metadata, - operands, - ) - - -@pytest.mark.L0 -def test_reference_wgrad_operands_follow_route_weights_and_dense_formulas(): - topk_idx = torch.tensor( - [[0, 2], [2, 0], [0, 2]], - dtype=torch.int32, - ) - topk_weights = torch.tensor( - [[0.5, 0.25], [0.0, 0.75], [1.0, 0.125]], - dtype=torch.float32, - ) - ( - activation, - fc2_weight, - grad_output, - fc1_c, - route_metadata, - operands, - ) = _reference_wgrad_case(topk_idx, topk_weights) - - assert isinstance(operands, WgradOperandsReference) - assert operands.expert_offsets.tolist() == [256, 256, 512] - assert operands.valid_route_counts.tolist() == [3, 0, 3] - assert operands.fc1_a.shape == (32, 512) - assert operands.fc1_b.shape == (512, 64) - assert operands.fc2_a.shape == (32, 512) - assert operands.fc2_b.shape == (512, 32) - - # Rebuild the staged values independently in compact metadata order. - staged_x = quantize_blockwise( - activation, - MoeFormat.MXFP8, - axis=1, - ).dequantize() - staged_dy = quantize_blockwise( - grad_output, - MoeFormat.MXFP8, - axis=1, - ).dequantize() - staged_w2 = quantize_blockwise( - fc2_weight.transpose(1, 2), - MoeFormat.MXFP8, - axis=1, - ).dequantize().transpose(1, 2) - compact_x = [] - compact_weighted_h = [] - compact_dy = [] - compact_dc = [] - for row, (expert, _, token, slot) in enumerate(route_metadata.tolist()): - c_gate, c_up = fc1_c[row].float().split(32) - sigmoid = torch.sigmoid(c_gate) - silu = c_gate * sigmoid - h = silu * c_up - p = topk_weights[token, slot] - dy = staged_dy[token] - dh = (dy @ staged_w2[expert].transpose(0, 1)) * p - dc = torch.cat( - ( - dh * c_up * sigmoid * (1 + c_gate * (1 - sigmoid)), - dh * silu, - ) - ) - compact_x.append(staged_x[token]) - compact_weighted_h.append(p * h) - compact_dy.append(dy) - compact_dc.append(dc) - - def padded(rows): - result = torch.zeros(512, rows[0].numel()) - result[:3] = torch.stack(rows[:3]) - result[256:259] = torch.stack(rows[3:]) - return result - - expected_x = quantize_blockwise( - padded(compact_x).transpose(0, 1), - MoeFormat.MXFP8, - axis=1, - ) - expected_h = quantize_blockwise( - padded(compact_weighted_h).transpose(0, 1), - MoeFormat.MXFP8, - axis=1, - ) - expected_pdy = quantize_blockwise( - padded(compact_dy), - MoeFormat.MXFP8, - axis=0, - ) - expected_dc = quantize_blockwise( - padded(compact_dc), - MoeFormat.MXFP8, - axis=0, - ) - for actual, expected in ( - (operands.fc1_a, expected_x), - (operands.fc1_b, expected_dc), - (operands.fc2_a, expected_h), - (operands.fc2_b, expected_pdy), - ): - assert torch.equal(actual.data, expected.data) - assert torch.equal(actual.scale, expected.scale) - - # The valid route with p=0 keeps x/dY but contributes zero to both wgrads. - zero_weight_row = 257 - assert operands.fc1_a.dequantize()[:, zero_weight_row].abs().sum() > 0 - assert operands.fc2_a.dequantize()[:, zero_weight_row].eq(0).all() - assert operands.fc1_b.dequantize()[zero_weight_row].eq(0).all() - assert operands.fc2_b.dequantize()[zero_weight_row].abs().sum() > 0 - - dw1, dw2 = operands.dense_wgrads() - a1, b1 = operands.fc1_a.dequantize(), operands.fc1_b.dequantize() - a2, b2 = operands.fc2_a.dequantize(), operands.fc2_b.dequantize() - torch.testing.assert_close(dw1[0], a1[:, :256] @ b1[:256]) - torch.testing.assert_close(dw1[1], torch.zeros_like(dw1[1])) - torch.testing.assert_close(dw1[2], a1[:, 256:] @ b1[256:]) - torch.testing.assert_close(dw2[0], a2[:, :256] @ b2[:256]) - torch.testing.assert_close(dw2[1], torch.zeros_like(dw2[1])) - torch.testing.assert_close(dw2[2], a2[:, 256:] @ b2[256:]) - decoded_dw1, decoded_dw2 = _dense_wgrads_from_operands( - _as_production_operands(operands) - ) - torch.testing.assert_close(decoded_dw1, dw1) - torch.testing.assert_close(decoded_dw2, dw2) - - -@pytest.mark.L0 -def test_reference_wgrad_empty_routes_padding_offsets_and_invalid_routes(): - topk_idx = torch.full((2, 2), -1, dtype=torch.int32) - topk_weights = torch.randn(2, 2) - *_, fc1_c, route_metadata, operands = _reference_wgrad_case( - topk_idx, - topk_weights, - ) - - assert fc1_c.shape == (0, 64) - assert route_metadata.shape == (0, 4) - assert operands.expert_offsets.tolist() == [0, 0, 0] - assert operands.valid_route_counts.tolist() == [0, 0, 0] - assert operands.fc1_a.shape == (32, 0) - assert operands.fc1_b.shape == (0, 64) - assert operands.fc2_a.shape == (32, 0) - assert operands.fc2_b.shape == (0, 32) - for dense in operands.dense_wgrads(): - assert dense.eq(0).all() - - invalid = torch.tensor([[0, -2], [3, -1]], dtype=torch.int32) - with pytest.raises(ValueError, match="out-of-range expert id"): - _reference_wgrad_case(invalid, topk_weights) - - -def _assemble_reference_scale( - tensor, - expert_offsets: torch.Tensor, -) -> torch.Tensor: - k_axis = tensor.axis - scale = tensor.scale - parts = [] - begin = 0 - for end_tensor in expert_offsets: - end = int(end_tensor.item()) - if k_axis == 1: - raw = scale[:, begin // 32 : end // 32] - else: - raw = scale[begin // 32 : end // 32].transpose(0, 1) - parts.append(_blocked_reference(raw.view(torch.uint8))) - begin = end - non_k = tensor.shape[1 - k_axis] - rounded_non_k = (non_k + 127) // 128 * 128 - scale_columns = (tensor.shape[k_axis] // 32 + 3) // 4 * 4 - if parts: - packed = torch.cat(parts) - else: - packed = torch.empty(0, dtype=torch.uint8, device=tensor.device) - return packed.reshape(rounded_non_k, scale_columns).view( - torch.float8_e8m0fnu - ) - - -def _as_production_operands(reference: WgradOperandsReference): - return MoeEpWgradOperands( - fc1_a=reference.fc1_a.data, - fc1_sfa=_assemble_reference_scale( - reference.fc1_a, - reference.expert_offsets, - ), - fc1_b=reference.fc1_b.data.transpose(0, 1).contiguous().transpose(0, 1), - fc1_sfb=_assemble_reference_scale( - reference.fc1_b, - reference.expert_offsets, - ), - fc2_a=reference.fc2_a.data, - fc2_sfa=_assemble_reference_scale( - reference.fc2_a, - reference.expert_offsets, - ), - fc2_b=reference.fc2_b.data.transpose(0, 1).contiguous().transpose(0, 1), - fc2_sfb=_assemble_reference_scale( - reference.fc2_b, - reference.expert_offsets, - ), - expert_offsets=reference.expert_offsets, - valid_route_counts=reference.valid_route_counts, - route_metadata=reference.route_metadata, - ) - - -@pytest.mark.L1 -@pytest.mark.gpu_exclusive -def test_returned_operand_abi_runs_direct_grouped_wgrad(): - if not torch.cuda.is_available(): - pytest.skip("grouped wgrad integration requires CUDA") - device = torch.device("cuda", 0) - major, minor = torch.cuda.get_device_capability(device) - if (major, minor) != (10, 0): - pytest.skip("the in-tree grouped wgrad integration kernel targets SM100") - - import cudnn - - generator = torch.Generator(device=device).manual_seed(20260821) - hidden, intermediate = 128, 256 - activation = torch.randn( - 3, - hidden, - generator=generator, - device=device, - ) / 8 - fc1_weight = torch.randn( - 2, - hidden, - 2 * intermediate, - generator=generator, - device=device, - ) / 8 - fc2_weight = torch.randn( - 2, - intermediate, - hidden, - generator=generator, - device=device, - ) / 8 - topk_idx = torch.tensor( - [[0, 1], [1, 0], [0, 1]], - dtype=torch.int32, - device=device, - ) - topk_weights = torch.tensor( - [[0.5, 0.25], [0.75, 0.5], [1.0, 0.125]], - device=device, - ) - grad_output = torch.randn( - 3, - hidden, - generator=generator, - device=device, - ) / 8 - reference = MoeEpReference( - num_experts=2, - hidden_size=hidden, - intermediate_size=intermediate, - top_k=2, - max_tokens_per_rank=3, - generate_c=True, - backward_wgrad_mode="operands", - token_padding_size=256, - ) - _, fc1_c, metadata, forward_stash = reference( - activation, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights, - ) - _, _, logical = reference.backward( - grad_output, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights, - fc1_c, - metadata, - wgrad_forward_stash=forward_stash, - ) - operands = _as_production_operands(logical) - - common = { - "offsets_tensor": operands.expert_offsets, - "output_mode": "dense", - "wgrad_dtype": torch.bfloat16, - "acc_dtype": torch.float32, - "mma_tiler_mn": (128, 128), - "cluster_shape_mn": (1, 1), - "sf_vec_size": 32, - } - fc1 = cudnn.grouped_gemm_wgrad_wrapper_sm100( - a_tensor=operands.fc1_a, - b_tensor=operands.fc1_b, - sfa_tensor=operands.fc1_sfa, - sfb_tensor=operands.fc1_sfb, - **common, - )["wgrad_tensor"] - fc2 = cudnn.grouped_gemm_wgrad_wrapper_sm100( - a_tensor=operands.fc2_a, - b_tensor=operands.fc2_b, - sfa_tensor=operands.fc2_sfa, - sfb_tensor=operands.fc2_sfb, - **common, - )["wgrad_tensor"] - torch.cuda.synchronize(device) - - expected_fc1, expected_fc2 = logical.dense_wgrads() - torch.testing.assert_close( - fc1.float(), - expected_fc1, - rtol=0.15, - atol=0.125, - ) - torch.testing.assert_close( - fc2.float(), - expected_fc2, - rtol=0.15, - atol=0.125, - ) - - -@pytest.mark.L1 -@pytest.mark.gpu_exclusive -@pytest.mark.parametrize("world_size", [1, 2, 4], ids=["ep1", "ep2", "ep4"]) -def test_production_wgrad_operands_run_end_to_end(world_size, tmp_path): - if world_size == 1: - _run_wgrad_operand_case( - device=_sm107_device(), - ep_group=None, - rank=0, - world_size=1, - ) - return - - _require_distributed_sm107(world_size) - os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") - init_file = tmp_path / f"mxfp8_wgrad_operands_ep{world_size}.init" - mp.spawn( - _distributed_wgrad_worker, - args=(world_size, str(init_file)), - nprocs=world_size, - join=True, - ) From 49468ce1796f448d60ed0bfbf3f5bf4b6d26f4f9 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Fri, 28 Aug 2026 07:58:15 -0700 Subject: [PATCH 13/31] Document the fixed-resource MoeEP contract Align the public and backend documentation with current training, CUDA Graph, overflow, topology, and validation constraints. --- docs/fe-oss-apis/moe_ep.md | 1337 ++++++----------- docs/fe-oss-apis/overview.md | 14 +- .../cudnn/moe_ep/_megamoe_backend/README.md | 253 ++-- 3 files changed, 599 insertions(+), 1005 deletions(-) diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md index 59df7de1a..0de402ffa 100644 --- a/docs/fe-oss-apis/moe_ep.md +++ b/docs/fe-oss-apis/moe_ep.md @@ -1,919 +1,516 @@ -# MoE + Expert Parallel API - -Status: public API, validated lazy backend seam, internal runtime/workspace -owners, and executable PyTorch reference. The current device target is the -Rubin training `fwd_glu` kernel plus the restricted `bwd_dglu` path on exactly -SM107 (compute capability 10.7). It accepts MXFP8 E4M3/E8M0 operands or plain -BF16/FP16/FP32 operands staged to MXFP8, supports BF16 or MXFP8 combine with -BF16 output, and any positive EP size for forward and backward. Unsupported -combinations fail explicitly rather than returning uncomputed storage. - -The design removes workspace pointers, peer pointer mappers, streams, and -individual scheduler knobs from the semantic runtime interface. Performance -tuning is encapsulated in the optional `MoeEpTuningConfig`. - -## 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`. -- In the public contract, `output_format` describes the - post-top-k-reduction `(T, H)` result and `combine_format` independently - describes each per-route FC2 contribution on the EP return path. -- The public API reserves BF16, MXFP8, and NVFP4 spellings for both choices. - Quantized output is a data-plus-scale object, never a scale-free PyTorch - tensor. This is a future contract, not the current device capability. -- The current Rubin backend accepts `combine_format="bf16"` or `"mxfp8"` and - requires `output_format="bf16"`. NVFP4 combine and quantized public output - are rejected before runtime initialization. - -The distinction between public output and combine traffic is intentional. A -future transport-only MXFP8/NVFP4 mode would set `combine_format` to that -format and keep `output_format="bf16"`. The current SM107 MXFP8 combine path is -the training kernel's direct `MXFP8(FP32 accumulator)` conversion. - -## Initial device-backend implementation scope - -The public contract below remains the target. Device support will be enabled -incrementally, and unsupported combinations must fail explicitly rather than -return uninitialized storage. - -The current implementation connects the complete EP-subgroup path through the -executable CuTeDSL backend. Its deliberately narrow capability is: - -- the device must report exactly compute capability `(10, 7)` (SM107). If - `CUTE_DSL_ARCH` is unset, the compile runner sets it to `sm_107a`; an - existing value must target `sm_107` or `sm_107a`; -- `activation`, `fc1_weight`, and `fc2_weight` may be MXFP8 - `BlockScaledTensor` objects using logical, unswizzled scales and `axis=1`, - or plain BF16/FP16/FP32 tensors that staging quantizes to MXFP8; -- MXFP8 payloads use FP8 E4M3 and scales use FP8 E8M0, with the shapes specified - in "Block-scaled representation"; -- `combine_format` may be `"bf16"` or `"mxfp8"` and `output_format` must be - `"bf16"`; both - `generate_c=False` and eager-only `generate_c=True` forward are supported; -- `generate_c=True` performs one lockstep route-count collective and returns - fresh compact `fc1_c`/`route_metadata` tensors. The default path uses - 128-row internal expert alignment. Opt-in - `backward_wgrad_mode="operands"` requires 256-row alignment and additionally - returns caller-owned MXFP8 grouped-wgrad operands. Training execution does - not support CUDA Graph capture; the restricted device backward returns - activation and router-weight gradients; backward hardware acceptance - currently covers EP1/EP2/EP4; -- forward and backward support optional `gate_up_clamp`; Rubin training - execution requires `apply_topk_in_fc1=True`; -- `max_tokens_per_rank` must be explicitly positive; `top_k <= 32`, - `H % 128 == 0`, and `I % 256 == 0`; -- `ep_group=None` remains explicit single-rank execution. Distributed execution - accepts any initialized `torch.distributed.ProcessGroup`, including - non-contiguous global-rank membership, with any positive EP size. The public - contract requires `top_k <= num_experts` and the device path additionally - requires `top_k <= 32`; `top_k` may exceed `experts_per_rank`. Expert - ownership, peer tables, and route metadata use dense group-relative EP ranks; - experts remain contiguous and equally partitioned; -- every subgroup rank must call forward, backward, warmup, and graph replay - (where supported) in lockstep, including zero-token and zero-valid-route - ranks. Teardown must also be coordinated, although `close()` does not insert - a process-group barrier. Validation is rank-local, and collective - participation is a caller contract rather than an extra host-synchronized - control collective; -- the first lazy subgroup forward launch performs a one-time readiness - rendezvous after staging and JIT: each rank synchronizes its current stream, - all-gathers the effective tuning signature and rejects a mismatch, then - enters a process-group barrier before peer metadata writes can begin. The - first distributed backward launch likewise synchronizes and barriers without - a second tuning gather; EP1 skips both rendezvous. Subsequent eager launches - and graph replays do not add this control collective; -- while the process-global NVSHMEM runtime is active, all backend instances in - that process share one CUDA device binding, one ordered EP-subgroup - membership, and one reference count; distributed deployment uses one process - per GPU. A different subgroup cannot become active in the same process until - the first is fully released. A non-WORLD subgroup does not attach to an - externally initialized NVSHMEM runtime whose membership cannot be verified; - a matching full-WORLD external runtime may be attached but is never finalized - by this backend; -- each operator owns its local workspace and its NVSHMEM-symmetric workspace, - plus transformed-weight and reduction scratch. Allocations are plan-scoped - and stable across eager calls; `generate_c=True` additionally owns the BF16 - high-watermark C buffer described above; -- the forward and backward compile paths instantiate the vendored Rubin - training `Sm107MegaMoEMxfp8GluKernel` and - `Sm107MegaMoEMxfp8DgluKernel`, respectively. They do not select the - Blackwell inference MegaMoE implementation. - -Plain BF16/FP16/FP32 operands do not select a separate floating-point kernel -specialization: staging quantizes them to MXFP8 E4M3/E8M0 before launch. -Pre-quantized MXFP8 payloads and logical scales are preserved rather than -dequantized and requantized. - -The private CuTeDSL source snapshot lives under -`python/cudnn/moe_ep/_megamoe_backend/cutedsl_src`. Its -`VENDOR_INFO.md` records the upstream source revisions, vendoring dates, -copied-file manifest, local import adaptations, and update procedure. The -Apache-2.0 headers and redistribution terms are packaged beside the source. -The current kernel entry points are vendored under -`kernel_src/rubin/training/mega/fwd_glu` and -`kernel_src/rubin/training/mega/bwd_dglu`; those packages must use -package-relative imports, must not depend on a sibling `cutedsl_megamoe` -checkout, and must not import `kernel_src.blackwell`. - -After installing the `moe_ep` optional dependencies, run the L0 device -validation with: +# Mixture of Experts with Expert Parallelism -```bash -python -m pytest \ - test/python/fe_api/moe_ep/test_moe_ep_forward.py \ - test/python/fe_api/moe_ep/test_moe_ep_backward.py \ - -m L0 -``` +`cudnn.moe_ep` provides a fused SwiGLU MoE implementation for Rubin SM107. +Experts are sharded contiguously across an optional expert-parallel process +group. + +## Supported configuration + +- CUDA execution on Rubin SM107 (compute capability 10.7) +- fused SwiGLU with contiguous expert sharding across `ep_group` +- BF16 output, including when the combine path uses MXFP8 +- BF16 or MXFP8 combine +- plain BF16, FP16, or FP32 inference operands, or MXFP8 + `BlockScaledTensor` operands +- `apply_topk_in_fc1=True` +- `hidden_size` divisible by 128 +- `intermediate_size` divisible by 256 +- `top_k <= min(32, num_experts)` +- `num_experts` divisible by the EP group size +- explicit positive `max_tokens_per_rank` + +`output_format` is currently executable only as `"bf16"`. NVFP4 is represented +by the public format types but NVFP4 operands, combine, and output are not +executable by the current MegaMoE backend. -These forward and backward core suites include public-contract and capability -checks, staging, runtime, reference, and supported single-rank numerical -coverage. On SM107, the L0 forward suite exercises the production compile and -launch path. Full numerical parity, stress, CUDA Graph replay, multi-rank -forward, and device backward acceptance are covered by L1 tests and the -hardware/container runner rather than by this L0 command alone. +The fixed-resource CUDA Graph path has hardware acceptance through EP32 within +one direct-P2P MNNVL peer-access domain. The Python capability layer does not +impose an EP-size ceiling; this statement describes validated hardware scope, +not support for cross-MNNVL execution. -The ordinary L1 hardware/container runner is single-node: its distributed -forward matrix uses `mp.spawn` for EP2, EP3, and EP4. EP7, EP12, EP15, and -EP16 use the torchrun-native multi-node suite instead. From the first node of -an existing Slurm allocation, start one torchrun agent per selected node with: +## Installation + +Install the dedicated optional dependencies: ```bash -torchrun \ - --nnodes="${NNODES}" \ - --node-rank="${NODE_RANK}" \ - --nproc-per-node="${NPROC_PER_NODE}" \ - --master-addr="${MASTER_ADDR}" \ - --master-port="${MASTER_PORT}" \ - -m pytest \ - test/python/fe_api/moe_ep/test_moe_ep_forward_multinode.py \ - -m "L1 and moe_ep_multinode" \ - -k "${PYTEST_FILTER}" +pip install nvidia-cudnn-frontend[moe_ep] ``` -NVSHMEM requires the same number of EP PEs on every participating node. Use -`NNODES=7`, `NPROC_PER_NODE=2`, and `PYTEST_FILTER=ep7-world14` for EP7; its -subgroup selects local rank zero on every node. Use `NNODES=5`, -`NPROC_PER_NODE=4`, and `PYTEST_FILTER=ep15-world20` for EP15; its subgroup -selects local ranks zero through two on every node. EP12 uses `NNODES=3`, -`NPROC_PER_NODE=4`, and `PYTEST_FILTER=ep12-world12`; EP16 uses `NNODES=4`, -`NPROC_PER_NODE=4`, and `PYTEST_FILTER=ep16-world16`. The remaining WORLD -ranks synchronize without constructing `MoeEp`. The Slurm harness must provide -a distinct `NODE_RANK` to each node and shared `MASTER_ADDR`/`MASTER_PORT` -values. Each case checks BF16 and MXFP8 combine against the executable -reference plus all-`-1` route behavior. Capability support alone is not a -hardware PASS: preserve the torchrun logs for acceptance evidence. - -PyTorch is a prerequisite of the `MoeEp` API and remains in the `moe_ep` -optional extra rather than becoming a base dependency of the entire -`nvidia-cudnn-frontend` distribution. The same extra selects the CUDA-13 -CuTeDSL stack and `nvshmem4py-cu13>=0.3.1`. Given an installed PyTorch, -ordinary `import cudnn` and `import cudnn.moe_ep` do not import CuTeDSL, -NVSHMEM4Py, or CUDA Python and do not initialize CUDA. Missing CuTeDSL or -NVSHMEM runtime components are reported as `BackendUnavailableError` only -when a supported device forward first needs the backend. - -The hardware, API, stress, and packaging runner defaults to -`MOE_EP_DEPENDENCY_MODE=rubin-internal`. In this mode it first removes every -installed `nvidia-cutlass-dsl*` distribution, then installs the latest -pre-release `nvidia-cutlass-dsl` from NVIDIA's Rubin-capable CUTLASS DSL master -index with PyPI as a dependency fallback. The runner verifies that -`cutlass.utils.rubin_helpers` is importable before compiling the kernel. -`MOE_EP_DEPENDENCY_MODE=latest` and `minimum` retain public-release acceptance; -the latter pins CUTLASS DSL, NVSHMEM4Py, and Apache TVM FFI to 4.8.0, 0.3.1, -and 0.1.11 respectively. The public 4.7.0 wheel does not contain -`rubin_helpers`. The packaging runner also validates an isolated wheel import -and the vendored Rubin import closure. - -## Public API - -The production class is `cudnn.moe_ep.MoeEp`. The test-only -`MoeEpReference` is the executable semantic and numerical oracle; it is not the -device backend. +The extra supplies the CuTeDSL and NVSHMEM Python dependencies. PyTorch with +CUDA support is also required. EP2+ additionally requires an initialized NCCL +process group and a usable NVSHMEM peer topology. + +## Public API and constructor + +The public surface exports `MoeEp`, `MoeEpTrainingWeights`, +`MoeEpTrainingResources`, `MoeEpTrainingSlot`, `MoeEpExecutionLane`, +`MoeEpTrainingWgradOperands`, `BlockScaledTensor`, `MoeFormat`, and +`MoeEpTuningConfig`. + +The `MoeEp` constructor accepts: + +| Parameter | Current contract | +| --- | --- | +| `num_experts` | Positive global expert count; divisible by EP size | +| `hidden_size` | Positive and divisible by 128 | +| `intermediate_size` | Positive and divisible by 256 | +| `top_k` | Positive and no larger than 32 or `num_experts` | +| `ep_group` | Optional initialized `torch.distributed.ProcessGroup`; `None` selects EP1 | +| `max_tokens_per_rank` | Required by the executable backend and must be positive | +| `max_recv_size_per_rank` | Optional positive receive-pool capacity | +| `drop_on_overflow` | `False` by default; selects fatal-assert versus reporting/drop policy | +| `output_format` | `"bf16"` only for current execution | +| `combine_format` | `"bf16"` or `"mxfp8"` | +| `apply_topk_in_fc1` | Must be `True` | +| `gate_up_clamp` | Optional finite clamp magnitude | +| `token_padding_size` | Positive; training fixed resources use 128 internally | +| `sf_padding_size` | Positive multiple of 128; training fixed resources use 128 internally | +| `tuning` | Optional `MoeEpTuningConfig`; must match on every EP rank | + +When `max_recv_size_per_rank` is omitted, the backend allocates for the +worst-case receive count: + +```text +ep_size * max_tokens_per_rank * top_k +``` + +An explicit value is capped at that same worst-case count. + +## Breaking training API migration + +This release removes the legacy dynamic compact training API: + +- `MoeEp.backward(...)` +- constructor arguments `generate_c` and `backward_wgrad_mode` +- forward returns containing compact `fc1_c` and `route_metadata` +- `MoeEpWgradForwardStash` and `MoeEpWgradOperands` + +Old: ```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, - max_recv_size_per_rank: Optional[int] = None, - drop_on_overflow: bool = False, - 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, - backward_wgrad_mode: Literal["none", "operands"] = "none", - token_padding_size: int = 128, - sf_padding_size: int = 128, - tuning: Optional[MoeEpTuningConfig] = None, - ) -> 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] - | tuple[ - Tensor | BlockScaledTensor, - Tensor, - Tensor, - MoeEpWgradForwardStash, - ] - ): ... - - def warmup( - self, - activation: Tensor | BlockScaledTensor, - fc1_weight: Tensor | BlockScaledTensor, - fc2_weight: Tensor | BlockScaledTensor, - topk_idx: Tensor, - topk_weights: Tensor, - ) -> None: ... - - def backward( - self, - grad_output: Tensor, - fc1_weight: Tensor | BlockScaledTensor, - fc2_weight: Tensor | BlockScaledTensor, - topk_idx: Tensor, - topk_weights: Tensor, - fc1_c: Tensor, - route_metadata: Tensor, - *, - wgrad_forward_stash: Optional[MoeEpWgradForwardStash] = None, - ) -> ( - tuple[Tensor, Tensor] - | tuple[Tensor, Tensor, MoeEpWgradOperands] - ): ... - - def close(self) -> None: ... - - def __enter__(self) -> "MoeEp": ... - - def __exit__(self, exc_type, exc_value, traceback) -> bool: ... +output, fc1_c, route_metadata = op( + activation, w1, w2, topk_idx, topk_weights +) +dx, dprob = op.backward( + grad_output, w1, w2, topk_idx, topk_weights, fc1_c, route_metadata +) ``` -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. - -`close()` is terminal and idempotent. Once device workspaces have been created, -the backend synchronizes outstanding CUDA work before releasing transformed -weights, workspace, and runtime ownership. It does not issue a process-group -barrier; distributed callers must coordinate teardown so one rank cannot -release symmetric storage while a peer is still launching or replaying. -The context-manager form is preferred when deterministic release matters. - -### 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 constructor, but the current device capability gate requires an explicit positive value on first execution. | -| `max_recv_size_per_rank` | Optional bound on routed rows received by one EP rank. The default is the conservative `ep_size * max_tokens_per_rank * top_k`; a smaller bound reduces workspace size. | -| `drop_on_overflow` | If `True`, truncate routed rows beyond `max_recv_size_per_rank`; if `False` (default), overflow traps instead of silently changing results. | -| `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 the combine-rounded FC2 route contribution in the standalone top-k reducer. | -| `gate_up_clamp` | If set, use `gate = clamp(gate, max=abs(limit))` and `up = clamp(up, min=-abs(limit), max=abs(limit))`. | -| `generate_c` | Training integration: additionally return `fc1_c` (the BF16-rounded pre-SwiGLU FC1 accumulator of every route this rank's experts processed) and its row-aligned `route_metadata`. | -| `backward_wgrad_mode` | `"none"` preserves the default API. `"operands"` opts into caller-owned MXFP8 operands for external grouped FC1/FC2 wgrad GEMMs; it requires `generate_c=True`, `token_padding_size=256`, and `sf_padding_size=128`. | -| `token_padding_size` | Token-dimension padding used by the Rubin execution plan; positive integer, default 128. | -| `sf_padding_size` | Scale-factor padding; positive multiple of 128, default 128. Operand mode currently requires exactly 128. | -| `tuning` | Optional `MoeEpTuningConfig`; every rank in an EP group must use the same effective configuration. | - -The constructor validates public static dimensions, format alignment, padding, -and tuning. MXFP8 format spellings require `H % 32 == 0`; NVFP4 format -spellings require `H % 16 == 0`. Device-specific requirements such as SM107, -`H % 128 == 0`, `I % 256 == 0`, `top_k <= 32`, and explicit positive capacity -are checked lazily before backend initialization. - -The wgrad mode is strictly opt-in. With the default -`backward_wgrad_mode="none"`, the constructor default, forward return, backward -signature, backward return, padding default, allocation behavior, and -documented numerical semantics are unchanged. - -Scheduler settings are exposed through the optional `MoeEpTuningConfig` -object. The current public knobs are `token_back_mode`, `epi_flag_batch`, -`token_in_flag_batch`, `group_hint`, and `reduce_topk_in_kernel`; they are -keyword configuration rather than positional runtime arguments. - -`reduce_topk_in_kernel=True` enables the in-kernel top-k reduction path and is -restricted to BF16 combine/output, `apply_topk_in_fc1=True`, and -`token_back_mode="epi_warps"`. BF16 reduction order can change rounding, so -this path is accepted against the documented numerical tolerance rather than -bitwise Form A equality. - -### 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/FP16/FP32 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. - -The first successful backend creation binds a `MoeEp` instance and its stable -workspace allocations to that device. A later call on another device raises -`ValueError`; callers must create one `MoeEp` instance per device. - -### 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. - -With `generate_c=True`, the call instead returns -`(output, fc1_c, route_metadata)`; see the training-integration sections -below. - -With both `generate_c=True` and `backward_wgrad_mode="operands"`, it returns -`(output, fc1_c, route_metadata, wgrad_forward_stash)`. The fourth item belongs -to that exact routed forward call and must be passed to its corresponding -backward call. - -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`; the -production entry point is `MoeEp.backward`, while `MoeEpReference.backward` -defines its executable semantic oracle. The Rubin MXFP8 device path supports -BF16/MXFP8 combine, BF16 output, any positive EP size, -`apply_topk_in_fc1=True`, and optional `gate_up_clamp`. 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)` | Any floating dtype on the request device; incoming gradient of the *dequantized* public output (all encodes are straight-through). | -| `fc1_weight`, `fc2_weight`, `topk_idx`, `topk_weights` | as in forward | the framework re-supplies the forward weights and routing 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. | -| `wgrad_forward_stash` | `MoeEpWgradForwardStash` | Keyword-only and required only in operand mode. It contains the forward `x.T` operand plus padded expert offsets/counts and exact route identity. Passing it in default mode is an error. | - -Default mode returns two 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_topk_weights` | `(T, K)` | per-route router-weight gradient; exact zero at `-1` slots. | - -Operand mode returns -`(grad_activation, grad_topk_weights, wgrad_operands)`. The first two values -retain the default meanings and dtypes; `wgrad_operands` is described below. - -On the Rubin device path, dGLU materializes `grad_activation` in BF16 and the -public wrapper widens it to FP32, so its numerical granularity remains BF16. -The semantic router-weight gradient is recomputed from the unquantized -floating `grad_output` and returned in FP32. - -Internally, the compact public `fc1_c` rows are lowered into an external -pool-layout BF16 `fc1_preact` tensor. The kernel also writes a pre-zeroed, -symmetric source-domain FP32 dprob plane with shape -`(max_tokens_per_rank, top_k)` for its peer-atomic ABI; that internal plane is -not the public return because the public straight-through semantics use the -FP32 recomputation described above. - -The save-set, recompute rules, and numerical conventions behind this -signature are specified in "Saved tensors for backward" below. - -### Grouped-wgrad operand contract - -`backward_wgrad_mode="operands"` exposes the operands needed by the existing -`GroupedGemmWgradSm100` / `grouped_gemm_wgrad_wrapper_sm100` Tensor2D ABI. It -does not launch those GEMMs and does not return dense weight gradients. - -For local expert `e`, let `R_e` be its valid route count, -`P_e = ceil(R_e / 256) * 256`, and `Kp = sum_e P_e`. Every expert occupies one -contiguous range in the shared K dimension. Valid rows precede zero padding; -an empty expert contributes zero extent, so cumulative offsets may repeat. -`expert_offsets[e] = sum_{j <= e} P_j`, and `valid_route_counts[e] = R_e`. - -Forward returns `MoeEpWgradForwardStash` with: - -| Field | Logical shape | Meaning | -|---|---:|---| -| `fc1_a`, `fc1_sfa` | `(H, Kp)`, `(round_up(H,128), round_up(Kp/32,4))` | MXFP8 `x.T` data and assembled E8M0 scales. | -| `expert_offsets` | `(E_local,)` | Int32 cumulative 256-padded expert ends. | -| `valid_route_counts` | `(E_local,)` | Int32 unpadded route counts. | -| `route_metadata` | `(local_routes, 4)` | The same compact route identity returned beside `fc1_c`. | - -Backward returns `MoeEpWgradOperands`. It carries those five fields plus: - -| Field | Logical shape | Meaning | -|---|---:|---| -| `fc1_b`, `fc1_sfb` | `(Kp, 2I)`, `(round_up(2I,128), round_up(Kp/32,4))` | MXFP8 `dC`, where `C = x @ W1`. | -| `fc2_a`, `fc2_sfa` | `(I, Kp)`, `(round_up(I,128), round_up(Kp/32,4))` | MXFP8 route-weighted recomputed `(p * h).T`, `h = silu(gate) * up`. | -| `fc2_b`, `fc2_sfb` | `(Kp, H)`, `(round_up(H,128), round_up(Kp/32,4))` | MXFP8 unweighted routed `dY`. | - -For each expert range, including its zero padding, the represented dense -operations are: +New: -```text -dW1[e] = fc1_a[e] @ fc1_b[e] = x[e].T @ dC[e] -dW2[e] = fc2_a[e] @ fc2_b[e] = (p[e] * h[e]).T @ dY[e] +```python +resources = op.prepare_training_resources( + training_weights, + slot_count=2, + lane_count=1, +) +slot = resources.slots[0] +lane = resources.lanes[0] +resources.refresh_weights() +output = resources.forward( + slot, lane, activation, topk_idx, topk_weights +) +dx, dprob, operands = resources.backward(slot, lane, grad_output) +overflow = resources.finalize_overflow((slot,), lane) ``` -The data operands are E4M3. Scale factors are E8M0 with logical 1x32 scaling -and grouped-wgrad's assembled physical 128x4 scale tiles. A operands have -unit K stride; B operands use the grouped-wgrad K-major view with unit K -stride. - -The Rubin staging/quantization order is part of the operand model: - -1. `x` is first staged to MXFP8 along H (plain inputs only), routed and padded, - then column-requantized along expert K to form `fc1_a`. -2. BF16 `fc1_c` is recomputed through clamp/SwiGLU without a router weight; - that `h` is column-requantized along K to form `fc2_a`. -3. `grad_output` is staged to MXFP8 along H before re-dispatch. The route - weight is then applied exactly once, and the result is - column-requantized along K to form `fc2_b`. -4. Staged `dY` and backward-staged `W2.T` produce `dH`; the route weight is - applied before the SwiGLU derivative, and the resulting `dC` is directly - column-requantized along K to form `fc1_b`. - -All returned stash and operand tensors are caller-owned allocations. They do -not alias reusable execution-plan workspace and are not overwritten by later -forward/backward calls. The caller must retain the forward stash through its -matching backward call and retain the returned operands until every external -grouped-wgrad consumer has completed. Route metadata and counts are validated -against the matching call; stashes are not interchangeable between different -routing inputs. - -### Example - -The following illustrates the reserved future quantized-output contract. It is -valid reference-level API semantics, but the current SM107 device backend -rejects `output_format="nvfp4"`. The shown `combine_format="mxfp8"` is supported -by the device backend when paired with `output_format="bf16"`. +`dprob` now follows the MXFP8-staged kernel numerical contract and relaxed +atomic accumulation order. Dynamic inputs use a trusted-caller contract. +Distributed graph support requires one direct-P2P MNNVL domain, and +distributed lanes must be ordered consistently across ranks with captured +events. The fixed-capacity WGrad result is a producer ABI; no specific grouped +WGrad consumer is guaranteed in this release. + +## Inference forward + +`MoeEp.__call__` is the inference-forward surface: ```python -moe = MoeEpReference( - num_experts=8, - hidden_size=4096, - intermediate_size=14336, - top_k=2, +from cudnn import MoeEp + +op = MoeEp( + num_experts=num_experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=top_k, ep_group=ep_group, - max_tokens_per_rank=2048, - combine_format="mxfp8", - output_format="nvfp4", + max_tokens_per_rank=max_tokens, + max_recv_size_per_rank=recv_capacity, ) -# 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) +output = op( + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, +) ``` -## Mathematical semantics +`activation`, `fc1_weight`, and `fc2_weight` may independently be plain BF16, +FP16, or FP32 tensors, or MXFP8 `BlockScaledTensor` values. The logical shapes +are: -For valid route `(t, k)` with global expert `e` and router weight `p[t, k]`: +- `activation`: `(T, H)` +- `fc1_weight`: `(E_local, H, 2I)` +- `fc2_weight`: `(E_local, I, H)` +- `topk_idx` and `topk_weights`: `(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 +MXFP8 operands are block-scaled along logical axis 1. The output always has +shape `(T, H)` and dtype `torch.bfloat16`. This surface does not return compact +FC1 or route metadata stashes and does not provide backward. -if apply_topk_in_fc1: - hidden[t,k] *= p[t,k] +### Inference CUDA Graph capture -expert[t,k] = hidden[t,k] @ fp32(W2[e]) -combine[t,k] = dequantize(quantize(expert[t,k], combine_format)) +Call `warmup` with the exact tensors that will be captured: -if not apply_topk_in_fc1: - combine[t,k] *= p[t,k] +```python +import torch + +op.warmup(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) +if ep_group is not None: + torch.distributed.barrier(group=ep_group) -result[t] = sum_k(combine[t,k]) -output = encode(result, output_format) +graph = torch.cuda.CUDAGraph() +with torch.cuda.graph(graph): + graph_output = op( + activation, fc1_weight, fc2_weight, topk_idx, topk_weights + ) ``` -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. +`warmup` completes runtime bootstrap, symmetric allocation, weight staging, +JIT compilation, and one real launch. It is collective by contract for EP2+ +but intentionally does not issue a process-group barrier. Captured inference +weights must expose usable PyTorch version counters and must match the warmed +weight cache. Replay may update captured tensor contents in place but may not +replace their storage. -For `combine_format="mxfp8"`, the current Rubin forward and backward paths -directly convert each FP32 route accumulator to MXFP8 before top-k reduction. +`MoeEp` supports explicit `close()` and context-manager use. One instance is +bound to one CUDA device. Do not close an operator while a stream is capturing +or while graph work using its resources remains outstanding. -Moving the router weight from the post-SwiGLU intermediate to the -combine-rounded FC2 contribution is algebraically equivalent only without the -intervening low-precision conversions. The option is therefore semantic and is -fixed in the constructor, matching the MegaMoE kernel. +## Fixed-resource training -## Block-scaled representation +Training uses `prepare_training_resources`: -The API uses logical, unswizzled scales. Backend-specific F8_128x4 swizzling is -an implementation detail performed while constructing tensor maps or staging -weights. +```python +from cudnn import MoeEpTrainingWeights -| Format | Payload | Scale | Block | Quantized axis | -|---|---|---|---:|---| -| BF16 | BF16 | none | n/a | n/a | -| MXFP8 | FP8 E4M3 | FP8 E8M0 | 32 | public `axis=1`: contraction axis for weights, feature/output axis otherwise | -| NVFP4 | packed FP4 E2M1, low nibble first | FP8 E4M3 | 16 | public `axis=1`: contraction axis for weights, feature/output axis otherwise | +weights = MoeEpTrainingWeights( + forward_fc1=forward_fc1_mxfp8, + forward_fc2=forward_fc2_mxfp8, + backward_w2_transpose=backward_w2t_mxfp8, + backward_w1_transpose=backward_w1t_mxfp8, +) -MXFP8 scale calculation is: +resources = op.prepare_training_resources( + weights, + slot_count=2, + lane_count=1, +) -```text -raw_scale = amax(block) / 448 -scale = 2 ** ceil(log2(raw_scale)) # E8M0 round toward +infinity -data = e4m3(clamp(block / scale, -448, 448)) +slot0, slot1 = resources.slots +lane0 = resources.lanes[0] + +# Required after each in-place source-weight update and before the first +# forward/backward that consumes that version. +resources.refresh_weights() + +y0 = resources.forward( + slot0, + lane0, + activation0, + topk_idx0, + topk_weights0, +) +dx0, dprob0, operands0 = resources.backward(slot0, lane0, grad_output0) +overflow = resources.finalize_overflow((slot0,), lane0) ``` -NVFP4 scale calculation is: +`prepare_training_resources` is collective across `ep_group` and must execute +outside CUDA Graph capture. All ranks must use matching static configuration, +slot/lane counts, and tuning. The training backend internally enables FC1 +preactivation generation and fixed-capacity WGrad operands, and fixes token and +scale-factor padding to 128. -```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 +`forward`, `backward`, and `finalize_overflow` enqueue the same device +operations in ordinary execution and in a caller-owned CUDA Graph. The caller +owns capture, replay, stream/event dependencies, slot reuse, and resource +lifetime. + +A `MoeEp` instance can own only one training-resource set. Closing that set is +terminal for the operator: replacing source-weight storage requires a new +operator, new resources, and new graph captures. + +### Slots and lanes + +A persistent slot owns state that survives from matching forward to backward: + +- routing indices and weights +- pool-native FC1 preactivation +- expert counts and padded offsets +- kernel dprob +- backward auxiliaries and outputs +- fixed-capacity WGrad operands +- per-slot overflow flags + +An execution lane owns mutable router, protocol, and kernel scratch. Multiple +active streams require distinct lanes. Distributed MegaMoE communication +kernels must be ordered identically on every rank with captured CUDA events; +the kernels cannot be launched in unordered concurrent lanes. + +All peer-visible regions are constructed in deterministic order. Their sizes +are validated and normalized across ranks before symmetric allocation so every +named region has the same peer offset. + +## Tensor contracts + +For inference through `MoeEp.__call__`, `activation` has: + +- shape `(T, H)` +- BF16, FP16, FP32, or MXFP8 block-scaled input + +Inference `fc1_weight` and `fc2_weight` accept the same plain or MXFP8 operand +families. MXFP8 activation and weights must be represented by +`BlockScaledTensor` with logical block axis 1. + +The inference `topk_idx` contract is: + +- shape `(T, K)` +- Int32 or Int64 +- each element is `-1` or a valid global expert ID + +The inference `topk_weights` contract is: + +- shape `(T, K)` +- floating point + +Fixed-resource training uses a narrower, graph-stable staging ABI: + +- contiguous BF16 or FP32 `activation` and `grad_output`, each shaped `(T, H)`; +- contiguous Int32 `topk_idx` shaped `(T, K)`; +- contiguous FP32 `topk_weights` shaped `(T, K)`; +- all tensors on one device and `T <= max_tokens_per_rank`. + +Expert IDs and finite dynamic values remain a trusted-caller replay contract; +they are not revalidated by host code after graph capture. + +`MoeEpTrainingWeights` contains four contiguous MXFP8 block-scaled tensors: + +- `forward_fc1`: `(E_local, H, 2I)` +- `forward_fc2`: `(E_local, I, H)` +- `backward_w2_transpose`: `(E_local, H, I)` +- `backward_w1_transpose`: `(E_local, 2I, H)` + +Each data and scale tensor must be contiguous, reside on one device, and use +logical block axis 1. Plain FP16 operands are accepted by inference staging, +but fixed-resource training accepts only BF16 or FP32 `activation` and +`grad_output`. + +Replacing weight storage requires preparing resources and capturing again. +Callers must establish stream/event ordering for in-place weight updates. + +### Explicit weight refresh contract + +`MoeEpTrainingWeights` uses a public contiguous MXFP8 layout, while the Rubin +kernels consume fixed-address K-major, gate/up-interleaved, and blocked-scale +layouts. `resources.refresh_weights()` enqueues the required device-only copies +and layout transforms into the internal kernel bindings. + +The caller must obey all of the following: + +- update both the data and scale tensors in place; their storage addresses, + shape, stride, dtype, device, and capacity must remain unchanged; +- call `resources.refresh_weights()` after every source-weight update and + before any forward or backward that consumes the new version; +- establish stream ordering from the weight update to the refresh and from the + refresh to the first consumer, using the same stream or CUDA events; +- do not refresh between a matching forward and backward; both operations must + observe the same four-tensor weight version; +- do not overlap a refresh with any forward/backward that reads the shared + internal weight bindings, including operations using another slot or lane; +- replace any source storage only by closing the existing operator, creating a + new `MoeEp` instance and resources, and recapturing every graph that + references them. Closed resources cannot be reopened or replaced on the same + operator. + +For CUDA Graph execution, capture the refresh at the appropriate update +boundary: + +```python +with torch.cuda.graph(graph, stream=stream): + # An optimizer or external producer must complete its in-place updates + # before this node. + resources.refresh_weights() + y = resources.forward(slot, lane, x, topk_idx, topk_weights) + dx, dprob, operands = resources.backward(slot, lane, grad_output) + overflow = resources.finalize_overflow((slot,), lane) ``` -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`. +Replay then executes the captured device refresh; Python is not called during +replay. Activation `x` has the same storage rule under CUDA Graph capture: +its contents may change in place, but replacing its captured storage requires +recapture. -Examples of scale shapes are: +## Backward outputs -| 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))` | +Fixed-resource backward returns: -NVFP4 payload shape replaces the quantized axis by `ceil(axis_extent / 2)`. +- `grad_activation`: fixed-slot `(T, H)` FP32 view +- `dprob`: source-order `(T, K)` kernel dprob +- `MoeEpTrainingWgradOperands` -## Expert-parallel execution +Kernel dprob follows the MXFP8-staged backward numerical contract and relaxed +atomic accumulation order. Bitwise determinism is not guaranteed. -The semantic data flow is: +`MoeEpTrainingWgradOperands` contains: -```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 +- FC1 operands: `fc1_a`, `fc1_sfa`, `fc1_b`, `fc1_sfb` +- FC2 operands: `fc2_a`, `fc2_sfa`, `fc2_b`, `fc2_sfb` +- `expert_offsets` and `valid_route_counts` + +These tensors have fixed addresses and fixed capacity. Device +`expert_offsets`/`valid_route_counts` describe the live expert segments and +padding rows are zero. This release guarantees the producer ABI only; a +specific grouped-WGrad consumer is future integration work. The result is not +a pair of dense gradients that can be passed directly to an optimizer. + +## Tuning + +`MoeEpTuningConfig` exposes semantic-preserving performance controls: + +- `token_back_mode`: `"epi_warps"`, `"standalone_warps"`, or + `"reuse_dispatch_warps"` +- `epi_flag_batch`: one of the validated `(M, N)` flag-batch pairs +- `token_in_flag_batch`: `1`, `2`, `4`, `8`, or `16` +- `group_hint`: `None`, `64`, `128`, `256`, `512`, `768`, or `1024` +- `reduce_topk_in_kernel`: Boolean + +Every rank in an EP group must use the same tuning configuration. +`reduce_topk_in_kernel=True` requires BF16 combine/output, +`apply_topk_in_fc1=True`, and `token_back_mode="epi_warps"`. + +## Capacity and overflow + +`max_recv_size_per_rank` defines bounded receive capacity. Resources cannot +grow during capture. A capacity change requires preparation and recapture. + +Inference checks its per-call overflow result after the fused launch. The +fixed-resource training transport deterministically truncates overflow so all +ranks complete the communication protocol. `finalize_overflow` combines the +forward and backward flags for all selected slots and performs one captured +scalar MAX all-reduce for EP2+. + +With `drop_on_overflow=True`, `finalize_overflow` returns a one-element Int32 +CUDA tensor: zero means no overflow and nonzero means truncation occurred. +Dropped routes contribute zero. With `drop_on_overflow=False`, overflow is a +fatal device assertion; this mode requires `torch._assert_async`, and EP2+ +requires an NCCL process group. The training transport still truncates first +to let every rank finish the protocol before the public policy is applied at +graph tail. + +## CUDA Graph execution + +For inference: + +1. all ranks call `MoeEp.warmup` with the exact capture bindings; +2. the caller aligns ranks after warmup; +3. each rank captures its forward graph; +4. graph execs are replayed in the same cross-rank order. + +For fixed-resource training: + +1. all ranks collectively call `prepare_training_resources`; +2. all ranks perform an ordinary + `refresh_weights -> forward -> backward -> finalize_overflow` warmup so + every staging, MegaMoE, and WGrad-export kernel is compiled; +3. each rank captures its outer graph; +4. ranks align after capture; +5. graph execs are submitted in lockstep without host synchronization inside + a replay burst; +6. all stream work completes before resources are closed. + +For EP2+: + +Distributed MegaMoE launches must have the same order on every rank. +Independent lane storage does not permit unordered collective-kernel overlap. +Use distinct lanes for simultaneously active streams and captured CUDA events +to impose the same cross-stream order on every rank. + +Each graph binds fixed tensor shapes, addresses, slots, lanes, and token +extent. Dynamic routing values may change in place; dynamic shapes may not. +Different token extents may share prepared resources, but each extent must be +warmed and captured as its own graph specialization. + +## NVSHMEM topology + +MegaMoE kernels use direct symmetric peer pointers. With +`NVSHMEM_REMOTE_TRANSPORT=none`, every EP rank must appear in the P2P connected +list and `NVSHMEM_TEAM_SHARED` must span the complete EP world. + +Selecting `ibrc` does not by itself make a non-P2P peer directly addressable. +Cross-MNNVL execution is not part of the current support matrix. + +Global expert `e` belongs to group-relative EP rank +`e // experts_per_rank`. Noncontiguous global-rank process groups are accepted +when their group-relative topology and direct peer access are valid. When +creating multiple subgroups, every world rank must create them in the same +order. + +## Validation coverage + +The following configurations are exercised by the current test and probe +suite. This list describes validation coverage and does not broaden the +topology contract beyond one direct-P2P MNNVL domain: + +- EP1 inference, fixed-resource training, overflow, and CUDA Graph replay +- single-node EP2/EP3/EP4 inference +- single-node EP2/EP4 training +- noncontiguous EP2 inference and training subgroups +- multi-node forward acceptance for EP4/EP6/EP12/EP16 +- multi-node backward acceptance for EP8/EP16/EP32 +- fixed-resource CUDA Graph launchers for EP8/EP16/EP32 + +## Validation + +Run host-side and local tests: + +```bash +python -m pytest \ + test/python/moe_ep/test_moe_ep_forward.py \ + test/python/moe_ep/test_moe_ep_backward.py \ + -m L0 ``` -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`. `max_recv_size_per_rank` selects a -smaller static workspace bound; callers must provide a router capacity contract -or accept `drop_on_overflow=True`. With `drop_on_overflow=False`, exceeding the -bound fails explicitly. - -## Mapping to the MegaMoE interface - -| concept | Public 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 | -| forward column-requant `x.T` | `MoeEpWgradForwardStash.fc1_a/fc1_sfa` in operand mode | -| dGLU column-requant `dC`, `(p*h).T`, `dY` | `MoeEpWgradOperands` returned by backward | - -### `fc1_c` and `route_metadata` (training integration) - -With `generate_c=True`, `__call__` returns `(output, fc1_c, route_metadata)`. -Operand mode appends `wgrad_forward_stash` as a fourth item. -`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 output gradients to the expert ranks, which is where the -stashed preactivations 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 -FC1 FP32 accumulator rounded to BF16 and captured **before** SwiGLU, before the -gate/up clamp, and without the router weight (which applies after SwiGLU when -`apply_topk_in_fc1=True`). Columns `[0:I)` are gate and `[I:2I)` are up. The -kernel's mode-dependent per-expert padding (128 rows by default, 256 in operand -mode) and internal gate/up-interleaved layout are absent from the logical -`fc1_c` 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 - -By default, `MoeEpReference.backward(grad_output, fc1_weight, fc2_weight, -topk_idx, topk_weights, fc1_c, route_metadata)` returns -`(grad_activation, grad_topk_weights)` and is the executable statement of the -save-set. Operand mode also takes `wgrad_forward_stash=` and appends a -`WgradOperandsReference` result. 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; it yields `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ᵀ`. | -| `topk_idx`, `topk_weights` | source rank (framework inputs) | `topk_idx` regenerates the exact dispatch plan; `topk_weights` scales the activation-gradient path 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). | -| `wgrad_forward_stash` | expert rank (operand mode only) | Caller-owned MXFP8 `x.T` plus 256-padded offsets/counts and route identity; supplies FC1 A and fixes the grouped K segmentation. | - -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. - -In the reference backward, input decode, `combine_format`, and `output_format` -round trips 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; 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 it is -`d_w = ⟨d_y, y_pre⟩`. The Rubin device dGLU path additionally MXFP8-stages -`grad_output` and transposed weights, modeled in parity tests by -`backward_operand_format="mxfp8"`, while recomputing `grad_topk_weights` from -the original floating `grad_output`. - -## 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. - - Bootstrap and allocate once outside CUDA graph capture; capture requires a - prior warmup whose stream has completed before capture begins. - - `MoeEp.warmup(...)` runs one complete eager forward and synchronizes its - CUDA device. For an EP subgroup, all member ranks must call it concurrently - and the caller must align ranks before capture; the method deliberately - does not issue a process-group barrier. - - Distributed graph capture is per-rank and replay is lockstep. Capture - records the cross-rank kernel without executing it; every replay must be - issued by every EP rank in the same iteration. - - Cache successful eager route-value validation by tensor identity and - version so unchanged routing avoids repeated host synchronization. - - 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. - - A captured graph has a static-storage contract: the `MoeEp` instance, - workspace, captured input/output storages, and transformed weights must - outlive every replay. Weights must not be modified after capture; replace - or modify them only after retiring the graph, then warm up and recapture. - - The same `MoeEp` instance must not be used concurrently by graph replay - and eager execution. Eager calls on different streams are serialized by a - completion event before shared workspace staging. - - Warmup validates route values. Capture/replay requires every route to - remain `-1` or a valid global expert ID; data-dependent host validation is - not capturable. - - Inference tensors without a PyTorch version counter are repacked on every - eager call and are not accepted as weights during graph capture. - - `generate_c=True` is eager-only because its exact `(local_routes, 2I)` - result requires data-dependent route counting and compaction. Capture is - rejected before the count collective, allocation, or kernel launch. - - Distributed callers must coordinate `close()` because symmetric allocation - release and an owned NVSHMEM finalization must not race a peer that is - still launching or replaying the operator. `close()` itself does not issue - a process-group barrier. - -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. -- 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. -- `WgradForwardStashReference` and `WgradOperandsReference` model logical - MXFP8 values, 256-route expert padding, and dense - `x.T @ dC` / `(p*h).T @ dY` results. Their scales are logical rather than - the production 128x4 physical assembly. -- It has no expert-capacity drop policy beyond explicit `-1` routes. - -Device parity tests configure the same `MoeEpReference` with -`intermediate_format="mxfp8"` to model the internal post-SwiGLU MXFP8 round -trip before FC2, and with `backward_operand_format="mxfp8"` to model dGLU -operand staging. These options are diagnostic backend approximations and are -deliberately not part of the public mathematical contract above. - -## Validation and test matrix - -Public-contract and reference tests remain broader than the current device -backend. They cover MXFP8/NVFP4 representation and quantization semantics, -both router-weight locations, quantized combine/output round trips, and the -explicit reference backward. Those checks do not imply kernel capability. - -The status below applies specifically to the production Rubin device backend. -“Hardware-validated” means that the case passed on SM107 (compute capability -10.7); CPU/reference tests, collection, and skips do not establish that status. - -### Hardware-validated on SM107 - -- EP1 forward has passed with MXFP8 E4M3/E8M0 operands and with mixed plain - BF16/FP16/FP32 activation/weights, BF16 or MXFP8 combine, BF16 output, - `apply_topk_in_fc1=True`, gate/up clamp, all-`-1` routes, and fresh output - ownership. INT32 and INT64 routing indices and supported shape/top-k cases - through `top_k=32` have also passed. -- Eager EP1 `generate_c=True` has passed with the Rubin-required - `apply_topk_in_fc1=True`. The checks cover compact BF16 `fc1_c`, INT32 - `route_metadata`, pre-clamp/unweighted values, repeated calls, and fresh - tensor ownership. -- EP1 warmup, non-default tuning, 100-call eager stress, and CUDA Graph replay - have passed. CUDA Graph replay with post-FC2 router weighting has also passed. -- Single-node WORLD-group eager forward parity has passed for EP2, EP3, and - EP4 with BF16 and MXFP8 combine, including all-`-1` route behavior. -- A WORLD4 eager test has passed for disjoint non-contiguous EP2 subgroups - `[0,2]` and `[1,3]`, including the case where subgroup rank zero is not global - rank zero. -- Multi-node eager forward parity has passed for EP12/WORLD12 (three nodes, - four PEs per node) and EP16/WORLD16 (four nodes, four PEs per node), with - BF16 and MXFP8 combine and all-`-1` route behavior. -- MXFP8 backward has passed for EP1, EP2, and EP4 with BF16 and - MXFP8 combine. It checks activation and router-weight gradients against the - executable reference; EP1 additionally covers repeated calls and explicitly - reordered forward stashes. -- The returned operand field/stride/scale ABI has passed direct execution - through both FC1 and FC2 grouped-wgrad GEMMs on SM100 using - reference-generated operands. This establishes consumer ABI integration, not - end-to-end Rubin operand production. - -### Supported but awaiting hardware validation - -- Forward capability accepts every positive EP size. EP5, EP6, EP8 through - EP11, EP13, EP14, and sizes above EP16 have no current hardware acceptance - case. Sizes above EP16 use the generated vector peer-offset path instead of - the fixed 128-byte by-value table used through EP16. -- EP7/WORLD14 is defined as a seven-node subgroup with one EP PE per node, and - EP15/WORLD20 is defined as a five-node subgroup with three EP PEs per node. - Both acceptance cases remain pending on allocations with the required node - counts. -- Forward CUDA Graph capture/replay and lifecycle contracts apply to - distributed EP groups, but current EP2+ acceptance runs cover eager parity - only; distributed stress and graph replay remain pending. -- Plain/mixed operands, gate/up clamp, post-FC2 router weighting, and explicit - `generate_c=True` output semantics are hardware-validated at EP1, but are not - separately validated across every supported distributed EP size. -- End-to-end Rubin forward/backward production of wgrad operands remains - awaiting SM107 hardware validation. The available SM100 test can execute the - grouped-wgrad consumer but cannot execute the SM107-only MegaMoE producer. - -### Currently unsupported by the device backend - -- Devices other than SM107, non-positive or unspecified `max_tokens_per_rank`, - `hidden_size` not divisible by 128, - `intermediate_size` not divisible by 256, and `top_k > 32`. -- Native NVFP4 operands or combine, any non-BF16 public output, and plain - operand dtypes other than BF16/FP16/FP32. MXFP8 operands and BF16/MXFP8 - combine with BF16 output are supported. -- Backward execution with `apply_topk_in_fc1=False`, non-BF16 output, or - backward CUDA Graph capture. -- Wgrad operand mode with padding other than 256, without `generate_c=True`, - or outside the restricted Rubin MXFP8 backward configuration. The mode - returns operands only; dense `dW1`/`dW2` computation remains an explicit - grouped-wgrad call by the integration layer. -- CUDA Graph capture with `generate_c=True`, same-process concurrent EP - subgroups, expert bias, shared/dense experts inside this operator, implicit - top-k normalization, capacity-factor routing, and implicit route drop. -- Dense weight-gradient returns and an integrated `torch.autograd.Function` - wrapper. - -Source/packaging tests separately require the vendored Rubin -`training/mega/fwd_glu` and `training/mega/bwd_dglu` packages, reject sibling -`cutedsl_megamoe` dependencies and `kernel_src.blackwell` imports, and validate -isolated-wheel imports. These checks validate packaging rather than additional -device capabilities. - -Future format families must add the same CUDA/NVSHMEM, stress, graph, and -isolated-package matrix before their capability gates are removed. - -## Current 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; -- backward semantics are fixed by `MoeEpReference.backward` (consuming the - `generate_c=True` stash); opt-in wgrad operands are exposed, while dense - weight-gradient returns and a `torch.autograd` wrapper are not part of this - API; -- logical scales at the Python boundary, backend swizzle internally. +Run SM107 single-node distributed tests from an exclusive GPU allocation: + +```bash +python -m pytest \ + test/python/moe_ep/test_moe_ep_forward.py \ + test/python/moe_ep/test_moe_ep_backward.py \ + -m L1 +``` + +`test/python/moe_ep/test_moe_ep_multinode.py` is torchrun-native and requires +the standard `LOCAL_RANK`, `LOCAL_WORLD_SIZE`, `RANK`, and `WORLD_SIZE` +environment. The multi-node fixture initializes NCCL and defaults +`NVIDIA_IMEX_CHANNELS=0`. + +Run distributed fixed-resource probes from an existing Slurm allocation: + +```bash +NVSHMEM_REMOTE_TRANSPORT=none \ +data/script/run_moe_ep_forward_multinode_slurm.sh backward-ep8 + +NVSHMEM_REMOTE_TRANSPORT=none \ +data/script/run_moe_ep_forward_multinode_slurm.sh graph-ep8 +``` + +The launcher provides `backward-ep16`, `backward-ep32`, `graph-ep16`, and +`graph-ep32` for the larger EP configurations. The fatal captured-overflow +assertion is a separate `graph-ep8-error` expected-failure task. + +The graph tasks invoke `test/python/moe_ep/probe_moe_ep_training_graph.py`, +which covers collective warmup, capture alignment, lockstep replay bursts, +dynamic routing/overflow recovery, ordered multi-lane execution, and +collective teardown. diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 5edf5a3bf..233caae86 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -36,7 +36,8 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d - [RMSNorm + RHT + Amax](rmsnorm_rht_amax.md) - [SDPA Backward (SM120)](attention/sdpa_bwd_sm120.md) - [RMSNorm + SiLU](rmsnorm_silu.md) -- [MoE + Expert Parallel API](moe_ep.md) +- [MoE + Expert Parallel API](moe_ep.md) — Rubin SM107 fused SwiGLU with + fixed-resource training and CUDA Graph support ## Installation and setup @@ -57,12 +58,23 @@ MoE + Expert Parallel requires its dedicated optional dependencies: pip install nvidia-cudnn-frontend[moe_ep] ``` +MoeEP is currently CUDA/PyTorch-only and targets Rubin SM107. EP2+ execution +also requires NCCL, NVSHMEM, and a direct-P2P MNNVL peer-access domain. See the +[MoeEP support matrix and tensor contracts](moe_ep.md#supported-configuration) +before integrating it. + After installation, you can import the APIs directly from the `cudnn` package, i.e. `from cudnn import {your_operation}` ## API Usage Each operation exposes two APIs: +MoeEP is an exception to the generic wrapper/kernel pattern below. It exposes +an object API: `MoeEp.__call__` for inference and +`MoeEp.prepare_training_resources` for fixed-resource training. See +[MoE + Expert Parallel API](moe_ep.md) for its lifecycle and CUDA Graph +contract. + ### 1. High-level wrapper - Single pythonic function call diff --git a/python/cudnn/moe_ep/_megamoe_backend/README.md b/python/cudnn/moe_ep/_megamoe_backend/README.md index 0b622fb14..a55986ad7 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/README.md +++ b/python/cudnn/moe_ep/_megamoe_backend/README.md @@ -1,139 +1,124 @@ -# MegaMoE backend capabilities - -This private backend implements the public `cudnn.moe_ep` contract with Rubin -SM107 CuTeDSL products. Public validation and backend capability checks are -separate: a request may be valid for `MoeEp` but unavailable in this backend. - -## Implemented forward paths - -- MXFP8 inputs use the training MegaMoE forward GLU product. -- Plain BF16, FP16, and FP32 operands are quantized into the same logical - MXFP8 representation before launch. -- Block-scaled operands must use the public MXFP8 representation. Native NVFP4 - operands are part of the public contract but are not executable in this - backend. -- Combine format may be BF16 or MXFP8. Final output format is BF16. -- MXFP8 combine quantizes each FP32 route accumulator directly. -- Rubin training execution requires `apply_topk_in_fc1=True`. -- Forward with `generate_c=False` supports CUDA Graph capture after warmup; - `generate_c=True` is eager-only. - -## Explicit Rubin limits - -- CUDA compute capability must be 10.7. -- `max_tokens_per_rank` must be positive and explicit. -- `hidden_size` must be divisible by 128, and `intermediate_size` must be - divisible by 256. -- `top_k` must not exceed 32. -- EP sizes above 16 use a generated vector peer-offset table; EP sizes through - 16 use the fixed 128-byte by-value table. -- `max_recv_size_per_rank` may bound per-rank routed rows below the conservative - `EP * max_tokens_per_rank * top_k` workspace size. Overflow traps by default - or truncates only when `drop_on_overflow=True`. - -These are backend limits, not additional public `MoeEp` semantics. They remain -precise, product-specific capability gates rather than hidden padding or a -silent numerical fallback. - -## Backward status - -`MoeEp.backward` has a validated backend seam and requires a forward stash from -`generate_c=True`. In the default `backward_wgrad_mode="none"`, the restricted -Rubin MXFP8 path returns -`(grad_activation, grad_topk_weights)` for any positive EP size with BF16 or -MXFP8 combine, BF16 output, `apply_topk_in_fc1=True`, optional -`gate_up_clamp`, and eager execution. It uses `fc1_c` and `route_metadata` to reconstruct an -external pool-layout `fc1_preact` tensor and converts `grad_output` to FP32 -before re-dispatching it for semantic dprob. The kernel's source-domain dprob -plane is symmetric and reset before every launch; the public router-weight -gradient remains an FP32 semantic recomputation. Default mode does not accept -or retain the forward activation, does not produce FC1/FC2 wgrad operands, and -does not depend on "most recent forward" state. - -### Opt-in grouped-wgrad operands - -Constructing with `backward_wgrad_mode="operands"` requires -`generate_c=True`, `token_padding_size=256`, and `sf_padding_size=128`. -Forward then returns -`(output, fc1_c, route_metadata, wgrad_forward_stash)`. The fourth value is a -`MoeEpWgradForwardStash` for that exact routed call: MXFP8 `x.T` data/scales, -cumulative padded expert offsets, valid route counts, and the same route -metadata. - -Backward takes the stash by keyword: +# MegaMoE backend + +The private MegaMoE backend provides Rubin SM107 MXFP8 execution for +`cudnn.moe_ep`. + +## Executable capability + +- CUDA Rubin SM107 (compute capability 10.7) +- BF16 output with BF16 or MXFP8 combine +- `hidden_size % 128 == 0` +- `intermediate_size % 256 == 0` +- `top_k <= min(32, num_experts)` +- explicit positive `max_tokens_per_rank` +- `apply_topk_in_fc1=True` + +Inference accepts plain BF16/FP16/FP32 or MXFP8 operands and stages plain +operands to MXFP8. Fixed-resource training accepts contiguous BF16/FP32 +activation and grad-output tensors, contiguous Int32 routing indices, +contiguous FP32 routing weights, and four contiguous MXFP8 training-weight +packs. NVFP4 operands, non-BF16 output, and `apply_topk_in_fc1=False` are not +executable. + +## Public execution paths + +`MoeEp.__call__` is the inference-forward surface. It returns only the fused +BF16 `(T, H)` output and does not expose a compact training stash or backward. +Inference CUDA Graph capture requires `MoeEp.warmup` with the exact capture +bindings before capture. EP ranks must align after warmup and replay in the +same cross-rank order. + +Training uses fixed resources: ```python -grad_activation, grad_topk_weights, operands = op.backward( - grad_output, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights, - fc1_c, - route_metadata, - wgrad_forward_stash=wgrad_forward_stash, +resources = op.prepare_training_resources( + weights, + slot_count=2, + lane_count=1, ) +slot0, slot1 = resources.slots +lane0 = resources.lanes[0] + +resources.refresh_weights() +y0 = resources.forward(slot0, lane0, x0, topk_idx0, topk_weights0) +dx0, dprob0, wgrad0 = resources.backward(slot0, lane0, grad0) +overflow = resources.finalize_overflow((slot0,), lane0) ``` -`MoeEpWgradOperands` is directly shaped for the grouped-wgrad Tensor2D ABI: - -- `fc1_a=(H,Kp)`, `fc1_b=(Kp,2I)` represent - `dW1 = x.T @ dC`; -- `fc2_a=(I,Kp)`, `fc2_b=(Kp,H)` represent the upstream factorization - `dW2 = (p * h).T @ dY`; -- every local expert's valid rows precede zero padding to 256 routes; - `expert_offsets` contains cumulative padded ends and may repeat for empty - experts, while `valid_route_counts` excludes padding; -- E4M3 data uses unit stride on K. E8M0 scales represent logical 1x32 blocks - assembled into the grouped kernel's physical 128x4 layout. - -The device order is deliberate. Forward first MXFP8-stages `x` along H, then -column-requantizes routed/padded rows along K. Backward follows the upstream -FC2-gradient factorization: the recomputed `h` export carries the route weight, -while the token-axis `grad_y2` export is the unweighted routed `dY`. Their -grouped product is therefore `dW2 = (p*h).T @ dY`. Staged `dY` and `W2.T` -produce `dH`, the route weight is applied before the SwiGLU derivative, and -`dC` is directly column-requantized along K. All three backward auxiliary -scale outputs use the upstream MN-major 128-column by 4-token-block atom -layout. - -The forward stash and backward operand tensors are caller-owned fresh -allocations, not views of reusable execution-plan workspace. Callers must keep -the forward stash alive through its matching backward call and keep returned -operands alive until external grouped-wgrad work completes. Later operator -calls do not overwrite them. Route identity/count validation prevents mixing -stashes from different forwards. - -This mode only produces operands; it does not launch grouped wgrad or return -dense `dW1`/`dW2`. It remains eager-only and inherits the restricted Rubin -MXFP8 backward gates (SM107, BF16 output, -`apply_topk_in_fc1=True`, and BF16/MXFP8 combine). End-to-end operand -production still requires SM107 acceptance. Direct FC1/FC2 consumer execution -has been validated separately on SM100 with reference-generated operands. - -The dGLU product emits BF16 `grad_activation`; the backend converts it to FP32 -for the public return. This is a documented BF16-rounded numerical limitation, -not strict FP32 dgrad parity. `apply_topk_in_fc1=False`, NVFP4 operands or -combine, non-BF16 output, and backward CUDA Graph capture remain -capability-gated. - -## Validation boundary - -L0 tests cover public validation, plain-to-MXFP8 staging, compile/cache keys, -workspace sizing, combine semantics, overflow audit behavior, and backward -layout/dispatch capability gates. Current CUDA Graph acceptance covers EP1 -forward. The single-node eager forward suite defines WORLD EP2/EP3/EP4, and -the current SM107 L1 hardware run establishes PASS for all three EP sizes. -The torchrun-native multi-node suite balances NVSHMEM PEs across participating -nodes: EP7 uses a WORLD14 subgroup over seven nodes with two workers per node, -EP12 uses WORLD12 over three nodes with four workers per node, EP15 uses a -WORLD20 subgroup over five nodes with four workers per node, and EP16 uses -WORLD16 over four nodes with four workers per node. Multi-node collection or -skip results do not establish a hardware PASS. Current hardware runs establish -PASS for EP12 and EP16; EP7 and EP15 remain pending. - -End-to-end device forward/backward parity requires SM107 hardware and the -`moe_ep` optional runtime dependencies, including a CuTeDSL installation that -provides `cutlass.utils.rubin_helpers`. Backward acceptance remains limited to -EP1/EP2/EP4; larger EP sizes are enabled but not covered by current backward -hardware acceptance. +`prepare_training_resources` is collective over the EP group, executes outside +capture, and fixes the training kernel to FC1-preactivation generation, +fixed-capacity WGrad operands, and token/scale-factor padding 128. + +The same methods execute ordinarily during warmup and enqueue identical nodes +inside a caller-owned outer CUDA Graph. MoeEP does not own or wrap graph +replay. The ordinary warmup must cover +`refresh_weights -> forward -> backward -> finalize_overflow` so staging, +forward, backward, and WGrad-export kernels are compiled before capture. + +## Fixed resource model + +- A persistent slot owns one microbatch's routing snapshot, pool-native FC1 + preactivation, kernel dprob, outputs, backward auxiliaries, overflow flags, + and fixed-capacity WGrad operands. +- An execution lane owns mutable router, barrier, and kernel scratch. +- Every symmetric region is built in deterministic order and its size is + normalized by name across EP ranks before allocation. +- Multiple streams require distinct lanes. Distributed MegaMoE kernels must be + ordered consistently on every rank with captured CUDA events; independent + lane storage does not permit unordered communication overlap. +- `max_recv_size_per_rank` bounds allocation. Capacity never grows during + capture; changing it requires new resources and graph capture. + +## Weights and WGrad outputs + +`MoeEpTrainingWeights` contains four address-stable MXFP8 block-scaled tensors: +forward W1/W2 and independently quantized backward W2-transpose/W1-transpose. +Their public layout differs from the K-major, gate/up-interleaved, and +blocked-scale kernel bindings. After every in-place data+scale update, the +caller must enqueue `resources.refresh_weights()` before the first consumer, +with explicit stream/event ordering. A matching forward/backward pair must use +one version; refresh cannot overlap any consumer on another slot/lane. Replacing +source storage requires closing the old operator, creating a new `MoeEp` +instance and resources, and capturing a new graph. Closed resources are +terminal and cannot be replaced on the same operator. Capturing the refresh +turns these transforms into fixed-address graph nodes, so replay does not call +Python. + +Backward returns kernel dprob directly. It follows the MXFP8-staged numerical +contract and relaxed atomic accumulation order. + +`MoeEpTrainingWgradOperands` is a fixed-capacity producer ABI. Device +`expert_offsets` and `valid_route_counts` describe the current valid K extent; +padding is zeroed. No specific downstream grouped-WGrad consumer is guaranteed +by this milestone. + +## Overflow policy + +`max_recv_size_per_rank` bounds the fixed receive pool. When omitted, it uses +the worst-case `ep_size * max_tokens_per_rank * top_k`; an explicit value is +capped at that count. + +The fixed-resource transport truncates deterministically so every rank +completes its communication protocol. `finalize_overflow` aggregates the +selected slots and performs a scalar MAX all-reduce for EP2+. With +`drop_on_overflow=True`, it returns a one-element Int32 status tensor and +dropped routes contribute zero. With `drop_on_overflow=False`, the graph tail +uses `torch._assert_async`; EP2+ error mode requires NCCL. + +## Distributed support + +Hardware acceptance covers EP1, EP2/4, EP8, EP16, and EP32 on one MNNVL +peer-access domain. The Python capability layer has no hard EP-size ceiling; +the listed sizes are validated scope rather than cross-MNNVL support. + +Current tests additionally cover single-node EP3 inference, noncontiguous EP2 +subgroups, multi-node EP4/6/12/16 forward, multi-node EP8/16/32 backward, and +EP8/16/32 fixed-resource graph launchers. EP2+ probes perform collective +warmup, independent capture, capture alignment, diagnostic replay, lockstep +production-like replay bursts, overflow/recovery, ordered multi-lane +execution, and collective teardown. + +The kernels use direct peer pointers obtained from NVSHMEM symmetric tensors. +`NVSHMEM_REMOTE_TRANSPORT=none` is valid only when every EP rank is directly +P2P-accessible (`NVSHMEM_TEAM_SHARED` spans the EP world). IBRC initialization +alone does not make non-P2P peers directly addressable by these kernels. From 6799b9dc060e2f957e59f1a30c0fa6ad88aad674 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Fri, 28 Aug 2026 09:07:33 -0700 Subject: [PATCH 14/31] style: format MoeEP sources for repository checks Apply the current Black contract to non-vendored MoeEP code so the rebased branch satisfies the all-files style check. --- .../moe_ep/_megamoe_backend/_capability.py | 75 ++---- python/cudnn/moe_ep/_megamoe_backend/_comm.py | 76 ++---- python/cudnn/moe_ep/_megamoe_backend/_plan.py | 34 +-- .../cudnn/moe_ep/_megamoe_backend/_runtime.py | 148 +++--------- .../moe_ep/_megamoe_backend/_workspace.py | 68 ++---- .../moe_ep/_megamoe_backend/mxfp8/_adapter.py | 146 +++-------- .../mxfp8/_backward_launch.py | 4 +- .../moe_ep/_megamoe_backend/mxfp8/_config.py | 43 +--- .../_megamoe_backend/mxfp8/_fingerprint.py | 40 ++- .../moe_ep/_megamoe_backend/mxfp8/_formats.py | 1 - .../moe_ep/_megamoe_backend/mxfp8/_launch.py | 36 +-- .../mxfp8/_training_execute.py | 18 +- .../mxfp8/_training_resources.py | 228 +++++------------- .../_megamoe_backend/mxfp8/_training_stage.py | 61 +---- .../mxfp8/_training_stage_kernel.py | 17 +- .../mxfp8/_training_weights.py | 97 ++++---- .../_megamoe_backend/mxfp8/_training_wgrad.py | 5 +- .../mxfp8/_training_wgrad_kernel.py | 81 ++----- python/cudnn/moe_ep/_tuning.py | 60 +---- python/cudnn/moe_ep/_types.py | 117 +++------ python/cudnn/moe_ep/_validation.py | 106 ++------ python/cudnn/moe_ep/api.py | 116 ++------- .../moe_ep/moe_ep_distributed_workers.py | 15 +- test/python/moe_ep/moe_ep_reference.py | 101 +++----- .../moe_ep/probe_moe_ep_training_graph.py | 90 +++---- test/python/moe_ep/test_moe_ep_multinode.py | 47 +--- 26 files changed, 444 insertions(+), 1386 deletions(-) diff --git a/python/cudnn/moe_ep/_megamoe_backend/_capability.py b/python/cudnn/moe_ep/_megamoe_backend/_capability.py index d34f74038..993d3c621 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/_capability.py +++ b/python/cudnn/moe_ep/_megamoe_backend/_capability.py @@ -24,34 +24,22 @@ def _validate_operand(name, tensor) -> None: if isinstance(tensor, BlockScaledTensor): if tensor.format is MoeFormat.MXFP8: return - raise NotImplementedError( - "MoeEp training MegaMoE supports only MXFP8 BlockScaledTensor " - f"inputs; {name} has format={tensor.format.value!r}" - ) + raise NotImplementedError("MoeEp training MegaMoE supports only MXFP8 BlockScaledTensor " f"inputs; {name} has format={tensor.format.value!r}") if tensor.dtype not in { torch.bfloat16, torch.float16, torch.float32, }: - raise NotImplementedError( - f"MoeEp MegaMoE {name} staging supports BF16, FP16, " - f"or FP32 plain tensors, got {tensor.dtype}" - ) + raise NotImplementedError(f"MoeEp MegaMoE {name} staging supports BF16, FP16, " f"or FP32 plain tensors, got {tensor.dtype}") def _validate_device(device: torch.device) -> None: if device.type != "cuda": - raise NotImplementedError( - f"MoeEp MegaMoE backend requires a CUDA device, got {device}" - ) + raise NotImplementedError(f"MoeEp MegaMoE backend requires a CUDA device, got {device}") major, minor = torch.cuda.get_device_capability(device) if (major, minor) != (10, 7): - raise NotImplementedError( - "MoeEp MegaMoE backend requires Rubin SM107 " - "(compute capability 10.7); " - f"found compute capability {major}.{minor}" - ) + raise NotImplementedError("MoeEp MegaMoE backend requires Rubin SM107 " "(compute capability 10.7); " f"found compute capability {major}.{minor}") def _is_cuda_stream_capturing(device: torch.device) -> bool: @@ -63,25 +51,14 @@ def _is_cuda_stream_capturing(device: torch.device) -> bool: def _validate_wgrad_config(config: ForwardConfig) -> None: if config.backward_wgrad_mode not in ("none", "operands"): - raise ValueError( - "unsupported backward_wgrad_mode " - f"{config.backward_wgrad_mode!r}" - ) + raise ValueError("unsupported backward_wgrad_mode " f"{config.backward_wgrad_mode!r}") if config.backward_wgrad_mode == "operands": if not config.generate_c: - raise ValueError( - "backward_wgrad_mode='operands' requires generate_c=True" - ) + raise ValueError("backward_wgrad_mode='operands' requires generate_c=True") if config.token_padding_size != 128: - raise ValueError( - "backward_wgrad_mode='operands' requires " - "token_padding_size=128" - ) + raise ValueError("backward_wgrad_mode='operands' requires " "token_padding_size=128") if config.sf_padding_size != 128: - raise ValueError( - "backward_wgrad_mode='operands' requires " - "sf_padding_size=128" - ) + raise ValueError("backward_wgrad_mode='operands' requires " "sf_padding_size=128") def validate_config(config: ForwardConfig) -> None: @@ -89,45 +66,25 @@ def validate_config(config: ForwardConfig) -> None: _validate_wgrad_config(config) if config.output_format != MoeFormat.BF16.value: - raise NotImplementedError( - "MoeEp training MegaMoE supports output_format='bf16' only" - ) + raise NotImplementedError("MoeEp training MegaMoE supports output_format='bf16' only") supported_combine_formats = { MoeFormat.BF16.value, MoeFormat.MXFP8.value, } if config.combine_format not in supported_combine_formats: - raise NotImplementedError( - "MoeEp training MegaMoE supports combine_format='bf16' " - "or 'mxfp8'" - ) + raise NotImplementedError("MoeEp training MegaMoE supports combine_format='bf16' " "or 'mxfp8'") if config.max_tokens_per_rank is None: - raise NotImplementedError( - "MoeEp MegaMoE backend requires an explicit max_tokens_per_rank" - ) + raise NotImplementedError("MoeEp MegaMoE backend requires an explicit max_tokens_per_rank") if config.max_tokens_per_rank == 0: - raise NotImplementedError( - "MoeEp SM107 MXFP8 execution requires " - "max_tokens_per_rank to be positive" - ) + raise NotImplementedError("MoeEp SM107 MXFP8 execution requires " "max_tokens_per_rank to be positive") if config.hidden_size % 128: - raise NotImplementedError( - "MoeEp SM107 MXFP8 kernel currently requires hidden_size " - f"to be divisible by 128, got {config.hidden_size}" - ) + raise NotImplementedError("MoeEp SM107 MXFP8 kernel currently requires hidden_size " f"to be divisible by 128, got {config.hidden_size}") if config.intermediate_size % 256: - raise NotImplementedError( - "MoeEp SM107 MXFP8 kernel currently requires intermediate_size " - f"to be divisible by 256, got {config.intermediate_size}" - ) + raise NotImplementedError("MoeEp SM107 MXFP8 kernel currently requires intermediate_size " f"to be divisible by 256, got {config.intermediate_size}") if config.top_k > 32: - raise NotImplementedError( - "MoeEp SM107 MXFP8 dispatch currently requires top_k <= 32" - ) + raise NotImplementedError("MoeEp SM107 MXFP8 dispatch currently requires top_k <= 32") if not config.apply_topk_in_fc1: - raise NotImplementedError( - "MoeEp Rubin training MegaMoE requires apply_topk_in_fc1=True" - ) + raise NotImplementedError("MoeEp Rubin training MegaMoE requires apply_topk_in_fc1=True") def validate_request(request: ValidatedForwardRequest) -> None: diff --git a/python/cudnn/moe_ep/_megamoe_backend/_comm.py b/python/cudnn/moe_ep/_megamoe_backend/_comm.py index d837558e7..b4f67ad04 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/_comm.py +++ b/python/cudnn/moe_ep/_megamoe_backend/_comm.py @@ -51,9 +51,7 @@ def _core(): try: import nvshmem.core as core except (ImportError, OSError) as exc: - raise RuntimeUnavailableError( - "symmetric workspace requires nvshmem4py and NVSHMEM libraries" - ) from exc + raise RuntimeUnavailableError("symmetric workspace requires nvshmem4py and NVSHMEM libraries") from exc return core def allocate(self, nbytes: int, device: torch.device) -> torch.Tensor: @@ -75,9 +73,7 @@ def allocate(self, nbytes: int, device: torch.device) -> torch.Tensor: error=repr(exc), elapsed_seconds=f"{time.monotonic() - started_at:.3f}", ) - raise RuntimeUnavailableError( - f"failed to allocate {nbytes} bytes from the NVSHMEM symmetric heap" - ) from exc + raise RuntimeUnavailableError(f"failed to allocate {nbytes} bytes from the NVSHMEM symmetric heap") from exc _runtime_debug( "symmetric.allocate.end", nbytes=nbytes, @@ -102,9 +98,7 @@ def free(self, tensor: torch.Tensor) -> None: error=repr(exc), elapsed_seconds=f"{time.monotonic() - started_at:.3f}", ) - raise RuntimeUnavailableError( - "failed to free the NVSHMEM symmetric root slab" - ) from exc + raise RuntimeUnavailableError("failed to free the NVSHMEM symmetric root slab") from exc _runtime_debug( "symmetric.free.end", elapsed_seconds=f"{time.monotonic() - started_at:.3f}", @@ -127,9 +121,7 @@ def peer_address(self, tensor: torch.Tensor, peer: int) -> int: error=repr(exc), elapsed_seconds=f"{time.monotonic() - started_at:.3f}", ) - raise RuntimeUnavailableError( - f"failed to map symmetric root slab for peer {peer}" - ) from exc + raise RuntimeUnavailableError(f"failed to map symmetric root slab for peer {peer}") from exc peer_pointer = int(peer_tensor.data_ptr()) _runtime_debug( "symmetric.peer-map.end", @@ -152,13 +144,9 @@ def __post_init__(self) -> None: if len(self.offsets) == 0: raise ValueError("peer mapping requires at least one rank") if self.rank < 0 or self.rank >= len(self.offsets): - raise ValueError( - f"peer mapping rank {self.rank} is outside {len(self.offsets)} ranks" - ) + raise ValueError(f"peer mapping rank {self.rank} is outside {len(self.offsets)} ranks") if self.offsets[self.rank] != 0: - raise ValueError( - f"local peer offset must be zero, got {self.offsets[self.rank]}" - ) + raise ValueError(f"local peer offset must be zero, got {self.offsets[self.rank]}") @property def world_size(self) -> int: @@ -195,26 +183,18 @@ def __init__( self._runtime = runtime self._nbytes = nbytes - self._provider = provider or ( - _NvshmemMemoryProvider() - if runtime.nvshmem_enabled - else _TorchMemoryProvider() - ) + self._provider = provider or (_NvshmemMemoryProvider() if runtime.nvshmem_enabled else _TorchMemoryProvider()) self._root: Optional[torch.Tensor] = None self._mapping: Optional[PeerMapping] = None self._cleanup_required = False def ensure_allocated(self) -> None: if self._cleanup_required: - raise RuntimeError( - "symmetric slab requires cleanup before allocation" - ) + raise RuntimeError("symmetric slab requires cleanup before allocation") if self._root is not None and self._mapping is not None: return if self._root is not None: - raise RuntimeError( - "symmetric slab has an allocation pending cleanup" - ) + raise RuntimeError("symmetric slab has an allocation pending cleanup") _runtime_debug( "symmetric-slab.ensure.begin", @@ -224,21 +204,13 @@ def ensure_allocated(self) -> None: ) root = self._provider.allocate(self._nbytes, self._runtime.device) if not isinstance(root, torch.Tensor): - raise TypeError( - "symmetric memory provider must return a torch.Tensor" - ) + raise TypeError("symmetric memory provider must return a torch.Tensor") self._root = root try: if root.dtype is not torch.uint8 or root.numel() < self._nbytes: - raise ValueError( - "symmetric root must be a uint8 tensor with at least " - f"{self._nbytes} elements" - ) + raise ValueError("symmetric root must be a uint8 tensor with at least " f"{self._nbytes} elements") if root.device != self._runtime.device: - raise ValueError( - "symmetric root device does not match runtime device: " - f"root={root.device}, runtime={self._runtime.device}" - ) + raise ValueError("symmetric root device does not match runtime device: " f"root={root.device}, runtime={self._runtime.device}") if not root.is_contiguous(): raise ValueError("symmetric root tensor must be contiguous") except Exception: @@ -260,21 +232,14 @@ def ensure_allocated(self) -> None: if peer == self._runtime.rank: offsets.append(0) continue - offsets.append( - self._provider.peer_address(root, peer) - base_address - ) + offsets.append(self._provider.peer_address(root, peer) - base_address) mapping = PeerMapping( base_address=base_address, offsets=tuple(offsets), rank=self._runtime.rank, ) - if ( - mapping.world_size != self._runtime.world_size - or mapping.rank != self._runtime.rank - ): - raise RuntimeError( - "symmetric peer mapping does not match the EP subgroup" - ) + if mapping.world_size != self._runtime.world_size or mapping.rank != self._runtime.rank: + raise RuntimeError("symmetric peer mapping does not match the EP subgroup") except Exception: self._cleanup_required = True raise @@ -297,11 +262,7 @@ def closed(self) -> bool: @property def allocated(self) -> bool: - return ( - not self._cleanup_required - and self._root is not None - and self._mapping is not None - ) + return not self._cleanup_required and self._root is not None and self._mapping is not None @property def mapping(self) -> PeerMapping: @@ -321,10 +282,7 @@ def root(self) -> torch.Tensor: def byte_view(self, offset: int, nbytes: int) -> torch.Tensor: if offset < 0 or nbytes < 0 or offset + nbytes > self._nbytes: - raise ValueError( - f"byte view [{offset}, {offset + nbytes}) exceeds " - f"symmetric slab size {self._nbytes}" - ) + raise ValueError(f"byte view [{offset}, {offset + nbytes}) exceeds " f"symmetric slab size {self._nbytes}") return self.root.narrow(0, offset, nbytes) def close(self) -> None: diff --git a/python/cudnn/moe_ep/_megamoe_backend/_plan.py b/python/cudnn/moe_ep/_megamoe_backend/_plan.py index c78b880cb..aad99cfc4 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/_plan.py +++ b/python/cudnn/moe_ep/_megamoe_backend/_plan.py @@ -44,9 +44,7 @@ def __init__( local_provider: Optional[LocalMemoryProvider] = None, ) -> None: if config.max_tokens_per_rank != requirements.max_tokens_per_rank: - raise ValueError( - "workspace capacity must match ForwardConfig.max_tokens_per_rank" - ) + raise ValueError("workspace capacity must match ForwardConfig.max_tokens_per_rank") self.config = config self.device = torch.device(device) self.requirements = requirements @@ -61,12 +59,7 @@ def __init__( @property def prepared(self) -> bool: - return ( - not self._cleanup_required - and self._runtime is not None - and self._workspace is not None - and self._workspace.allocated - ) + return not self._cleanup_required and self._runtime is not None and self._workspace is not None and self._workspace.allocated @property def cleanup_required(self) -> bool: @@ -84,28 +77,15 @@ def prepare( if self._closed: raise RuntimeError("MegaMoE execution plan is closed") if self._cleanup_required: - raise RuntimeError( - "MegaMoE execution plan requires cleanup before prepare" - ) + raise RuntimeError("MegaMoE execution plan requires cleanup before prepare") if request.config is not self.config: raise ValueError("request does not belong to this static plan") if torch.device(request.device) != self.device: - raise ValueError( - f"execution plan is bound to {self.device}, got {request.device}" - ) + raise ValueError(f"execution plan is bound to {self.device}, got {request.device}") if request.token_count > self.requirements.max_tokens_per_rank: - raise ValueError( - f"token count {request.token_count} exceeds " - f"max_tokens_per_rank={self.requirements.max_tokens_per_rank}" - ) - if ( - not self.prepared - and torch.cuda.is_current_stream_capturing() - ): - raise RuntimeError( - "MegaMoE runtime/workspace must be warmed up before " - "CUDA graph capture" - ) + raise ValueError(f"token count {request.token_count} exceeds " f"max_tokens_per_rank={self.requirements.max_tokens_per_rank}") + if not self.prepared and torch.cuda.is_current_stream_capturing(): + raise RuntimeError("MegaMoE runtime/workspace must be warmed up before " "CUDA graph capture") if not self.prepared: runtime = self._runtime_manager.acquire(self.config, self.device) diff --git a/python/cudnn/moe_ep/_megamoe_backend/_runtime.py b/python/cudnn/moe_ep/_megamoe_backend/_runtime.py index 126431d85..434ba1d59 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/_runtime.py +++ b/python/cudnn/moe_ep/_megamoe_backend/_runtime.py @@ -47,8 +47,7 @@ def _runtime_debug(event: str, **details: object) -> None: **details, } print( - "[moe-ep-runtime] " - + " ".join(f"{name}={value}" for name, value in fields.items()), + "[moe-ep-runtime] " + " ".join(f"{name}={value}" for name, value in fields.items()), file=sys.stderr, flush=True, ) @@ -71,9 +70,7 @@ def __init__(self, event: str) -> None: self._event = event self._stopped = threading.Event() try: - self._interval = float( - os.environ.get("MOE_EP_RUNTIME_WATCHDOG_SECONDS", "30") - ) + self._interval = float(os.environ.get("MOE_EP_RUNTIME_WATCHDOG_SECONDS", "30")) except ValueError: self._interval = 30.0 self._thread: Optional[threading.Thread] = None @@ -116,9 +113,7 @@ def _run(self) -> None: except OSError as exc: wait_channels.append(f"{task_dir.name}:unavailable({exc.errno})") else: - wait_channels.append( - f"{task_dir.name}:{thread_name}:{wait_channel or '-'}" - ) + wait_channels.append(f"{task_dir.name}:{thread_name}:{wait_channel or '-'}") _runtime_debug( "watchdog", blocked_event=self._event, @@ -174,40 +169,23 @@ def finalize(self) -> None: ... def _resolve_world(config: ForwardConfig) -> RuntimeWorld: if config.ep_group is None: - if ( - config.ep_size != 1 - or config.ep_rank != 0 - or config.ep_global_ranks - ): - raise ValueError( - "ep_group=None requires ep_size=1, ep_rank=0, and no " - "distributed rank membership" - ) + if config.ep_size != 1 or config.ep_rank != 0 or config.ep_global_ranks: + raise ValueError("ep_group=None requires ep_size=1, ep_rank=0, and no " "distributed rank membership") return RuntimeWorld(rank=0, size=1, group=None, global_ranks=()) if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError( - "distributed MegaMoE runtime requires torch.distributed to be initialized" - ) + raise RuntimeError("distributed MegaMoE runtime requires torch.distributed to be initialized") group = config.ep_group rank = dist.get_rank(group) size = dist.get_world_size(group) - global_ranks = tuple( - dist.get_global_rank(group, group_rank) - for group_rank in range(size) - ) + global_ranks = tuple(dist.get_global_rank(group, group_rank) for group_rank in range(size)) if (rank, size) != (config.ep_rank, config.ep_size): raise RuntimeError( - "ForwardConfig EP geometry does not match its process group: " - f"config=({config.ep_rank}, {config.ep_size}), " - f"runtime=({rank}, {size})" + "ForwardConfig EP geometry does not match its process group: " f"config=({config.ep_rank}, {config.ep_size}), " f"runtime=({rank}, {size})" ) if global_ranks != config.ep_global_ranks: - raise RuntimeError( - "ForwardConfig EP membership does not match its process group: " - f"config={config.ep_global_ranks}, runtime={global_ranks}" - ) + raise RuntimeError("ForwardConfig EP membership does not match its process group: " f"config={config.ep_global_ranks}, runtime={global_ranks}") return RuntimeWorld( rank=rank, size=size, @@ -237,9 +215,7 @@ def _load_nvshmem_core(): try: import nvshmem.core as core except (ImportError, OSError) as exc: - raise RuntimeUnavailableError( - "MegaMoE distributed runtime requires nvshmem4py and NVSHMEM libraries" - ) from exc + raise RuntimeUnavailableError("MegaMoE distributed runtime requires nvshmem4py and NVSHMEM libraries") from exc return core @@ -249,9 +225,7 @@ def _normalize_nvshmem_init_state(status) -> RuntimeInitState: name = getattr(status, "name", "") if name.endswith("NOT_INITIALIZED"): return RuntimeInitState.NOT_INITIALIZED - if name.endswith("IS_INITIALIZED") or name.endswith( - ("LIMITED_MPG", "FULL_MPG") - ): + if name.endswith("IS_INITIALIZED") or name.endswith(("LIMITED_MPG", "FULL_MPG")): return RuntimeInitState.INITIALIZED if name.endswith("IS_BOOTSTRAPPED"): return RuntimeInitState.PARTIAL @@ -275,9 +249,7 @@ def initialization_state(self) -> RuntimeInitState: try: status = core.init_status() except Exception as exc: - raise RuntimeUnavailableError( - "failed to query NVSHMEM initialization status" - ) from exc + raise RuntimeUnavailableError("failed to query NVSHMEM initialization status") from exc return _normalize_nvshmem_init_state(status) def initialize(self, device: torch.device, world: RuntimeWorld) -> None: @@ -313,16 +285,11 @@ def initialize(self, device: torch.device, world: RuntimeWorld) -> None: uid_bytes = uid._data.view(np.uint8).copy() uid_tensor = torch.from_numpy(uid_bytes) group_backend = dist.get_backend(world.group) - if ( - group_backend == dist.Backend.NCCL - or str(group_backend).lower() == "nccl" - ): + if group_backend == dist.Backend.NCCL or str(group_backend).lower() == "nccl": uid_tensor = uid_tensor.to(device=device) root_global_rank = dist.get_global_rank(world.group, 0) if root_global_rank != world.global_ranks[0]: - raise RuntimeError( - "EP subgroup root changed during NVSHMEM bootstrap" - ) + raise RuntimeError("EP subgroup root changed during NVSHMEM bootstrap") _runtime_debug( "initialize.uid-broadcast.begin", backend=group_backend, @@ -368,9 +335,7 @@ def initialize(self, device: torch.device, world: RuntimeWorld) -> None: error=repr(exc), elapsed_seconds=f"{time.monotonic() - started_at:.3f}", ) - raise RuntimeUnavailableError( - "failed to initialize the NVSHMEM EP subgroup runtime" - ) from exc + raise RuntimeUnavailableError("failed to initialize the NVSHMEM EP subgroup runtime") from exc def rank(self) -> int: try: @@ -382,9 +347,7 @@ def world_size(self) -> int: try: return int(_load_nvshmem_core().n_pes()) except Exception as exc: - raise RuntimeUnavailableError( - "failed to query the NVSHMEM PE world size" - ) from exc + raise RuntimeUnavailableError("failed to query the NVSHMEM PE world size") from exc def device(self) -> torch.device: try: @@ -395,9 +358,7 @@ def device(self) -> torch.device: raise RuntimeError("NVSHMEM cached device is empty") return torch.device("cuda", int(cached.device_id)) except Exception as exc: - raise RuntimeUnavailableError( - "failed to query the NVSHMEM initialization device" - ) from exc + raise RuntimeUnavailableError("failed to query the NVSHMEM initialization device") from exc def finalize(self) -> None: core = _load_nvshmem_core() @@ -501,9 +462,7 @@ class RuntimeManager: def __init__( self, *, - provider_factory: Callable[[], NvshmemRuntimeProvider] = ( - _DefaultNvshmemRuntimeProvider - ), + provider_factory: Callable[[], NvshmemRuntimeProvider] = (_DefaultNvshmemRuntimeProvider), world_resolver: Callable[[ForwardConfig], RuntimeWorld] = _resolve_world, keep_alive: bool = False, ) -> None: @@ -541,19 +500,11 @@ def acquire( cleanup_required=active.cleanup_required, ) if active.cleanup_required: - raise RuntimeError( - "MegaMoE process runtime requires cleanup before reacquire" - ) + raise RuntimeError("MegaMoE process runtime requires cleanup before reacquire") if active.device != device: - raise ValueError( - f"MegaMoE process runtime is bound to {active.device}; " - f"cannot acquire it for {device}" - ) + raise ValueError(f"MegaMoE process runtime is bound to {active.device}; " f"cannot acquire it for {device}") if active.world.identity != world.identity: - raise RuntimeError( - "MegaMoE process runtime is already bound to a different " - "EP subgroup" - ) + raise RuntimeError("MegaMoE process runtime is already bound to a different " "EP subgroup") active.ref_count += 1 _runtime_debug( "manager.acquire-reuse.end", @@ -583,13 +534,8 @@ def acquire( init_status=status.value, ) if status is RuntimeInitState.PARTIAL: - raise RuntimeError( - "cannot attach to a partially initialized NVSHMEM runtime" - ) - if ( - status is RuntimeInitState.INITIALIZED - and not _spans_default_distributed_world(world) - ): + raise RuntimeError("cannot attach to a partially initialized NVSHMEM runtime") + if status is RuntimeInitState.INITIALIZED and not _spans_default_distributed_world(world): raise RuntimeError( "cannot safely attach an externally initialized NVSHMEM " "runtime to a non-WORLD EP subgroup because its ordered " @@ -629,20 +575,11 @@ def acquire( device, world, RuntimeError( - "NVSHMEM initialization device does not match " - f"the requested device: nvshmem={provider_device}, " - f"requested={device}" + "NVSHMEM initialization device does not match " f"the requested device: nvshmem={provider_device}, " f"requested={device}" ), ) - ownership = ( - "owned" - if owns_runtime - else "externally initialized" - ) - raise RuntimeError( - f"{ownership} NVSHMEM runtime is bound to " - f"{provider_device}, not the requested device {device}" - ) + ownership = "owned" if owns_runtime else "externally initialized" + raise RuntimeError(f"{ownership} NVSHMEM runtime is bound to " f"{provider_device}, not the requested device {device}") if (provider_rank, provider_size) != (world.rank, world.size): if owns_runtime: self._cleanup_owned_runtime_after_error( @@ -702,9 +639,7 @@ def _rollback_failed_initialization( provider.finalize() except Exception as cleanup_error: cls._mark_cleanup_required(provider, device, world) - raise RuntimeError( - "NVSHMEM initialization failed and rollback requires retry" - ) from cleanup_error + raise RuntimeError("NVSHMEM initialization failed and rollback requires retry") from cleanup_error _logger.debug( "rolled back failed NVSHMEM initialization: %s", initialization_error, @@ -722,9 +657,7 @@ def _cleanup_owned_runtime_after_error( provider.finalize() except Exception as cleanup_error: cls._mark_cleanup_required(provider, device, world) - raise RuntimeError( - "NVSHMEM validation failed and cleanup requires retry" - ) from cleanup_error + raise RuntimeError("NVSHMEM validation failed and cleanup requires retry") from cleanup_error _logger.debug( "finalized owned NVSHMEM after validation failure: %s", original_error, @@ -738,13 +671,9 @@ def retry_cleanup(self) -> None: if active is None: return if not active.cleanup_required or active.ref_count != 0: - raise RuntimeError( - "MegaMoE process runtime does not have retryable cleanup" - ) + raise RuntimeError("MegaMoE process runtime does not have retryable cleanup") if active.provider is None: - raise RuntimeError( - "retryable MegaMoE runtime cleanup has no provider" - ) + raise RuntimeError("retryable MegaMoE runtime cleanup has no provider") active.provider.finalize() _PROCESS_RUNTIME_REGISTRY.active = None @@ -756,13 +685,8 @@ def shutdown(self) -> None: if active is None: return if active.ref_count != 0: - raise RuntimeError( - "cannot shut down the MegaMoE process runtime while " - f"{active.ref_count} handles remain active" - ) - if active.provider is not None and ( - active.owns_runtime or active.cleanup_required - ): + raise RuntimeError("cannot shut down the MegaMoE process runtime while " f"{active.ref_count} handles remain active") + if active.provider is not None and (active.owns_runtime or active.cleanup_required): _runtime_debug( "manager.shutdown-finalize.begin", cleanup_required=active.cleanup_required, @@ -779,9 +703,7 @@ def _release(self, token: object) -> None: _runtime_debug("manager.release-stale") return if active.ref_count <= 0: - raise RuntimeError( - "MegaMoE process runtime has invalid release state" - ) + raise RuntimeError("MegaMoE process runtime has invalid release state") _runtime_debug( "manager.release.begin", @@ -791,9 +713,7 @@ def _release(self, token: object) -> None: ) if active.cleanup_required: if active.ref_count != 1 or active.provider is None: - raise RuntimeError( - "MegaMoE process runtime has invalid retry state" - ) + raise RuntimeError("MegaMoE process runtime has invalid retry state") active.provider.finalize() _PROCESS_RUNTIME_REGISTRY.active = None _runtime_debug("manager.release.cleanup-retry.end") diff --git a/python/cudnn/moe_ep/_megamoe_backend/_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/_workspace.py index ad65394e9..24a24d71d 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/_workspace.py +++ b/python/cudnn/moe_ep/_megamoe_backend/_workspace.py @@ -45,13 +45,9 @@ def __post_init__(self) -> None: if not self.name: raise ValueError("workspace region name must not be empty") if self.nbytes < 0: - raise ValueError( - f"workspace region {self.name!r} has negative size {self.nbytes}" - ) + raise ValueError(f"workspace region {self.name!r} has negative size {self.nbytes}") if self.alignment <= 0 or self.alignment & (self.alignment - 1): - raise ValueError( - f"workspace region {self.name!r} alignment must be a power of two" - ) + raise ValueError(f"workspace region {self.name!r} alignment must be a power of two") @dataclass(frozen=True) @@ -122,10 +118,7 @@ def __post_init__(self) -> None: local_names = {region.name for region in self.local_regions} duplicates = symmetric_names & local_names if duplicates: - raise ValueError( - "workspace region names must be unique across roots: " - f"{sorted(duplicates)}" - ) + raise ValueError("workspace region names must be unique across roots: " f"{sorted(duplicates)}") @classmethod def for_mxfp8( @@ -156,9 +149,7 @@ def for_mxfp8( if value < 0: raise ValueError(f"{name} must be non-negative, got {value}") if bool(col_quant_data_bytes) != bool(col_quant_sf_bytes): - raise ValueError( - "column requant data and scale workspace must be enabled together" - ) + raise ValueError("column requant data and scale workspace must be enabled together") backward_sizes = ( backward_fc1_preact_bytes, backward_dprob_bytes, @@ -166,21 +157,14 @@ def for_mxfp8( backward_aux_scale_bytes, ) if any(backward_sizes) and not all(backward_sizes): - raise ValueError( - "backward preactivation, dprob, data, and scale workspace " - "must be enabled together" - ) + raise ValueError("backward preactivation, dprob, data, and scale workspace " "must be enabled together") tokens = config.max_tokens_per_rank hidden = config.hidden_size top_k = config.top_k kernel_sf_columns = padded_mxfp8_scale_columns(hidden) - backward_symmetric_regions = ( - (BufferRegion("backward_dprob", backward_dprob_bytes),) - if backward_dprob_bytes - else () - ) + backward_symmetric_regions = (BufferRegion("backward_dprob", backward_dprob_bytes),) if backward_dprob_bytes else () symmetric_regions = ( BufferRegion("activation_data", tokens * hidden), BufferRegion("activation_scale", tokens * kernel_sf_columns), @@ -238,6 +222,7 @@ def for_mxfp8( local_regions=local_regions, ) + class LocalMemoryProvider(Protocol): """Injectable local allocation boundary.""" @@ -269,15 +254,9 @@ def __init__( if not isinstance(root, torch.Tensor): raise TypeError("local memory provider must return a torch.Tensor") if root.dtype is not torch.uint8 or root.numel() < nbytes: - raise ValueError( - "local root must be a uint8 tensor with at least " - f"{nbytes} elements" - ) + raise ValueError("local root must be a uint8 tensor with at least " f"{nbytes} elements") if root.device != device: - raise ValueError( - "local root device does not match runtime device: " - f"root={root.device}, runtime={device}" - ) + raise ValueError("local root device does not match runtime device: " f"root={root.device}, runtime={device}") if not root.is_contiguous(): raise ValueError("local root tensor must be contiguous") _runtime_debug("local-slab.zero.begin", nbytes=nbytes) @@ -297,10 +276,7 @@ def root(self) -> torch.Tensor: def byte_view(self, offset: int, nbytes: int) -> torch.Tensor: if offset < 0 or nbytes < 0 or offset + nbytes > self._nbytes: - raise ValueError( - f"byte view [{offset}, {offset + nbytes}) exceeds " - f"local slab size {self._nbytes}" - ) + raise ValueError(f"byte view [{offset}, {offset + nbytes}) exceeds " f"local slab size {self._nbytes}") return self.root.narrow(0, offset, nbytes) def close(self) -> None: @@ -334,9 +310,7 @@ def __init__( ) -> None: self.requirements = requirements self.runtime = runtime - self.symmetric_layout = BufferLayout.build( - requirements.symmetric_regions - ) + self.symmetric_layout = BufferLayout.build(requirements.symmetric_regions) self.local_layout = BufferLayout.build(requirements.local_regions) if self.symmetric_layout.total_bytes <= 0: raise ValueError("workspace requires at least one symmetric byte") @@ -353,12 +327,7 @@ def __init__( @property def allocated(self) -> bool: - return ( - not self._cleanup_required - and self._symmetric is not None - and self._symmetric.allocated - and self._local is not None - ) + return not self._cleanup_required and self._symmetric is not None and self._symmetric.allocated and self._local is not None @property def cleanup_required(self) -> bool: @@ -373,9 +342,7 @@ def ensure_allocated(self) -> None: if self._closed: raise RuntimeError("workspace owner is closed") if self._cleanup_required: - raise RuntimeError( - "workspace owner requires cleanup before allocation" - ) + raise RuntimeError("workspace owner requires cleanup before allocation") if self.allocated: return self.runtime.ensure_open() @@ -419,14 +386,9 @@ def ensure_allocated(self) -> None: def views(self, token_count: int) -> WorkspaceViews: with self._lock: if token_count < 0: - raise ValueError( - f"token_count must be non-negative, got {token_count}" - ) + raise ValueError(f"token_count must be non-negative, got {token_count}") if token_count > self.requirements.max_tokens_per_rank: - raise ValueError( - f"token count {token_count} exceeds " - f"max_tokens_per_rank={self.requirements.max_tokens_per_rank}" - ) + raise ValueError(f"token count {token_count} exceeds " f"max_tokens_per_rank={self.requirements.max_tokens_per_rank}") self.ensure_allocated() assert self._symmetric is not None assert self._local is not None diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py index 34ec485f1..12057847a 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py @@ -62,15 +62,8 @@ def _quantize_plain_mxfp8( scale_for_math.reciprocal(), 0.0, ) - normalized = ( - blocks * reciprocal.unsqueeze(-1) - ).clamp(-448.0, 448.0) - data = ( - normalized.to(_MXFP8_DATA_DTYPE) - .reshape(*moved.shape)[..., :logical_extent] - .movedim(-1, axis) - .contiguous() - ) + normalized = (blocks * reciprocal.unsqueeze(-1)).clamp(-448.0, 448.0) + data = normalized.to(_MXFP8_DATA_DTYPE).reshape(*moved.shape)[..., :logical_extent].movedim(-1, axis).contiguous() return BlockScaledTensor( data=data, scale=scale.movedim(-1, axis).contiguous(), @@ -83,10 +76,7 @@ def _quantize_plain_mxfp8( def _as_mxfp8(tensor: torch.Tensor | BlockScaledTensor) -> BlockScaledTensor: if isinstance(tensor, BlockScaledTensor): if tensor.format is not MoeFormat.MXFP8: - raise NotImplementedError( - "MXFP8 staging cannot convert " - f"{tensor.format.value!r} block-scaled input" - ) + raise NotImplementedError("MXFP8 staging cannot convert " f"{tensor.format.value!r} block-scaled input") return tensor return _quantize_plain_mxfp8(tensor) @@ -101,10 +91,7 @@ def _typed_view( expected_bytes *= extent expected_bytes *= dtype.itemsize if byte_tensor.numel() != expected_bytes: - raise ValueError( - f"byte region has {byte_tensor.numel()} bytes, " - f"expected {expected_bytes} for shape={shape}, dtype={dtype}" - ) + raise ValueError(f"byte region has {byte_tensor.numel()} bytes, " f"expected {expected_bytes} for shape={shape}, dtype={dtype}") return byte_tensor.view(dtype).reshape(shape) @@ -135,14 +122,8 @@ def _validate_int32_downcast(tensor: torch.Tensor) -> None: if tensor.dtype is torch.int32: return if tensor.dtype is not torch.int64: - raise TypeError( - "topk_idx staging requires torch.int32 or torch.int64, " - f"got {tensor.dtype}" - ) - capturing = ( - tensor.device.type == "cuda" - and torch.cuda.is_current_stream_capturing() - ) + raise TypeError("topk_idx staging requires torch.int32 or torch.int64, " f"got {tensor.dtype}") + capturing = tensor.device.type == "cuda" and torch.cuda.is_current_stream_capturing() if capturing or tensor.numel() == 0: # The public validator checked the same tensor before capture. During # replay callers must preserve its documented expert-id invariant. @@ -160,9 +141,7 @@ def _zero_workspace_prefix( name: str, ) -> None: if nbytes < 0 or nbytes > workspace.numel(): - raise ValueError( - f"{name} zero prefix {nbytes} exceeds {workspace.numel()} bytes" - ) + raise ValueError(f"{name} zero prefix {nbytes} exceeds {workspace.numel()} bytes") workspace[:nbytes].zero_() @@ -174,10 +153,7 @@ def _zero_workspace_range( name: str, ) -> None: if offset < 0 or nbytes < 0 or offset + nbytes > workspace.numel(): - raise ValueError( - f"{name} byte range [{offset}, {offset + nbytes}) exceeds " - f"{workspace.numel()} bytes" - ) + raise ValueError(f"{name} byte range [{offset}, {offset + nbytes}) exceeds " f"{workspace.numel()} bytes") workspace.narrow(0, offset, nbytes).zero_() @@ -188,15 +164,9 @@ def _interleave_gate_up_rows( """Convert gate-half/up-half rows to 32-row gate/up pairs.""" if intermediate % _GATE_UP_INTERLEAVE: - raise ValueError( - "MXFP8 gate/up interleave requires intermediate_size to be " - f"divisible by {_GATE_UP_INTERLEAVE}, got {intermediate}" - ) + raise ValueError("MXFP8 gate/up interleave requires intermediate_size to be " f"divisible by {_GATE_UP_INTERLEAVE}, got {intermediate}") if tensor.ndim != 3 or tensor.shape[1] != 2 * intermediate: - raise ValueError( - f"expected (experts, {2 * intermediate}, K) tensor, " - f"got {tuple(tensor.shape)}" - ) + raise ValueError(f"expected (experts, {2 * intermediate}, K) tensor, " f"got {tuple(tensor.shape)}") experts, _gate_up, reduction = tensor.shape pairs = intermediate // _GATE_UP_INTERLEAVE @@ -212,11 +182,7 @@ def _interleave_gate_up_rows( _GATE_UP_INTERLEAVE, reduction, ) - return ( - torch.stack((gate, up), dim=2) - .reshape(experts, 2 * intermediate, reduction) - .contiguous() - ) + return torch.stack((gate, up), dim=2).reshape(experts, 2 * intermediate, reduction).contiguous() def _to_blocked_bytes(scale_2d: torch.Tensor) -> torch.Tensor: @@ -245,12 +211,7 @@ def _to_blocked_bytes(scale_2d: torch.Tensor) -> torch.Tensor: 1, 3, ) - return ( - blocks.reshape(-1, 4, 32, 4) - .transpose(1, 2) - .reshape(-1, 32, 16) - .flatten() - ) + return blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16).flatten() def _stack_blocked_scales(raw_scales: torch.Tensor) -> torch.Tensor: @@ -381,16 +342,8 @@ def weights_have_version_counters( def _request_weight_key( request: ValidatedForwardRequest, ) -> tuple | None: - fc1 = ( - _block_scaled_fingerprint(request.fc1_weight) - if isinstance(request.fc1_weight, BlockScaledTensor) - else _tensor_fingerprint(request.fc1_weight) - ) - fc2 = ( - _block_scaled_fingerprint(request.fc2_weight) - if isinstance(request.fc2_weight, BlockScaledTensor) - else _tensor_fingerprint(request.fc2_weight) - ) + fc1 = _block_scaled_fingerprint(request.fc1_weight) if isinstance(request.fc1_weight, BlockScaledTensor) else _tensor_fingerprint(request.fc1_weight) + fc2 = _block_scaled_fingerprint(request.fc2_weight) if isinstance(request.fc2_weight, BlockScaledTensor) else _tensor_fingerprint(request.fc2_weight) if fc1 is None or fc2 is None: return None return fc1, fc2 @@ -422,16 +375,8 @@ def _prepare_weights( # Retain the source storages while this entry is cached so allocator # pointer reuse cannot produce a false cache hit. self._weight_sources = ( - *( - (request.fc1_weight.data, request.fc1_weight.scale) - if isinstance(request.fc1_weight, BlockScaledTensor) - else (request.fc1_weight,) - ), - *( - (request.fc2_weight.data, request.fc2_weight.scale) - if isinstance(request.fc2_weight, BlockScaledTensor) - else (request.fc2_weight,) - ), + *((request.fc1_weight.data, request.fc1_weight.scale) if isinstance(request.fc1_weight, BlockScaledTensor) else (request.fc1_weight,)), + *((request.fc2_weight.data, request.fc2_weight.scale) if isinstance(request.fc2_weight, BlockScaledTensor) else (request.fc2_weight,)), ) self._weight_refresh_count += 1 return weights @@ -465,10 +410,7 @@ def stage( or fc1_c.shape[1] != config.fc1_out or not fc1_c.is_contiguous() ): - raise ValueError( - "fc1_c buffer must be contiguous BF16 on the request " - f"device with shape (capacity, {config.fc1_out})" - ) + raise ValueError("fc1_c buffer must be contiguous BF16 on the request " f"device with shape (capacity, {config.fc1_out})") elif fc1_c is not None: raise ValueError("generate_c=False must not receive an fc1_c buffer") hidden_sf_columns = (config.hidden + 31) // 32 @@ -514,9 +456,7 @@ def stage( ) if config.enable_col_quant: if col_quant_data_rows <= 0 or col_quant_sf_elements <= 0: - raise ValueError( - "enabled column requant requires positive output capacities" - ) + raise ValueError("enabled column requant requires positive output capacities") col_quant_data = _typed_k_major_view( local["col_quant_data"], _MXFP8_DATA_DTYPE, @@ -529,9 +469,7 @@ def stage( ) else: if col_quant_data_rows != 0 or col_quant_sf_elements != 0: - raise ValueError( - "disabled column requant must not reserve output capacity" - ) + raise ValueError("disabled column requant must not reserve output capacity") col_quant_data = None col_quant_sf = None local_workspace = local["kernel_local_workspace"] @@ -539,13 +477,9 @@ def stage( staged_activation = _as_mxfp8(request.activation) _as_bytes(activation).zero_() - _as_bytes(activation[:token_count]).copy_( - _as_bytes(staged_activation.data) - ) + _as_bytes(activation[:token_count]).copy_(_as_bytes(staged_activation.data)) _as_bytes(activation_sf).zero_() - _as_bytes( - activation_sf[:token_count, :hidden_sf_columns] - ).copy_(_as_bytes(staged_activation.scale)) + _as_bytes(activation_sf[:token_count, :hidden_sf_columns]).copy_(_as_bytes(staged_activation.scale)) _validate_int32_downcast(request.topk_idx) topk_indices.fill_(-1) topk_indices[:token_count].copy_(request.topk_idx) @@ -584,21 +518,12 @@ def stage( or pre_reduced_activation_sf_offset is not None or pre_reduced_activation_sf_bytes_per_token != 0 ): - raise ValueError( - "in-kernel top-k reduction must not receive a " - "standalone pre-reduced activation workspace" - ) + raise ValueError("in-kernel top-k reduction must not receive a " "standalone pre-reduced activation workspace") # output_data is the in-kernel REDG accumulation base and was # cleared above. else: - if ( - pre_reduced_activation_offset is None - or pre_reduced_activation_bytes_per_token <= 0 - ): - raise ValueError( - "standalone top-k reduction requires a pre-reduced " - "activation workspace" - ) + if pre_reduced_activation_offset is None or pre_reduced_activation_bytes_per_token <= 0: + raise ValueError("standalone top-k reduction requires a pre-reduced " "activation workspace") # The kernel writes only valid routes into this persistent combine # plane. Clear the active token rows so dropped routes cannot reuse # contributions from a previous launch. @@ -610,29 +535,16 @@ def stage( ) quantized_combine = config.combine_format != "bf16" if quantized_combine: - if ( - pre_reduced_activation_sf_offset is None - or pre_reduced_activation_sf_bytes_per_token <= 0 - ): - raise ValueError( - "quantized standalone top-k reduction requires a " - "pre-reduced scale workspace" - ) + if pre_reduced_activation_sf_offset is None or pre_reduced_activation_sf_bytes_per_token <= 0: + raise ValueError("quantized standalone top-k reduction requires a " "pre-reduced scale workspace") _zero_workspace_range( shared_workspace, pre_reduced_activation_sf_offset, - token_count - * pre_reduced_activation_sf_bytes_per_token, + token_count * pre_reduced_activation_sf_bytes_per_token, name="pre-reduced activation scale workspace", ) - elif ( - pre_reduced_activation_sf_offset is not None - or pre_reduced_activation_sf_bytes_per_token != 0 - ): - raise ValueError( - "BF16 standalone top-k reduction must not receive a " - "pre-reduced scale workspace" - ) + elif pre_reduced_activation_sf_offset is not None or pre_reduced_activation_sf_bytes_per_token != 0: + raise ValueError("BF16 standalone top-k reduction must not receive a " "pre-reduced scale workspace") weights = self._prepare_weights(request, config) return Mxfp8LaunchInputs( diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py index 8680e37b0..9e4fe27ed 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py @@ -40,9 +40,7 @@ def launch_backward_dglu( _check_overflow(inputs.overflow_flag) return Mxfp8DgluResult( - grad_activation=inputs.output_activation[ - : inputs.token_count - ].float(), + grad_activation=inputs.output_activation[: inputs.token_count].float(), # The dGLU epilogue has already returned source-order dprob through # the symmetric token-communication plane. Own the public result so a # later launch cannot overwrite it. diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py index e6a6040df..3f09f1d6f 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py @@ -52,10 +52,7 @@ class Mxfp8KernelConfig: fc2_tma_stages: int | None = None def __post_init__(self) -> None: - if ( - self.max_recv_size_per_rank is not None - and self.max_recv_size_per_rank <= 0 - ): + if self.max_recv_size_per_rank is not None and self.max_recv_size_per_rank <= 0: raise ValueError("max_recv_size_per_rank must be positive") if self.col_quant_num_ctas <= 0: raise ValueError("col_quant_num_ctas must be positive") @@ -65,19 +62,11 @@ def from_forward_config(cls, config: ForwardConfig) -> "Mxfp8KernelConfig": if config.ep_size < 1: raise ValueError("MXFP8 execution requires a positive EP size") if config.ep_rank < 0 or config.ep_rank >= config.ep_size: - raise ValueError( - f"ep_rank {config.ep_rank} is outside EP size {config.ep_size}" - ) + raise ValueError(f"ep_rank {config.ep_rank} is outside EP size {config.ep_size}") if config.max_tokens_per_rank is None: raise ValueError("MXFP8 execution requires max_tokens_per_rank") - worst_case_recv_size = ( - config.ep_size * config.max_tokens_per_rank * config.top_k - ) - max_recv_size_per_rank = ( - worst_case_recv_size - if config.max_recv_size_per_rank is None - else min(config.max_recv_size_per_rank, worst_case_recv_size) - ) + worst_case_recv_size = config.ep_size * config.max_tokens_per_rank * config.top_k + max_recv_size_per_rank = worst_case_recv_size if config.max_recv_size_per_rank is None else min(config.max_recv_size_per_rank, worst_case_recv_size) if max_recv_size_per_rank <= 0: raise ValueError("max_recv_size_per_rank must be positive") return cls( @@ -94,22 +83,16 @@ def from_forward_config(cls, config: ForwardConfig) -> "Mxfp8KernelConfig": max_recv_size_per_rank=max_recv_size_per_rank, drop_on_overflow=config.drop_on_overflow, combine_format=combine_wire_format(config.combine_format), - enable_col_quant=( - config.backward_wgrad_mode == "operands" - ), + enable_col_quant=(config.backward_wgrad_mode == "operands"), token_padding_block=( - config.token_padding_size - if config.backward_wgrad_mode == "operands" - else 128 if config.generate_c else config.token_padding_size + config.token_padding_size if config.backward_wgrad_mode == "operands" else 128 if config.generate_c else config.token_padding_size ), sf_padding_block=config.sf_padding_size, group_hint=config.tuning.group_hint, token_back_mode=config.tuning.token_back_mode, epi_flag_batch=config.tuning.epi_flag_batch, flag_batch=config.tuning.token_in_flag_batch, - fc2_in_kernel_topk_reduce=( - config.tuning.reduce_topk_in_kernel - ), + fc2_in_kernel_topk_reduce=(config.tuning.reduce_topk_in_kernel), ) @property @@ -126,11 +109,7 @@ def tuning_signature( ) -> tuple[str, tuple[int, int], int, int, bool]: """Return the effective rank-independent transport/scheduler knobs.""" - group_hint = ( - launch_cluster_count - if self.group_hint is None - else self.group_hint - ) + group_hint = launch_cluster_count if self.group_hint is None else self.group_hint return ( self.token_back_mode, self.epi_flag_batch, @@ -142,11 +121,7 @@ def tuning_signature( def effective_config(self, launch_cluster_count: int) -> dict[str, object]: """Return the complete JSON-safe compile-time configuration.""" - effective_group_hint = ( - launch_cluster_count - if self.group_hint is None - else self.group_hint - ) + effective_group_hint = launch_cluster_count if self.group_hint is None else self.group_hint return { "num_experts_per_rank": self.num_experts, "world_size": self.world_size, diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_fingerprint.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_fingerprint.py index d7600d7b4..1dd1dbf31 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_fingerprint.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_fingerprint.py @@ -12,7 +12,6 @@ from pathlib import Path from typing import Any - FINGERPRINT_SCHEMA_VERSION = 1 _KERNEL_IDENTITY_FIELDS = ( "kernel_name", @@ -38,12 +37,7 @@ def canonical_json_sha256(value: object) -> str: def kernel_identity_sha256(fingerprint: dict[str, Any]) -> str: """Hash fields that must match across AOT and in-process JIT paths.""" - return canonical_json_sha256( - { - field: fingerprint.get(field) - for field in _KERNEL_IDENTITY_FIELDS - } - ) + return canonical_json_sha256({field: fingerprint.get(field) for field in _KERNEL_IDENTITY_FIELDS}) def source_tree_sha256(root: Path) -> str: @@ -53,9 +47,7 @@ def source_tree_sha256(root: Path) -> str: if not resolved.is_dir(): raise RuntimeError(f"kernel source tree does not exist: {resolved}") digest = hashlib.sha256() - sources = sorted( - path for path in resolved.rglob("*.py") if path.is_file() - ) + sources = sorted(path for path in resolved.rglob("*.py") if path.is_file()) if not sources: raise RuntimeError(f"kernel source tree contains no Python files: {resolved}") for path in sources: @@ -84,13 +76,15 @@ def _cutlass_version() -> str: def _json_layout_signature(signature: tuple) -> list[object]: return [ - None - if entry is None - else { - "shape": list(entry[0]), - "stride": list(entry[1]), - "dtype": str(entry[2]), - } + ( + None + if entry is None + else { + "shape": list(entry[0]), + "stride": list(entry[1]), + "dtype": str(entry[2]), + } + ) for entry in signature ] @@ -105,9 +99,7 @@ def build_kernel_fingerprint( kernel = prepared.kernel source_root = Path(__file__).resolve().parents[1] / "cutedsl_src" - effective_config = prepared.config.effective_config( - prepared.launch_cluster_count - ) + effective_config = prepared.config.effective_config(prepared.launch_cluster_count) launch_geometry = { "grid": [ prepared.config.cluster_shape_mnk[0], @@ -117,9 +109,7 @@ def build_kernel_fingerprint( "block": [int(kernel.threads_per_cta), 1, 1], "cluster": list(prepared.config.cluster_shape_mnk), "min_blocks_per_mp": int(getattr(kernel, "occupancy", 1)), - "dynamic_shared_memory_bytes": int( - getattr(kernel, "smem_capacity", 0) - ), + "dynamic_shared_memory_bytes": int(getattr(kernel, "smem_capacity", 0)), } layout = _json_layout_signature(layout_signature) fingerprint = { @@ -134,9 +124,7 @@ def build_kernel_fingerprint( "layout_signature_sha256": canonical_json_sha256(layout), "compiled_binary_sha256": compiled_binary_sha256, } - fingerprint["kernel_identity_sha256"] = kernel_identity_sha256( - fingerprint - ) + fingerprint["kernel_identity_sha256"] = kernel_identity_sha256(fingerprint) fingerprint["fingerprint_sha256"] = canonical_json_sha256(fingerprint) return fingerprint diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_formats.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_formats.py index c3d4df4f4..6529ee4f3 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_formats.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_formats.py @@ -7,7 +7,6 @@ from ..._types import MoeFormat, parse_format - _COMBINE_WIRE_FORMATS = { MoeFormat.BF16: "bf16", MoeFormat.MXFP8: "32e4m3xe8m0", diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_launch.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_launch.py index 98d9c24e1..a5dfbb032 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_launch.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_launch.py @@ -28,9 +28,7 @@ def _to_cute( ) if not dynamic_layout: return cute_tensor - return cute_tensor.mark_layout_dynamic( - leading_dim=cutlass_torch.get_leading_dim(tensor) - ) + return cute_tensor.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(tensor)) def _to_cute_ptr(tensor: torch.Tensor, assumed_align: int = 128): @@ -42,10 +40,7 @@ def _to_cute_ptr(tensor: torch.Tensor, assumed_align: int = 128): address = int(tensor.data_ptr()) if address % assumed_align: - raise ValueError( - f"Rubin workspace address {address:#x} is not " - f"{assumed_align}-byte aligned" - ) + raise ValueError(f"Rubin workspace address {address:#x} is not " f"{assumed_align}-byte aligned") return make_ptr( cutlass.Uint8, address, @@ -71,11 +66,7 @@ def build_runtime_kwargs( "fc1_weight_sf": _to_cute(weights.fc1_weight_sf), "fc2_weight": _to_cute(weights.fc2_weight), "fc2_weight_sf": _to_cute(weights.fc2_weight_sf), - "fc1_c": ( - None - if inputs.fc1_c is None - else _to_cute(inputs.fc1_c, dynamic_layout=False) - ), + "fc1_c": (None if inputs.fc1_c is None else _to_cute(inputs.fc1_c, dynamic_layout=False)), "output_activation": _to_cute(inputs.output_data), "col_quant_data": ( None @@ -101,9 +92,7 @@ def build_runtime_kwargs( ), "local_workspace": _to_cute_ptr(inputs.local_workspace), "shared_workspace": _to_cute_ptr(inputs.shared_workspace), - "peer_rank_ptr_mapper_host": ( - resources.workspace.peer_mapping.to_sym_buffer_host() - ), + "peer_rank_ptr_mapper_host": (resources.workspace.peer_mapping.to_sym_buffer_host()), "stream": cuda.CUstream(stream.cuda_stream), } return kwargs @@ -127,28 +116,17 @@ def layout_signature(inputs: Mxfp8LaunchInputs) -> tuple: inputs.local_workspace, inputs.shared_workspace, ) - return tuple( - None - if tensor is None - else (tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype) - for tensor in tensors - ) + return tuple(None if tensor is None else (tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype) for tensor in tensors) def _check_overflow(overflow_flag: torch.Tensor) -> None: - message = ( - "Rubin MegaMoE receive route-pool overflow; the output is invalid for " - "this routing distribution" - ) + message = "Rubin MegaMoE receive route-pool overflow; the output is invalid for " "this routing distribution" assert_async = getattr(torch, "_assert_async", None) if assert_async is not None: assert_async(overflow_flag == 0, message) return if torch.cuda.is_current_stream_capturing(): - raise NotImplementedError( - "CUDA graph capture requires torch._assert_async to surface " - "Rubin MegaMoE overflow" - ) + raise NotImplementedError("CUDA graph capture requires torch._assert_async to surface " "Rubin MegaMoE overflow") # Compatibility fallback for PyTorch builds without a device-side assert. value = int(overflow_flag.item()) if value != 0: diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py index f90890b3d..0a48f81a9 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py @@ -54,11 +54,7 @@ def _activation_views( capacity: int, hidden: int, ) -> tuple[torch.Tensor, torch.Tensor]: - workspace = ( - execution.backward.workspace - if backward - else execution.forward.workspace - ) + workspace = execution.backward.workspace if backward else execution.forward.workspace return ( _typed_view( workspace.symmetric["activation_data"], @@ -79,9 +75,7 @@ def _write_expert_offsets( ) -> None: snapshot = execution.forward_expert_size_snapshot if snapshot is None: - raise RuntimeError( - "training forward requires the persistent expert-size snapshot" - ) + raise RuntimeError("training forward requires the persistent expert-size snapshot") counts = execution.slot.valid_route_counts offsets = execution.slot.expert_offsets counts.copy_(snapshot) @@ -159,9 +153,7 @@ def launch_training_forward( ) _runtime_debug("training-forward.compile.end", slot=execution.slot.index) _runtime_debug("training-forward.launch.begin", slot=execution.slot.index) - compiled.callable( - **build_runtime_kwargs(inputs, execution.forward) - ) + compiled.callable(**build_runtime_kwargs(inputs, execution.forward)) _runtime_debug("training-forward.launch.end", slot=execution.slot.index) _runtime_debug("training-forward.offsets.begin", slot=execution.slot.index) _write_expert_offsets(execution, config.token_padding_block) @@ -256,9 +248,7 @@ def launch_training_backward( ) _runtime_debug("training-backward.compile.end", slot=execution.slot.index) _runtime_debug("training-backward.launch.begin", slot=execution.slot.index) - compiled.callable( - **build_backward_runtime_kwargs(inputs, execution.backward) - ) + compiled.callable(**build_backward_runtime_kwargs(inputs, execution.backward)) _runtime_debug("training-backward.launch.end", slot=execution.slot.index) slot.grad_activation.copy_(slot.backward_output) _runtime_debug("training-backward.wgrad-export.begin", slot=execution.slot.index) diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py index 873e08e68..4a5e1427a 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py @@ -42,22 +42,15 @@ from ._training_weights import Mxfp8TrainingWeightBindings from ._training_wgrad import Mxfp8TrainingWgradExporter - _DATA_DTYPE = torch.float8_e4m3fn _SCALE_DTYPE = torch.float8_e8m0fnu _ROUTING_SYMMETRIC = frozenset({"topk_weights"}) _ROUTING_LOCAL = frozenset({"topk_idx"}) _FORWARD_SLOT_SYMMETRIC = frozenset({"output_data", *_ROUTING_SYMMETRIC}) -_FORWARD_SLOT_LOCAL = frozenset( - {"overflow_flag", "col_quant_data", "col_quant_sf", *_ROUTING_LOCAL} -) -_BACKWARD_SLOT_SYMMETRIC = frozenset( - {"output_data", "backward_dprob", *_ROUTING_SYMMETRIC} -) -_BACKWARD_SLOT_LOCAL = frozenset( - {"overflow_flag", "backward_aux_data", "backward_aux_scale", *_ROUTING_LOCAL} -) +_FORWARD_SLOT_LOCAL = frozenset({"overflow_flag", "col_quant_data", "col_quant_sf", *_ROUTING_LOCAL}) +_BACKWARD_SLOT_SYMMETRIC = frozenset({"output_data", "backward_dprob", *_ROUTING_SYMMETRIC}) +_BACKWARD_SLOT_LOCAL = frozenset({"overflow_flag", "backward_aux_data", "backward_aux_scale", *_ROUTING_LOCAL}) def _round_up(value: int, multiple: int) -> int: @@ -106,11 +99,7 @@ def _region_map( requirements: WorkspaceRequirements, space: str, ) -> dict[str, BufferRegion]: - regions = ( - requirements.symmetric_regions - if space == "symmetric" - else requirements.local_regions - ) + regions = requirements.symmetric_regions if space == "symmetric" else requirements.local_regions return {region.name: region for region in regions} @@ -122,9 +111,7 @@ def _required_region( try: return _region_map(requirements, space)[name] except KeyError as exc: - raise ValueError( - f"{space} workspace requirements do not contain {name!r}" - ) from exc + raise ValueError(f"{space} workspace requirements do not contain {name!r}") from exc def _add_lane_regions( @@ -136,17 +123,11 @@ def _add_lane_regions( space: str, slot_names: frozenset[str], ) -> None: - regions = ( - requirements.symmetric_regions - if space == "symmetric" - else requirements.local_regions - ) + regions = requirements.symmetric_regions if space == "symmetric" else requirements.local_regions for region in regions: if region.name in slot_names: continue - if phase == "backward" and space == "local" and region.name == ( - "backward_fc1_preact" - ): + if phase == "backward" and space == "local" and region.name == ("backward_fc1_preact"): # The graph path aliases forward's raw receiver pool directly. continue output.append( @@ -173,11 +154,7 @@ def build_training_workspace_requirements( if not config.generate_c: raise ValueError("training resources require generate_c=True") if forward.pool_token_capacity != backward.pool_token_capacity: - raise ValueError( - "forward/backward pool capacities must match, got " - f"{forward.pool_token_capacity} and " - f"{backward.pool_token_capacity}" - ) + raise ValueError("forward/backward pool capacities must match, got " f"{forward.pool_token_capacity} and " f"{backward.pool_token_capacity}") forward_requirements = forward.workspace_requirements backward_requirements = backward.workspace_requirements @@ -246,10 +223,7 @@ def build_training_workspace_requirements( forward_local = _region_map(forward_requirements, "local") backward_symmetric = _region_map(backward_requirements, "symmetric") backward_local = _region_map(backward_requirements, "local") - fc1_c_shape = tuple( - int(extent) - for extent in forward.kernel.get_aux_output_shapes()["fc1_c"] - ) + fc1_c_shape = tuple(int(extent) for extent in forward.kernel.get_aux_output_shapes()["fc1_c"]) fc1_c_bytes = math.prod(fc1_c_shape) * torch.bfloat16.itemsize backward_preact = _required_region( backward_requirements, @@ -257,14 +231,8 @@ def build_training_workspace_requirements( "backward_fc1_preact", ) if fc1_c_bytes != backward_preact.nbytes: - raise ValueError( - "forward fc1_c and backward preactivation byte sizes differ: " - f"{fc1_c_bytes} != {backward_preact.nbytes}" - ) - aux_shapes = { - name: tuple(int(extent) for extent in shape) - for name, shape in backward.kernel.get_aux_output_shapes().items() - } + raise ValueError("forward fc1_c and backward preactivation byte sizes differ: " f"{fc1_c_bytes} != {backward_preact.nbytes}") + aux_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in backward.kernel.get_aux_output_shapes().items()} aux_dtypes = { "fc1_recompute": _DATA_DTYPE, "fc1_recompute_sf": _SCALE_DTYPE, @@ -349,9 +317,7 @@ def build_training_workspace_requirements( ), BufferRegion( _custom_slot_name(slot, "routing_topk_idx"), - int(config.max_tokens_per_rank) - * config.top_k - * torch.int32.itemsize, + int(config.max_tokens_per_rank) * config.top_k * torch.int32.itemsize, alignment=16, ), BufferRegion( @@ -366,9 +332,7 @@ def build_training_workspace_requirements( ), BufferRegion( _custom_slot_name(slot, "grad_activation"), - int(config.max_tokens_per_rank) - * config.hidden_size - * torch.float32.itemsize, + int(config.max_tokens_per_rank) * config.hidden_size * torch.float32.itemsize, alignment=16, ), ) @@ -376,9 +340,7 @@ def build_training_workspace_requirements( symmetric_regions.append( BufferRegion( _custom_slot_symmetric_name(slot, "routing_topk_weights"), - int(config.max_tokens_per_rank) - * config.top_k - * torch.float32.itemsize, + int(config.max_tokens_per_rank) * config.top_k * torch.float32.itemsize, alignment=16, ) ) @@ -423,14 +385,9 @@ def _harmonize_symmetric_regions( dist.all_reduce(minimum_count, op=dist.ReduceOp.MIN, group=runtime.group) dist.all_reduce(maximum_count, op=dist.ReduceOp.MAX, group=runtime.group) if int(minimum_count.item()) != int(maximum_count.item()): - raise RuntimeError( - "symmetric workspace region counts differ across EP ranks: " - f"min={int(minimum_count.item())}, max={int(maximum_count.item())}" - ) + raise RuntimeError("symmetric workspace region counts differ across EP ranks: " f"min={int(minimum_count.item())}, max={int(maximum_count.item())}") - metadata = "\0".join( - f"{region.name}:{region.alignment}" for region in regions - ).encode() + metadata = "\0".join(f"{region.name}:{region.alignment}" for region in regions).encode() signature_value = int.from_bytes( hashlib.blake2b(metadata, digest_size=8).digest(), "little", @@ -470,9 +427,7 @@ def _harmonize_symmetric_regions( dist.all_reduce(maximum_sizes, op=dist.ReduceOp.MAX, group=runtime.group) harmonized_sizes = tuple(int(value) for value in maximum_sizes.cpu().tolist()) changes = tuple( - f"{region.name}:{region.nbytes}->{harmonized_size}" - for region, harmonized_size in zip(regions, harmonized_sizes) - if region.nbytes != harmonized_size + f"{region.name}:{region.nbytes}->{harmonized_size}" for region, harmonized_size in zip(regions, harmonized_sizes) if region.nbytes != harmonized_size ) _runtime_debug( "training-resources.symmetric-layout-harmonized", @@ -537,9 +492,7 @@ def _prepared_kernel_abi(prepared) -> dict[str, object]: return { "name": str(kernel.name()), "architecture": list(prepared.architecture), - "effective_config": prepared.config.effective_config( - prepared.launch_cluster_count - ), + "effective_config": prepared.config.effective_config(prepared.launch_cluster_count), "launch": { "cluster_count": int(prepared.launch_cluster_count), "threads_per_cta": int(kernel.threads_per_cta), @@ -590,9 +543,7 @@ def _build_training_abi_facts( "intermediate": int(config.intermediate_size), "top_k": int(config.top_k), "max_tokens_per_rank": int(config.max_tokens_per_rank), - "max_recv_size_per_rank": int( - forward.config.max_recv_size_per_rank - ), + "max_recv_size_per_rank": int(forward.config.max_recv_size_per_rank), }, "policy": { "drop_on_overflow": bool(config.drop_on_overflow), @@ -633,9 +584,7 @@ def _verify_training_abi_across_ranks( rank_digests: list[Any] = [None] * runtime.world_size dist.all_gather_object(rank_digests, digest, group=runtime.group) raise RuntimeError( - "MoeEp training ABI differs across expert-parallel ranks before " - "workspace allocation: " - f"digests={rank_digests}, local_facts={facts}" + "MoeEp training ABI differs across expert-parallel ranks before " "workspace allocation: " f"digests={rank_digests}, local_facts={facts}" ) @@ -735,12 +684,7 @@ def __init__( @property def prepared(self) -> bool: - return ( - not self._closed - and self._runtime is not None - and self._workspace is not None - and self._workspace.allocated - ) + return not self._closed and self._runtime is not None and self._workspace is not None and self._workspace.allocated def prepare(self) -> None: with self._lock: @@ -749,21 +693,17 @@ def prepare(self) -> None: if self.prepared: return if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "training resources must be prepared before CUDA graph capture" - ) + raise RuntimeError("training resources must be prepared before CUDA graph capture") _runtime_debug( "training-resources.prepare.begin", slot_count=self.slot_count, lane_count=self.lane_count, - local_bytes=self.requirements.local_layout.total_bytes - if hasattr(self.requirements, "local_layout") - else sum( - region.nbytes for region in self.requirements.local_regions - ), - symmetric_bytes=sum( - region.nbytes for region in self.requirements.symmetric_regions + local_bytes=( + self.requirements.local_layout.total_bytes + if hasattr(self.requirements, "local_layout") + else sum(region.nbytes for region in self.requirements.local_regions) ), + symmetric_bytes=sum(region.nbytes for region in self.requirements.symmetric_regions), ) _runtime_debug("training-resources.runtime-acquire.begin") runtime = self._runtime_manager.acquire(self.config, self.device) @@ -777,13 +717,9 @@ def prepare(self) -> None: ) self._runtime = runtime try: - layout_watchdog = _RuntimeWatchdog( - "training-resources.symmetric-layout-harmonize" - ) + layout_watchdog = _RuntimeWatchdog("training-resources.symmetric-layout-harmonize") layout_watchdog.start() - _runtime_debug( - "training-resources.symmetric-layout-harmonize.begin" - ) + _runtime_debug("training-resources.symmetric-layout-harmonize.begin") try: self.requirements = _harmonize_symmetric_regions( self.requirements, @@ -792,13 +728,9 @@ def prepare(self) -> None: ) finally: layout_watchdog.close() - _runtime_debug( - "training-resources.symmetric-layout-harmonize.end" - ) + _runtime_debug("training-resources.symmetric-layout-harmonize.end") if runtime.world_size > 1: - abi_watchdog = _RuntimeWatchdog( - "training-resources.abi-handshake" - ) + abi_watchdog = _RuntimeWatchdog("training-resources.abi-handshake") abi_watchdog.start() _runtime_debug("training-resources.abi-handshake.begin") try: @@ -811,12 +743,10 @@ def prepare(self) -> None: slot_count=self.slot_count, lane_count=self.lane_count, ) - self._abi_fingerprint = ( - _verify_training_abi_across_ranks( - abi_facts, - runtime, - self.device, - ) + self._abi_fingerprint = _verify_training_abi_across_ranks( + abi_facts, + runtime, + self.device, ) finally: abi_watchdog.close() @@ -837,9 +767,7 @@ def prepare(self) -> None: symmetric_bytes=workspace.symmetric_layout.total_bytes, ) self._workspace = workspace - allocation_watchdog = _RuntimeWatchdog( - "training-resources.workspace-allocate" - ) + allocation_watchdog = _RuntimeWatchdog("training-resources.workspace-allocate") allocation_watchdog.start() try: workspace.ensure_allocated() @@ -850,22 +778,16 @@ def prepare(self) -> None: # Symmetric-root zeroing is asynchronous. No rank may # enter the first device barrier until every peer has # completed allocation and root initialization. - stream_watchdog = _RuntimeWatchdog( - "training-resources.stream-synchronize" - ) + stream_watchdog = _RuntimeWatchdog("training-resources.stream-synchronize") stream_watchdog.start() - _runtime_debug( - "training-resources.stream-synchronize.begin" - ) + _runtime_debug("training-resources.stream-synchronize.begin") try: torch.cuda.current_stream(self.device).synchronize() finally: stream_watchdog.close() _runtime_debug("training-resources.stream-synchronize.end") - barrier_watchdog = _RuntimeWatchdog( - "training-resources.rank-barrier" - ) + barrier_watchdog = _RuntimeWatchdog("training-resources.rank-barrier") barrier_watchdog.start() _runtime_debug("training-resources.rank-barrier.begin") try: @@ -898,44 +820,24 @@ def _phase_workspace( ) -> WorkspaceViews: symmetric = {} local = {} - slot_symmetric = ( - _FORWARD_SLOT_SYMMETRIC - if phase == "forward" - else _BACKWARD_SLOT_SYMMETRIC - ) - slot_local = ( - _FORWARD_SLOT_LOCAL - if phase == "forward" - else _BACKWARD_SLOT_LOCAL - ) + slot_symmetric = _FORWARD_SLOT_SYMMETRIC if phase == "forward" else _BACKWARD_SLOT_SYMMETRIC + slot_local = _FORWARD_SLOT_LOCAL if phase == "forward" else _BACKWARD_SLOT_LOCAL for region in requirements.symmetric_regions: if region.name in _ROUTING_SYMMETRIC: - symmetric[region.name] = flat.symmetric[ - _custom_slot_symmetric_name(slot, "routing_topk_weights") - ] + symmetric[region.name] = flat.symmetric[_custom_slot_symmetric_name(slot, "routing_topk_weights")] continue scope_name = ( - _slot_name(slot, phase, "symmetric", region.name) - if region.name in slot_symmetric - else _lane_name(lane, phase, "symmetric", region.name) + _slot_name(slot, phase, "symmetric", region.name) if region.name in slot_symmetric else _lane_name(lane, phase, "symmetric", region.name) ) symmetric[region.name] = flat.symmetric[scope_name] for region in requirements.local_regions: if region.name in _ROUTING_LOCAL: - local[region.name] = flat.local[ - _custom_slot_name(slot, "routing_topk_idx") - ] + local[region.name] = flat.local[_custom_slot_name(slot, "routing_topk_idx")] continue if phase == "backward" and region.name == "backward_fc1_preact": - local[region.name] = flat.local[ - _custom_slot_name(slot, "fc1_preact") - ] + local[region.name] = flat.local[_custom_slot_name(slot, "fc1_preact")] continue - scope_name = ( - _slot_name(slot, phase, "local", region.name) - if region.name in slot_local - else _lane_name(lane, phase, "local", region.name) - ) + scope_name = _slot_name(slot, phase, "local", region.name) if region.name in slot_local else _lane_name(lane, phase, "local", region.name) local[region.name] = flat.local[scope_name] return WorkspaceViews( token_count=flat.token_count, @@ -951,17 +853,9 @@ def _slot_views( ) -> Mxfp8TrainingSlotViews: config = self.config capacity = int(config.max_tokens_per_rank) - fwd_shapes = { - name: tuple(int(extent) for extent in shape) - for name, shape in self.forward_prepared.kernel.get_aux_output_shapes().items() - } - bwd_shapes = { - name: tuple(int(extent) for extent in shape) - for name, shape in self.backward_prepared.kernel.get_aux_output_shapes().items() - } - scale_columns = _align_scale_columns( - self.forward_prepared.pool_token_capacity - ) + fwd_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in self.forward_prepared.kernel.get_aux_output_shapes().items()} + bwd_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in self.backward_prepared.kernel.get_aux_output_shapes().items()} + scale_columns = _align_scale_columns(self.forward_prepared.pool_token_capacity) def local_bytes(name: str) -> torch.Tensor: return flat.local[_custom_slot_name(slot, name)] @@ -1175,13 +1069,9 @@ def views( ) -> Mxfp8TrainingExecutionViews: with self._lock: if slot < 0 or slot >= self.slot_count: - raise ValueError( - f"slot {slot} is outside [0, {self.slot_count})" - ) + raise ValueError(f"slot {slot} is outside [0, {self.slot_count})") if lane < 0 or lane >= self.lane_count: - raise ValueError( - f"lane {lane} is outside [0, {self.lane_count})" - ) + raise ValueError(f"lane {lane} is outside [0, {self.lane_count})") flat = self._flat_views(token_count) forward_workspace = self._phase_workspace( flat, @@ -1199,9 +1089,7 @@ def views( ) snapshot = None if self.forward_prepared.col_quant_sizes_offset is not None: - snapshot_bytes = forward_workspace.local[ - "kernel_local_workspace" - ].narrow( + snapshot_bytes = forward_workspace.local["kernel_local_workspace"].narrow( 0, self.forward_prepared.col_quant_sizes_offset, self.forward_prepared.col_quant_sizes_bytes, @@ -1259,9 +1147,7 @@ def finalize_overflow( raise ValueError("finalize_overflow slots must be unique") for slot in slots: if slot < 0 or slot >= self.slot_count: - raise ValueError( - f"slot {slot} is outside [0, {self.slot_count})" - ) + raise ValueError(f"slot {slot} is outside [0, {self.slot_count})") if lane < 0 or lane >= self.lane_count: raise ValueError(f"lane {lane} is outside [0, {self.lane_count})") flat = self._flat_views(0) @@ -1303,10 +1189,7 @@ def finalize_overflow( if not self.config.drop_on_overflow: assert_async = getattr(torch, "_assert_async", None) if assert_async is None: - raise RuntimeError( - "drop_on_overflow=False training resources require " - "torch._assert_async" - ) + raise RuntimeError("drop_on_overflow=False training resources require " "torch._assert_async") overflow_ok = _typed_view( flat.local[ _lane_name( @@ -1322,8 +1205,7 @@ def finalize_overflow( torch.eq(global_overflow, 0, out=overflow_ok) assert_async( overflow_ok, - "Rubin MegaMoE receive route-pool overflow; " - "the fixed-slot outputs are invalid", + "Rubin MegaMoE receive route-pool overflow; " "the fixed-slot outputs are invalid", ) return global_overflow diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py index 16a05501d..a410e1740 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py @@ -32,14 +32,9 @@ def _validate( output_topk_weights: torch.Tensor, ) -> int: if source.dtype not in (torch.bfloat16, torch.float32): - raise TypeError( - "training staging source must be BF16 or FP32, " - f"got {source.dtype}" - ) + raise TypeError("training staging source must be BF16 or FP32, " f"got {source.dtype}") if source.ndim != 2 or source.shape[1] != self.hidden: - raise ValueError( - f"training staging source must have shape (T, {self.hidden})" - ) + raise ValueError(f"training staging source must have shape (T, {self.hidden})") if not source.is_contiguous(): raise ValueError("training staging source must be contiguous") token_count = int(source.shape[0]) @@ -49,27 +44,12 @@ def _validate( raise TypeError("training staging topk_idx must be contiguous Int32") if topk_weights.shape != topk_idx.shape: raise ValueError("training staging topk_weights shape mismatch") - if ( - topk_weights.dtype is not torch.float32 - or not topk_weights.is_contiguous() - ): - raise TypeError( - "training staging topk_weights must be contiguous FP32" - ) - if ( - output.dtype is not torch.float8_e4m3fn - or output.ndim != 2 - or output.shape[1] != self.hidden - or not output.is_contiguous() - ): - raise ValueError( - "training staging output must be contiguous E4M3 " - f"(capacity, {self.hidden})" - ) + if topk_weights.dtype is not torch.float32 or not topk_weights.is_contiguous(): + raise TypeError("training staging topk_weights must be contiguous FP32") + if output.dtype is not torch.float8_e4m3fn or output.ndim != 2 or output.shape[1] != self.hidden or not output.is_contiguous(): + raise ValueError("training staging output must be contiguous E4M3 " f"(capacity, {self.hidden})") if token_count > output.shape[0]: - raise ValueError( - f"token count {token_count} exceeds capacity {output.shape[0]}" - ) + raise ValueError(f"token count {token_count} exceeds capacity {output.shape[0]}") logical_sf_columns = self.hidden // 32 if ( output_sf.dtype is not torch.float8_e8m0fnu @@ -83,11 +63,7 @@ def _validate( ("output_topk_idx", output_topk_idx, torch.int32), ("output_topk_weights", output_topk_weights, torch.float32), ): - if ( - tensor.shape != (output.shape[0], self.top_k) - or tensor.dtype is not dtype - or not tensor.is_contiguous() - ): + if tensor.shape != (output.shape[0], self.top_k) or tensor.dtype is not dtype or not tensor.is_contiguous(): raise ValueError(f"training staging {name} has an invalid ABI") devices = { source.device, @@ -124,20 +100,10 @@ def stage( output_topk_weights, ) output_sf.zero_() - routing_in_place = ( - topk_idx.data_ptr() == output_topk_idx.data_ptr() - and topk_weights.data_ptr() == output_topk_weights.data_ptr() - ) - routing_partially_aliased = ( - topk_idx.data_ptr() == output_topk_idx.data_ptr() - ) != ( - topk_weights.data_ptr() == output_topk_weights.data_ptr() - ) + routing_in_place = topk_idx.data_ptr() == output_topk_idx.data_ptr() and topk_weights.data_ptr() == output_topk_weights.data_ptr() + routing_partially_aliased = (topk_idx.data_ptr() == output_topk_idx.data_ptr()) != (topk_weights.data_ptr() == output_topk_weights.data_ptr()) if routing_partially_aliased: - raise ValueError( - "training staging routing inputs must either both alias " - "their outputs or neither alias" - ) + raise ValueError("training staging routing inputs must either both alias " "their outputs or neither alias") if not routing_in_place: output_topk_idx.fill_(-1) output_topk_weights.zero_() @@ -182,10 +148,7 @@ def stage( compiled = self._compiled.get(key) if compiled is None: if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "MXFP8 training stager must be compiled before " - "CUDA graph capture" - ) + raise RuntimeError("MXFP8 training stager must be compiled before " "CUDA graph capture") import cutlass.cute as cute from ._training_stage_kernel import Mxfp8TrainingStageKernel diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py index c92cf6035..f570b9b3d 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py @@ -25,14 +25,9 @@ def __init__(self, hidden: int, top_k: int) -> None: self.hidden = int(hidden) self.top_k = int(top_k) if self.hidden <= 0 or self.hidden % self._sf_vec: - raise ValueError( - "MXFP8 training stage requires hidden divisible by 32" - ) + raise ValueError("MXFP8 training stage requires hidden divisible by 32") if self.top_k <= 0 or self.top_k > self._threads_per_cta: - raise ValueError( - "MXFP8 training stage requires " - f"1 <= top_k <= {self._threads_per_cta}" - ) + raise ValueError("MXFP8 training stage requires " f"1 <= top_k <= {self._threads_per_cta}") @cute.jit def __call__( @@ -78,9 +73,7 @@ def _kernel( sf_vec: cutlass.Constexpr[int] = self._sf_vec threads: cutlass.Constexpr[int] = self._threads_per_cta block_count: cutlass.Constexpr[int] = hidden // sf_vec - rounds: cutlass.Constexpr[int] = ( - block_count + threads - 1 - ) // threads + rounds: cutlass.Constexpr[int] = (block_count + threads - 1) // threads for block_round in cutlass.range_constexpr(rounds): block = tid + Int32(block_round * threads) @@ -126,9 +119,7 @@ def _kernel( if tid < Int32(self.top_k): output_topk_idx[token, tid] = Int32(topk_idx[token, tid]) - output_topk_weights[token, tid] = Float32( - topk_weights[token, tid] - ) + output_topk_weights[token, tid] = Float32(topk_weights[token, tid]) __all__ = ["Mxfp8TrainingStageKernel"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py index afb693534..623bf50e6 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py @@ -19,9 +19,7 @@ def _round_up(value: int, multiple: int) -> int: def _empty_k_major_like(tensor: torch.Tensor) -> torch.Tensor: if tensor.ndim != 3: - raise ValueError( - f"K-major training weight must be rank 3, got {tensor.ndim}" - ) + raise ValueError(f"K-major training weight must be rank 3, got {tensor.ndim}") experts, reduction, output = tensor.shape return torch.empty_strided( tensor.shape, @@ -52,9 +50,7 @@ def _copy_k_major( source: torch.Tensor, ) -> None: if target.shape != source.shape: - raise ValueError( - f"K-major copy shape mismatch: {target.shape} != {source.shape}" - ) + raise ValueError(f"K-major copy shape mismatch: {target.shape} != {source.shape}") target.copy_(source) @@ -131,28 +127,25 @@ def _copy_blocked_scales_plain( experts = source.shape[0] if tuple(source.shape) != (experts, raw_columns, raw_rows): - raise ValueError( - "plain training scale shape mismatch: " - f"{tuple(source.shape)} != " - f"{(experts, raw_columns, raw_rows)}" - ) + raise ValueError("plain training scale shape mismatch: " f"{tuple(source.shape)} != " f"{(experts, raw_columns, raw_rows)}") if raw_rows % 128 or raw_columns % 4: - raise ValueError( - "training scale pack requires rows divisible by 128 and " - "columns divisible by 4" - ) + raise ValueError("training scale pack requires rows divisible by 128 and " "columns divisible by 4") row_blocks = raw_rows // 128 column_blocks = raw_columns // 4 - source_view = source.view( - torch.uint8, - ).view( - experts, - column_blocks, - 4, - row_blocks, - 4, - 32, - ).permute(0, 3, 1, 5, 4, 2) + source_view = ( + source.view( + torch.uint8, + ) + .view( + experts, + column_blocks, + 4, + row_blocks, + 4, + 32, + ) + .permute(0, 3, 1, 5, 4, 2) + ) target.view(torch.uint8).view( experts, row_blocks, @@ -178,21 +171,22 @@ def _copy_blocked_scales_gate_up_rows( if tuple(source.shape) != (experts, raw_columns, raw_rows): raise ValueError("forward FC1 training scale shape mismatch") if intermediate % 64 or raw_columns % 4: - raise ValueError( - "forward FC1 scale pack requires intermediate divisible by 64 " - "and reduction blocks divisible by 4" - ) + raise ValueError("forward FC1 scale pack requires intermediate divisible by 64 " "and reduction blocks divisible by 4") row_blocks = raw_rows // 128 column_blocks = raw_columns // 4 - source_view = source.view(torch.uint8).view( - experts, - column_blocks, - 4, - 2, - row_blocks, - 2, - 32, - ).permute(0, 4, 1, 6, 5, 3, 2) + source_view = ( + source.view(torch.uint8) + .view( + experts, + column_blocks, + 4, + 2, + row_blocks, + 2, + 32, + ) + .permute(0, 4, 1, 6, 5, 3, 2) + ) target.view(torch.uint8).view( experts, row_blocks, @@ -222,21 +216,22 @@ def _copy_blocked_scales_gate_up_columns( ): raise ValueError("backward W1-transpose training scale shape mismatch") if output % 128 or reduction_blocks % 2: - raise ValueError( - "backward W1-transpose scale pack requires output divisible by " - "128 and intermediate divisible by 64" - ) + raise ValueError("backward W1-transpose scale pack requires output divisible by " "128 and intermediate divisible by 64") row_blocks = output // 128 column_blocks = reduction_blocks // 2 - source_view = source.view(torch.uint8).view( - experts, - 2, - column_blocks, - 2, - row_blocks, - 4, - 32, - ).permute(0, 4, 2, 6, 5, 3, 1) + source_view = ( + source.view(torch.uint8) + .view( + experts, + 2, + column_blocks, + 2, + row_blocks, + 4, + 32, + ) + .permute(0, 4, 2, 6, 5, 3, 1) + ) target.view(torch.uint8).view( experts, row_blocks, diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py index 0a1fd851e..9f6687bb9 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py @@ -91,10 +91,7 @@ def _expand_scales( compiled = self._compiled.get(key) if compiled is None: if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "WGrad scale expansion must be compiled before " - "CUDA graph capture" - ) + raise RuntimeError("WGrad scale expansion must be compiled before " "CUDA graph capture") import cutlass.cute as cute from ._training_wgrad_kernel import ( diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py index c30a4f4f3..3e2909b31 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py @@ -30,26 +30,14 @@ def __init__( self.non_k_size = int(non_k_size) self.expert_count = int(expert_count) self.source_sf_padding = int(source_sf_padding) - self.deinterleave_gate_up = ( - None - if deinterleave_gate_up is None - else int(deinterleave_gate_up) - ) + self.deinterleave_gate_up = None if deinterleave_gate_up is None else int(deinterleave_gate_up) if self.non_k_size <= 0 or self.non_k_size % 128: raise ValueError("WGrad scale non-K size must be divisible by 128") if self.expert_count <= 0: raise ValueError("WGrad scale expansion requires experts") - if ( - self.source_sf_padding <= 0 - or self.source_sf_padding % 128 - ): - raise ValueError( - "WGrad source SF padding must be a positive multiple of 128" - ) - if ( - self.deinterleave_gate_up is not None - and self.non_k_size != 2 * self.deinterleave_gate_up - ): + if self.source_sf_padding <= 0 or self.source_sf_padding % 128: + raise ValueError("WGrad source SF padding must be a positive multiple of 128") + if self.deinterleave_gate_up is not None and self.non_k_size != 2 * self.deinterleave_gate_up: raise ValueError("gate/up scale deinterleave size mismatch") @cute.jit @@ -86,10 +74,7 @@ def _kernel( expert_offsets: cute.Tensor, output: cute.Tensor, ) -> None: - linear = ( - cute.arch.block_idx()[0] * Int32(self._threads) - + cute.arch.thread_idx()[0] - ) + linear = cute.arch.block_idx()[0] * Int32(self._threads) + cute.arch.thread_idx()[0] atom_bytes: cutlass.Constexpr[int] = self._atom_bytes non_k_atoms: cutlass.Constexpr[int] = self.non_k_size // 128 @@ -103,20 +88,11 @@ def _kernel( for expert in cutlass.range_constexpr(self.expert_count): end = Int32(expert_offsets[expert]) target_token_atoms = (end - previous_end) // Int32(128) - source_token_atoms = ( - ( - Int32(valid_counts[expert]) - + Int32(self.source_sf_padding - 1) - ) - // Int32(self.source_sf_padding) - ) * Int32(self.source_sf_padding // 128) - target_atom_count = ( - Int32(non_k_atoms) * target_token_atoms - ) - in_expert = ( - (atom >= target_atom_base) - & (atom < target_atom_base + target_atom_count) + source_token_atoms = ((Int32(valid_counts[expert]) + Int32(self.source_sf_padding - 1)) // Int32(self.source_sf_padding)) * Int32( + self.source_sf_padding // 128 ) + target_atom_count = Int32(non_k_atoms) * target_token_atoms + in_expert = (atom >= target_atom_base) & (atom < target_atom_base + target_atom_count) if in_expert & (target_token_atoms > Int32(0)): relative_atom = atom - target_atom_base hidden_atom = relative_atom // target_token_atoms @@ -124,49 +100,24 @@ def _kernel( if token_atom < source_token_atoms: source_hidden_atom = hidden_atom source_byte = byte_in_atom - if cutlass.const_expr( - self.deinterleave_gate_up is not None - ): + if cutlass.const_expr(self.deinterleave_gate_up is not None): lane = byte_in_atom // Int32(16) byte_tail = byte_in_atom % Int32(16) group = byte_tail // Int32(4) column_lane = byte_tail % Int32(4) - feature = ( - hidden_atom * Int32(128) - + group * Int32(32) - + lane - ) + feature = hidden_atom * Int32(128) + group * Int32(32) + lane intermediate = Int32(self.deinterleave_gate_up) source_feature = Int32(0) if feature < intermediate: - source_feature = ( - (feature // Int32(32)) * Int32(64) - + feature % Int32(32) - ) + source_feature = (feature // Int32(32)) * Int32(64) + feature % Int32(32) else: up_feature = feature - intermediate - source_feature = ( - (up_feature // Int32(32)) * Int32(64) - + Int32(32) - + up_feature % Int32(32) - ) + source_feature = (up_feature // Int32(32)) * Int32(64) + Int32(32) + up_feature % Int32(32) source_hidden_atom = source_feature // Int32(128) source_feature_in_atom = source_feature % Int32(128) - source_byte = ( - (source_feature_in_atom % Int32(32)) - * Int32(16) - + (source_feature_in_atom // Int32(32)) - * Int32(4) - + column_lane - ) - source_atom = ( - source_atom_base - + source_hidden_atom * source_token_atoms - + token_atom - ) - value = source[ - source_atom * Int32(atom_bytes) + source_byte - ] + source_byte = (source_feature_in_atom % Int32(32)) * Int32(16) + (source_feature_in_atom // Int32(32)) * Int32(4) + column_lane + source_atom = source_atom_base + source_hidden_atom * source_token_atoms + token_atom + value = source[source_atom * Int32(atom_bytes) + source_byte] target_atom_base += target_atom_count source_atom_base += Int32(non_k_atoms) * source_token_atoms previous_end = end diff --git a/python/cudnn/moe_ep/_tuning.py b/python/cudnn/moe_ep/_tuning.py index d9393c481..1d7bfa779 100644 --- a/python/cudnn/moe_ep/_tuning.py +++ b/python/cudnn/moe_ep/_tuning.py @@ -8,7 +8,6 @@ from dataclasses import dataclass from typing import Literal - TokenBackMode = Literal[ "epi_warps", "standalone_warps", @@ -57,53 +56,18 @@ class MoeEpTuningConfig: reduce_topk_in_kernel: bool = False def __post_init__(self) -> None: - if ( - not isinstance(self.token_back_mode, str) - or self.token_back_mode not in _TOKEN_BACK_MODES - ): - raise ValueError( - "token_back_mode must be one of " - f"{tuple(sorted(_TOKEN_BACK_MODES))}, got " - f"{self.token_back_mode!r}" - ) - if ( - not isinstance(self.epi_flag_batch, tuple) - or self.epi_flag_batch not in _EPI_FLAG_BATCHES - ): - raise ValueError( - "epi_flag_batch must be one of " - f"{tuple(sorted(_EPI_FLAG_BATCHES))}, got " - f"{self.epi_flag_batch!r}" - ) - if ( - isinstance(self.token_in_flag_batch, bool) - or self.token_in_flag_batch not in _TOKEN_IN_FLAG_BATCHES - ): - raise ValueError( - "token_in_flag_batch must be one of " - f"{tuple(sorted(_TOKEN_IN_FLAG_BATCHES))}, got " - f"{self.token_in_flag_batch!r}" - ) - if self.group_hint is not None and ( - isinstance(self.group_hint, bool) - or self.group_hint not in _GROUP_HINTS - ): - raise ValueError( - "group_hint must be None or one of " - f"{tuple(sorted(_GROUP_HINTS))}, got {self.group_hint!r}" - ) + if not isinstance(self.token_back_mode, str) or self.token_back_mode not in _TOKEN_BACK_MODES: + raise ValueError("token_back_mode must be one of " f"{tuple(sorted(_TOKEN_BACK_MODES))}, got " f"{self.token_back_mode!r}") + if not isinstance(self.epi_flag_batch, tuple) or self.epi_flag_batch not in _EPI_FLAG_BATCHES: + raise ValueError("epi_flag_batch must be one of " f"{tuple(sorted(_EPI_FLAG_BATCHES))}, got " f"{self.epi_flag_batch!r}") + if isinstance(self.token_in_flag_batch, bool) or self.token_in_flag_batch not in _TOKEN_IN_FLAG_BATCHES: + raise ValueError("token_in_flag_batch must be one of " f"{tuple(sorted(_TOKEN_IN_FLAG_BATCHES))}, got " f"{self.token_in_flag_batch!r}") + if self.group_hint is not None and (isinstance(self.group_hint, bool) or self.group_hint not in _GROUP_HINTS): + raise ValueError("group_hint must be None or one of " f"{tuple(sorted(_GROUP_HINTS))}, got {self.group_hint!r}") if not isinstance(self.reduce_topk_in_kernel, bool): - raise ValueError( - "reduce_topk_in_kernel must be a bool, got " - f"{self.reduce_topk_in_kernel!r}" - ) - if ( - self.reduce_topk_in_kernel - and self.token_back_mode != "epi_warps" - ): - raise ValueError( - "reduce_topk_in_kernel requires " - "token_back_mode='epi_warps'" - ) + raise ValueError("reduce_topk_in_kernel must be a bool, got " f"{self.reduce_topk_in_kernel!r}") + if self.reduce_topk_in_kernel and self.token_back_mode != "epi_warps": + raise ValueError("reduce_topk_in_kernel requires " "token_back_mode='epi_warps'") + __all__ = ["MoeEpTuningConfig"] diff --git a/python/cudnn/moe_ep/_types.py b/python/cudnn/moe_ep/_types.py index e1bfb48de..4ae2287ec 100644 --- a/python/cudnn/moe_ep/_types.py +++ b/python/cudnn/moe_ep/_types.py @@ -30,9 +30,7 @@ def parse_format(value: Union[MoeFormat, str]) -> MoeFormat: 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 + raise ValueError(f"unsupported format {value!r}; expected one of: {choices}") from exc def _normalize_axis(axis: int, ndim: int) -> int: @@ -60,53 +58,34 @@ class BlockScaledTensor: def __post_init__(self) -> None: if not isinstance(self.data, torch.Tensor): - raise ValueError( - f"data must be a torch.Tensor, got {type(self.data).__name__}" - ) + raise ValueError(f"data must be a torch.Tensor, got {type(self.data).__name__}") if not isinstance(self.scale, torch.Tensor): - raise ValueError( - f"scale must be a torch.Tensor, got {type(self.scale).__name__}" - ) + raise ValueError(f"scale must be a torch.Tensor, got {type(self.scale).__name__}") fmt = parse_format(self.format) if fmt is MoeFormat.BF16: raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") if self.data.device != self.scale.device: - raise ValueError( - f"data device {self.data.device} does not match " - f"scale device {self.scale.device}" - ) + raise ValueError(f"data device {self.data.device} does not match " f"scale device {self.scale.device}") try: raw_logical_shape = tuple(self.logical_shape) except TypeError as exc: - raise ValueError( - "logical_shape must be an iterable of integers" - ) from exc + raise ValueError("logical_shape must be an iterable of integers") from exc logical_shape = [] for dim in raw_logical_shape: if isinstance(dim, bool): - raise ValueError( - f"logical_shape dimensions must be integers, got {dim!r}" - ) + raise ValueError(f"logical_shape dimensions must be integers, got {dim!r}") try: dim = operator.index(dim) except TypeError as exc: - raise ValueError( - f"logical_shape dimensions must be integers, got {dim!r}" - ) from exc + raise ValueError(f"logical_shape dimensions must be integers, got {dim!r}") from exc if dim < 0: - raise ValueError( - f"logical_shape dimensions must be non-negative, got {dim}" - ) + raise ValueError(f"logical_shape dimensions must be non-negative, got {dim}") logical_shape.append(dim) normalized_shape = tuple(logical_shape) axis = _normalize_axis(self.axis, len(normalized_shape)) logical_extent = normalized_shape[axis] block_size = 32 if fmt is MoeFormat.MXFP8 else 16 - payload_extent = ( - logical_extent - if fmt is MoeFormat.MXFP8 - else (logical_extent + 1) // 2 - ) + payload_extent = logical_extent if fmt is MoeFormat.MXFP8 else (logical_extent + 1) // 2 scale_extent = (logical_extent + block_size - 1) // block_size expected_data_shape = list(normalized_shape) expected_data_shape[axis] = payload_extent @@ -115,15 +94,9 @@ def __post_init__(self) -> None: expected_data_shape = tuple(expected_data_shape) expected_scale_shape = tuple(expected_scale_shape) if tuple(self.data.shape) != expected_data_shape: - raise ValueError( - f"{fmt.value} data shape must be {expected_data_shape}, " - f"got {tuple(self.data.shape)}" - ) + raise ValueError(f"{fmt.value} data shape must be {expected_data_shape}, " f"got {tuple(self.data.shape)}") if tuple(self.scale.shape) != expected_scale_shape: - raise ValueError( - f"{fmt.value} scale shape must be {expected_scale_shape}, " - f"got {tuple(self.scale.shape)}" - ) + raise ValueError(f"{fmt.value} scale shape must be {expected_scale_shape}, " f"got {tuple(self.scale.shape)}") e4m3_dtype = getattr(torch, "float8_e4m3fn", None) if e4m3_dtype is None: raise RuntimeError("this PyTorch build does not provide torch.float8_e4m3fn") @@ -131,22 +104,14 @@ def __post_init__(self) -> None: expected_data_dtype = e4m3_dtype expected_scale_dtype = getattr(torch, "float8_e8m0fnu", None) if expected_scale_dtype is None: - raise RuntimeError( - "this PyTorch build does not provide torch.float8_e8m0fnu" - ) + raise RuntimeError("this PyTorch build does not provide torch.float8_e8m0fnu") else: expected_data_dtype = torch.uint8 expected_scale_dtype = e4m3_dtype if self.data.dtype is not expected_data_dtype: - raise ValueError( - f"{fmt.value} data must have dtype {expected_data_dtype}, " - f"got {self.data.dtype}" - ) + raise ValueError(f"{fmt.value} data must have dtype {expected_data_dtype}, " f"got {self.data.dtype}") if self.scale.dtype is not expected_scale_dtype: - raise ValueError( - f"{fmt.value} scale must have dtype {expected_scale_dtype}, " - f"got {self.scale.dtype}" - ) + raise ValueError(f"{fmt.value} scale must have dtype {expected_scale_dtype}, " f"got {self.scale.dtype}") object.__setattr__(self, "format", fmt) object.__setattr__(self, "logical_shape", normalized_shape) object.__setattr__(self, "axis", axis) @@ -179,9 +144,7 @@ def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: packed = self.data.movedim(self.axis, -1) low = packed & 0x0F high = packed >> 4 - codes = torch.stack((low, high), dim=-1).flatten(-2)[ - ..., :logical_extent - ] + codes = torch.stack((low, high), dim=-1).flatten(-2)[..., :logical_extent] table = torch.tensor( [ 0.0, @@ -277,14 +240,8 @@ def __init__( self._resource_token = object() self.weights = weights self.device = torch.device(device) - self.slots = tuple( - MoeEpTrainingSlot(index, self._resource_token) - for index in range(slot_count) - ) - self.lanes = tuple( - MoeEpExecutionLane(index, self._resource_token) - for index in range(lane_count) - ) + self.slots = tuple(MoeEpTrainingSlot(index, self._resource_token) for index in range(slot_count)) + self.lanes = tuple(MoeEpExecutionLane(index, self._resource_token) for index in range(lane_count)) self._closed = False @property @@ -301,15 +258,9 @@ def _check_binding( raise RuntimeError("MoeEp training resources are closed") if self._operator_token is not operator_token: raise ValueError("training resources belong to another MoeEp instance") - if ( - slot._resource_token is not self._resource_token - or slot not in self.slots - ): + if slot._resource_token is not self._resource_token or slot not in self.slots: raise ValueError("training slot does not belong to these resources") - if ( - lane._resource_token is not self._resource_token - or lane not in self.lanes - ): + if lane._resource_token is not self._resource_token or lane not in self.lanes: raise ValueError("execution lane does not belong to these resources") def refresh_weights(self) -> None: @@ -379,12 +330,10 @@ def backward( launch_training_backward, ) - grad_activation, grad_topk_weights, operands = ( - launch_training_backward( - self._owner, - execution, - grad_output, - ) + grad_activation, grad_topk_weights, operands = launch_training_backward( + self._owner, + execution, + grad_output, ) return grad_activation, grad_topk_weights, operands @@ -399,24 +348,12 @@ def finalize_overflow( raise RuntimeError("MoeEp training resources are closed") if lane is None: lane = self.lanes[0] - if ( - not isinstance(lane, MoeEpExecutionLane) - or lane._resource_token is not self._resource_token - or lane not in self.lanes - ): - raise ValueError( - "overflow execution lane does not belong to these resources" - ) + if not isinstance(lane, MoeEpExecutionLane) or lane._resource_token is not self._resource_token or lane not in self.lanes: + raise ValueError("overflow execution lane does not belong to these resources") slot_indices = [] for slot in slots: - if ( - not isinstance(slot, MoeEpTrainingSlot) - or slot._resource_token is not self._resource_token - or slot not in self.slots - ): - raise ValueError( - "overflow slot does not belong to these resources" - ) + if not isinstance(slot, MoeEpTrainingSlot) or slot._resource_token is not self._resource_token or slot not in self.slots: + raise ValueError("overflow slot does not belong to these resources") slot_indices.append(slot.index) return self._owner.finalize_overflow( tuple(slot_indices), diff --git a/python/cudnn/moe_ep/_validation.py b/python/cudnn/moe_ep/_validation.py index f2b725001..e7ec4067f 100644 --- a/python/cudnn/moe_ep/_validation.py +++ b/python/cudnn/moe_ep/_validation.py @@ -44,9 +44,7 @@ def _logical_shape(tensor: MoeTensor) -> Tuple[int, ...]: return tensor.logical_shape if isinstance(tensor, torch.Tensor): return tuple(tensor.shape) - raise ValueError( - f"expected torch.Tensor or BlockScaledTensor, got {type(tensor).__name__}" - ) + raise ValueError(f"expected torch.Tensor or BlockScaledTensor, got {type(tensor).__name__}") def _tensor_device(tensor: MoeTensor) -> torch.device: @@ -54,9 +52,7 @@ def _tensor_device(tensor: MoeTensor) -> torch.device: return tensor.device if isinstance(tensor, torch.Tensor): return tensor.device - raise ValueError( - f"expected torch.Tensor or BlockScaledTensor, got {type(tensor).__name__}" - ) + raise ValueError(f"expected torch.Tensor or BlockScaledTensor, got {type(tensor).__name__}") def _validate_strided(name: str, tensor: torch.Tensor) -> None: @@ -71,20 +67,14 @@ def _validate_tensor_representation( ) -> None: logical_shape = _logical_shape(tensor) if logical_shape != expected_logical_shape: - raise ValueError( - f"{name} logical shape must be {expected_logical_shape}, " - f"got {logical_shape}" - ) + raise ValueError(f"{name} logical shape must be {expected_logical_shape}, " f"got {logical_shape}") if isinstance(tensor, torch.Tensor): _validate_strided(name, tensor) if not tensor.is_floating_point(): raise ValueError(f"{name} must be floating point, got {tensor.dtype}") return if not isinstance(tensor, BlockScaledTensor): - raise ValueError( - f"{name} must be a torch.Tensor or BlockScaledTensor, " - f"got {type(tensor).__name__}" - ) + raise ValueError(f"{name} must be a torch.Tensor or BlockScaledTensor, " f"got {type(tensor).__name__}") if tensor.axis != 1: raise ValueError(f"{name} block-scaled axis must be 1, got {tensor.axis}") _validate_strided(f"{name}.data", tensor.data) @@ -112,25 +102,13 @@ def _validate_tensor_representation( _ceil_div(logical_extent, block_size), ) if tuple(tensor.data.shape) != expected_data_shape: - raise ValueError( - f"{name}.data shape must be {expected_data_shape}, " - f"got {tuple(tensor.data.shape)}" - ) + raise ValueError(f"{name}.data shape must be {expected_data_shape}, " f"got {tuple(tensor.data.shape)}") if tuple(tensor.scale.shape) != expected_scale_shape: - raise ValueError( - f"{name}.scale shape must be {expected_scale_shape}, " - f"got {tuple(tensor.scale.shape)}" - ) + raise ValueError(f"{name}.scale shape must be {expected_scale_shape}, " f"got {tuple(tensor.scale.shape)}") if tensor.data.dtype != expected_data_dtype: - raise ValueError( - f"{name}.data must have dtype {expected_data_dtype}, " - f"got {tensor.data.dtype}" - ) + raise ValueError(f"{name}.data must have dtype {expected_data_dtype}, " f"got {tensor.data.dtype}") if tensor.scale.dtype != expected_scale_dtype: - raise ValueError( - f"{name}.scale must have dtype {expected_scale_dtype}, " - f"got {tensor.scale.dtype}" - ) + raise ValueError(f"{name}.scale must have dtype {expected_scale_dtype}, " f"got {tensor.scale.dtype}") def _validate_expert_ids( @@ -139,12 +117,7 @@ def _validate_expert_ids( ) -> None: valid_experts = topk_idx.reshape(-1) valid_experts = valid_experts[valid_experts != -1] - if valid_experts.numel() > 0 and bool( - ( - (valid_experts < 0) - | (valid_experts >= config.num_experts) - ).any().item() - ): + if valid_experts.numel() > 0 and bool(((valid_experts < 0) | (valid_experts >= config.num_experts)).any().item()): raise ValueError("topk_idx contains out-of-range expert ids") @@ -157,43 +130,22 @@ def _validate_routes( validate_expert_ids: bool, ) -> None: if not isinstance(topk_idx, torch.Tensor): - raise ValueError( - f"topk_idx must be a torch.Tensor, got {type(topk_idx).__name__}" - ) + raise ValueError(f"topk_idx must be a torch.Tensor, got {type(topk_idx).__name__}") if not isinstance(topk_weights, torch.Tensor): - raise ValueError( - "topk_weights must be a torch.Tensor, " - f"got {type(topk_weights).__name__}" - ) + raise ValueError("topk_weights must be a torch.Tensor, " f"got {type(topk_weights).__name__}") _validate_strided("topk_idx", topk_idx) _validate_strided("topk_weights", topk_weights) route_shape = (token_count, config.top_k) if tuple(topk_idx.shape) != route_shape: - raise ValueError( - f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.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}, " - f"got {tuple(topk_weights.shape)}" - ) + raise ValueError(f"topk_weights shape must be {route_shape}, " f"got {tuple(topk_weights.shape)}") if topk_idx.dtype not in (torch.int32, torch.int64): - raise ValueError( - "topk_idx must have dtype torch.int32 or torch.int64, " - f"got {topk_idx.dtype}" - ) + raise ValueError("topk_idx must have dtype torch.int32 or torch.int64, " f"got {topk_idx.dtype}") if not topk_weights.is_floating_point(): - raise ValueError( - f"topk_weights must be floating point, got {topk_weights.dtype}" - ) - if ( - config.max_tokens_per_rank is not None - and token_count > config.max_tokens_per_rank - ): - raise ValueError( - f"token count {token_count} exceeds " - f"max_tokens_per_rank={config.max_tokens_per_rank}" - ) + raise ValueError(f"topk_weights must be floating point, got {topk_weights.dtype}") + if config.max_tokens_per_rank is not None and token_count > config.max_tokens_per_rank: + raise ValueError(f"token count {token_count} exceeds " f"max_tokens_per_rank={config.max_tokens_per_rank}") if validate_expert_ids: _validate_expert_ids(config, topk_idx) @@ -212,10 +164,7 @@ def validate_forward( activation_shape = _logical_shape(activation) if len(activation_shape) != 2 or activation_shape[1] != config.hidden_size: - raise ValueError( - f"activation logical shape must be (T, {config.hidden_size}), " - f"got {activation_shape}" - ) + raise ValueError(f"activation logical shape must be (T, {config.hidden_size}), " f"got {activation_shape}") token_count = activation_shape[0] _validate_tensor_representation("activation", activation, activation_shape) _validate_tensor_representation( @@ -279,10 +228,7 @@ def validate_training_weights( """Validate fixed MXFP8 weight bindings used by training resources.""" if not isinstance(weights, MoeEpTrainingWeights): - raise TypeError( - "weights must be a MoeEpTrainingWeights, " - f"got {type(weights).__name__}" - ) + raise TypeError("weights must be a MoeEpTrainingWeights, " f"got {type(weights).__name__}") expected = ( ( "weights.forward_fc1", @@ -324,19 +270,11 @@ def validate_training_weights( for name, tensor, shape in expected: _validate_tensor_representation(name, tensor, shape) if not isinstance(tensor, BlockScaledTensor): - raise TypeError( - f"{name} must be an MXFP8 BlockScaledTensor for " - "fixed training resources" - ) + raise TypeError(f"{name} must be an MXFP8 BlockScaledTensor for " "fixed training resources") if tensor.format is not MoeFormat.MXFP8: - raise NotImplementedError( - f"{name} must use format='mxfp8', got {tensor.format.value!r}" - ) + raise NotImplementedError(f"{name} must use format='mxfp8', got {tensor.format.value!r}") if not tensor.data.is_contiguous() or not tensor.scale.is_contiguous(): - raise ValueError( - f"{name} data and scale must be contiguous for fixed " - "training weight binding" - ) + raise ValueError(f"{name} data and scale must be contiguous for fixed " "training weight binding") device = weights.forward_fc1.device for name, tensor, _shape in expected[1:]: if tensor.device != device: diff --git a/python/cudnn/moe_ep/api.py b/python/cudnn/moe_ep/api.py index 03273859a..317870433 100644 --- a/python/cudnn/moe_ep/api.py +++ b/python/cudnn/moe_ep/api.py @@ -34,6 +34,7 @@ ) from ._validation import validate_forward, validate_training_weights + def _resolve_ep_topology( ep_group: Optional[dist.ProcessGroup], ) -> tuple[int, int, tuple[int, ...]]: @@ -42,25 +43,18 @@ def _resolve_ep_topology( if ep_group is None: return 1, 0, () if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError( - "ep_group requires an initialized torch.distributed process group" - ) + 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 ep_size <= 0 or ep_rank < 0 or ep_rank >= ep_size: raise ValueError("the current process must be a member of ep_group") - ep_global_ranks = tuple( - dist.get_global_rank(ep_group, group_rank) - for group_rank in range(ep_size) - ) + ep_global_ranks = tuple(dist.get_global_rank(ep_group, group_rank) for group_rank in range(ep_size)) if len(set(ep_global_ranks)) != ep_size: raise RuntimeError("ep_group returned duplicate global ranks") if ep_global_ranks[ep_rank] != dist.get_rank(): - raise RuntimeError( - "ep_group rank mapping is inconsistent with the current global rank" - ) + raise RuntimeError("ep_group rank mapping is inconsistent with the current global rank") return ep_size, ep_rank, ep_global_ranks @@ -70,18 +64,12 @@ def _validate_training_assert_capability(config: ForwardConfig) -> None: if config.drop_on_overflow: return if not callable(getattr(torch, "_assert_async", None)): - raise RuntimeError( - "drop_on_overflow=False training resources require callable " - "torch._assert_async before CUDA Graph capture" - ) + raise RuntimeError("drop_on_overflow=False training resources require callable " "torch._assert_async before CUDA Graph capture") if config.ep_size <= 1: return backend = dist.get_backend(config.ep_group) if backend != dist.Backend.NCCL and str(backend).lower() != "nccl": - raise NotImplementedError( - "drop_on_overflow=False EP2+ training resources require an NCCL " - "process group for the captured scalar global overflow OR" - ) + raise NotImplementedError("drop_on_overflow=False EP2+ training resources require an NCCL " "process group for the captured scalar global overflow OR") class MoeEp: @@ -141,22 +129,12 @@ def __init__( 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 ( - isinstance(max_tokens_per_rank, bool) - or not isinstance(max_tokens_per_rank, int) - or max_tokens_per_rank < 0 - ): - raise ValueError( - "max_tokens_per_rank must be a non-negative integer or None" - ) + if max_tokens_per_rank is not None and (isinstance(max_tokens_per_rank, bool) or not isinstance(max_tokens_per_rank, int) or max_tokens_per_rank < 0): + raise ValueError("max_tokens_per_rank must be a non-negative integer or None") if max_recv_size_per_rank is not None and ( - isinstance(max_recv_size_per_rank, bool) - or not isinstance(max_recv_size_per_rank, int) - or max_recv_size_per_rank <= 0 + isinstance(max_recv_size_per_rank, bool) or not isinstance(max_recv_size_per_rank, int) or max_recv_size_per_rank <= 0 ): - raise ValueError( - "max_recv_size_per_rank must be a positive integer or None" - ) + raise ValueError("max_recv_size_per_rank must be a positive integer or None") if not isinstance(drop_on_overflow, bool): raise ValueError("drop_on_overflow must be a bool") if not isinstance(apply_topk_in_fc1, bool): @@ -166,19 +144,11 @@ def __init__( ("sf_padding_size", sf_padding_size), ): if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise ValueError( - f"{name} must be a positive integer, got {value!r}" - ) + raise ValueError(f"{name} must be a positive integer, got {value!r}") if sf_padding_size % 128: - raise ValueError( - "sf_padding_size must be a positive multiple of 128, " - f"got {sf_padding_size}" - ) + raise ValueError("sf_padding_size must be a positive multiple of 128, " f"got {sf_padding_size}") if tuning is not None and not isinstance(tuning, MoeEpTuningConfig): - raise TypeError( - "tuning must be a MoeEpTuningConfig or None, " - f"got {type(tuning).__name__}" - ) + raise TypeError("tuning must be a MoeEpTuningConfig or None, " f"got {type(tuning).__name__}") if gate_up_clamp is not None: if isinstance(gate_up_clamp, bool) or not isinstance(gate_up_clamp, Real): raise ValueError("gate_up_clamp must be a finite real number or None") @@ -187,10 +157,7 @@ def __init__( raise ValueError("gate_up_clamp must be a finite real number or None") if ep_group is not None and not isinstance(ep_group, dist.ProcessGroup): - raise ValueError( - f"ep_group must be a torch.distributed.ProcessGroup or None, " - f"got {type(ep_group).__name__}" - ) + raise ValueError(f"ep_group must be a torch.distributed.ProcessGroup or None, " f"got {type(ep_group).__name__}") ep_size, ep_rank, ep_global_ranks = _resolve_ep_topology(ep_group) if num_experts % ep_size != 0: raise ValueError(f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})") @@ -215,14 +182,9 @@ def __init__( self.sf_padding_size = sf_padding_size self.tuning = MoeEpTuningConfig() if tuning is None else tuning if self.tuning.reduce_topk_in_kernel and ( - self.combine_format is not MoeFormat.BF16 - or self.output_format is not MoeFormat.BF16 - or not self.apply_topk_in_fc1 + self.combine_format is not MoeFormat.BF16 or self.output_format is not MoeFormat.BF16 or not self.apply_topk_in_fc1 ): - raise ValueError( - "reduce_topk_in_kernel requires BF16 combine/output and " - "apply_topk_in_fc1=True" - ) + raise ValueError("reduce_topk_in_kernel requires BF16 combine/output and " "apply_topk_in_fc1=True") for name, fmt in ( ("output_format", self.output_format), @@ -280,14 +242,8 @@ def _get_backend(self, request): raise RuntimeError("MoeEp is closed") from . import _backend - if ( - self._forward_backend is not None - and request.device != self._forward_backend_device - ): - raise ValueError( - f"MoeEp backend is bound to {self._forward_backend_device}; " - f"create a separate MoeEp instance for {request.device}" - ) + if self._forward_backend is not None and request.device != self._forward_backend_device: + raise ValueError(f"MoeEp backend is bound to {self._forward_backend_device}; " f"create a separate MoeEp instance for {request.device}") _backend.validate_config(self._forward_config) _backend.validate_request(request) @@ -322,11 +278,7 @@ def __call__( if self._closed: raise RuntimeError("MoeEp is closed") topk_version = self._tensor_version(topk_idx) - validate_expert_ids = not ( - self._validated_topk_idx is topk_idx - and topk_version is not None - and topk_version == self._validated_topk_version - ) + validate_expert_ids = not (self._validated_topk_idx is topk_idx and topk_version is not None and topk_version == self._validated_topk_version) request = validate_forward( self._forward_config, activation, @@ -337,10 +289,7 @@ def __call__( validate_expert_ids=validate_expert_ids, ) version_after_validation = self._tensor_version(topk_idx) - if ( - topk_version is not None - and topk_version == version_after_validation - ): + if topk_version is not None and topk_version == version_after_validation: self._validated_topk_idx = topk_idx self._validated_topk_version = topk_version else: @@ -406,21 +355,12 @@ def prepare_training_resources( ("slot_count", slot_count), ("lane_count", lane_count), ): - if ( - isinstance(value, bool) - or not isinstance(value, int) - or value <= 0 - ): - raise ValueError( - f"{name} must be a positive integer, got {value!r}" - ) + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") if self._training_resources is not None: if not self._training_resources.closed: raise RuntimeError("MoeEp training resources already exist") - raise RuntimeError( - "MoeEp training resources were closed; create a new " - "MoeEp instance before preparing replacement weights" - ) + raise RuntimeError("MoeEp training resources were closed; create a new " "MoeEp instance before preparing replacement weights") device = validate_training_weights( self._forward_config, @@ -430,14 +370,8 @@ def prepare_training_resources( from . import _backend _backend.validate_config(self._forward_config) - if ( - self._forward_backend is not None - and device != self._forward_backend_device - ): - raise ValueError( - f"MoeEp backend is bound to " - f"{self._forward_backend_device}; got {device}" - ) + if self._forward_backend is not None and device != self._forward_backend_device: + raise ValueError(f"MoeEp backend is bound to " f"{self._forward_backend_device}; got {device}") if self._forward_backend is None: self._forward_backend = _backend.create_backend( self._forward_config, diff --git a/test/python/moe_ep/moe_ep_distributed_workers.py b/test/python/moe_ep/moe_ep_distributed_workers.py index 0b0442b73..3e5923d5b 100644 --- a/test/python/moe_ep/moe_ep_distributed_workers.py +++ b/test/python/moe_ep/moe_ep_distributed_workers.py @@ -137,18 +137,12 @@ def _distributed_subgroup_output_worker( ) try: subgroup_memberships = ((0, 2), (1, 3)) - subgroups = [ - dist.new_group(list(members), backend="nccl") - for members in subgroup_memberships - ] + subgroups = [dist.new_group(list(members), backend="nccl") for members in subgroup_memberships] subgroup_index = global_rank % 2 ep_group = subgroups[subgroup_index] ep_rank = dist.get_rank(ep_group) ep_size = dist.get_world_size(ep_group) - actual_global_ranks = tuple( - dist.get_global_rank(ep_group, group_rank) - for group_rank in range(ep_size) - ) + actual_global_ranks = tuple(dist.get_global_rank(ep_group, group_rank) for group_rank in range(ep_size)) _run_forward_output_case( device=device, @@ -373,10 +367,7 @@ def _distributed_subgroup_backward_reference_worker( ep_group = subgroups[subgroup_index] ep_rank = dist.get_rank(ep_group) ep_size = dist.get_world_size(ep_group) - actual_global_ranks = tuple( - dist.get_global_rank(ep_group, group_rank) - for group_rank in range(ep_size) - ) + actual_global_ranks = tuple(dist.get_global_rank(ep_group, group_rank) for group_rank in range(ep_size)) assert ep_size == len(expected_global_ranks) assert ep_rank == expected_global_ranks.index(global_rank) assert actual_global_ranks == expected_global_ranks diff --git a/test/python/moe_ep/moe_ep_reference.py b/test/python/moe_ep/moe_ep_reference.py index 3c56db8b9..0358edd66 100644 --- a/test/python/moe_ep/moe_ep_reference.py +++ b/test/python/moe_ep/moe_ep_reference.py @@ -426,13 +426,9 @@ def _padded_expert_rows( as_tuple=False, ).flatten() if int(positions.numel()) != int(count): - raise ValueError( - f"expert {expert} has {positions.numel()} rows, expected {count}" - ) + raise ValueError(f"expert {expert} has {positions.numel()} rows, expected {count}") if count: - padded[begin : begin + count].copy_( - rows.index_select(0, positions) - ) + padded[begin : begin + count].copy_(rows.index_select(0, positions)) begin = int(end) return padded @@ -484,23 +480,13 @@ def __init__( 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 backward_wgrad_mode not in ("none", "operands"): - raise ValueError( - "backward_wgrad_mode must be 'none' or 'operands'" - ) + raise ValueError("backward_wgrad_mode must be 'none' or 'operands'") if backward_wgrad_mode == "operands" and not generate_c: - raise ValueError( - "backward_wgrad_mode='operands' requires generate_c=True" - ) + raise ValueError("backward_wgrad_mode='operands' requires generate_c=True") if not isinstance(token_padding_size, int) or token_padding_size <= 0: raise ValueError("token_padding_size must be a positive integer") - if ( - backward_wgrad_mode == "operands" - and token_padding_size != 256 - ): - raise ValueError( - "backward_wgrad_mode='operands' requires " - "token_padding_size=256" - ) + if backward_wgrad_mode == "operands" and token_padding_size != 256: + raise ValueError("backward_wgrad_mode='operands' requires " "token_padding_size=256") if ep_group is None: ep_size, ep_rank = 1, 0 @@ -523,14 +509,8 @@ def __init__( self.max_tokens_per_rank = max_tokens_per_rank self.output_format = _parse_format(output_format) self.combine_format = _parse_format(combine_format) - self.intermediate_format = ( - None if intermediate_format is None else _parse_format(intermediate_format) - ) - self.backward_operand_format = ( - None - if backward_operand_format is None - else _parse_format(backward_operand_format) - ) + self.intermediate_format = None if intermediate_format is None else _parse_format(intermediate_format) + self.backward_operand_format = None if backward_operand_format is None else _parse_format(backward_operand_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) @@ -787,11 +767,7 @@ def __call__( send_token_idx = plan.send_token_idx send_slot_idx = plan.send_slot_idx send_counts, recv_counts = plan.send_counts, plan.recv_counts - forward_activation_float = ( - wgrad_activation_float - if wgrad_activation_float is not None - else activation_float - ) + forward_activation_float = wgrad_activation_float if wgrad_activation_float is not None else activation_float send_tokens = forward_activation_float.index_select( 0, send_token_idx, @@ -855,15 +831,20 @@ def __call__( for value in torch.bincount( recv_expert, minlength=self.experts_per_rank, - ).cpu().tolist() + ) + .cpu() + .tolist() ) padded_ends = [] total = 0 for count in valid_counts: - total += _ceil_div( - count, - self.token_padding_size, - ) * self.token_padding_size + total += ( + _ceil_div( + count, + self.token_padding_size, + ) + * self.token_padding_size + ) padded_ends.append(total) ordered_tokens = recv_wgrad_tokens.index_select( 0, @@ -908,9 +889,7 @@ def backward( fc1_c: torch.Tensor, route_metadata: torch.Tensor, *, - wgrad_forward_stash: Optional[ - WgradForwardStashReference - ] = None, + wgrad_forward_stash: Optional[WgradForwardStashReference] = None, ) -> Union[ Tuple[torch.Tensor, torch.Tensor], Tuple[ @@ -943,22 +922,14 @@ def backward( wgrad_forward_stash, WgradForwardStashReference, ): - raise TypeError( - "wgrad_forward_stash must be a " - "WgradForwardStashReference" - ) + raise TypeError("wgrad_forward_stash must be a " "WgradForwardStashReference") if not torch.equal( wgrad_forward_stash.route_metadata, route_metadata, ): - raise ValueError( - "wgrad_forward_stash route identity does not match " - "route_metadata" - ) + raise ValueError("wgrad_forward_stash route identity does not match " "route_metadata") elif wgrad_forward_stash is not None: - raise ValueError( - "wgrad_forward_stash is only accepted in operands mode" - ) + raise ValueError("wgrad_forward_stash is only accepted in operands mode") 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)}") @@ -981,10 +952,7 @@ def backward( ) semantic_fc2_float = fc2_float effective_backward_format = self.backward_operand_format - if ( - effective_backward_format is None - and self.backward_wgrad_mode == "operands" - ): + if effective_backward_format is None and self.backward_wgrad_mode == "operands": effective_backward_format = MoeFormat.MXFP8 if effective_backward_format is not None: # The dGLU adapter requantizes both transposed weights along the @@ -1086,16 +1054,11 @@ def backward( d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) if self.apply_topk_in_fc1: d_h = d_h_fc2 * w - semantic_d_h = ( - semantic_d_y - @ semantic_fc2_float[expert].transpose(0, 1) - ) + semantic_d_h = semantic_d_y @ semantic_fc2_float[expert].transpose(0, 1) d_w_rows[positions] = (semantic_d_h * h).sum(dim=-1) else: d_h = d_h_fc2 - d_w_rows[positions] = ( - semantic_d_y * (h @ semantic_fc2_float[expert]) - ).sum(dim=-1) + d_w_rows[positions] = (semantic_d_y * (h @ semantic_fc2_float[expert])).sum(dim=-1) d_g = d_h * u * (sig * (1 + g * (1 - sig))) d_u = d_h * s @@ -1128,14 +1091,8 @@ def backward( ) if self.backward_wgrad_mode == "operands": stash = wgrad_forward_stash - padded_ends = tuple( - int(value) - for value in stash.expert_offsets.cpu().tolist() - ) - valid_counts = tuple( - int(value) - for value in stash.valid_route_counts.cpu().tolist() - ) + padded_ends = tuple(int(value) for value in stash.expert_offsets.cpu().tolist()) + valid_counts = tuple(int(value) for value in stash.valid_route_counts.cpu().tolist()) padded_dc = _padded_expert_rows( dc_rows, expert_rows, @@ -1189,4 +1146,4 @@ def backward( "backward_combine_round_trip", "forward_combine_round_trip", "quantize_blockwise", -] \ No newline at end of file +] diff --git a/test/python/moe_ep/probe_moe_ep_training_graph.py b/test/python/moe_ep/probe_moe_ep_training_graph.py index e3eea1c4d..556fb0c1c 100644 --- a/test/python/moe_ep/probe_moe_ep_training_graph.py +++ b/test/python/moe_ep/probe_moe_ep_training_graph.py @@ -46,9 +46,7 @@ def _debug_phase(rank: int, phase: str) -> None: if os.environ.get("MOE_EP_DEBUG_RUNTIME", "0") != "1": return print( - "[moe-ep-probe] " - f"time={time.monotonic():.6f} host={socket.gethostname()} " - f"pid={os.getpid()} rank={rank} phase={phase}", + "[moe-ep-probe] " f"time={time.monotonic():.6f} host={socket.gethostname()} " f"pid={os.getpid()} rank={rank} phase={phase}", flush=True, ) @@ -85,28 +83,19 @@ def _parse_args() -> argparse.Namespace: "--multistream-replays", type=int, default=10, - help=( - "two-lane cross-stream graph replays; use a larger value such as " - "100 for dedicated stress runs" - ), + help=("two-lane cross-stream graph replays; use a larger value such as " "100 for dedicated stress runs"), ) parser.add_argument( "--max-recv-size-per-rank", type=int, default=1, - help=( - "bounded receive capacity; must remain below the forced-overflow " - "route count so the probe retains overflow coverage" - ), + help=("bounded receive capacity; must remain below the forced-overflow " "route count so the probe retains overflow coverage"), ) parser.add_argument( "--cycles", type=int, default=2, - help=( - "create/capture/destroy cycles; the first is exhaustive and later " - "cycles use a minimal replay to verify teardown/re-init" - ), + help=("create/capture/destroy cycles; the first is exhaustive and later " "cycles use a minimal replay to verify teardown/re-init"), ) parser.add_argument("--timeout-seconds", type=int, default=600) parser.add_argument( @@ -117,10 +106,7 @@ def _parse_args() -> argparse.Namespace: parser.add_argument( "--expect-overflow-assert", action="store_true", - help=( - "run only the fatal drop_on_overflow=False graph assertion probe; " - "success requires every rank to observe the expected CUDA error" - ), + help=("run only the fatal drop_on_overflow=False graph assertion probe; " "success requires every rank to observe the expected CUDA error"), ) return parser.parse_args() @@ -147,9 +133,7 @@ def _assert_replay_tensor( } if actual.dtype in low_precision: if not torch.equal(actual, expected): - raise AssertionError( - f"{name} is not bitwise equal after graph replay" - ) + raise AssertionError(f"{name} is not bitwise equal after graph replay") return torch.testing.assert_close( actual, @@ -434,20 +418,20 @@ def _run_training_resource_probe( for name, tensor in zip( comparison_names, ( - y0, - y1, - dx0, - dx1, - dp0, - dp1, - operands0.fc1_a, - operands0.fc1_b, - operands0.fc2_a, - operands0.fc2_b, - operands1.fc1_a, - operands1.fc1_b, - operands1.fc2_a, - operands1.fc2_b, + y0, + y1, + dx0, + dx1, + dp0, + dp1, + operands0.fc1_a, + operands0.fc1_b, + operands0.fc2_a, + operands0.fc2_b, + operands1.fc1_a, + operands1.fc1_b, + operands1.fc2_a, + operands1.fc2_b, ), ) } @@ -542,9 +526,7 @@ def _run_training_resource_probe( stream.synchronize() dist.barrier(group=group) if int(graph_overflow.item()) != 0: - raise AssertionError( - "fixed-resource diagnostic replay overflowed" - ) + raise AssertionError("fixed-resource diagnostic replay overflowed") torch.testing.assert_close( graph_dp0, dprob_reference, @@ -604,8 +586,7 @@ def _run_training_resource_probe( mode = "full" if full_probe else "reinit" effective_burst = burst_replays if full_probe else 0 print( - f"MOE_EP_EP{world_size}_TRAINING_RESOURCES_GRAPH_PASS " - f"mode={mode} burst={effective_burst}", + f"MOE_EP_EP{world_size}_TRAINING_RESOURCES_GRAPH_PASS " f"mode={mode} burst={effective_burst}", flush=True, ) finally: @@ -650,9 +631,7 @@ def _run_multistream_resource_probe( resources.refresh_weights() with _debug_phase_scope(rank, "multistream.lane0-forward"): - eager_y0 = resources.forward( - slot0, lane0, args0[0], args0[3], args0[4] - ) + eager_y0 = resources.forward(slot0, lane0, args0[0], args0[3], args0[4]) with _debug_phase_scope(rank, "multistream.lane0-backward"): eager_dx0, eager_dp0, _ = resources.backward( slot0, @@ -670,9 +649,7 @@ def _run_multistream_resource_probe( dist.barrier(group=group) with _debug_phase_scope(rank, "multistream.lane1-forward"): - eager_y1 = resources.forward( - slot1, lane1, args1[0], args1[3], args1[4] - ) + eager_y1 = resources.forward(slot1, lane1, args1[0], args1[3], args1[4]) with _debug_phase_scope(rank, "multistream.lane1-backward"): eager_dx1, eager_dp1, _ = resources.backward( slot1, @@ -767,9 +744,7 @@ def _run_multistream_resource_probe( with torch.cuda.stream(capture_stream): for _ in range(replays): graph.replay() - replay_watchdog = _RuntimeWatchdog( - "multistream.replay-synchronize" - ) + replay_watchdog = _RuntimeWatchdog("multistream.replay-synchronize") replay_watchdog.start() with _debug_phase_scope( rank, @@ -800,8 +775,7 @@ def _run_multistream_resource_probe( ) if rank == 0: print( - f"MOE_EP_EP{world_size}_MULTISTREAM_GRAPH_PASS " - f"replays={replays}", + f"MOE_EP_EP{world_size}_MULTISTREAM_GRAPH_PASS " f"replays={replays}", flush=True, ) finally: @@ -884,8 +858,7 @@ def _run_error_mode_assert_probe( stream.synchronize() except BaseException as exc: print( - f"MOE_EP_EP{world_size}_ERROR_MODE_ASSERT_PASS " - f"rank={rank} error={type(exc).__name__}", + f"MOE_EP_EP{world_size}_ERROR_MODE_ASSERT_PASS " f"rank={rank} error={type(exc).__name__}", flush=True, ) # CUDA device assertions poison the process context. Do not run Python @@ -914,9 +887,7 @@ def main() -> None: rank = int(os.environ.get("RANK", "0")) local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) if world_size < 2: - raise RuntimeError( - f"this probe requires WORLD_SIZE >= 2, got {world_size}" - ) + raise RuntimeError(f"this probe requires WORLD_SIZE >= 2, got {world_size}") forced_overflow_routes = world_size * 8 * 2 if args.max_recv_size_per_rank >= forced_overflow_routes: raise ValueError( @@ -929,10 +900,7 @@ def main() -> None: torch.cuda.set_device(device) capability = torch.cuda.get_device_capability(device) if capability != (10, 7): - raise RuntimeError( - "this probe requires Rubin SM107; " - f"rank {rank} found compute capability {capability}" - ) + raise RuntimeError("this probe requires Rubin SM107; " f"rank {rank} found compute capability {capability}") os.environ.setdefault("CUTE_DSL_ARCH", "sm_107a") dist.init_process_group( diff --git a/test/python/moe_ep/test_moe_ep_multinode.py b/test/python/moe_ep/test_moe_ep_multinode.py index 26b65636e..b0da9fe99 100644 --- a/test/python/moe_ep/test_moe_ep_multinode.py +++ b/test/python/moe_ep/test_moe_ep_multinode.py @@ -22,7 +22,6 @@ make_distributed_forward_inputs, ) - pytestmark = [ pytest.mark.L1, pytest.mark.gpu_exclusive, @@ -59,10 +58,7 @@ class _TorchrunWorld: def _require_torchrun_environment() -> tuple[int, int, int, int]: missing = [name for name in _TORCHRUN_ENV if name not in os.environ] if missing: - pytest.skip( - "multi-node MoE EP tests require torchrun environment variables: " - + ", ".join(missing) - ) + pytest.skip("multi-node MoE EP tests require torchrun environment variables: " + ", ".join(missing)) return ( int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"]), @@ -76,20 +72,13 @@ def torchrun_world(): if not dist.is_available() or not dist.is_nccl_available(): pytest.skip("multi-node Rubin MXFP8 tests require NCCL") - rank, world_size, local_rank, local_world_size = ( - _require_torchrun_environment() - ) + rank, world_size, local_rank, local_world_size = _require_torchrun_environment() if local_rank < 0 or local_rank >= torch.cuda.device_count(): - pytest.skip( - f"torchrun LOCAL_RANK={local_rank} is not backed by a visible GPU" - ) + pytest.skip(f"torchrun LOCAL_RANK={local_rank} is not backed by a visible GPU") device = torch.device("cuda", local_rank) if torch.cuda.get_device_capability(device) != (10, 7): - pytest.skip( - "multi-node Rubin MXFP8 tests require exactly SM107 " - "(compute capability 10.7) on every rank" - ) + pytest.skip("multi-node Rubin MXFP8 tests require exactly SM107 " "(compute capability 10.7) on every rank") try: import nvshmem.core # noqa: F401 except (ImportError, OSError): @@ -99,9 +88,7 @@ def torchrun_world(): torch.cuda.set_device(device) if dist.is_initialized(): if dist.get_rank() != rank or dist.get_world_size() != world_size: - raise RuntimeError( - "existing process group does not match torchrun RANK/WORLD_SIZE" - ) + raise RuntimeError("existing process group does not match torchrun RANK/WORLD_SIZE") else: dist.init_process_group( backend="nccl", @@ -181,10 +168,7 @@ def test_mxfp8_forward_multinode_matches_reference( combine_format, ): world = torchrun_world - if ( - world.world_size != required_world_size - or world.local_world_size != required_local_world_size - ): + if world.world_size != required_world_size or world.local_world_size != required_local_world_size: pytest.skip( f"EP{ep_size} requires torchrun WORLD_SIZE={required_world_size}, " f"LOCAL_WORLD_SIZE={required_local_world_size}; got " @@ -275,10 +259,7 @@ def test_fixed_training_resources_multinode_match_independent_reference( ): world = torchrun_world if world.world_size != required_world_size: - pytest.skip( - f"EP{ep_size} requires torchrun WORLD_SIZE={required_world_size}; " - f"got WORLD_SIZE={world.world_size}" - ) + pytest.skip(f"EP{ep_size} requires torchrun WORLD_SIZE={required_world_size}; " f"got WORLD_SIZE={world.world_size}") _run_backward_reference_case( device=world.device, @@ -332,9 +313,7 @@ def test_training_prepare_multinode_rejects_rank_abi_mismatch( ) caught_error = None try: - lane_count = ( - rank_zero_lane_count if world.rank == 0 else other_lane_count - ) + lane_count = rank_zero_lane_count if world.rank == 0 else other_lane_count try: op.prepare_training_resources( weights, @@ -352,14 +331,8 @@ def test_training_prepare_multinode_rejects_rank_abi_mismatch( device_ids=[world.local_rank], ) - assert isinstance(caught_error, RuntimeError), ( - f"rank {world.rank} expected RuntimeError from collective prepare, " - f"got {caught_error!r}" - ) + assert isinstance(caught_error, RuntimeError), f"rank {world.rank} expected RuntimeError from collective prepare, " f"got {caught_error!r}" message = str(caught_error) - assert ( - "symmetric workspace region counts differ" in message - or "ABI differs" in message - ), f"rank {world.rank} got unexpected prepare error: {message}" + assert "symmetric workspace region counts differ" in message or "ABI differs" in message, f"rank {world.rank} got unexpected prepare error: {message}" finally: op.close() From 0ecef6da64e13aefac22f3b080d85447d73db2f7 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Fri, 28 Aug 2026 09:07:44 -0700 Subject: [PATCH 15/31] fix: align MoeEP packaging with backend boundaries Use reusable communication dependencies, gate Rubin kernels on CUTLASS DSL 4.8, and make vendored source licensing and provenance explicit. --- .pre-commit-config.yaml | 4 +- pyproject.toml | 11 +- python/cudnn/__init__.py | 18 +-- .../_megamoe_backend/cutedsl_src/VENDOR.md | 70 ++++++++ .../cutedsl_src/VENDOR_INFO.md | 139 ---------------- .../_megamoe_backend/cutedsl_src/__init__.py | 3 + .../_megamoe_backend/cutedsl_src/api.py | 3 + .../cutedsl_src/communication/__init__.py | 3 + .../communication/nvlink_domain/__init__.py | 3 + .../nvlink_domain/symmetric_buffer.py | 3 + .../communication/nvlink_domain/token_comm.py | 3 + .../nvlink_domain/token_comm_deterministic.py | 3 + .../cutedsl_src/helpers/__init__.py | 3 + .../cutedsl_src/helpers/constants.py | 3 + .../cutedsl_src/helpers/cute_py_helpers.py | 3 + .../cutedsl_src/helpers/device_workspace.py | 3 + .../cutedsl_src/helpers/dsl_helpers.py | 3 + .../cutedsl_src/helpers/flag_batch.py | 3 + .../cutedsl_src/helpers/iket_compat.py | 3 + .../cutedsl_src/helpers/ptx_helpers.py | 3 + .../cutedsl_src/helpers/smem_workspace.py | 3 + .../cutedsl_src/helpers/software_sync.py | 3 + .../cutedsl_src/helpers/utils.py | 3 + .../cutedsl_src/kernel_src/__init__.py | 3 + .../block_scaled_swap_ab_fc12_extension.py | 3 + .../cutedsl_src/kernel_src/rubin/__init__.py | 3 + .../kernel_src/rubin/training/__init__.py | 3 + .../kernel_src/schedulers/__init__.py | 3 + .../cutedsl_src/kernel_src/schedulers/base.py | 3 + .../kernel_src/schedulers/fc12_mapping.py | 3 + .../kernel_src/schedulers/fc12_scheduler.py | 3 + .../schedulers/non_clc_mixed_cga.py | 3 + .../kernel_src/schedulers/work_id_claim.py | 3 + .../moe_ep/_megamoe_backend/mxfp8/_backend.py | 85 +++------- .../mxfp8/_backward_compile.py | 118 +++++--------- .../moe_ep/_megamoe_backend/mxfp8/_compile.py | 150 +++++------------- .../moe_ep/_megamoe_backend/mxfp8/_cutedsl.py | 56 +++++++ test/python/moe_ep/test_moe_ep_cutedsl.py | 66 ++++++++ 38 files changed, 383 insertions(+), 418 deletions(-) create mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR.md delete mode 100644 python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR_INFO.md create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_cutedsl.py create mode 100644 test/python/moe_ep/test_moe_ep_cutedsl.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5a87729d3..bcfcb0b7c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,5 +1,5 @@ -# Vendored third-party sources are kept byte-for-byte as upstream ships them. -exclude: '^include/cudnn_frontend/thirdparty/' +# Vendored third-party source bodies retain their upstream formatting. +exclude: '^(include/cudnn_frontend/thirdparty/|python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/)' repos: - repo: https://github.com/pre-commit/mirrors-clang-format diff --git a/pyproject.toml b/pyproject.toml index d40aa7987..a8710d393 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,13 +69,10 @@ cutedsl = [ "cuda-python", "apache-tvm-ffi>=0.1.11", ] -moe_ep = [ - "nvidia-cutlass-dsl[cu13]>=4.8.0", +comm = [ + # Communication runtimes are grouped by backend so future distributed + # operation graphs can reuse the same installation extra. "nvshmem4py-cu13>=0.3.1", - "cuda-python", - "torch", - "apache-tvm-ffi>=0.1.11", - "torch-c-dlpack-ext", ] cutile = [ # The cuTile linear-attention engines. Base cuda-tile only -- its [tileiras] @@ -130,5 +127,5 @@ version = {attr = "cudnn.__version__"} include = ["**/*"] "cudnn.moe_ep._megamoe_backend.cutedsl_src" = [ "LICENSE.Apache-2.0", - "VENDOR_INFO.md", + "VENDOR.md", ] diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 395228ae0..54277ba6c 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -305,12 +305,8 @@ def _dlopen_cudnn(): ) __all__ = [*_EAGER_PUBLIC_NAMES, "Graph", "wrapper"] -_CUTEDSL_INSTALL_HINT = ( - "Install with 'pip install nvidia-cudnn-frontend[cutedsl]'" -) -_MOE_EP_INSTALL_HINT = ( - "Install with 'pip install nvidia-cudnn-frontend[moe_ep]'" -) +_CUTEDSL_INSTALL_HINT = "Install with 'pip install nvidia-cudnn-frontend[cutedsl]'" +_MOE_EP_INSTALL_HINT = "Install with 'pip install " '"nvidia-cudnn-frontend[cutedsl,comm]" torch torch-c-dlpack-ext\'' _MOE_EP_OPTIONAL_IMPORTS = { "moe_ep", "BlockScaledTensor", @@ -418,14 +414,8 @@ def _load_optional_symbol(name: str) -> Any: module = importlib.import_module(module_name, package=__name__) value = module if attr_name is None else getattr(module, attr_name) except Exception as e: - install_hint = ( - _MOE_EP_INSTALL_HINT - if name in _MOE_EP_OPTIONAL_IMPORTS - else _CUTEDSL_INSTALL_HINT - ) - raise ImportError( - f"{name} requires optional dependencies. {install_hint}: {e}" - ) from e + install_hint = _MOE_EP_INSTALL_HINT if name in _MOE_EP_OPTIONAL_IMPORTS else _CUTEDSL_INSTALL_HINT + raise ImportError(f"{name} requires optional dependencies. {install_hint}: {e}") from e globals()[name] = value return value diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR.md b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR.md new file mode 100644 index 000000000..a4e366e2e --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR.md @@ -0,0 +1,70 @@ +# Vendoring record: cutedsl_megamoe + +This file records provenance and synchronization state for the CuTeDSL +MegaMoE source snapshot. Runtime behavior and integration details are +documented in the parent backend `README.md`. + +## Upstream + +- **Project**: `cutedsl_megamoe` (NVIDIA-internal repository; URL omitted). +- **Source tree**: `cutedsl_megamoe/next/sources`. +- **Current synchronized commit**: + `5b89819cb16069dfe20a1a0ba0778d35cb428352`. +- **Earlier import points**: + - base forward: + `882c83e2ce4086c3cd4211fc5a2296143c5e2aea`; + - selected forward updates and backward dGLU: + `92dd334af2eeedb36087834354b58ace08e880c6`. +- **Last synced**: 2026-08-28. Earlier imports occurred on 2026-08-11, + 2026-08-17, 2026-08-20, and 2026-08-24. +- **Vendored subset**: the recursive Python import closure required by Rubin + SM107 training MegaMoE forward GLU, optional forward MXFP8 column + requantization, and backward dGLU. + +Complete kernel products, runners, tests, repository scaffolding, and +unrelated Blackwell and Rubin inference sources are excluded. Three +architecture-neutral Blackwell donor modules are retained because the Rubin +`topk_reduce.py` and `tmem_transpose.py` source-copy shims import them. + +## Policy + +- Vendored Python source bodies track the corresponding upstream paths at the + synchronized commit. +- Repository-required copyright and BSD-3-Clause SPDX headers may be added + where the upstream snapshot did not carry them. +- Integration behavior belongs in the parent `_megamoe_backend` package, not + in the vendored source bodies. +- Local kernel fixes should go upstream first and then be synchronized here. + Any unavoidable local source difference must be listed below. +- Snapshot updates must preserve the minimal recursive import closure and + verify source-body equality while ignoring repository-added header lines. + +The synchronized Python sources use BSD-3-Clause SPDX identifiers. +`LICENSE.Apache-2.0` is retained as historical snapshot metadata. + +## Local differences from upstream + +- `kernel_src/rubin/training/__init__.py` is reduced to a package marker. This + avoids importing the unused traditional-wgrad product. +- Repository-required copyright and BSD-3-Clause SPDX headers are added to + source files that lacked explicit headers. + +No other vendored Python source-body differences are expected. + +## Integration boundary + +Public API validation, symmetric-workspace ownership, overflow reporting, +input and weight staging, CUDA Graph handling, dprob materialization, and +grouped-WGrad layout conversion live in the parent `_megamoe_backend` +package. + +The vendored Rubin sources require a CUTLASS DSL distribution that provides +`cutlass.utils.rubin_helpers`. The executable backend enforces +`nvidia-cutlass-dsl>=4.8.0` before importing these kernels. + +## Consumers + +- `_megamoe_backend/mxfp8/_compile.py`: Rubin MXFP8 forward preparation and + compilation. +- `_megamoe_backend/mxfp8/_backward_compile.py`: Rubin MXFP8 backward dGLU + preparation and compilation. diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR_INFO.md b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR_INFO.md deleted file mode 100644 index 6c293ef91..000000000 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR_INFO.md +++ /dev/null @@ -1,139 +0,0 @@ -# Vendored CuTeDSL MegaMoE sources - -## Provenance - -- Source project: `cutedsl_megamoe` -- Source tree: `cutedsl_megamoe/next/sources` -- Base forward upstream revision: - `882c83e2ce4086c3cd4211fc5a2296143c5e2aea` -- Selected forward updates and backward dGLU upstream revision: - `92dd334af2eeedb36087834354b58ace08e880c6` -- Latest synchronized upstream revision: - `5b89819cb16069dfe20a1a0ba0778d35cb428352` -- Vendoring dates: 2026-08-11 (base), 2026-08-17 (selected updates), and - 2026-08-20, 2026-08-24, and 2026-08-28 (latest synchronizations). -- On 2026-08-28 every vendored Python source except the intentionally minimal - `kernel_src/rubin/training/__init__.py` was synchronized byte-for-byte with - the revision above. Other integration-specific behavior lives outside this - directory. -- License: synchronized Python sources retain their upstream BSD-3-Clause - SPDX identifiers. `LICENSE.Apache-2.0` remains as historical snapshot - metadata. - -The source repository URL is intentionally omitted because it is an internal -development location. The revisions above identify the upstream baselines; -the manifest and local-modification notes below describe the packaged snapshot. - -## Scope - -This directory contains the recursive Python import closures for these Rubin -SM107 products: - -- training MegaMoE forward GLU; -- optional forward MXFP8 column requantization (disabled by default); -- training MegaMoE backward dGLU. - -It preserves the `next/sources` package hierarchy and includes the shared API, -quantization, workspace, synchronization, NVLink token communication, -schedulers, TopK reduction, and Rubin helper modules required by those roots. -The Rubin training initializer remains a minimal package marker so importing -the MegaMoE products does not pull in the unused traditional-wgrad product. -The public backend exposes the backward dGLU product through a restricted Rubin -MXFP8 dgrad/dprob path. Unsupported formats and semantics retain explicit -capability gates; see the backend README. - -Complete Blackwell kernel products, runners, tests, and repository-only tooling -are excluded. Three architecture-neutral Blackwell donor modules remain because -the upstream Rubin source-copy shims import them. Imports of external CUTLASS -utility modules remain because they are CUTLASS helpers, not vendored kernel -support. - -## Manifest - -```text -LICENSE.Apache-2.0 -VENDOR_INFO.md -__init__.py -api.py -communication/__init__.py -communication/nvlink_domain/__init__.py -communication/nvlink_domain/symmetric_buffer.py -communication/nvlink_domain/token_comm.py -communication/nvlink_domain/token_comm_deterministic.py -communication/token_protocol.py -helpers/__init__.py -helpers/constants.py -helpers/cute_py_helpers.py -helpers/device_workspace.py -helpers/dsl_helpers.py -helpers/flag_batch.py -helpers/iket_compat.py -helpers/ptx_helpers.py -helpers/software_sync.py -helpers/smem_workspace.py -helpers/utils.py -kernel_src/__init__.py -kernel_src/function_mapping.py -kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py -kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py -kernel_src/blackwell/inference/mega/topk_reduce.py -kernel_src/rubin/__init__.py -kernel_src/rubin/training/__init__.py -kernel_src/rubin/training/mega/__init__.py -kernel_src/rubin/training/mega/bwd_dglu/__init__.py -kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_epilogue.py -kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_extension.py -kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py -kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py -kernel_src/rubin/training/mega/fwd_glu/__init__.py -kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py -kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_epilogue.py -kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_extension.py -kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_fc12_kernel.py -kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_mega_moe_kernel.py -kernel_src/rubin/training/mega/helpers/__init__.py -kernel_src/rubin/training/mega/helpers/constants.py -kernel_src/rubin/training/mega/helpers/utils.py -kernel_src/rubin/training/mega/tmem_transpose.py -kernel_src/rubin/training/mega/topk_reduce.py -kernel_src/schedulers/__init__.py -kernel_src/schedulers/base.py -kernel_src/schedulers/fc12_mapping.py -kernel_src/schedulers/fc12_scheduler.py -kernel_src/schedulers/non_clc_mixed_cga.py -kernel_src/schedulers/work_id_claim.py -quant_def.py -``` - -## Integration boundary - -- Every `.py` file listed in the manifest except - `kernel_src/rubin/training/__init__.py` is a byte-for-byte copy of the same - relative path at revision `5b89819cb16069dfe20a1a0ba0778d35cb428352`. -- `kernel_src/rubin/training/__init__.py` is intentionally reduced to a - package marker. This avoids vendoring and eagerly importing the unused Rubin - traditional-wgrad product. -- `helpers/software_sync.py` replaces the former integration-only - `communication/nvlink_domain/software_sync.py` path. -- The upstream Rubin `topk_reduce.py` and `tmem_transpose.py` source-copy shims - require the three architecture-neutral Blackwell donor modules in the - manifest. Unrelated Blackwell and Rubin inference products remain excluded. -- Public API validation, symmetric-workspace ownership, overflow reporting, - staging, CUDA Graph handling, dprob materialization, and grouped-wgrad layout - conversion live in the parent `_megamoe_backend` package. -- The synchronized Rubin sources require a CUTLASS DSL distribution that - provides `cutlass.utils.rubin_helpers`. - -## Updating the snapshot - -1. Review source changes from the revision above. -2. Recompute recursive relative imports from all Rubin SM107 product roots, - exact package initializers, and shared product entry points. -3. Include only the exact donor and eager-import dependencies required by that - closure; do not add unrelated products. -4. Copy every selected Python source without modification and remove paths no - longer present in the selected upstream closure. -5. Update the revision, date, manifest, and integration-boundary notes. -6. Assert byte equality for every vendored Python source, then run compileall, - relative-import closure validation, package import smoke tests, and the - MoeEP regression suite. diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/__init__.py index 8feb6089f..6d10a7e6e 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/__init__.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/__init__.py @@ -1 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Integration-ready CuTeDSL MegaMoE sources.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/api.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/api.py index c3fd0af46..1935befbb 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/api.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/api.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Stable construction API for composable kernel implementations.""" import types diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/__init__.py index cc1ef5158..a3fa301fa 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/__init__.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/__init__.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Cross-rank communication protocols and implementations.""" from ..quant_def import CombineFormat, QuantKind diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/__init__.py index c72d7918a..962122907 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/__init__.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/__init__.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """NVLink-domain pointer mapping and token communication.""" from ...helpers.software_sync import NvlinkBarrier, SoftwareGridSync diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py index 0fd98b9fb..c166b817c 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Symmetric-heap peer pointer mapping carried in kernel arguments.""" from dataclasses import dataclass diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm.py index 793ea2167..a24a0997f 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Metadata-push routing and fused non-deterministic token communication.""" import dataclasses diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm_deterministic.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm_deterministic.py index 05381406f..972d26b79 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm_deterministic.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/token_comm_deterministic.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Metadata-push routing with a fixed token sequence and fused communication.""" import dataclasses diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/__init__.py index 0e7d0e526..780c03c78 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/__init__.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/__init__.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Low-level workspace, synchronization, and PTX helpers.""" from .device_workspace import DeviceWorkspace diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/constants.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/constants.py index 7870293a0..839ba88d6 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/constants.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/constants.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Specification-defined numeric constants shared across kernel components.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/cute_py_helpers.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/cute_py_helpers.py index 8cf623735..13736388e 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/cute_py_helpers.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/cute_py_helpers.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + from dataclasses import dataclass from math import gcd from typing import Optional, Tuple, Union diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py index 4956e6bd2..252adfc16 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Declarative GMEM workspace shared by host layout and device access.""" import dataclasses diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/dsl_helpers.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/dsl_helpers.py index 807d85c27..7221b2954 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/dsl_helpers.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/dsl_helpers.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """General-purpose CuTe DSL helpers.""" from typing import Callable, Literal, Optional diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/flag_batch.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/flag_batch.py index 28088ec8f..e69d60d29 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/flag_batch.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/flag_batch.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Lane- and warp-distributed GPU release-counter batching.""" import dataclasses diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/iket_compat.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/iket_compat.py index d4cfd4c97..a076067ae 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/iket_compat.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/iket_compat.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Compatibility wrapper for optional in-kernel event tracing support.""" try: diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py index 56d81bc87..9f56235e9 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Minimal inline-PTX primitives required by the greenfield NVFP4 kernels.""" from typing import Optional diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py index e6fe04c53..db5cdb729 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/smem_workspace.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Pure-Python SMEM declarations with lifetime-aware overlay placement.""" import dataclasses diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/software_sync.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/software_sync.py index 9b1ef4010..af282e016 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/software_sync.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/software_sync.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Software grid and NVLink synchronization for persistent kernels.""" import cutlass diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py index b14dae2cd..59d78e743 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Small integer and layout helpers shared by workspace implementations.""" from typing import Iterable, List, Tuple, Union diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/__init__.py index 337547020..1b80fdf34 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/__init__.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/__init__.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Kernel source modules independent of repository-only runners.""" from .schedulers import ( diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py index fc47fe0e3..c45dc66cf 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Block-scaled SwapAb adapter between FC12 scheduling and kernel tensor views.""" import dataclasses diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/__init__.py index b84f4d01b..f763de2fc 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/__init__.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/__init__.py @@ -1 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Rubin kernel implementations.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/__init__.py index 4e1908109..6152b881d 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/__init__.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/__init__.py @@ -1 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Rubin SM107 training kernel package.""" diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/__init__.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/__init__.py index 6b1dde777..19945fd80 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/__init__.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/__init__.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Scheduler protocols and implementations.""" from .base import SchedulerBase, SchedulerConsumer, SchedulerWorkTileBase, WorkIdAcquisitionMode diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/base.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/base.py index f4a1f40ef..ceffeb665 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/base.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/base.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Scheduler façade and architecture-independent work-tile transport.""" from abc import ABC, abstractmethod diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py index 2afe2ba83..af270cf14 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """FC12 work-tile ABI and grouped or phase-interleaved task mapping.""" import dataclasses diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_scheduler.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_scheduler.py index ffedfb5c0..a0edfbc87 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_scheduler.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_scheduler.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Composable grouped and phase-interleaved FC12 schedulers.""" import math diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/non_clc_mixed_cga.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/non_clc_mixed_cga.py index c0f0fa8a5..c1df89496 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/non_clc_mixed_cga.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/non_clc_mixed_cga.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Reusable preferred/fallback cluster scheduling without hardware CLC.""" import dataclasses diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/work_id_claim.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/work_id_claim.py index 93a9be875..222885b64 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/work_id_claim.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/work_id_claim.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + """Blackwell persistent work-ID claim backends.""" import dataclasses diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py index c3e67cf87..f32978734 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py @@ -69,8 +69,7 @@ def _ensure_prepared_kernel(self) -> PreparedMxfp8Kernel: ) except (ImportError, OSError) as exc: raise BackendUnavailableError( - "MoeEp MXFP8 backend requires the 'moe_ep' optional " - "dependencies and their shared libraries" + "MoeEp MXFP8 backend requires the 'cutedsl' and 'comm' " "optional dependencies and their shared libraries" ) from exc return self._prepared_kernel @@ -83,27 +82,16 @@ def _ensure_ep_launch_ready(self, resources, stream) -> None: # those writes. stream.synchronize() if resources.runtime.group is None: - raise RuntimeError( - "distributed MXFP8 launch requires a " - "torch.distributed process group" - ) - tuning_signature = self.kernel_config.tuning_signature( - self._ensure_prepared_kernel().launch_cluster_count - ) + raise RuntimeError("distributed MXFP8 launch requires a " "torch.distributed process group") + tuning_signature = self.kernel_config.tuning_signature(self._ensure_prepared_kernel().launch_cluster_count) rank_tuning_signatures = [None] * resources.runtime.world_size dist.all_gather_object( rank_tuning_signatures, tuning_signature, group=resources.runtime.group, ) - if any( - signature != rank_tuning_signatures[0] - for signature in rank_tuning_signatures[1:] - ): - raise RuntimeError( - "MoeEp tuning must match on every expert-parallel rank; " - f"effective signatures by rank: {rank_tuning_signatures}" - ) + if any(signature != rank_tuning_signatures[0] for signature in rank_tuning_signatures[1:]): + raise RuntimeError("MoeEp tuning must match on every expert-parallel rank; " f"effective signatures by rank: {rank_tuning_signatures}") dist.barrier(group=resources.runtime.group) self._ep_launch_ready = True @@ -112,32 +100,18 @@ def forward(self, request: ValidatedForwardRequest): if self._closed: raise RuntimeError("MoeEp MXFP8 backend is closed") if request.device != self.device: - raise ValueError( - f"MoeEp MXFP8 backend is bound to {self.device}, " - f"got {request.device}" - ) + raise ValueError(f"MoeEp MXFP8 backend is bound to {self.device}, " f"got {request.device}") with torch.cuda.device(self.device): capturing = torch.cuda.is_current_stream_capturing() - if ( - capturing - and not self._adapter.weights_have_version_counters( - request - ) - ): + if capturing and not self._adapter.weights_have_version_counters(request): raise NotImplementedError( "CUDA graph capture does not support inference tensor " "weights without version counters; eager calls remain " "supported and repack those weights on every call" ) - if capturing and ( - not self._warmed_up - or not self._adapter.has_cached_weights(request) - ): - raise RuntimeError( - "MoeEp MXFP8 backend and weights must be warmed up " - "before CUDA graph capture" - ) + if capturing and (not self._warmed_up or not self._adapter.has_cached_weights(request)): + raise RuntimeError("MoeEp MXFP8 backend and weights must be warmed up " "before CUDA graph capture") stream = torch.cuda.current_stream(self.device) if self._device_work_may_be_pending: @@ -167,24 +141,12 @@ def forward(self, request: ValidatedForwardRequest): request, resources, self.kernel_config, - local_workspace_zero_bytes=( - prepared.local_workspace_zero_bytes - ), - shared_workspace_zero_bytes=( - prepared.shared_workspace_zero_bytes - ), - pre_reduced_activation_offset=( - prepared.pre_reduced_activation_offset - ), - pre_reduced_activation_bytes_per_token=( - prepared.pre_reduced_activation_bytes_per_token - ), - pre_reduced_activation_sf_offset=( - prepared.pre_reduced_activation_sf_offset - ), - pre_reduced_activation_sf_bytes_per_token=( - prepared.pre_reduced_activation_sf_bytes_per_token - ), + local_workspace_zero_bytes=(prepared.local_workspace_zero_bytes), + shared_workspace_zero_bytes=(prepared.shared_workspace_zero_bytes), + pre_reduced_activation_offset=(prepared.pre_reduced_activation_offset), + pre_reduced_activation_bytes_per_token=(prepared.pre_reduced_activation_bytes_per_token), + pre_reduced_activation_sf_offset=(prepared.pre_reduced_activation_sf_offset), + pre_reduced_activation_sf_bytes_per_token=(prepared.pre_reduced_activation_sf_bytes_per_token), col_quant_data_rows=prepared.col_quant_data_rows, col_quant_sf_elements=prepared.col_quant_sf_elements, fc1_c=None, @@ -202,8 +164,7 @@ def forward(self, request: ValidatedForwardRequest): ) except (ImportError, OSError) as exc: raise BackendUnavailableError( - "MoeEp MXFP8 backend requires the 'moe_ep' optional " - "dependencies and their shared libraries" + "MoeEp MXFP8 backend requires the 'cutedsl' and 'comm' " "optional dependencies and their shared libraries" ) from exc finally: if device_work_attempted and not capturing: @@ -240,9 +201,7 @@ def prepare_training_resources( token_padding_size=128, sf_padding_size=128, ) - training_kernel_config = Mxfp8KernelConfig.from_forward_config( - training_config - ) + training_kernel_config = Mxfp8KernelConfig.from_forward_config(training_config) # Graph transport must complete its cross-rank protocol before the # frontend applies the public trap/drop policy at graph tail. graph_kernel_config = replace( @@ -290,14 +249,8 @@ def close(self) -> None: return with torch.cuda.device(self.device): if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "MoeEp MXFP8 backend cannot be closed during " - "CUDA graph capture" - ) - if ( - self._plan is not None - or self._training_resource_owner is not None - ): + raise RuntimeError("MoeEp MXFP8 backend cannot be closed during " "CUDA graph capture") + if self._plan is not None or self._training_resource_owner is not None: torch.cuda.synchronize(self.device) self._adapter.close() if self._training_resource_owner is not None: diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py index 931f1d55b..190647e9d 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py @@ -21,6 +21,7 @@ _pre_reduced_workspace_metadata, ) from ._config import Mxfp8KernelConfig +from ._cutedsl import require_rubin_cutedsl from ._formats import combine_wire_format from ._launch import _to_cute, _to_cute_ptr @@ -88,20 +89,16 @@ def prepare_backward_kernel( ) -> PreparedMxfp8BackwardKernel: """Instantiate the fixed Rubin dGLU specialization.""" + require_rubin_cutedsl() torch.cuda.set_device(device) architecture = torch.cuda.get_device_capability(device) if architecture != (10, 7): - raise RuntimeError( - "Rubin MXFP8 backward requires compute capability (10, 7), " - f"got {architecture}" - ) + raise RuntimeError("Rubin MXFP8 backward requires compute capability (10, 7), " f"got {architecture}") configured_architecture = os.environ.get("CUTE_DSL_ARCH") if configured_architecture is None: os.environ["CUTE_DSL_ARCH"] = "sm_107a" elif configured_architecture not in ("sm_107", "sm_107a"): - raise RuntimeError( - "CUTE_DSL_ARCH must target SM107 for the Rubin MXFP8 backward" - ) + raise RuntimeError("CUTE_DSL_ARCH must target SM107 for the Rubin MXFP8 backward") import cutlass import cutlass.utils as utils @@ -110,18 +107,10 @@ def prepare_backward_kernel( ) from ..cutedsl_src.quant_def import CombineFormat - launch_cluster_count = int( - utils.HardwareInfo().get_max_active_clusters(config.cluster_size) - ) + launch_cluster_count = int(utils.HardwareInfo().get_max_active_clusters(config.cluster_size)) if launch_cluster_count <= 0: - raise RuntimeError( - "hardware occupancy query returned no launchable Rubin clusters" - ) - group_hint = ( - launch_cluster_count - if config.group_hint is None - else config.group_hint - ) + raise RuntimeError("hardware occupancy query returned no launchable Rubin clusters") + group_hint = launch_cluster_count if config.group_hint is None else config.group_hint operands_mode = forward_config.backward_wgrad_mode == "operands" dfc2_recompute = operands_mode dfc2_col_output = operands_mode @@ -156,9 +145,7 @@ def prepare_backward_kernel( token_back_mode="epi_warps", epi_flag_batch=config.epi_flag_batch, flag_batch=config.flag_batch, - combine_format=CombineFormat.parse( - combine_wire_format(forward_config.combine_format) - ), + combine_format=CombineFormat.parse(combine_wire_format(forward_config.combine_format)), act_func=config.act_func, gate_up_clamp=config.gate_up_clamp, dfc2_recompute=dfc2_recompute, @@ -170,36 +157,32 @@ def prepare_backward_kernel( local_zero, shared_zero = kernel.require_zero_workspace_leading_bytes device_workspace = kernel._mega_device_workspace pool_capacity = int(kernel.pool_token_capacity) - fc1_preact_shape = tuple( - int(extent) for extent in kernel.get_fc1_preact_shape() - ) + fc1_preact_shape = tuple(int(extent) for extent in kernel.get_fc1_preact_shape()) expected_preact_shape = ( pool_capacity, 2 * config.intermediate, ) if fc1_preact_shape != expected_preact_shape: - raise RuntimeError( - "Rubin dGLU fc1_preact shape mismatch: " - f"{fc1_preact_shape} != {expected_preact_shape}" + raise RuntimeError("Rubin dGLU fc1_preact shape mismatch: " f"{fc1_preact_shape} != {expected_preact_shape}") + aux_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in kernel.get_aux_output_shapes().items()} + fc1_preact_bytes = math.prod(fc1_preact_shape) * torch.bfloat16.itemsize + dprob_bytes = math.prod(aux_shapes["dprob"]) * torch.float32.itemsize + aux_data_bytes = ( + max( + math.prod(aux_shapes["fc1_recompute"]), + math.prod(aux_shapes["fc1_col_output"]), + math.prod(aux_shapes["grad_y2"]), ) - aux_shapes = { - name: tuple(int(extent) for extent in shape) - for name, shape in kernel.get_aux_output_shapes().items() - } - fc1_preact_bytes = ( - math.prod(fc1_preact_shape) * torch.bfloat16.itemsize + * torch.float8_e4m3fn.itemsize + ) + aux_scale_bytes = ( + max( + math.prod(aux_shapes["fc1_recompute_sf"]), + math.prod(aux_shapes["fc1_col_output_sf"]), + math.prod(aux_shapes["grad_y2_sf"]), + ) + * torch.float8_e8m0fnu.itemsize ) - dprob_bytes = math.prod(aux_shapes["dprob"]) * torch.float32.itemsize - aux_data_bytes = max( - math.prod(aux_shapes["fc1_recompute"]), - math.prod(aux_shapes["fc1_col_output"]), - math.prod(aux_shapes["grad_y2"]), - ) * torch.float8_e4m3fn.itemsize - aux_scale_bytes = max( - math.prod(aux_shapes["fc1_recompute_sf"]), - math.prod(aux_shapes["fc1_col_output_sf"]), - math.prod(aux_shapes["grad_y2_sf"]), - ) * torch.float8_e8m0fnu.itemsize requirements = WorkspaceRequirements.for_mxfp8( forward_config, kernel_local_workspace_bytes=local_bytes, @@ -209,23 +192,17 @@ def prepare_backward_kernel( backward_aux_data_bytes=aux_data_bytes, backward_aux_scale_bytes=aux_scale_bytes, ) - pre_reduced_offset, pre_reduced_bytes_per_token = ( - _pre_reduced_workspace_metadata( - device_workspace, - config, - shared_bytes, - ) + pre_reduced_offset, pre_reduced_bytes_per_token = _pre_reduced_workspace_metadata( + device_workspace, + config, + shared_bytes, ) if pre_reduced_offset is None or pre_reduced_bytes_per_token <= 0: - raise RuntimeError( - "Rubin MXFP8 backward requires standalone pre-reduced activation" - ) - pre_reduced_sf_offset, pre_reduced_sf_bytes_per_token = ( - _pre_reduced_sf_workspace_metadata( - device_workspace, - config, - shared_bytes, - ) + raise RuntimeError("Rubin MXFP8 backward requires standalone pre-reduced activation") + pre_reduced_sf_offset, pre_reduced_sf_bytes_per_token = _pre_reduced_sf_workspace_metadata( + device_workspace, + config, + shared_bytes, ) return PreparedMxfp8BackwardKernel( config=config, @@ -238,9 +215,7 @@ def prepare_backward_kernel( pre_reduced_activation_offset=pre_reduced_offset, pre_reduced_activation_bytes_per_token=pre_reduced_bytes_per_token, pre_reduced_activation_sf_offset=pre_reduced_sf_offset, - pre_reduced_activation_sf_bytes_per_token=( - pre_reduced_sf_bytes_per_token - ), + pre_reduced_activation_sf_bytes_per_token=(pre_reduced_sf_bytes_per_token), local_workspace_zero_bytes=int(local_zero), shared_workspace_zero_bytes=int(shared_zero), dfc2_recompute=dfc2_recompute, @@ -250,15 +225,8 @@ def prepare_backward_kernel( def _layout_signature(inputs: Mxfp8BackwardLaunchInputs) -> tuple: - tensors = tuple( - value - for value in inputs.__dict__.values() - if isinstance(value, torch.Tensor) - ) - return tuple( - (tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype) - for tensor in tensors - ) + tensors = tuple(value for value in inputs.__dict__.values() if isinstance(value, torch.Tensor)) + return tuple((tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype) for tensor in tensors) def build_backward_runtime_kwargs( @@ -321,9 +289,7 @@ def build_backward_runtime_kwargs( ), "local_workspace": _to_cute_ptr(inputs.local_workspace), "shared_workspace": _to_cute_ptr(inputs.shared_workspace), - "peer_rank_ptr_mapper_host": ( - resources.workspace.peer_mapping.to_sym_buffer_host() - ), + "peer_rank_ptr_mapper_host": (resources.workspace.peer_mapping.to_sym_buffer_host()), "stream": cuda.CUstream(stream.cuda_stream), } @@ -349,9 +315,7 @@ def compile_backward_or_get( if cached is not None: return cached if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "MXFP8 backward kernel must be compiled before capture" - ) + raise RuntimeError("MXFP8 backward kernel must be compiled before capture") import cutlass.cute as cute runtime_kwargs = build_backward_runtime_kwargs(inputs, resources) diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py index 6b4a3bd3e..26769c201 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py @@ -17,6 +17,7 @@ from .._workspace import WorkspaceRequirements from ._adapter import Mxfp8LaunchInputs from ._config import Mxfp8KernelConfig +from ._cutedsl import require_rubin_cutedsl from ._fingerprint import build_kernel_fingerprint from ._launch import build_runtime_kwargs, layout_signature @@ -54,20 +55,15 @@ class CompiledMxfp8Kernel: _COMPILE_LOCK = threading.RLock() _COMPILE_CACHE: dict[tuple, CompiledMxfp8Kernel] = {} _TOKEN_SRC_METADATA_REGION = "nvlink.token_comm.token_src_metadata" -_COL_QUANT_SIZES_REGION = ( - "rubin.glu_mxfp8.mega.col_quant_expert_token_sizes" -) -_PRE_REDUCED_ACTIVATION_REGION = ( - "nvlink.token_comm.pre_reduced_activation" -) -_PRE_REDUCED_ACTIVATION_SF_REGION = ( - "nvlink.token_comm.pre_reduced_activation_sf" -) +_COL_QUANT_SIZES_REGION = "rubin.glu_mxfp8.mega.col_quant_expert_token_sizes" +_PRE_REDUCED_ACTIVATION_REGION = "nvlink.token_comm.pre_reduced_activation" +_PRE_REDUCED_ACTIVATION_SF_REGION = "nvlink.token_comm.pre_reduced_activation_sf" def _compile_kernel(kernel: Any, compile_kwargs: dict[str, Any]) -> Any: """Import CuTeDSL only on a cache miss and compile one callable.""" + require_rubin_cutedsl() import cutlass.cute as cute return cute.compile(kernel, **compile_kwargs) @@ -84,15 +80,9 @@ def _pre_reduced_workspace_metadata( region = device_workspace.region(_PRE_REDUCED_ACTIVATION_REGION) if region.buffer_space != "shared": - raise RuntimeError( - "Rubin pre_reduced_activation must reside in shared workspace" - ) - offset = int( - device_workspace.offset(_PRE_REDUCED_ACTIVATION_REGION) - ) - nbytes = int( - device_workspace.nbytes(_PRE_REDUCED_ACTIVATION_REGION) - ) + raise RuntimeError("Rubin pre_reduced_activation must reside in shared workspace") + offset = int(device_workspace.offset(_PRE_REDUCED_ACTIVATION_REGION)) + nbytes = int(device_workspace.nbytes(_PRE_REDUCED_ACTIVATION_REGION)) wire_bits_per_element = { "bf16": 16, "32e4m3xe8m0": 8, @@ -100,9 +90,7 @@ def _pre_reduced_workspace_metadata( try: element_bits = wire_bits_per_element[config.combine_format] except KeyError as exc: - raise ValueError( - f"unsupported combine wire format {config.combine_format!r}" - ) from exc + raise ValueError(f"unsupported combine wire format {config.combine_format!r}") from exc wire_bits_per_token = config.top_k * config.hidden * element_bits if wire_bits_per_token % 8: raise RuntimeError("combine wire row is not byte aligned") @@ -110,14 +98,10 @@ def _pre_reduced_workspace_metadata( expected_bytes = config.max_tokens_per_rank * bytes_per_token if nbytes != expected_bytes: raise RuntimeError( - "Rubin pre_reduced_activation size does not match " - f"combine_format={config.combine_format!r}: {nbytes} bytes, " - f"expected {expected_bytes}" + "Rubin pre_reduced_activation size does not match " f"combine_format={config.combine_format!r}: {nbytes} bytes, " f"expected {expected_bytes}" ) if offset + nbytes > shared_bytes: - raise RuntimeError( - "Rubin pre_reduced_activation region exceeds shared workspace" - ) + raise RuntimeError("Rubin pre_reduced_activation region exceeds shared workspace") return offset, bytes_per_token @@ -132,19 +116,13 @@ def _pre_reduced_sf_workspace_metadata( region = device_workspace.region(_PRE_REDUCED_ACTIVATION_SF_REGION) if region.buffer_space != "shared": - raise RuntimeError( - "Rubin pre_reduced_activation_sf must reside in shared workspace" - ) + raise RuntimeError("Rubin pre_reduced_activation_sf must reside in shared workspace") offset = int(device_workspace.offset(_PRE_REDUCED_ACTIVATION_SF_REGION)) nbytes = int(device_workspace.nbytes(_PRE_REDUCED_ACTIVATION_SF_REGION)) if nbytes % config.max_tokens_per_rank: - raise RuntimeError( - "Rubin pre_reduced_activation_sf size is not token aligned" - ) + raise RuntimeError("Rubin pre_reduced_activation_sf size is not token aligned") if offset + nbytes > shared_bytes: - raise RuntimeError( - "Rubin pre_reduced_activation_sf region exceeds shared workspace" - ) + raise RuntimeError("Rubin pre_reduced_activation_sf region exceeds shared workspace") return offset, nbytes // config.max_tokens_per_rank @@ -155,21 +133,16 @@ def prepare_kernel( ) -> PreparedMxfp8Kernel: """Instantiate the kernel and derive exact allocation requirements.""" + require_rubin_cutedsl() torch.cuda.set_device(device) architecture = torch.cuda.get_device_capability(device) if architecture != (10, 7): - raise RuntimeError( - "Rubin MXFP8 kernel preparation requires compute capability " - f"(10, 7), got {architecture}" - ) + raise RuntimeError("Rubin MXFP8 kernel preparation requires compute capability " f"(10, 7), got {architecture}") configured_architecture = os.environ.get("CUTE_DSL_ARCH") if configured_architecture is None: os.environ["CUTE_DSL_ARCH"] = "sm_107a" elif configured_architecture not in ("sm_107", "sm_107a"): - raise RuntimeError( - "CUTE_DSL_ARCH must target SM107 for the Rubin MXFP8 backend, " - f"got {configured_architecture!r}" - ) + raise RuntimeError("CUTE_DSL_ARCH must target SM107 for the Rubin MXFP8 backend, " f"got {configured_architecture!r}") import cutlass import cutlass.utils as utils @@ -179,18 +152,10 @@ def prepare_kernel( ) from ..cutedsl_src.quant_def import CombineFormat - launch_cluster_count = int( - utils.HardwareInfo().get_max_active_clusters(config.cluster_size) - ) + launch_cluster_count = int(utils.HardwareInfo().get_max_active_clusters(config.cluster_size)) if launch_cluster_count <= 0: - raise RuntimeError( - "hardware occupancy query returned no launchable Rubin clusters" - ) - group_hint = ( - launch_cluster_count - if config.group_hint is None - else config.group_hint - ) + raise RuntimeError("hardware occupancy query returned no launchable Rubin clusters") + group_hint = launch_cluster_count if config.group_hint is None else config.group_hint kernel_kwargs = dict( mma_tiler_mnk=config.mma_tiler_mnk, cluster_shape_mnk=config.cluster_shape_mnk, @@ -235,74 +200,37 @@ def prepare_kernel( ) kernel = Sm107MegaMoEMxfp8GluKernel.from_kwargs(**kernel_kwargs) local_bytes, shared_bytes = kernel.get_workspace_sizes() - local_zero_bytes, shared_zero_bytes = ( - kernel.require_zero_workspace_leading_bytes - ) + local_zero_bytes, shared_zero_bytes = kernel.require_zero_workspace_leading_bytes for name, zero_bytes, total_bytes in ( ("local", local_zero_bytes, local_bytes), ("shared", shared_zero_bytes, shared_bytes), ): if zero_bytes < 0 or zero_bytes > total_bytes: - raise RuntimeError( - f"Rubin kernel {name} zero prefix {zero_bytes} exceeds " - f"workspace size {total_bytes}" - ) + raise RuntimeError(f"Rubin kernel {name} zero prefix {zero_bytes} exceeds " f"workspace size {total_bytes}") device_workspace = kernel._mega_device_workspace metadata_region = device_workspace.region(_TOKEN_SRC_METADATA_REGION) if metadata_region.buffer_space != "shared": - raise RuntimeError( - "Rubin token_src_metadata must reside in shared workspace" - ) - token_src_metadata_offset = int( - device_workspace.offset(_TOKEN_SRC_METADATA_REGION) - ) - token_src_metadata_bytes = int( - device_workspace.nbytes(_TOKEN_SRC_METADATA_REGION) - ) + raise RuntimeError("Rubin token_src_metadata must reside in shared workspace") + token_src_metadata_offset = int(device_workspace.offset(_TOKEN_SRC_METADATA_REGION)) + token_src_metadata_bytes = int(device_workspace.nbytes(_TOKEN_SRC_METADATA_REGION)) if token_src_metadata_offset + token_src_metadata_bytes > shared_bytes: - raise RuntimeError( - "Rubin token_src_metadata region exceeds shared workspace" - ) + raise RuntimeError("Rubin token_src_metadata region exceeds shared workspace") pool_token_capacity = int(kernel.pool_token_capacity) if token_src_metadata_bytes != pool_token_capacity * 8: - raise RuntimeError( - "Rubin token_src_metadata must contain one Int64 per pool token" - ) - col_quant_data_rows = ( - pool_token_capacity if config.enable_col_quant else 0 - ) - col_quant_sf_elements = ( - int(kernel.token_comm.worst_case_sf_token_count) - * (config.hidden // config.sf_vec_size) - if config.enable_col_quant - else 0 - ) + raise RuntimeError("Rubin token_src_metadata must contain one Int64 per pool token") + col_quant_data_rows = pool_token_capacity if config.enable_col_quant else 0 + col_quant_sf_elements = int(kernel.token_comm.worst_case_sf_token_count) * (config.hidden // config.sf_vec_size) if config.enable_col_quant else 0 if config.enable_col_quant: - col_quant_sizes_region = device_workspace.region( - _COL_QUANT_SIZES_REGION - ) + col_quant_sizes_region = device_workspace.region(_COL_QUANT_SIZES_REGION) if col_quant_sizes_region.buffer_space != "local": - raise RuntimeError( - "Rubin col-quant expert-size snapshot must reside in " - "local workspace" - ) - col_quant_sizes_offset = int( - device_workspace.offset(_COL_QUANT_SIZES_REGION) - ) - col_quant_sizes_bytes = int( - device_workspace.nbytes(_COL_QUANT_SIZES_REGION) - ) + raise RuntimeError("Rubin col-quant expert-size snapshot must reside in " "local workspace") + col_quant_sizes_offset = int(device_workspace.offset(_COL_QUANT_SIZES_REGION)) + col_quant_sizes_bytes = int(device_workspace.nbytes(_COL_QUANT_SIZES_REGION)) expected_sizes_bytes = config.num_experts * torch.int32.itemsize if col_quant_sizes_bytes != expected_sizes_bytes: - raise RuntimeError( - "Rubin col-quant expert-size snapshot has " - f"{col_quant_sizes_bytes} bytes, expected " - f"{expected_sizes_bytes}" - ) + raise RuntimeError("Rubin col-quant expert-size snapshot has " f"{col_quant_sizes_bytes} bytes, expected " f"{expected_sizes_bytes}") if col_quant_sizes_offset + col_quant_sizes_bytes > local_bytes: - raise RuntimeError( - "Rubin col-quant expert-size snapshot exceeds local workspace" - ) + raise RuntimeError("Rubin col-quant expert-size snapshot exceeds local workspace") else: col_quant_sizes_offset = None col_quant_sizes_bytes = 0 @@ -344,13 +272,9 @@ def prepare_kernel( col_quant_sizes_offset=col_quant_sizes_offset, col_quant_sizes_bytes=col_quant_sizes_bytes, pre_reduced_activation_offset=pre_reduced_activation_offset, - pre_reduced_activation_bytes_per_token=( - pre_reduced_activation_bytes_per_token - ), + pre_reduced_activation_bytes_per_token=(pre_reduced_activation_bytes_per_token), pre_reduced_activation_sf_offset=pre_reduced_activation_sf_offset, - pre_reduced_activation_sf_bytes_per_token=( - pre_reduced_activation_sf_bytes_per_token - ), + pre_reduced_activation_sf_bytes_per_token=(pre_reduced_activation_sf_bytes_per_token), local_workspace_zero_bytes=int(local_zero_bytes), shared_workspace_zero_bytes=int(shared_zero_bytes), ) diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_cutedsl.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_cutedsl.py new file mode 100644 index 000000000..e8fa3b48c --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_cutedsl.py @@ -0,0 +1,56 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CUTLASS DSL compatibility gate for Rubin MegaMoE kernels.""" + +from __future__ import annotations + +import importlib.metadata + +RUBIN_CUTEDSL_MIN_VERSION = (4, 8, 0) + + +def _public_cutedsl_version() -> str | None: + """Return public-wheel metadata without importing CUTLASS DSL.""" + + try: + return importlib.metadata.version("nvidia-cutlass-dsl") + except importlib.metadata.PackageNotFoundError: + return None + + +def _parse_version(version: str) -> tuple[int, int, int] | None: + """Parse the numeric release prefix and tolerate prerelease suffixes.""" + + parts = version.split("+", 1)[0].split(".") + parsed = [] + try: + for part in parts[:3]: + digits = "" + for character in part: + if not character.isdigit(): + break + digits += character + if not digits: + return None + parsed.append(int(digits)) + except (TypeError, ValueError): + return None + return tuple(parsed) if len(parsed) == 3 else None + + +def require_rubin_cutedsl() -> None: + """Reject public CUTLASS DSL wheels older than Rubin kernel support.""" + + version = _public_cutedsl_version() + parsed = None if version is None else _parse_version(version) + if parsed is not None and parsed < RUBIN_CUTEDSL_MIN_VERSION: + raise RuntimeError( + "Rubin MegaMoE MXFP8 kernels require " + "nvidia-cutlass-dsl>=4.8.0; found " + f"{version}. Other cuDNN Frontend APIs remain available with " + "the package minimum of 4.5.0" + ) + + +__all__ = ["RUBIN_CUTEDSL_MIN_VERSION", "require_rubin_cutedsl"] diff --git a/test/python/moe_ep/test_moe_ep_cutedsl.py b/test/python/moe_ep/test_moe_ep_cutedsl.py new file mode 100644 index 000000000..2b9c58536 --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_cutedsl.py @@ -0,0 +1,66 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CUTLASS DSL version-gate tests for Rubin MegaMoE.""" + +import pytest + + +@pytest.mark.L0 +@pytest.mark.parametrize("version", ["4.5.0", "4.6.1", "4.7.0"]) +def test_rubin_cutedsl_gate_rejects_public_wheels_below_4_8( + monkeypatch, + version, +): + from cudnn.moe_ep._megamoe_backend.mxfp8 import _cutedsl + + monkeypatch.setattr(_cutedsl, "_public_cutedsl_version", lambda: version) + + with pytest.raises( + RuntimeError, + match=r"nvidia-cutlass-dsl>=4\.8\.0", + ): + _cutedsl.require_rubin_cutedsl() + + +@pytest.mark.L0 +@pytest.mark.parametrize("version", ["4.8.0", "4.8.0rc1", "4.9.0"]) +def test_rubin_cutedsl_gate_accepts_4_8_or_newer(monkeypatch, version): + from cudnn.moe_ep._megamoe_backend.mxfp8 import _cutedsl + + monkeypatch.setattr(_cutedsl, "_public_cutedsl_version", lambda: version) + + _cutedsl.require_rubin_cutedsl() + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "module_name, function_name", + [ + ( + "cudnn.moe_ep._megamoe_backend.mxfp8._compile", + "prepare_kernel", + ), + ( + "cudnn.moe_ep._megamoe_backend.mxfp8._backward_compile", + "prepare_backward_kernel", + ), + ], +) +def test_rubin_prepare_gates_before_cuda_initialization( + monkeypatch, + module_name, + function_name, +): + module = __import__(module_name, fromlist=[function_name]) + + class GateReached(RuntimeError): + pass + + def reject(): + raise GateReached + + monkeypatch.setattr(module, "require_rubin_cutedsl", reject) + + with pytest.raises(GateReached): + getattr(module, function_name)(None, None, None) From ab810941750bd735fc693e547caa09e5ba8558ef Mon Sep 17 00:00:00 2001 From: zhibinz Date: Fri, 28 Aug 2026 09:07:54 -0700 Subject: [PATCH 16/31] test: focus MoeEP tuning coverage on contracts Remove redundant tuning value matrices while retaining public wiring, cache, distributed-consistency, and semantic coverage. --- test/python/moe_ep/test_moe_ep_backward.py | 48 ++---- test/python/moe_ep/test_moe_ep_forward.py | 175 +++------------------ 2 files changed, 43 insertions(+), 180 deletions(-) diff --git a/test/python/moe_ep/test_moe_ep_backward.py b/test/python/moe_ep/test_moe_ep_backward.py index a451db390..6de24fe51 100644 --- a/test/python/moe_ep/test_moe_ep_backward.py +++ b/test/python/moe_ep/test_moe_ep_backward.py @@ -50,7 +50,6 @@ from cudnn.moe_ep._megamoe_backend.mxfp8._training_weights import ( Mxfp8TrainingWeightBindings, ) -from cudnn.moe_ep._tuning import MoeEpTuningConfig from moe_ep.moe_ep_reference import ( MoeEpReference, ) @@ -641,9 +640,12 @@ def unexpected_call(*args, **kwargs): monkeypatch.setattr(backend_seam, "create_backend", unexpected_call) counts = {"slot_count": 1, "lane_count": 1, field: value} - with _operator() as operator, pytest.raises( - ValueError, - match=rf"{field} must be a positive integer", + with ( + _operator() as operator, + pytest.raises( + ValueError, + match=rf"{field} must be a positive integer", + ), ): operator.prepare_training_resources( SimpleNamespace(mock_training_weights=True), @@ -864,7 +866,11 @@ def all_reduce(tensor, *, op, group): "route-contiguity": (lambda t: t.update(topk_idx=t["topk_idx"].t().contiguous().t()), TypeError, "contiguous Int32"), "weight-dtype": (lambda t: t.update(topk_weights=t["topk_weights"].to(torch.bfloat16)), TypeError, "contiguous FP32"), "weight-contiguity": (lambda t: t.update(topk_weights=t["topk_weights"].t().contiguous().t()), TypeError, "contiguous FP32"), - "capacity": (lambda t: t.update(**{name: value[:4] for name, value in t.items() if name.startswith("output")}), ValueError, "token count 5 exceeds capacity 4"), + "capacity": ( + lambda t: t.update(**{name: value[:4] for name, value in t.items() if name.startswith("output")}), + ValueError, + "token count 5 exceeds capacity 4", + ), "device": (lambda t: t.update(source=torch.empty_like(t["source"], device="meta")), ValueError, "must share one device"), } @@ -872,10 +878,7 @@ def all_reduce(tensor, *, op, group): @pytest.mark.L1 @pytest.mark.parametrize( ("mutator", "error_type", "message"), - [ - pytest.param(*case, id=name) - for name, case in _STAGER_FAILURES.items() - ], + [pytest.param(*case, id=name) for name, case in _STAGER_FAILURES.items()], ) def test_training_stager_rejects_invalid_inputs(mutator, error_type, message): tensors = _training_staging_tensors() @@ -895,28 +898,14 @@ def test_training_stager_rejects_invalid_inputs(mutator, error_type, message): "combine_format", "gate_up_clamp", "top_k", - "tuning", "all_dropped", ), [ - pytest.param("fixed", "bf16", None, 2, MoeEpTuningConfig(), False, id="bf16-unclamped"), - pytest.param("fixed", "mxfp8", 0.5, 2, MoeEpTuningConfig(), False, id="mxfp8-clamp-0.5"), - pytest.param("routed", "bf16", None, 1, MoeEpTuningConfig(), False, id="topk1-default-tuning"), - pytest.param( - "routed", - "bf16", - None, - 2, - MoeEpTuningConfig( - token_back_mode="reuse_dispatch_warps", - epi_flag_batch=(2, 2), - token_in_flag_batch=2, - group_hint=128, - ), - False, - id="topk2-nondefault-tuning", - ), - pytest.param("routed", "bf16", None, 2, MoeEpTuningConfig(), True, id="topk2-all-dropped"), + pytest.param("fixed", "bf16", None, 2, False, id="bf16-unclamped"), + pytest.param("fixed", "mxfp8", 0.5, 2, False, id="mxfp8-clamp-0.5"), + pytest.param("routed", "bf16", None, 1, False, id="topk1"), + pytest.param("routed", "bf16", None, 2, False, id="topk2"), + pytest.param("routed", "bf16", None, 2, True, id="topk2-all-dropped"), ], ) def test_fixed_training_resources_ep1_matches_independent_reference( @@ -924,7 +913,6 @@ def test_fixed_training_resources_ep1_matches_independent_reference( combine_format, gate_up_clamp, top_k, - tuning, all_dropped, ): device = _sm107_device() @@ -950,7 +938,6 @@ def test_fixed_training_resources_ep1_matches_independent_reference( grad_output, combine_format=combine_format, gate_up_clamp=gate_up_clamp, - tuning=tuning, ) with MoeEp( @@ -963,7 +950,6 @@ def test_fixed_training_resources_ep1_matches_independent_reference( drop_on_overflow=True, combine_format=combine_format, gate_up_clamp=gate_up_clamp, - tuning=tuning, ) as op: resources = op.prepare_training_resources( _fixed_training_weights(args), diff --git a/test/python/moe_ep/test_moe_ep_forward.py b/test/python/moe_ep/test_moe_ep_forward.py index 1ab9cd23a..9eebbf5ad 100644 --- a/test/python/moe_ep/test_moe_ep_forward.py +++ b/test/python/moe_ep/test_moe_ep_forward.py @@ -42,7 +42,6 @@ ) - def _public_nvfp4(data, scale, logical_shape): from cudnn import BlockScaledTensor @@ -91,76 +90,6 @@ def close(self): operator._closed = True -@pytest.mark.L0 -@pytest.mark.parametrize( - ("field", "value"), - [ - *( - ("token_back_mode", value) - for value in ( - "epi_warps", - "standalone_warps", - "reuse_dispatch_warps", - ) - ), - *( - ("epi_flag_batch", value) - for value in ( - (4, 2), - (1, 1), - (1, 2), - (1, 4), - (2, 1), - (2, 2), - (2, 4), - (4, 4), - ) - ), - *( - ("token_in_flag_batch", value) - for value in (1, 2, 4, 8, 16) - ), - *( - ("group_hint", value) - for value in (None, 64, 128, 256, 512, 768, 1024) - ), - ("reduce_topk_in_kernel", False), - ("reduce_topk_in_kernel", True), - ], -) -def test_moe_ep_tuning_accepts_candidate_values(field, value): - from cudnn import MoeEpTuningConfig - - tuning = MoeEpTuningConfig(**{field: value}) - assert getattr(tuning, field) == value - - -@pytest.mark.L0 -@pytest.mark.parametrize( - "kwargs", - [ - {"token_back_mode": "unknown"}, - {"token_back_mode": []}, - {"epi_flag_batch": [1, 1]}, - {"epi_flag_batch": (3, 3)}, - {"token_in_flag_batch": 3}, - {"token_in_flag_batch": True}, - {"group_hint": 0}, - {"group_hint": True}, - {"reduce_topk_in_kernel": 1}, - { - "token_back_mode": "standalone_warps", - "reduce_topk_in_kernel": True, - }, - ], -) -def test_moe_ep_tuning_rejects_unvalidated_values(kwargs): - from cudnn import MoeEpTuningConfig - - with pytest.raises(ValueError): - MoeEpTuningConfig(**kwargs) - - @pytest.mark.L0 def test_moe_ep_tuning_public_contract_mapping_and_cache_key(): from cudnn import MoeEp, MoeEpTuningConfig @@ -173,66 +102,32 @@ def test_moe_ep_tuning_public_contract_mapping_and_cache_key(): assert PackageMoeEpTuningConfig is MoeEpTuningConfig tuning = MoeEpTuningConfig( - token_back_mode="epi_warps", + token_back_mode="standalone_warps", epi_flag_batch=(4, 2), token_in_flag_batch=4, group_hint=768, - reduce_topk_in_kernel=True, ) - with MoeEp( - **_forward_config(), - token_padding_size=64, - sf_padding_size=256, - tuning=tuning, - ) as op: + with MoeEp(**_forward_config(), tuning=tuning) as op: assert op.tuning is tuning - kernel_config = Mxfp8KernelConfig.from_forward_config( - op._forward_config - ) + kernel_config = Mxfp8KernelConfig.from_forward_config(op._forward_config) - assert kernel_config.token_back_mode == "epi_warps" - assert kernel_config.epi_flag_batch == (4, 2) - assert kernel_config.flag_batch == 4 - assert kernel_config.group_hint == 768 - assert kernel_config.token_padding_block == 64 - assert kernel_config.sf_padding_block == 256 assert kernel_config.tuning_signature(123) == ( - "epi_warps", + "standalone_warps", (4, 2), 4, 768, - True, - ) - effective = kernel_config.effective_config(123) - assert effective["token_padding_block"] == 64 - assert effective["sf_padding_block"] == 256 - assert effective["effective_group_hint"] == 768 - assert effective["fc2_in_kernel_topk_reduce"] is True - assert effective["launch_cluster_count"] == 123 - assert effective["drop_on_overflow"] is False - assert effective["enable_col_quant"] is False - assert "output_format" not in effective - - with MoeEp(**_forward_config()) as default_op: - default_config = Mxfp8KernelConfig.from_forward_config( - default_op._forward_config - ) - assert default_config.tuning_signature(123) == ( - "epi_warps", - (1, 1), - 1, - 123, False, ) + + with MoeEp(**_forward_config()) as default_op: + default_config = Mxfp8KernelConfig.from_forward_config(default_op._forward_config) key_args = ( torch.device("cuda", 0), (10, 7), 123, (), ) - assert kernel_config.compile_key(*key_args) != default_config.compile_key( - *key_args - ) + assert kernel_config.compile_key(*key_args) != default_config.compile_key(*key_args) @pytest.mark.L0 @@ -252,11 +147,7 @@ def test_internal_column_requant_config_is_disabled_by_default_and_cache_distinc ) assert default_config.enable_col_quant is False - assert default_config.max_recv_size_per_rank == ( - default_forward.ep_size - * default_forward.max_tokens_per_rank - * default_forward.top_k - ) + assert default_config.max_recv_size_per_rank == (default_forward.ep_size * default_forward.max_tokens_per_rank * default_forward.top_k) assert enabled_config.enable_col_quant is True assert enabled_config.col_quant_num_ctas == 512 with pytest.raises(ValueError, match="max_recv_size_per_rank"): @@ -264,9 +155,7 @@ def test_internal_column_requant_config_is_disabled_by_default_and_cache_distinc with pytest.raises(ValueError, match="col_quant_num_ctas"): replace(default_config, col_quant_num_ctas=0) key_args = (torch.device("cuda", 0), (10, 7), 123, ()) - assert default_config.compile_key(*key_args) != enabled_config.compile_key( - *key_args - ) + assert default_config.compile_key(*key_args) != enabled_config.compile_key(*key_args) @pytest.mark.L0 @@ -306,12 +195,8 @@ def test_combine_format_maps_to_contract_wire(public_format, wire_format): Mxfp8KernelConfig, ) - with MoeEp( - **_forward_config(combine_format=public_format) - ) as op: - kernel_config = Mxfp8KernelConfig.from_forward_config( - op._forward_config - ) + with MoeEp(**_forward_config(combine_format=public_format)) as op: + kernel_config = Mxfp8KernelConfig.from_forward_config(op._forward_config) assert kernel_config.combine_format == wire_format @@ -501,17 +386,21 @@ def nbytes(self, _name): workspace = StandaloneWorkspace() if combine_format == "bf16": + class NoScaleWorkspace: def region(self, _name): raise AssertionError("BF16 combine must not query scale region") workspace = NoScaleWorkspace() - assert _pre_reduced_sf_workspace_metadata( - workspace, - config, - shared_bytes=128 + config.max_tokens_per_rank * 64, - ) == expected + assert ( + _pre_reduced_sf_workspace_metadata( + workspace, + config, + shared_bytes=128 + config.max_tokens_per_rank * 64, + ) + == expected + ) @pytest.mark.L0 @@ -581,9 +470,7 @@ def test_distributed_launch_rejects_mismatched_tuning_before_barrier( backend._ep_launch_ready = False stream = SimpleNamespace(synchronize=lambda: None) - resources = SimpleNamespace( - runtime=SimpleNamespace(group=object(), world_size=2) - ) + resources = SimpleNamespace(runtime=SimpleNamespace(group=object(), world_size=2)) prepared = SimpleNamespace(launch_cluster_count=123) monkeypatch.setattr( backend, @@ -806,9 +693,7 @@ def test_nondefault_moe_ep_tuning_matches_reference_and_reuses_plan(): assert backend._compiled is compiled assert backend._plan._workspace is workspace - assert backend.kernel_config.tuning_signature( - backend._prepared_kernel.launch_cluster_count - ) == ("standalone_warps", (4, 2), 4, 64, False) + assert backend.kernel_config.tuning_signature(backend._prepared_kernel.launch_cluster_count) == ("standalone_warps", (4, 2), 4, 64, False) _assert_matches_reference(first, expected) _assert_matches_reference(second, expected) @@ -1094,9 +979,7 @@ def test_intermediate_requires_full_mma_n_tile(): index_dtype=torch.int32, weight_dtype=torch.bfloat16, ) - with MoeEp( - **_forward_config(intermediate_size=128, max_tokens_per_rank=3) - ) as op: + with MoeEp(**_forward_config(intermediate_size=128, max_tokens_per_rank=3)) as op: with pytest.raises( NotImplementedError, match=r"intermediate_size .*divisible by 256", @@ -1123,11 +1006,7 @@ def test_activation_scale_rows_are_padded_to_16_bytes(): kernel_local_workspace_bytes=128, kernel_shared_workspace_bytes=128, ) - activation_scale = next( - region - for region in requirements.symmetric_regions - if region.name == "activation_scale" - ) + activation_scale = next(region for region in requirements.symmetric_regions if region.name == "activation_scale") assert activation_scale.nbytes == 5 * 16 @@ -1153,9 +1032,7 @@ def test_column_requant_workspace_is_allocated_only_when_enabled(): disabled_names = {region.name for region in disabled.local_regions} assert "col_quant_data" not in disabled_names assert "col_quant_sf" not in disabled_names - enabled_sizes = { - region.name: region.nbytes for region in enabled.local_regions - } + enabled_sizes = {region.name: region.nbytes for region in enabled.local_regions} assert enabled_sizes["col_quant_data"] == 640 assert enabled_sizes["col_quant_sf"] == 80 From f4d38ec9367b7e723d2880143048983ad78d8e95 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Fri, 28 Aug 2026 09:08:01 -0700 Subject: [PATCH 17/31] docs: separate MoeEP operation and API guidance Keep Python lifecycle documentation in the FE OSS guide while giving architecture, formats, and topology a dedicated operation reference. --- docs/fe-oss-apis/moe_ep.md | 41 ++++---------- docs/fe-oss-apis/overview.md | 9 +-- docs/operations/MoeEp.md | 104 +++++++++++++++++++++++++++++++++++ llms.txt | 2 + 4 files changed, 122 insertions(+), 34 deletions(-) create mode 100644 docs/operations/MoeEp.md diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md index 0de402ffa..689431c8b 100644 --- a/docs/fe-oss-apis/moe_ep.md +++ b/docs/fe-oss-apis/moe_ep.md @@ -2,43 +2,23 @@ `cudnn.moe_ep` provides a fused SwiGLU MoE implementation for Rubin SM107. Experts are sharded contiguously across an optional expert-parallel process -group. - -## Supported configuration - -- CUDA execution on Rubin SM107 (compute capability 10.7) -- fused SwiGLU with contiguous expert sharding across `ep_group` -- BF16 output, including when the combine path uses MXFP8 -- BF16 or MXFP8 combine -- plain BF16, FP16, or FP32 inference operands, or MXFP8 - `BlockScaledTensor` operands -- `apply_topk_in_fc1=True` -- `hidden_size` divisible by 128 -- `intermediate_size` divisible by 256 -- `top_k <= min(32, num_experts)` -- `num_experts` divisible by the EP group size -- explicit positive `max_tokens_per_rank` - -`output_format` is currently executable only as `"bf16"`. NVFP4 is represented -by the public format types but NVFP4 operands, combine, and output are not -executable by the current MegaMoE backend. - -The fixed-resource CUDA Graph path has hardware acceptance through EP32 within -one direct-P2P MNNVL peer-access domain. The Python capability layer does not -impose an EP-size ceiling; this statement describes validated hardware scope, -not support for cross-MNNVL execution. +group. This page documents the Python API and lifecycle. See the +[MoeEP operation reference](../operations/MoeEp.md) for supported +architectures, data formats, tensor contracts, and expert-parallel topology. ## Installation -Install the dedicated optional dependencies: +Install the reusable CuTeDSL and communication extras, then the PyTorch +integration dependencies: ```bash -pip install nvidia-cudnn-frontend[moe_ep] +pip install "nvidia-cudnn-frontend[cutedsl,comm]" torch torch-c-dlpack-ext ``` -The extra supplies the CuTeDSL and NVSHMEM Python dependencies. PyTorch with -CUDA support is also required. EP2+ additionally requires an initialized NCCL -process group and a usable NVSHMEM peer topology. +The package keeps the general CuTeDSL installation floor at 4.5.0. Rubin +MegaMoE execution checks for `nvidia-cutlass-dsl>=4.8.0` when preparing its +kernels. EP2+ additionally requires an initialized NCCL process group and a +usable NVSHMEM peer topology. ## Public API and constructor @@ -477,6 +457,7 @@ Run host-side and local tests: ```bash python -m pytest \ + test/python/moe_ep/test_moe_ep_cutedsl.py \ test/python/moe_ep/test_moe_ep_forward.py \ test/python/moe_ep/test_moe_ep_backward.py \ -m L0 diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 233caae86..6bfb12cbb 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -37,7 +37,8 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d - [SDPA Backward (SM120)](attention/sdpa_bwd_sm120.md) - [RMSNorm + SiLU](rmsnorm_silu.md) - [MoE + Expert Parallel API](moe_ep.md) — Rubin SM107 fused SwiGLU with - fixed-resource training and CUDA Graph support + fixed-resource training and CUDA Graph support; see the + [MoeEP operation reference](../operations/MoeEp.md) for support details ## Installation and setup @@ -53,14 +54,14 @@ pip install --group jax # jax >= 0.5 (XLA entry points via cutlass.jax, ship ``` (For the published wheel, `pip install torch torch-c-dlpack-ext` or `pip install "jax>=0.5"` directly.) -MoE + Expert Parallel requires its dedicated optional dependencies: +MoE + Expert Parallel composes the reusable CuTeDSL and communication extras: ```bash -pip install nvidia-cudnn-frontend[moe_ep] +pip install "nvidia-cudnn-frontend[cutedsl,comm]" torch torch-c-dlpack-ext ``` MoeEP is currently CUDA/PyTorch-only and targets Rubin SM107. EP2+ execution also requires NCCL, NVSHMEM, and a direct-P2P MNNVL peer-access domain. See the -[MoeEP support matrix and tensor contracts](moe_ep.md#supported-configuration) +[MoeEP support matrix and tensor contracts](../operations/MoeEp.md) before integrating it. After installation, you can import the APIs directly from the `cudnn` package, i.e. `from cudnn import {your_operation}` diff --git a/docs/operations/MoeEp.md b/docs/operations/MoeEp.md new file mode 100644 index 000000000..20fc5437d --- /dev/null +++ b/docs/operations/MoeEp.md @@ -0,0 +1,104 @@ +# Mixture of Experts with Expert Parallelism + +The MoeEP operation fuses token routing, expert SwiGLU computation, and +expert-parallel communication. Global experts are sharded contiguously across +the ranks of an expert-parallel process group. + +For token \(x_t\), selected expert \(e_{t,k}\), and routing weight \(p_{t,k}\), +the operation computes: + +\[ +y_t = \sum_{k=0}^{K-1} p_{t,k} + \left(\operatorname{SiLU}(x_t W^{gate}_{e_{t,k}}) + \odot (x_t W^{up}_{e_{t,k}})\right) + W^{down}_{e_{t,k}} +\] + +The current implementation is exposed by the frontend-only Python +[`cudnn.moe_ep.MoeEp`](../fe-oss-apis/moe_ep.md) API. It is distinct from the +cuDNN graph [MoE Grouped Matmul](MoeGroupedMatmul.md) operation. + +## Execution support + +- NVIDIA Rubin SM107 GPUs (compute capability 10.7). +- CUDA and PyTorch execution. +- `nvidia-cutlass-dsl>=4.8.0` for the Rubin kernels. The package-wide + `cutedsl` extra retains its 4.5.0 installation floor so other cuDNN Frontend + operations remain usable with older compatible DSL versions. +- Fused SwiGLU with contiguous expert sharding. +- `apply_topk_in_fc1=True`. +- `hidden_size` divisible by 128. +- `intermediate_size` divisible by 256. +- `top_k <= min(32, num_experts)`. +- `num_experts` divisible by the expert-parallel group size. +- An explicit positive `max_tokens_per_rank`. + +The fixed-resource CUDA Graph path has hardware acceptance through EP32 when +all ranks are in one direct-P2P MNNVL peer-access domain. The Python capability +layer does not impose an EP-size ceiling; cross-MNNVL execution is not part of +the validated support surface. + +## Data formats + +Inference activation and expert weights accept: + +- BF16, FP16, or FP32 plain tensors, staged internally to MXFP8; or +- MXFP8 `BlockScaledTensor` values with logical block axis 1. + +The current executable output format is BF16. The expert-combine path accepts +BF16 or MXFP8. NVFP4 types are represented by the public API but native NVFP4 +operands, combine, and output are not executable by this backend. + +Fixed-resource training narrows dynamic activation and gradient inputs to +contiguous BF16 or FP32 tensors. Training weights are contiguous MXFP8 +block-scaled tensors. + +## Tensor contracts + +Let: + +- \(T\) be the local token count; +- \(H\) be `hidden_size`; +- \(I\) be `intermediate_size`; +- \(K\) be `top_k`; +- \(E_{local}\) be the local expert count. + +Inference uses: + +- `activation`: `(T, H)`; +- `topk_idx`: `(T, K)`, Int32 or Int64, containing `-1` or a valid global + expert ID; +- `topk_weights`: `(T, K)`, floating point; +- FC1 weights: `(E_local, H, 2I)`; +- FC2 weights: `(E_local, I, H)`; +- output: `(T, H)`, BF16. + +Fixed-resource training additionally binds transposed backward weights with +shapes `(E_local, H, I)` and `(E_local, 2I, H)`. Dynamic tensors must share one +device and satisfy `T <= max_tokens_per_rank`. + +## Expert-parallel communication + +EP2+ execution requires: + +- an initialized NCCL process group; +- `nvshmem4py` and usable NVSHMEM libraries; +- direct peer access among every pair of participating ranks; and +- consistent rank ordering, resource sizes, tuning, slot selection, and lane + ordering across the group. + +`max_recv_size_per_rank` bounds receive capacity. When omitted, it defaults to +the worst-case route count: + +```text +ep_size * max_tokens_per_rank * top_k +``` + +Resources cannot grow during CUDA Graph replay. Capacity or storage changes +require resource preparation and graph capture again. + +## API reference + +See [MoE + Expert Parallel API](../fe-oss-apis/moe_ep.md) for installation, +constructor arguments, inference and training lifecycles, tuning, overflow +handling, and CUDA Graph usage. diff --git a/llms.txt b/llms.txt index 7081689f6..e2fd61536 100644 --- a/llms.txt +++ b/llms.txt @@ -18,6 +18,7 @@ Published documentation: https://docs.nvidia.com/deeplearning/cudnn/latest/devel - [Convolutions](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Convolutions.md) - [Normalizations (LayerNorm, RMSNorm, BatchNorm, InstanceNorm)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Normalizations.md) - [MoE Grouped Matmul](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/MoeGroupedMatmul.md) +- [MoE with Expert Parallelism](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/MoeEp.md) - [Pointwise](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Pointwise.md) - [Block Scaling (MXFP8/NVFP4 quantization)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/BlockScaling.md) - [RoPE](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/RoPE.md) @@ -31,6 +32,7 @@ Published documentation: https://docs.nvidia.com/deeplearning/cudnn/latest/devel - [FE OSS APIs overview — full catalog and usage pattern](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/overview.md) - [Block-sparse attention (BSA)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/bsa.md), [DeepSeek sparse attention (DSA)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/dsa.md), [Native sparse attention (NSA)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/nsa.md) - [GEMM fusions (amax, SwiGLU, sReLU, grouped/discrete MoE variants)](https://github.com/NVIDIA/cudnn-frontend/tree/main/docs/fe-oss-apis/gemm_fusions) +- [MoE + Expert Parallel Python API](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/moe_ep.md) - [RMSNorm + RHT + Amax](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/rmsnorm_rht_amax.md), [RMSNorm + SiLU](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/fe-oss-apis/rmsnorm_silu.md) ## How-to guides From 9a23d95657618eaba5c3be50debf3c51210da995 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Fri, 28 Aug 2026 18:00:22 -0700 Subject: [PATCH 18/31] fix: isolate grouped WGrad graph workspaces Key explicit dense outputs by address so same-signature CUDA Graph call sites do not reuse mutable descriptor state. --- python/cudnn/gemm/cutedsl/grouped/wgrad/api.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py index 948d7bc39..1aa90d76c 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py @@ -294,6 +294,23 @@ def grouped_gemm_wgrad_wrapper_sm100( ) if framework == "jax" and backend is GroupedGemmBackend.BLOCK_SCALED: raise ValueError(_BLOCK_SCALED_JAX_ERROR) + explicit_dense_output_identity = None + if ( + backend is GroupedGemmBackend.BLOCK_SCALED + and framework == "torch" + and output_mode == "dense" + and wgrad_tensor is not None + ): + # Temporary workaround: + # 1. Problem behavior: Multiple same-signature launches with explicit + # outputs can corrupt later results when captured in one CUDA graph + # and backed by one cached API instance. + # 2. Possible root cause: The launches alias mutable TMA descriptor + # workspace owned by the cached API instance. + # 3. Possible long-term fix: Cache compiled kernels by shape, but own + # descriptor workspace independently per graph call site. Then + # remove output identity from the cache key. + explicit_dense_output_identity = int(wgrad_tensor.data_ptr()) if wgrad_tensor is None and wgrad_ptrs is None: wgrad_shape = (expert_cnt, hidden, intermediate) if framework == "torch": @@ -332,6 +349,7 @@ def grouped_gemm_wgrad_wrapper_sm100( accumulate_on_output, input_order, int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")), + explicit_dense_output_identity, ) op = _cache_of_GroupedGemmWgradSm100Objects.get(cache_key) if op is None: From 60e29a7fa853d47862a05bd5d46e8e6391156650 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Fri, 28 Aug 2026 18:00:33 -0700 Subject: [PATCH 19/31] test: cover MoeEP production grouped WGrad Validate FC1/FC2 WGrad across EP1-32 against independent MXFP8 references, including accumulation and CUDA Graph replay scenarios. --- .../moe_ep/moe_ep_distributed_workers.py | 33 + test/python/moe_ep/moe_ep_test_support.py | 156 ++++- .../moe_ep/probe_moe_ep_training_graph.py | 589 +++++++++++++++++- test/python/moe_ep/test_moe_ep_backward.py | 214 ++++++- test/python/moe_ep/test_moe_ep_multinode.py | 6 + 5 files changed, 966 insertions(+), 32 deletions(-) diff --git a/test/python/moe_ep/moe_ep_distributed_workers.py b/test/python/moe_ep/moe_ep_distributed_workers.py index 3e5923d5b..99cca10a6 100644 --- a/test/python/moe_ep/moe_ep_distributed_workers.py +++ b/test/python/moe_ep/moe_ep_distributed_workers.py @@ -12,8 +12,11 @@ from moe_ep.moe_ep_test_support import ( _assert_backward_matches, + _assert_grouped_wgrads_match_reference, _assert_matches_reference, _assert_wgrads_match_reference, + _dense_wgrads_from_grouped_kernel, + _dense_wgrads_from_operands, _fixed_training_reference, _fixed_training_weights, _forward_config, @@ -296,6 +299,7 @@ def _run_backward_reference_case( grad_output, ) overflow = resources.finalize_overflow((slot,), lane) + grouped_wgrads = _dense_wgrads_from_grouped_kernel(actual_wgrads) torch.cuda.synchronize(device) # No rank may enter a local assertion while a peer is still inside a @@ -324,6 +328,35 @@ def _run_backward_reference_case( expected_wgrads, expected_dense=expected_dense_wgrads, ) + expected_offsets = torch.cumsum( + torch.div( + actual_wgrads.valid_route_counts + 127, + 128, + rounding_mode="floor", + ) + * 128, + dim=0, + dtype=actual_wgrads.expert_offsets.dtype, + ) + torch.testing.assert_close( + actual_wgrads.expert_offsets, + expected_offsets, + rtol=0, + atol=0, + ) + _assert_grouped_wgrads_match_reference( + grouped_wgrads, + expected_dense_wgrads, + reference_name="the independent PyTorch MXFP8 reference", + ) + _assert_grouped_wgrads_match_reference( + grouped_wgrads, + _dense_wgrads_from_operands(actual_wgrads), + reference_name="the decoded production operand bundle", + close_kwargs={"rtol": 0.1, "atol": 0.1}, + ) + assert grouped_wgrads[0][1].eq(0).all() + assert grouped_wgrads[1][1].eq(0).all() except BaseException as error: assertion_error = error dist.barrier(group=ep_group) diff --git a/test/python/moe_ep/moe_ep_test_support.py b/test/python/moe_ep/moe_ep_test_support.py index 52f6d47cb..efb90718e 100644 --- a/test/python/moe_ep/moe_ep_test_support.py +++ b/test/python/moe_ep/moe_ep_test_support.py @@ -24,9 +24,11 @@ ) __all__ = [ + "_allocate_dense_grouped_wgrad_outputs", "_assert_backward_matches", "_assert_fixed_training_drop_overflow_result", "_assert_fixed_training_matches_reference", + "_assert_grouped_wgrads_match_reference", "_assert_matches_reference", "_assert_training_graph_tails_are_reset", "_assert_training_weight_sources_changed", @@ -34,6 +36,7 @@ "_capture_fixed_training_batch", "_copy_training_weight_sources_", "_dense_wgrads_from_operands", + "_dense_wgrads_from_grouped_kernel", "_expected_backward", "_fixed_training_case", "_fixed_training_drop_overflow_case", @@ -51,6 +54,7 @@ "_replay_cuda_graph", "_require_distributed_sm107", "_run_fixed_training_batch", + "_run_grouped_wgrad_kernel", "_sm107_device", "_stress_backend_reuse", "_training_public_pointers", @@ -72,6 +76,32 @@ ] +def _allocate_dense_grouped_wgrad_outputs( + operands, + *, + fill_value=None, +): + """Allocate fixed-address dense BF16 outputs for FC1 and FC2 WGrad.""" + + expert_count = operands.expert_offsets.numel() + outputs = tuple( + torch.empty( + ( + expert_count, + getattr(operands, f"{prefix}_a").shape[0], + getattr(operands, f"{prefix}_b").shape[1], + ), + dtype=torch.bfloat16, + device=operands.expert_offsets.device, + ) + for prefix in ("fc1", "fc2") + ) + if fill_value is not None: + for output in outputs: + output.fill_(fill_value) + return outputs + + # Data @@ -825,6 +855,7 @@ def _replay_cuda_graph( {"rtol": 0.15, "atol": 0.125}, # router-weight gradient. ) _WGRAD_CLOSE_KWARGS = {"rtol": 0.2, "atol": 0.25} +_GROUPED_WGRAD_CLOSE_KWARGS = {"rtol": 0.1, "atol": 0.1} def _round_up(value: int, multiple: int) -> int: @@ -970,6 +1001,108 @@ def _dense_wgrads_from_operands(operands): return torch.stack(fc1_parts), torch.stack(fc2_parts) +def _run_grouped_wgrad_kernel( + operands, + prefix: str, + *, + wgrad_tensor=None, + accumulate_on_output: bool = False, + current_stream=None, +): + """Run one fixed-capacity operand bundle through production WGrad.""" + + import cudnn + + if prefix not in ("fc1", "fc2"): + raise ValueError(f"prefix must be 'fc1' or 'fc2', got {prefix!r}") + # Graph callers provide one persistent output per training slot. This is + # currently also the isolation key for a temporary production-WGrad + # workaround: an EP2 graph with two same-signature calls produced correct + # operands but corrupted the second WGrad when both calls shared the + # cached API object's mutable TMA descriptor workspace. Distinct fixed + # outputs make the calls use distinct workspaces. The production fix + # should instead share the compiled kernel while owning descriptor + # workspace per graph call site, after which output identity must no + # longer participate in the compile cache key. + return cudnn.grouped_gemm_wgrad_wrapper_sm100( + a_tensor=getattr(operands, f"{prefix}_a"), + b_tensor=getattr(operands, f"{prefix}_b"), + sfa_tensor=getattr(operands, f"{prefix}_sfa"), + sfb_tensor=getattr(operands, f"{prefix}_sfb"), + offsets_tensor=operands.expert_offsets, + output_mode="dense", + wgrad_tensor=wgrad_tensor, + wgrad_dtype=torch.bfloat16, + acc_dtype=torch.float32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=32, + accumulate_on_output=accumulate_on_output, + input_order="tensor2d", + current_stream=current_stream, + )["wgrad_tensor"] + + +def _dense_wgrads_from_grouped_kernel( + operands, + *, + wgrad_tensors=None, + accumulate_on_output: bool = False, + current_stream=None, +): + """Run both fixed-capacity operand bundles through production WGrad.""" + + if wgrad_tensors is None: + wgrad_tensors = (None, None) + if len(wgrad_tensors) != 2: + raise ValueError("wgrad_tensors must contain FC1 and FC2 outputs") + return tuple( + _run_grouped_wgrad_kernel( + operands, + prefix, + wgrad_tensor=output, + accumulate_on_output=accumulate_on_output, + current_stream=current_stream, + ) + for prefix, output in zip(("fc1", "fc2"), wgrad_tensors) + ) + + +def _assert_grouped_wgrads_match_reference( + actual, + expected, + *, + reference_name: str, + close_kwargs=None, +) -> None: + """Compare grouped-kernel FC1/FC2 outputs and report useful error maxima.""" + + if close_kwargs is None: + close_kwargs = _GROUPED_WGRAD_CLOSE_KWARGS + for name, actual_dw, expected_dw in zip( + ("grad_fc1_weight", "grad_fc2_weight"), + actual, + expected, + ): + actual_fp32 = actual_dw.float() + expected_fp32 = expected_dw.float() + absolute_error = (actual_fp32 - expected_fp32).abs() + max_absolute_error = absolute_error.max().item() + max_relative_error = ( + (absolute_error / expected_fp32.abs().clamp_min(1.0e-6)).max().item() + ) + torch.testing.assert_close( + actual_fp32, + expected_fp32, + msg=lambda default, name=name: ( + f"{name} does not match {reference_name}; " + f"max_abs_error={max_absolute_error:.6g}, " + f"max_rel_error={max_relative_error:.6g}\n{default}" + ), + **close_kwargs, + ) + + def _reference_backward(config) -> MoeEpReference: options = dict(config) for production_only in ( @@ -1212,16 +1345,37 @@ def _run_fixed_training_batch(resources, lane, cases): ) -def _capture_fixed_training_batch(resources, lane, cases, capture_stream): +def _capture_fixed_training_batch( + resources, + lane, + cases, + capture_stream, + *, + grouped_wgrad_outputs=None, +): """Capture the shared fixed-training sequence for one or more slots.""" + if grouped_wgrad_outputs is not None and len(grouped_wgrad_outputs) != len(cases): + raise ValueError("grouped_wgrad_outputs must match the captured case count") graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=capture_stream): actuals = _run_fixed_training_batch(resources, lane, cases) + grouped_wgrads = ( + None + if grouped_wgrad_outputs is None + else tuple( + _dense_wgrads_from_grouped_kernel( + actual.wgrads, + wgrad_tensors=outputs, + ) + for actual, outputs in zip(actuals, grouped_wgrad_outputs) + ) + ) capture_stream.synchronize() return SimpleNamespace( graph=graph, actuals=actuals, + grouped_wgrads=grouped_wgrads, public_pointers=tuple(_training_public_pointers(actual) for actual in actuals), ) diff --git a/test/python/moe_ep/probe_moe_ep_training_graph.py b/test/python/moe_ep/probe_moe_ep_training_graph.py index 556fb0c1c..537c25998 100644 --- a/test/python/moe_ep/probe_moe_ep_training_graph.py +++ b/test/python/moe_ep/probe_moe_ep_training_graph.py @@ -16,7 +16,7 @@ The probe exercises only the public fixed-resource ordinary/capture path, including fixed-address staging/reset operations, forward/backward CuTeDSL -callables, and a one-scalar NCCL overflow OR. +callables, FC1/FC2 production grouped WGrad, and a one-scalar NCCL overflow OR. """ from __future__ import annotations @@ -40,6 +40,11 @@ _RuntimeWatchdog, get_runtime_manager, ) +from moe_ep.moe_ep_test_support import ( + _allocate_dense_grouped_wgrad_outputs, + _dense_wgrads_from_grouped_kernel, + _dense_wgrads_from_operands, +) def _debug_phase(rank: int, phase: str) -> None: @@ -103,6 +108,15 @@ def _parse_args() -> argparse.Namespace: action="store_true", help="skip the two-lane ordered cross-stream resource probe", ) + parser.add_argument( + "--wgrad-capture-mode", + choices=("both", "slot0", "slot1"), + default="both", + help=( + "capture both grouped-WGrad calls, or isolate exactly one slot to " + "diagnose same-signature graph reuse" + ), + ) parser.add_argument( "--expect-overflow-assert", action="store_true", @@ -133,15 +147,126 @@ def _assert_replay_tensor( } if actual.dtype in low_precision: if not torch.equal(actual, expected): + _report_tensor_difference(name, actual, expected) raise AssertionError(f"{name} is not bitwise equal after graph replay") return - torch.testing.assert_close( - actual, - expected, - rtol=1e-5, - atol=1e-6, - msg=f"{name} differs after graph replay", + try: + torch.testing.assert_close( + actual, + expected, + rtol=1e-5, + atol=1e-6, + msg=f"{name} differs after graph replay", + ) + except AssertionError: + _report_tensor_difference(name, actual, expected) + raise + + +def _report_tensor_difference( + name: str, + actual: torch.Tensor, + expected: torch.Tensor, +) -> None: + """Print actionable mismatch statistics without changing pass criteria.""" + + rank = dist.get_rank() if dist.is_initialized() else 0 + if actual.shape != expected.shape or actual.dtype != expected.dtype: + print( + "MOE_EP_GRAPH_TENSOR_DIAGNOSTIC " + f"rank={rank} name={name} " + f"actual_shape={tuple(actual.shape)} expected_shape={tuple(expected.shape)} " + f"actual_dtype={actual.dtype} expected_dtype={expected.dtype}", + flush=True, + ) + return + + actual_fp32 = actual.float() + expected_fp32 = expected.float() + finite = torch.isfinite(actual_fp32) & torch.isfinite(expected_fp32) + absolute_error = (actual_fp32 - expected_fp32).abs() + relative_error = absolute_error / expected_fp32.abs().clamp_min(1.0e-6) + finite_absolute = absolute_error.masked_select(finite) + finite_relative = relative_error.masked_select(finite) + max_absolute = ( + float(finite_absolute.max().item()) if finite_absolute.numel() else float("nan") + ) + max_relative = ( + float(finite_relative.max().item()) if finite_relative.numel() else float("nan") + ) + exact_mismatch = actual.view(torch.uint8).ne(expected.view(torch.uint8)) + close_mismatch = ~torch.isclose( + actual_fp32, + expected_fp32, + rtol=1.0e-5, + atol=1.0e-6, + equal_nan=True, + ) + logical_mismatch = actual.ne(expected) + first_indices = logical_mismatch.nonzero() + first_description = "none" + if first_indices.numel(): + first_index = tuple(int(value) for value in first_indices[0].tolist()) + first_description = ( + f"index={first_index},actual={float(actual[first_index].float().item()):.9g}," + f"expected={float(expected[first_index].float().item()):.9g}" + ) + + print( + "MOE_EP_GRAPH_TENSOR_DIAGNOSTIC " + f"rank={rank} name={name} dtype={actual.dtype} shape={tuple(actual.shape)} " + f"byte_mismatches={int(exact_mismatch.sum().item())} " + f"logical_mismatches={int(logical_mismatch.sum().item())} " + f"close_mismatches={int(close_mismatch.sum().item())} " + f"max_abs={max_absolute:.9g} max_rel={max_relative:.9g} " + f"actual_nonfinite={int((~torch.isfinite(actual_fp32)).sum().item())} " + f"expected_nonfinite={int((~torch.isfinite(expected_fp32)).sum().item())} " + f"first_mismatch={first_description}", + flush=True, ) + if actual.ndim == 3: + expert_dims = (1, 2) + expert_max_absolute = absolute_error.amax(dim=expert_dims) + expert_max_relative = relative_error.amax(dim=expert_dims) + expert_close_mismatches = close_mismatch.sum(dim=expert_dims) + print( + "MOE_EP_GRAPH_EXPERT_DIAGNOSTIC " + f"rank={rank} name={name} " + f"max_abs={expert_max_absolute.detach().cpu().tolist()} " + f"max_rel={expert_max_relative.detach().cpu().tolist()} " + f"close_mismatches={expert_close_mismatches.detach().cpu().tolist()}", + flush=True, + ) + + +def _report_grouped_wgrad_operand_consistency( + slot_name: str, + operands, + grouped_wgrads, +) -> None: + """Report whether replayed WGrad agrees with its replayed operand bundle.""" + + rank = dist.get_rank() if dist.is_initialized() else 0 + try: + decoded = _dense_wgrads_from_operands(operands) + except BaseException as error: + print( + "MOE_EP_GRAPH_OPERAND_DECODE_ERROR " + f"rank={rank} slot={slot_name} " + f"error={type(error).__name__}:{error}", + flush=True, + ) + return + for prefix, actual, expected in zip( + ("fc1", "fc2"), + grouped_wgrads, + decoded, + ): + _report_tensor_difference( + f"{slot_name}.{prefix}_wgrad_vs_decoded_operands", + actual, + expected.to(actual.dtype), + ) def _make_inputs( @@ -335,6 +460,204 @@ def _close_probe_operator( dist.barrier(group=group) +def _run_single_wgrad_slot_capture_probe( + *, + rank: int, + world_size: int, + device: torch.device, + group, + max_recv_size_per_rank: int, + slot_index: int, +) -> None: + """Capture the full training chain with one grouped-WGrad invocation.""" + + args0, grad0, args1, grad1, _, _ = _make_two_slot_inputs( + rank, + world_size, + device, + ) + op = _make_operator( + world_size=world_size, + group=group, + max_recv_size_per_rank=max_recv_size_per_rank, + drop_on_overflow=True, + ) + graph = None + try: + resources = op.prepare_training_resources( + _make_training_weights(args0), + slot_count=2, + lane_count=1, + ) + slot0, slot1 = resources.slots + lane = resources.lanes[0] + + resources.refresh_weights() + eager_y0 = resources.forward(slot0, lane, args0[0], args0[3], args0[4]) + eager_y1 = resources.forward(slot1, lane, args1[0], args1[3], args1[4]) + eager_dx0, eager_dp0, eager_operands0 = resources.backward( + slot0, + lane, + grad0, + ) + eager_dx1, eager_dp1, eager_operands1 = resources.backward( + slot1, + lane, + grad1, + ) + eager_operands = (eager_operands0, eager_operands1) + grouped_outputs = tuple( + _allocate_dense_grouped_wgrad_outputs(operands) + for operands in eager_operands + ) + eager_grouped_wgrads = tuple( + _dense_wgrads_from_grouped_kernel( + operands, + wgrad_tensors=outputs, + ) + for operands, outputs in zip(eager_operands, grouped_outputs) + ) + eager_overflow = resources.finalize_overflow((slot0, slot1), lane) + torch.cuda.synchronize(device) + dist.barrier(group=group) + if int(eager_overflow.item()) != 0: + raise AssertionError("single-slot WGrad eager warmup overflowed") + + operand_fields = ( + "expert_offsets", + "valid_route_counts", + "fc1_a", + "fc1_sfa", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + ) + eager_common = tuple( + tensor.clone() + for tensor in ( + eager_y0, + eager_y1, + eager_dx0, + eager_dx1, + eager_dp0, + eager_dp1, + ) + ) + eager_operand_snapshot = { + field: getattr(eager_operands[slot_index], field).clone() + for field in operand_fields + } + eager_wgrad_snapshot = tuple( + tensor.clone() for tensor in eager_grouped_wgrads[slot_index] + ) + selected_outputs = grouped_outputs[slot_index] + selected_output_pointers = tuple( + output.data_ptr() for output in selected_outputs + ) + + stream = torch.cuda.Stream(device=device) + stream.wait_stream(torch.cuda.current_stream(device)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + resources.refresh_weights() + graph_y0 = resources.forward( + slot0, + lane, + args0[0], + args0[3], + args0[4], + ) + graph_y1 = resources.forward( + slot1, + lane, + args1[0], + args1[3], + args1[4], + ) + graph_dx0, graph_dp0, graph_operands0 = resources.backward( + slot0, + lane, + grad0, + ) + graph_dx1, graph_dp1, graph_operands1 = resources.backward( + slot1, + lane, + grad1, + ) + graph_operands = (graph_operands0, graph_operands1)[slot_index] + graph_grouped_wgrads = _dense_wgrads_from_grouped_kernel( + graph_operands, + wgrad_tensors=selected_outputs, + ) + graph_overflow = resources.finalize_overflow((slot0, slot1), lane) + dist.barrier(group=group) + + with torch.cuda.stream(stream): + graph.replay() + stream.synchronize() + dist.barrier(group=group) + if int(graph_overflow.item()) != 0: + raise AssertionError("single-slot WGrad graph replay overflowed") + if ( + tuple(output.data_ptr() for output in graph_grouped_wgrads) + != selected_output_pointers + ): + raise AssertionError("single-slot grouped WGrad output addresses changed") + + graph_common = ( + graph_y0, + graph_y1, + graph_dx0, + graph_dx1, + graph_dp0, + graph_dp1, + ) + for name, actual, expected in zip( + ("y0", "y1", "dx0", "dx1", "dprob0", "dprob1"), + graph_common, + eager_common, + ): + _assert_replay_tensor(name, actual, expected) + for field in operand_fields: + _assert_replay_tensor( + f"slot{slot_index}.{field}", + getattr(graph_operands, field), + eager_operand_snapshot[field], + ) + try: + for prefix, actual, expected in zip( + ("fc1", "fc2"), + graph_grouped_wgrads, + eager_wgrad_snapshot, + ): + _assert_replay_tensor( + f"slot{slot_index}.{prefix}_wgrad", + actual, + expected, + ) + except BaseException: + _report_grouped_wgrad_operand_consistency( + f"slot{slot_index}", + graph_operands, + graph_grouped_wgrads, + ) + raise + + if rank == 0: + print( + f"MOE_EP_EP{world_size}_SINGLE_WGRAD_SLOT_GRAPH_PASS " + f"slot=slot{slot_index}", + flush=True, + ) + finally: + if graph is not None: + del graph + _close_probe_operator(device=device, group=group, op=op) + + def _run_training_resource_probe( *, rank: int, @@ -391,6 +714,19 @@ def _run_training_resource_probe( ) dx0, dp0, operands0 = resources.backward(slot0, lane0, grad0) dx1, dp1, operands1 = resources.backward(slot1, lane0, grad1) + grouped_outputs0 = _allocate_dense_grouped_wgrad_outputs(operands0) + grouped_outputs1 = _allocate_dense_grouped_wgrad_outputs(operands1) + grouped_output_pointers = tuple( + output.data_ptr() for output in (*grouped_outputs0, *grouped_outputs1) + ) + grouped_wgrads0 = _dense_wgrads_from_grouped_kernel( + operands0, + wgrad_tensors=grouped_outputs0, + ) + grouped_wgrads1 = _dense_wgrads_from_grouped_kernel( + operands1, + wgrad_tensors=grouped_outputs1, + ) overflow_status = resources.finalize_overflow((slot0, slot1)) torch.cuda.synchronize(device) dist.barrier(group=group) @@ -404,14 +740,30 @@ def _run_training_resource_probe( "dx1", "dprob0", "dprob1", + "slot0.expert_offsets", + "slot0.valid_route_counts", "slot0.fc1_a", + "slot0.fc1_sfa", "slot0.fc1_b", + "slot0.fc1_sfb", "slot0.fc2_a", + "slot0.fc2_sfa", "slot0.fc2_b", + "slot0.fc2_sfb", + "slot1.expert_offsets", + "slot1.valid_route_counts", "slot1.fc1_a", + "slot1.fc1_sfa", "slot1.fc1_b", + "slot1.fc1_sfb", "slot1.fc2_a", + "slot1.fc2_sfa", "slot1.fc2_b", + "slot1.fc2_sfb", + "slot0.fc1_wgrad", + "slot0.fc2_wgrad", + "slot1.fc1_wgrad", + "slot1.fc2_wgrad", ) ordinary = { name: tensor.clone() @@ -424,14 +776,30 @@ def _run_training_resource_probe( dx1, dp0, dp1, + operands0.expert_offsets, + operands0.valid_route_counts, operands0.fc1_a, + operands0.fc1_sfa, operands0.fc1_b, + operands0.fc1_sfb, operands0.fc2_a, + operands0.fc2_sfa, operands0.fc2_b, + operands0.fc2_sfb, + operands1.expert_offsets, + operands1.valid_route_counts, operands1.fc1_a, + operands1.fc1_sfa, operands1.fc1_b, + operands1.fc1_sfb, operands1.fc2_a, + operands1.fc2_sfa, operands1.fc2_b, + operands1.fc2_sfb, + grouped_wgrads0[0], + grouped_wgrads0[1], + grouped_wgrads1[0], + grouped_wgrads1[1], ), ) } @@ -439,6 +807,10 @@ def _run_training_resource_probe( operands0.expert_offsets.clone(), operands1.expert_offsets.clone(), ) + ordinary_route_counts = ( + operands0.valid_route_counts.clone(), + operands1.valid_route_counts.clone(), + ) stream = torch.cuda.Stream(device=device) stream.wait_stream(torch.cuda.current_stream(device)) @@ -469,6 +841,14 @@ def _run_training_resource_probe( lane0, grad1, ) + graph_grouped_wgrads0 = _dense_wgrads_from_grouped_kernel( + graph_operands0, + wgrad_tensors=grouped_outputs0, + ) + graph_grouped_wgrads1 = _dense_wgrads_from_grouped_kernel( + graph_operands1, + wgrad_tensors=grouped_outputs1, + ) graph_overflow = resources.finalize_overflow((slot0, slot1)) dist.barrier(group=group) @@ -490,19 +870,59 @@ def _run_training_resource_probe( graph_dx1, graph_dp0, graph_dp1, + graph_operands0.expert_offsets, + graph_operands0.valid_route_counts, graph_operands0.fc1_a, + graph_operands0.fc1_sfa, graph_operands0.fc1_b, + graph_operands0.fc1_sfb, graph_operands0.fc2_a, + graph_operands0.fc2_sfa, graph_operands0.fc2_b, + graph_operands0.fc2_sfb, + graph_operands1.expert_offsets, + graph_operands1.valid_route_counts, graph_operands1.fc1_a, + graph_operands1.fc1_sfa, graph_operands1.fc1_b, + graph_operands1.fc1_sfb, graph_operands1.fc2_a, + graph_operands1.fc2_sfa, graph_operands1.fc2_b, + graph_operands1.fc2_sfb, + graph_grouped_wgrads0[0], + graph_grouped_wgrads0[1], + graph_grouped_wgrads1[0], + graph_grouped_wgrads1[1], ), ) } - for name in comparison_names: - _assert_replay_tensor(name, captured[name], ordinary[name]) + if ( + tuple( + output.data_ptr() + for output in ( + *graph_grouped_wgrads0, + *graph_grouped_wgrads1, + ) + ) + != grouped_output_pointers + ): + raise AssertionError("captured grouped WGrad output addresses changed") + try: + for name in comparison_names: + _assert_replay_tensor(name, captured[name], ordinary[name]) + except BaseException: + _report_grouped_wgrad_operand_consistency( + "slot0", + graph_operands0, + graph_grouped_wgrads0, + ) + _report_grouped_wgrad_operand_consistency( + "slot1", + graph_operands1, + graph_grouped_wgrads1, + ) + raise torch.testing.assert_close( graph_operands0.expert_offsets, ordinary_offsets[0], @@ -515,6 +935,18 @@ def _run_training_resource_probe( rtol=0, atol=0, ) + torch.testing.assert_close( + graph_operands0.valid_route_counts, + ordinary_route_counts[0], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + graph_operands1.valid_route_counts, + ordinary_route_counts[1], + rtol=0, + atol=0, + ) if full_probe: # Diagnostic mode aligns ranks after every replay and verifies that @@ -533,6 +965,28 @@ def _run_training_resource_probe( rtol=1e-5, atol=1e-6, ) + for name in ( + "slot0.fc1_wgrad", + "slot0.fc2_wgrad", + "slot1.fc1_wgrad", + "slot1.fc2_wgrad", + ): + _assert_replay_tensor(name, captured[name], ordinary[name]) + for index, graph_operands in enumerate( + (graph_operands0, graph_operands1) + ): + torch.testing.assert_close( + graph_operands.expert_offsets, + ordinary_offsets[index], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + graph_operands.valid_route_counts, + ordinary_route_counts[index], + rtol=0, + atol=0, + ) # Production-like burst: no synchronization or host collective in # the loop. The graph contains the captured scalar overflow OR. @@ -549,6 +1003,13 @@ def _run_training_resource_probe( rtol=1e-5, atol=1e-6, ) + for name in ( + "slot0.fc1_wgrad", + "slot0.fc2_wgrad", + "slot1.fc1_wgrad", + "slot1.fc2_wgrad", + ): + _assert_replay_tensor(name, captured[name], ordinary[name]) # Overflow both slots, then restore their distinct valid patterns. overflow = _route_pattern("overflow", rank, world_size, device) @@ -581,6 +1042,13 @@ def _run_training_resource_probe( f"slot1_routing_restored=" f"{torch.equal(args1[3], remote[0])}" ) + for name in ( + "slot0.fc1_wgrad", + "slot0.fc2_wgrad", + "slot1.fc1_wgrad", + "slot1.fc2_wgrad", + ): + _assert_replay_tensor(name, captured[name], ordinary[name]) if rank == 0: mode = "full" if full_probe else "reinit" @@ -589,6 +1057,11 @@ def _run_training_resource_probe( f"MOE_EP_EP{world_size}_TRAINING_RESOURCES_GRAPH_PASS " f"mode={mode} burst={effective_burst}", flush=True, ) + print( + f"MOE_EP_EP{world_size}_GROUPED_WGRAD_GRAPH_PASS " + f"mode={mode} burst={effective_burst}", + flush=True, + ) finally: if graph is not None: del graph @@ -633,11 +1106,16 @@ def _run_multistream_resource_probe( with _debug_phase_scope(rank, "multistream.lane0-forward"): eager_y0 = resources.forward(slot0, lane0, args0[0], args0[3], args0[4]) with _debug_phase_scope(rank, "multistream.lane0-backward"): - eager_dx0, eager_dp0, _ = resources.backward( + eager_dx0, eager_dp0, eager_operands0 = resources.backward( slot0, lane0, grad0, ) + grouped_outputs0 = _allocate_dense_grouped_wgrad_outputs(eager_operands0) + eager_grouped_wgrads0 = _dense_wgrads_from_grouped_kernel( + eager_operands0, + wgrad_tensors=grouped_outputs0, + ) with _debug_phase_scope(rank, "multistream.lane0-finalize"): resources.finalize_overflow((slot0,), lane0) _synchronize_with_watchdog( @@ -651,11 +1129,19 @@ def _run_multistream_resource_probe( with _debug_phase_scope(rank, "multistream.lane1-forward"): eager_y1 = resources.forward(slot1, lane1, args1[0], args1[3], args1[4]) with _debug_phase_scope(rank, "multistream.lane1-backward"): - eager_dx1, eager_dp1, _ = resources.backward( + eager_dx1, eager_dp1, eager_operands1 = resources.backward( slot1, lane1, grad1, ) + grouped_outputs1 = _allocate_dense_grouped_wgrad_outputs(eager_operands1) + eager_grouped_wgrads1 = _dense_wgrads_from_grouped_kernel( + eager_operands1, + wgrad_tensors=grouped_outputs1, + ) + grouped_output_pointers = tuple( + output.data_ptr() for output in (*grouped_outputs0, *grouped_outputs1) + ) with _debug_phase_scope(rank, "multistream.lane1-finalize"): resources.finalize_overflow((slot1,), lane1) _synchronize_with_watchdog( @@ -674,6 +1160,10 @@ def _run_multistream_resource_probe( eager_y1, eager_dx1, eager_dp1, + eager_grouped_wgrads0[0], + eager_grouped_wgrads0[1], + eager_grouped_wgrads1[0], + eager_grouped_wgrads1[1], ) ) @@ -709,11 +1199,15 @@ def _run_multistream_resource_probe( args0[3], args0[4], ) - graph_dx0, graph_dp0, _ = resources.backward( + graph_dx0, graph_dp0, graph_operands0 = resources.backward( slot0, lane0, grad0, ) + graph_grouped_wgrads0 = _dense_wgrads_from_grouped_kernel( + graph_operands0, + wgrad_tensors=grouped_outputs0, + ) done_event0.record(lane_stream0) lane_stream1.wait_event(done_event0) with torch.cuda.stream(lane_stream1): @@ -724,11 +1218,15 @@ def _run_multistream_resource_probe( args1[3], args1[4], ) - graph_dx1, graph_dp1, _ = resources.backward( + graph_dx1, graph_dp1, graph_operands1 = resources.backward( slot1, lane1, grad1, ) + graph_grouped_wgrads1 = _dense_wgrads_from_grouped_kernel( + graph_operands1, + wgrad_tensors=grouped_outputs1, + ) done_event1.record(lane_stream1) capture_stream.wait_event(done_event1) graph_overflow = resources.finalize_overflow( @@ -766,7 +1264,22 @@ def _run_multistream_resource_probe( graph_y1, graph_dx1, graph_dp1, + graph_grouped_wgrads0[0], + graph_grouped_wgrads0[1], + graph_grouped_wgrads1[0], + graph_grouped_wgrads1[1], ) + if ( + tuple( + output.data_ptr() + for output in ( + *graph_grouped_wgrads0, + *graph_grouped_wgrads1, + ) + ) + != grouped_output_pointers + ): + raise AssertionError("multistream grouped WGrad output addresses changed") for index, (value, reference) in enumerate(zip(actual, expected)): _assert_replay_tensor( f"multistream[{index}]", @@ -919,34 +1432,50 @@ def main() -> None: max_recv_size_per_rank=args.max_recv_size_per_rank, ) raise AssertionError("fatal overflow assertion probe returned") - for cycle in range(args.cycles): + if args.wgrad_capture_mode == "both": + for cycle in range(args.cycles): + with _debug_phase_scope( + rank, + f"training-resources-cycle-{cycle}", + ): + _run_training_resource_probe( + rank=rank, + world_size=world_size, + device=device, + group=dist.group.WORLD, + diagnostic_replays=args.diagnostic_replays, + burst_replays=args.burst_replays, + max_recv_size_per_rank=args.max_recv_size_per_rank, + full_probe=cycle == 0, + ) + if not args.skip_multistream: + with _debug_phase_scope(rank, "multistream"): + _run_multistream_resource_probe( + rank=rank, + world_size=world_size, + device=device, + group=dist.group.WORLD, + replays=args.multistream_replays, + max_recv_size_per_rank=args.max_recv_size_per_rank, + ) + else: + slot_index = int(args.wgrad_capture_mode[-1]) with _debug_phase_scope( rank, - f"training-resources-cycle-{cycle}", + f"single-wgrad-slot{slot_index}", ): - _run_training_resource_probe( - rank=rank, - world_size=world_size, - device=device, - group=dist.group.WORLD, - diagnostic_replays=args.diagnostic_replays, - burst_replays=args.burst_replays, - max_recv_size_per_rank=args.max_recv_size_per_rank, - full_probe=cycle == 0, - ) - if not args.skip_multistream: - with _debug_phase_scope(rank, "multistream"): - _run_multistream_resource_probe( + _run_single_wgrad_slot_capture_probe( rank=rank, world_size=world_size, device=device, group=dist.group.WORLD, - replays=args.multistream_replays, max_recv_size_per_rank=args.max_recv_size_per_rank, + slot_index=slot_index, ) if rank == 0: print( - f"MOE_EP_EP{world_size}_CUDA_GRAPH_PROBE_PASS", + f"MOE_EP_EP{world_size}_CUDA_GRAPH_PROBE_PASS " + f"wgrad_capture_mode={args.wgrad_capture_mode}", flush=True, ) finally: diff --git a/test/python/moe_ep/test_moe_ep_backward.py b/test/python/moe_ep/test_moe_ep_backward.py index 6de24fe51..efcda1532 100644 --- a/test/python/moe_ep/test_moe_ep_backward.py +++ b/test/python/moe_ep/test_moe_ep_backward.py @@ -58,12 +58,15 @@ _distributed_subgroup_backward_reference_worker, ) from moe_ep.moe_ep_test_support import ( + _allocate_dense_grouped_wgrad_outputs, _assert_fixed_training_drop_overflow_result, _assert_fixed_training_matches_reference, + _assert_grouped_wgrads_match_reference, _assert_training_graph_tails_are_reset, _assert_training_weight_sources_changed, _capture_fixed_training_batch, _copy_training_weight_sources_, + _dense_wgrads_from_grouped_kernel, _dense_wgrads_from_operands, _fixed_training_case, _fixed_training_drop_overflow_case, @@ -990,6 +993,186 @@ def test_fixed_training_resources_ep1_matches_independent_reference( assert all(tensor.eq(0).all() for tensor in zero_tensors) +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_fixed_training_resources_ep1_grouped_wgrad_matches_independent_reference(): + device = _sm107_device() + base_args = make_forward_inputs(device) + args = ( + base_args[0].dequantize(torch.bfloat16), + base_args[1], + base_args[2], + base_args[3], + base_args[4].float(), + ) + grad_output = _grad_output(device, args[0].shape[0], seed=20260831) + expected = _fixed_training_reference( + args, + grad_output, + combine_format="bf16", + gate_up_clamp=None, + ) + + with MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=args[0].shape[0], + max_recv_size_per_rank=args[0].shape[0] * args[3].shape[1], + drop_on_overflow=True, + combine_format="bf16", + ) as op: + resources = op.prepare_training_resources( + _fixed_training_weights(args), + slot_count=1, + lane_count=1, + ) + actual = _run_fixed_training_batch( + resources, + resources.lanes[0], + ((resources.slots[0], args, grad_output),), + )[0] + grouped_wgrads = _dense_wgrads_from_grouped_kernel(actual.wgrads) + torch.cuda.synchronize(device) + + assert actual.overflow.eq(0).all() + _assert_fixed_training_matches_reference( + (actual.y, actual.dx, actual.dprob, actual.wgrads), + expected, + args[3], + ) + torch.testing.assert_close( + actual.wgrads.valid_route_counts, + expected[3].valid_route_counts, + rtol=0, + atol=0, + ) + assert actual.wgrads.valid_route_counts.gt(0).all() + expected_offsets = torch.cumsum( + torch.div( + actual.wgrads.valid_route_counts + 127, + 128, + rounding_mode="floor", + ) + * 128, + dim=0, + dtype=actual.wgrads.expert_offsets.dtype, + ) + torch.testing.assert_close( + actual.wgrads.expert_offsets, + expected_offsets, + rtol=0, + atol=0, + ) + + expected_wgrads = expected[3].dense_wgrads() + _assert_grouped_wgrads_match_reference( + grouped_wgrads, + expected_wgrads, + reference_name="the independent PyTorch MXFP8 reference", + ) + _assert_grouped_wgrads_match_reference( + grouped_wgrads, + _dense_wgrads_from_operands(actual.wgrads), + reference_name="the decoded production operand bundle", + close_kwargs={"rtol": 0.1, "atol": 0.1}, + ) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_fixed_training_resources_ep1_grouped_wgrad_accumulates_two_microbatches(): + device = _sm107_device() + base_args = make_forward_inputs(device) + args0 = ( + base_args[0].dequantize(torch.bfloat16), + base_args[1], + base_args[2], + base_args[3], + base_args[4].float(), + ) + args1 = ( + args0[0].mul(-0.5), + args0[1], + args0[2], + args0[3].roll(1, dims=0), + args0[4].roll(1, dims=0), + ) + grad_outputs = ( + _grad_output(device, args0[0].shape[0], seed=20260902), + _grad_output(device, args1[0].shape[0], seed=20260903), + ) + references = tuple( + _fixed_training_reference( + args, + grad_output, + combine_format="bf16", + gate_up_clamp=None, + ) + for args, grad_output in zip((args0, args1), grad_outputs) + ) + + with MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=args0[0].shape[0], + max_recv_size_per_rank=args0[0].shape[0] * args0[3].shape[1], + drop_on_overflow=True, + combine_format="bf16", + ) as op: + resources = op.prepare_training_resources( + _fixed_training_weights(args0), + slot_count=2, + lane_count=1, + ) + batch = tuple( + (slot, args, grad_output) + for slot, args, grad_output in zip( + resources.slots, + (args0, args1), + grad_outputs, + ) + ) + actuals = _run_fixed_training_batch(resources, resources.lanes[0], batch) + accumulated = _allocate_dense_grouped_wgrad_outputs( + actuals[0].wgrads, + fill_value=0, + ) + output_pointers = tuple(output.data_ptr() for output in accumulated) + for actual in actuals: + returned = _dense_wgrads_from_grouped_kernel( + actual.wgrads, + wgrad_tensors=accumulated, + accumulate_on_output=True, + ) + assert tuple(output.data_ptr() for output in returned) == output_pointers + torch.cuda.synchronize(device) + + for actual, args, reference in zip(actuals, (args0, args1), references): + assert actual.overflow.eq(0).all() + _assert_fixed_training_matches_reference( + (actual.y, actual.dx, actual.dprob, actual.wgrads), + reference, + args[3], + ) + expected_accumulated = tuple( + reference0.float() + reference1.float() + for reference0, reference1 in zip( + references[0][3].dense_wgrads(), + references[1][3].dense_wgrads(), + ) + ) + _assert_grouped_wgrads_match_reference( + accumulated, + expected_accumulated, + reference_name="the sum of two independent PyTorch MXFP8 references", + close_kwargs={"rtol": 0.2, "atol": 0.25}, + ) + + @pytest.mark.L1 @pytest.mark.gpu_exclusive @pytest.mark.parametrize( @@ -1437,8 +1620,13 @@ def assert_result(actual, expected_overflow): batch = ((slot, args, grad_output),) # Compile the fixed T=1 specialization and validate overflow eagerly - # before capturing the same forward/backward/finalize sequence. + # before capturing the same forward/backward/finalize/WGrad sequence. warmup = _run_fixed_training_batch(resources, lane, batch)[0] + grouped_outputs = _allocate_dense_grouped_wgrad_outputs(warmup.wgrads) + _dense_wgrads_from_grouped_kernel( + warmup.wgrads, + wgrad_tensors=grouped_outputs, + ) torch.cuda.synchronize(device) assert_result(warmup, 1) @@ -1449,8 +1637,11 @@ def assert_result(actual, expected_overflow): lane, batch, capture_stream, + grouped_wgrad_outputs=(grouped_outputs,), ) graph_actual = captured.actuals[0] + graph_grouped_wgrads = captured.grouped_wgrads[0] + grouped_output_pointers = tuple(output.data_ptr() for output in grouped_outputs) for routing, expected_overflow in ( (overflow_routing, 1), @@ -1458,6 +1649,8 @@ def assert_result(actual, expected_overflow): (overflow_routing, 1), ): args[3].copy_(routing) + for output in grouped_outputs: + output.fill_(float("nan")) assert args[3].data_ptr() == routing_pointer expected = _fixed_training_drop_overflow_reference( args, @@ -1471,4 +1664,23 @@ def assert_result(actual, expected_overflow): *expected, expected_overflow=expected_overflow, ) + assert ( + tuple(output.data_ptr() for output in graph_grouped_wgrads) + == grouped_output_pointers + ) + _assert_grouped_wgrads_match_reference( + graph_grouped_wgrads, + expected[0][3].dense_wgrads(), + reference_name="the independent PyTorch MXFP8 graph reference", + ) + _assert_grouped_wgrads_match_reference( + graph_grouped_wgrads, + _dense_wgrads_from_operands(graph_actual.wgrads), + reference_name="the decoded captured production operand bundle", + close_kwargs={"rtol": 0.1, "atol": 0.1}, + ) + assert all(torch.isfinite(output).all() for output in graph_grouped_wgrads) + if expected_overflow: + assert graph_grouped_wgrads[0][1].eq(0).all() + assert graph_grouped_wgrads[1][1].eq(0).all() assert captured.public_pointers[0] == _training_public_pointers(graph_actual) diff --git a/test/python/moe_ep/test_moe_ep_multinode.py b/test/python/moe_ep/test_moe_ep_multinode.py index b0da9fe99..fe34e6171 100644 --- a/test/python/moe_ep/test_moe_ep_multinode.py +++ b/test/python/moe_ep/test_moe_ep_multinode.py @@ -249,6 +249,12 @@ def test_mxfp8_forward_multinode_matches_reference( "bf16", id="backward-ep32-world32-bf16-minimal", ), + pytest.param( + 32, + 32, + "mxfp8", + id="backward-ep32-world32-mxfp8", + ), ], ) def test_fixed_training_resources_multinode_match_independent_reference( From e3b437c8a2a3a7c23d2cb1a9e6f306c2d85190f8 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Fri, 28 Aug 2026 18:24:44 -0700 Subject: [PATCH 20/31] docs: clarify MoeEP formulas and usage --- docs/fe-oss-apis/moe_ep.md | 2 +- docs/fe-oss-apis/overview.md | 4 +- docs/operations/MoeEp.md | 104 ----------------- docs/operations/moe_ep.md | 212 +++++++++++++++++++++++++++++++++++ llms.txt | 2 +- 5 files changed, 216 insertions(+), 108 deletions(-) delete mode 100644 docs/operations/MoeEp.md create mode 100644 docs/operations/moe_ep.md diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md index 689431c8b..6a3b7ff0e 100644 --- a/docs/fe-oss-apis/moe_ep.md +++ b/docs/fe-oss-apis/moe_ep.md @@ -3,7 +3,7 @@ `cudnn.moe_ep` provides a fused SwiGLU MoE implementation for Rubin SM107. Experts are sharded contiguously across an optional expert-parallel process group. This page documents the Python API and lifecycle. See the -[MoeEP operation reference](../operations/MoeEp.md) for supported +[MoeEP operation reference](../operations/moe_ep.md) for supported architectures, data formats, tensor contracts, and expert-parallel topology. ## Installation diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 6bfb12cbb..0243082a2 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -38,7 +38,7 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d - [RMSNorm + SiLU](rmsnorm_silu.md) - [MoE + Expert Parallel API](moe_ep.md) — Rubin SM107 fused SwiGLU with fixed-resource training and CUDA Graph support; see the - [MoeEP operation reference](../operations/MoeEp.md) for support details + [MoeEP operation reference](../operations/moe_ep.md) for support details ## Installation and setup @@ -61,7 +61,7 @@ pip install "nvidia-cudnn-frontend[cutedsl,comm]" torch torch-c-dlpack-ext MoeEP is currently CUDA/PyTorch-only and targets Rubin SM107. EP2+ execution also requires NCCL, NVSHMEM, and a direct-P2P MNNVL peer-access domain. See the -[MoeEP support matrix and tensor contracts](../operations/MoeEp.md) +[MoeEP support matrix and tensor contracts](../operations/moe_ep.md) before integrating it. After installation, you can import the APIs directly from the `cudnn` package, i.e. `from cudnn import {your_operation}` diff --git a/docs/operations/MoeEp.md b/docs/operations/MoeEp.md deleted file mode 100644 index 20fc5437d..000000000 --- a/docs/operations/MoeEp.md +++ /dev/null @@ -1,104 +0,0 @@ -# Mixture of Experts with Expert Parallelism - -The MoeEP operation fuses token routing, expert SwiGLU computation, and -expert-parallel communication. Global experts are sharded contiguously across -the ranks of an expert-parallel process group. - -For token \(x_t\), selected expert \(e_{t,k}\), and routing weight \(p_{t,k}\), -the operation computes: - -\[ -y_t = \sum_{k=0}^{K-1} p_{t,k} - \left(\operatorname{SiLU}(x_t W^{gate}_{e_{t,k}}) - \odot (x_t W^{up}_{e_{t,k}})\right) - W^{down}_{e_{t,k}} -\] - -The current implementation is exposed by the frontend-only Python -[`cudnn.moe_ep.MoeEp`](../fe-oss-apis/moe_ep.md) API. It is distinct from the -cuDNN graph [MoE Grouped Matmul](MoeGroupedMatmul.md) operation. - -## Execution support - -- NVIDIA Rubin SM107 GPUs (compute capability 10.7). -- CUDA and PyTorch execution. -- `nvidia-cutlass-dsl>=4.8.0` for the Rubin kernels. The package-wide - `cutedsl` extra retains its 4.5.0 installation floor so other cuDNN Frontend - operations remain usable with older compatible DSL versions. -- Fused SwiGLU with contiguous expert sharding. -- `apply_topk_in_fc1=True`. -- `hidden_size` divisible by 128. -- `intermediate_size` divisible by 256. -- `top_k <= min(32, num_experts)`. -- `num_experts` divisible by the expert-parallel group size. -- An explicit positive `max_tokens_per_rank`. - -The fixed-resource CUDA Graph path has hardware acceptance through EP32 when -all ranks are in one direct-P2P MNNVL peer-access domain. The Python capability -layer does not impose an EP-size ceiling; cross-MNNVL execution is not part of -the validated support surface. - -## Data formats - -Inference activation and expert weights accept: - -- BF16, FP16, or FP32 plain tensors, staged internally to MXFP8; or -- MXFP8 `BlockScaledTensor` values with logical block axis 1. - -The current executable output format is BF16. The expert-combine path accepts -BF16 or MXFP8. NVFP4 types are represented by the public API but native NVFP4 -operands, combine, and output are not executable by this backend. - -Fixed-resource training narrows dynamic activation and gradient inputs to -contiguous BF16 or FP32 tensors. Training weights are contiguous MXFP8 -block-scaled tensors. - -## Tensor contracts - -Let: - -- \(T\) be the local token count; -- \(H\) be `hidden_size`; -- \(I\) be `intermediate_size`; -- \(K\) be `top_k`; -- \(E_{local}\) be the local expert count. - -Inference uses: - -- `activation`: `(T, H)`; -- `topk_idx`: `(T, K)`, Int32 or Int64, containing `-1` or a valid global - expert ID; -- `topk_weights`: `(T, K)`, floating point; -- FC1 weights: `(E_local, H, 2I)`; -- FC2 weights: `(E_local, I, H)`; -- output: `(T, H)`, BF16. - -Fixed-resource training additionally binds transposed backward weights with -shapes `(E_local, H, I)` and `(E_local, 2I, H)`. Dynamic tensors must share one -device and satisfy `T <= max_tokens_per_rank`. - -## Expert-parallel communication - -EP2+ execution requires: - -- an initialized NCCL process group; -- `nvshmem4py` and usable NVSHMEM libraries; -- direct peer access among every pair of participating ranks; and -- consistent rank ordering, resource sizes, tuning, slot selection, and lane - ordering across the group. - -`max_recv_size_per_rank` bounds receive capacity. When omitted, it defaults to -the worst-case route count: - -```text -ep_size * max_tokens_per_rank * top_k -``` - -Resources cannot grow during CUDA Graph replay. Capacity or storage changes -require resource preparation and graph capture again. - -## API reference - -See [MoE + Expert Parallel API](../fe-oss-apis/moe_ep.md) for installation, -constructor arguments, inference and training lifecycles, tuning, overflow -handling, and CUDA Graph usage. diff --git a/docs/operations/moe_ep.md b/docs/operations/moe_ep.md new file mode 100644 index 000000000..b9ba482ab --- /dev/null +++ b/docs/operations/moe_ep.md @@ -0,0 +1,212 @@ +# Mixture of Experts with Expert Parallelism + +The MoeEP operation fuses token routing, expert SwiGLU computation, and +expert-parallel communication. Global experts are sharded contiguously across +the ranks of an expert-parallel process group. + +## Operation + +Let $x_t \in \mathbb{R}^{H}$ be token $t$, $e_{t,k}$ its $k$-th +selected global expert, and $p_{t,k}$ the corresponding routing weight. For +each valid route, split the FC1 result into gate and up projections: + +$$ +\left[g_{t,k}, u_{t,k}\right] + = x_t W^{\mathrm{fc1}}_{e_{t,k}}, +\qquad +h_{t,k} + = p_{t,k}\left(\operatorname{SiLU}(g_{t,k}) \odot u_{t,k}\right), +\qquad +z_{t,k} + = h_{t,k} W^{\mathrm{fc2}}_{e_{t,k}}. +$$ + +The final token output is the sum over its selected experts: + +$$ +y_t = \sum_{\substack{0 \le k < K \\ e_{t,k} \ne -1}} z_{t,k}. +$$ + +When `gate_up_clamp=C`, the operation uses +$\min(g_{t,k}, C)$ for the gate and +$\operatorname{clip}(u_{t,k}, -C, C)$ for the up projection. A route whose +expert ID is `-1` contributes zero. Because the executable backend requires +`apply_topk_in_fc1=True`, it applies $p_{t,k}$ to the SwiGLU result before +FC2. The backend also stages plain inputs to MXFP8 and requantizes the routed +intermediate before FC2, so the equations describe the mathematical operation +rather than its finite-precision rounding. + +With $E$ global experts and an expert-parallel group of size $P$, each rank +stores $E_{\mathrm{local}}=E/P$ consecutive experts. Global expert $e$ is +owned by group-relative rank + +$$ +\operatorname{owner}(e) + = \left\lfloor \frac{e}{E_{\mathrm{local}}} \right\rfloor. +$$ + +## Python API + +The operation is exposed by the frontend-only `cudnn.MoeEp` object API. Static +model, parallelism, capacity, and format choices are set in the constructor: + +```python +from cudnn import MoeEp + +op = MoeEp( + num_experts=E, + hidden_size=H, + intermediate_size=I, + top_k=K, + ep_group=ep_group, # None for EP1 + max_tokens_per_rank=max_tokens, + max_recv_size_per_rank=None, # Defaults to P * max_tokens * K + drop_on_overflow=False, + output_format="bf16", + combine_format="bf16", # "bf16" or "mxfp8" + apply_topk_in_fc1=True, + gate_up_clamp=None, +) +``` + +`topk_idx` contains global expert IDs. Each rank passes its local tokens and +its contiguous shard of both expert-weight tensors: + +```python +output = op( + activation, # (T, H) + fc1_weight, # (E_local, H, 2I) + fc2_weight, # (E_local, I, H) + topk_idx, # (T, K), global expert IDs or -1 + topk_weights, # (T, K) +) # (T, H), BF16 +``` + +For inference CUDA Graph capture, call `op.warmup(...)` with the exact +bindings before capture. `MoeEp` supports `close()` and context-manager use. + +Fixed-resource training uses the same operator and binds graph-stable weights, +slots, and execution lanes: + +```python +from cudnn import MoeEpTrainingWeights + +weights = MoeEpTrainingWeights( + forward_fc1=forward_fc1_mxfp8, + forward_fc2=forward_fc2_mxfp8, + backward_w2_transpose=backward_w2t_mxfp8, + backward_w1_transpose=backward_w1t_mxfp8, +) +resources = op.prepare_training_resources(weights, slot_count=2, lane_count=1) +slot, lane = resources.slots[0], resources.lanes[0] + +resources.refresh_weights() +output = resources.forward(slot, lane, activation, topk_idx, topk_weights) +grad_activation, dprob, wgrad_operands = resources.backward( + slot, lane, grad_output +) +overflow = resources.finalize_overflow((slot,), lane) +``` + +The WGrad result is a fixed-capacity grouped-GEMM operand bundle, not dense +optimizer-ready weight gradients. See the detailed +[MoE + Expert Parallel API](../fe-oss-apis/moe_ep.md) reference for +installation, all constructor arguments, tensor formats, training resource +lifecycle, tuning, overflow handling, and CUDA Graph requirements. MoeEP is +distinct from the cuDNN graph [MoE Grouped Matmul](MoeGroupedMatmul.md) +operation. + +## Execution support + +- NVIDIA Rubin SM107 GPUs (compute capability 10.7). +- CUDA and PyTorch execution. +- `nvidia-cutlass-dsl>=4.8.0` for the Rubin kernels. The package-wide + `cutedsl` extra retains its 4.5.0 installation floor so other cuDNN Frontend + operations remain usable with older compatible DSL versions. +- Fused SwiGLU with contiguous expert sharding. +- `apply_topk_in_fc1=True`. +- `hidden_size` divisible by 128. +- `intermediate_size` divisible by 256. +- `top_k <= min(32, num_experts)`. +- `num_experts` divisible by the expert-parallel group size. +- An explicit positive `max_tokens_per_rank`. + +The fixed-resource CUDA Graph path has hardware acceptance through EP32 when +all ranks are in one direct-P2P MNNVL peer-access domain. The Python capability +layer does not impose an EP-size ceiling; cross-MNNVL execution is not part of +the validated support surface. + +## Data formats + +Inference activation and expert weights accept: + +- BF16, FP16, or FP32 plain tensors, staged internally to MXFP8; or +- MXFP8 `BlockScaledTensor` values with logical block axis 1. + +The current executable output format is BF16. The expert-combine path accepts +BF16 or MXFP8. NVFP4 types are represented by the public API but native NVFP4 +operands, combine, and output are not executable by this backend. + +Fixed-resource training narrows dynamic activation and gradient inputs to +contiguous BF16 or FP32 tensors. Training weights are contiguous MXFP8 +block-scaled tensors. + +## Tensor contracts + +Let: + +- $T$ be the local token count; +- $H$ be `hidden_size`; +- $I$ be `intermediate_size`; +- $K$ be `top_k`; +- $E_{\mathrm{local}}$ be the local expert count. + +Inference uses: + +- `activation`: `(T, H)`; +- `topk_idx`: `(T, K)`, Int32 or Int64, containing `-1` or a valid global + expert ID; +- `topk_weights`: `(T, K)`, floating point; +- FC1 weights: `(E_local, H, 2I)`; +- FC2 weights: `(E_local, I, H)`; +- output: `(T, H)`, BF16. + +All inference tensors must reside on one device, and the local token count must +satisfy `T <= max_tokens_per_rank`. + +Fixed-resource training uses a narrower graph-stable contract: + +- `activation` and `grad_output`: contiguous `(T, H)`, BF16 or FP32; +- `topk_idx`: contiguous `(T, K)`, Int32; +- `topk_weights`: contiguous `(T, K)`, FP32; +- forward FC1 and FC2 weights: contiguous MXFP8 block-scaled tensors with + shapes `(E_local, H, 2I)` and `(E_local, I, H)`; +- transposed backward weights: contiguous MXFP8 block-scaled tensors with + shapes `(E_local, H, I)` and `(E_local, 2I, H)`; +- forward output: `(T, H)`, BF16; +- `grad_activation`: fixed-slot `(T, H)`, FP32; +- `dprob`: source-order `(T, K)`, FP32; +- `wgrad_operands`: a fixed-capacity `MoeEpTrainingWgradOperands` bundle. + +All dynamic training tensors must reside on one device and satisfy +`T <= max_tokens_per_rank`. + +## Expert-parallel communication + +EP2+ execution requires: + +- an initialized NCCL process group; +- `nvshmem4py` and usable NVSHMEM libraries; +- direct peer access among every pair of participating ranks; and +- consistent rank ordering, resource sizes, tuning, slot selection, and lane + ordering across the group. + +`max_recv_size_per_rank` bounds receive capacity. When omitted, it defaults to +the worst-case route count: + +```text +ep_size * max_tokens_per_rank * top_k +``` + +Resources cannot grow during CUDA Graph replay. Capacity or storage changes +require resource preparation and graph capture again. diff --git a/llms.txt b/llms.txt index e2fd61536..9c6aef031 100644 --- a/llms.txt +++ b/llms.txt @@ -18,7 +18,7 @@ Published documentation: https://docs.nvidia.com/deeplearning/cudnn/latest/devel - [Convolutions](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Convolutions.md) - [Normalizations (LayerNorm, RMSNorm, BatchNorm, InstanceNorm)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Normalizations.md) - [MoE Grouped Matmul](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/MoeGroupedMatmul.md) -- [MoE with Expert Parallelism](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/MoeEp.md) +- [MoE with Expert Parallelism](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/moe_ep.md) - [Pointwise](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Pointwise.md) - [Block Scaling (MXFP8/NVFP4 quantization)](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/BlockScaling.md) - [RoPE](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/RoPE.md) From 377cf36538b513b0d40db2695df7645a5ddb9ae8 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Fri, 28 Aug 2026 18:31:42 -0700 Subject: [PATCH 21/31] docs: use GitHub-compatible MoeEP math macros --- docs/operations/moe_ep.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/operations/moe_ep.md b/docs/operations/moe_ep.md index b9ba482ab..6f0ad1cf7 100644 --- a/docs/operations/moe_ep.md +++ b/docs/operations/moe_ep.md @@ -15,7 +15,7 @@ $$ = x_t W^{\mathrm{fc1}}_{e_{t,k}}, \qquad h_{t,k} - = p_{t,k}\left(\operatorname{SiLU}(g_{t,k}) \odot u_{t,k}\right), + = p_{t,k}\left(\mathrm{SiLU}(g_{t,k}) \odot u_{t,k}\right), \qquad z_{t,k} = h_{t,k} W^{\mathrm{fc2}}_{e_{t,k}}. @@ -29,7 +29,7 @@ $$ When `gate_up_clamp=C`, the operation uses $\min(g_{t,k}, C)$ for the gate and -$\operatorname{clip}(u_{t,k}, -C, C)$ for the up projection. A route whose +$\mathrm{clip}(u_{t,k}, -C, C)$ for the up projection. A route whose expert ID is `-1` contributes zero. Because the executable backend requires `apply_topk_in_fc1=True`, it applies $p_{t,k}$ to the SwiGLU result before FC2. The backend also stages plain inputs to MXFP8 and requantizes the routed @@ -41,7 +41,7 @@ stores $E_{\mathrm{local}}=E/P$ consecutive experts. Global expert $e$ is owned by group-relative rank $$ -\operatorname{owner}(e) +\mathrm{owner}(e) = \left\lfloor \frac{e}{E_{\mathrm{local}}} \right\rfloor. $$ From 84dc7142d43aaae6d32d446b15f6b8350b256a75 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 30 Aug 2026 01:04:00 -0700 Subject: [PATCH 22/31] add the changes Signed-off-by: Varun Thumbe --- docs/fe-oss-apis/moe_ep.md | 23 +- .../cudnn/moe_ep/_megamoe_backend/README.md | 10 +- .../mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py | 9 +- .../bwd_dglu/dglu_mxfp8_mega_moe_kernel.py | 4 +- .../mxfp8/_training_weights.py | 263 ++++++++---------- python/cudnn/moe_ep/_types.py | 4 +- python/cudnn/moe_ep/_validation.py | 15 +- test/python/moe_ep/test_moe_ep_backward.py | 59 +++- 8 files changed, 201 insertions(+), 186 deletions(-) diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md index 6a3b7ff0e..551efb1cd 100644 --- a/docs/fe-oss-apis/moe_ep.md +++ b/docs/fe-oss-apis/moe_ep.md @@ -275,27 +275,30 @@ Fixed-resource training uses a narrower, graph-stable staging ABI: Expert IDs and finite dynamic values remain a trusted-caller replay contract; they are not revalidated by host code after graph capture. -`MoeEpTrainingWeights` contains four contiguous MXFP8 block-scaled tensors: +`MoeEpTrainingWeights` contains four MXFP8 block-scaled tensors: - `forward_fc1`: `(E_local, H, 2I)` - `forward_fc2`: `(E_local, I, H)` - `backward_w2_transpose`: `(E_local, H, I)` - `backward_w1_transpose`: `(E_local, 2I, H)` -Each data and scale tensor must be contiguous, reside on one device, and use -logical block axis 1. Plain FP16 operands are accepted by inference staging, -but fixed-resource training accepts only BF16 or FP32 `activation` and -`grad_output`. +Each data and scale tensor must be contiguous or use compact K-major strides +`(K*N, 1, K)`, reside on one device, and use logical block axis 1. The K-major +form lets framework integrations bind transposed weight views without an extra +copy. Plain FP16 operands are accepted by inference staging, but fixed-resource +training accepts only BF16 or FP32 `activation` and `grad_output`. Replacing weight storage requires preparing resources and capturing again. Callers must establish stream/event ordering for in-place weight updates. ### Explicit weight refresh contract -`MoeEpTrainingWeights` uses a public contiguous MXFP8 layout, while the Rubin -kernels consume fixed-address K-major, gate/up-interleaved, and blocked-scale -layouts. `resources.refresh_weights()` enqueues the required device-only copies -and layout transforms into the internal kernel bindings. +Rubin training recognizes compact K-major forward views plus +contiguous backward transposes. For that form, weight data is bound directly +and FC1 gate/up values must already use 32-element interleaving. +`resources.refresh_weights()` then swizzles only scales into fixed-address, +kernel-native buffers and never copies the weight payload. Existing contiguous +public packs retain the compatible data-and-scale staging path. The caller must obey all of the following: @@ -403,7 +406,7 @@ For fixed-resource training: 1. all ranks collectively call `prepare_training_resources`; 2. all ranks perform an ordinary `refresh_weights -> forward -> backward -> finalize_overflow` warmup so - every staging, MegaMoE, and WGrad-export kernel is compiled; + every scale-staging, MegaMoE, and WGrad-export kernel is compiled; 3. each rank captures its outer graph; 4. ranks align after capture; 5. graph execs are submitted in lockstep without host synchronization inside diff --git a/python/cudnn/moe_ep/_megamoe_backend/README.md b/python/cudnn/moe_ep/_megamoe_backend/README.md index a55986ad7..8da780fc6 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/README.md +++ b/python/cudnn/moe_ep/_megamoe_backend/README.md @@ -73,9 +73,13 @@ forward, backward, and WGrad-export kernels are compiled before capture. `MoeEpTrainingWeights` contains four address-stable MXFP8 block-scaled tensors: forward W1/W2 and independently quantized backward W2-transpose/W1-transpose. -Their public layout differs from the K-major, gate/up-interleaved, and -blocked-scale kernel bindings. After every in-place data+scale update, the -caller must enqueue `resources.refresh_weights()` before the first consumer, +When forward weights use compact K-major storage and backward transpose weights +use standard contiguous storage, the kernels alias weight data directly. In +this layout, W1 gate/up values are interleaved in 32-element strips and only +scales require kernel-native staging. When forward weights use standard +contiguous storage, weight data and scales are copied and reordered into +persistent kernel buffers. After every in-place data+scale update, the caller +must enqueue `resources.refresh_weights()` before the first consumer, with explicit stream/event ordering. A matching forward/backward pair must use one version; refresh cannot overlap any consumer on another slot/lane. Replacing source storage requires closing the old operator, creating a new `MoeEp` diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py index cff668b4a..031c753a5 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py @@ -857,8 +857,9 @@ def __call__( ), ) - # B_gemm (fc1 weights): (experts, hidden, intermediate_gateup) with hidden stride-1 (K-major) - # -> (N=intermediate_gateup, K=hidden, L=experts). + # B_gemm (W2T): reinterpret public C-contiguous (experts, hidden, inter_half) + # as (N=inter_half, K=hidden, L=experts). The stride permutation makes this + # an N-major GEMM operand without staging or moving data. experts, hidden_b, intermediate_gateup = fc1_weight.shape fc1_weight_gemm = cute.make_tensor( fc1_weight.iterator, @@ -924,7 +925,9 @@ def __call__( ), ) - # GEMM-domain transform for fc2 phase ── + # GEMM-domain transform for fc2 phase. W1T is public C-contiguous + # (experts, 2 * inter_half, hidden), with its reduction rows already in + # 32-wide gate/up order. Preserve that K ordering and expose hidden as N. experts2, intermediate_downproj_b2, hidden_b2 = fc2_weight.shape fc2_weight_gemm = cute.make_tensor( fc2_weight.iterator, diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py index a53770617..fb03d5ad3 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py @@ -143,9 +143,9 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): grad_out_sf=fake_tensor(sf_dtype, (tokens, self.token_comm.activation_sf_hidden_padded), (1, 0), {0}, 16), topk_idx=fake_tensor(cutlass.Int64, (tokens, self.num_topk), (1, 0), {0}, 16), topk_weights=fake_tensor(cutlass.Float32, (tokens, self.num_topk), (1, 0), {0}, 4), - fc1_weight=fake_tensor(self.ab_dtype, (experts, hidden, inter_half), (2, 0, 1), {0, 2}, 16), + fc1_weight=fake_tensor(self.ab_dtype, (experts, hidden, inter_half), (2, 1, 0), {0, 2}, 16), fc1_weight_sf=fake_tensor(sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16), - fc2_weight=fake_tensor(self.ab_dtype, (experts, gate_up, hidden), (2, 0, 1), {0, 2}, 16), + fc2_weight=fake_tensor(self.ab_dtype, (experts, gate_up, hidden), (2, 1, 0), {0, 2}, 16), fc2_weight_sf=fake_tensor(sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16), beta=fake_tensor(cutlass.Float32, (experts,), (0,), {0}, 4), fc1_preact=fake_tensor(cutlass.BFloat16, self.get_fc1_preact_shape(), (1, 0), set(), 128), diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py index 623bf50e6..8d900b3e5 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py @@ -18,8 +18,6 @@ def _round_up(value: int, multiple: int) -> int: def _empty_k_major_like(tensor: torch.Tensor) -> torch.Tensor: - if tensor.ndim != 3: - raise ValueError(f"K-major training weight must be rank 3, got {tensor.ndim}") experts, reduction, output = tensor.shape return torch.empty_strided( tensor.shape, @@ -45,13 +43,44 @@ def _empty_blocked_scales( ).view(dtype) -def _copy_k_major( +def _copy_blocked_scales_plain( target: torch.Tensor, source: torch.Tensor, + *, + raw_rows: int, + raw_columns: int, ) -> None: - if target.shape != source.shape: - raise ValueError(f"K-major copy shape mismatch: {target.shape} != {source.shape}") - target.copy_(source) + """Pack public ``(E,Kblocks,N)`` scales for a non-interleaved weight.""" + + experts = source.shape[0] + if tuple(source.shape) != (experts, raw_columns, raw_rows): + raise ValueError("plain training scale shape mismatch: " f"{tuple(source.shape)} != " f"{(experts, raw_columns, raw_rows)}") + if raw_rows % 128 or raw_columns % 4: + raise ValueError("training scale pack requires rows divisible by 128 and " "columns divisible by 4") + row_blocks = raw_rows // 128 + column_blocks = raw_columns // 4 + source_view = ( + source.view( + torch.uint8, + ) + .view( + experts, + column_blocks, + 4, + row_blocks, + 4, + 32, + ) + .permute(0, 3, 1, 5, 4, 2) + ) + target.view(torch.uint8).view( + experts, + row_blocks, + column_blocks, + 32, + 4, + 4, + ).copy_(source_view) def _copy_gate_up_interleaved_last( @@ -65,13 +94,7 @@ def _copy_gate_up_interleaved_last( if target.shape != source.shape or gate_up != 2 * intermediate: raise ValueError("forward FC1 training weight shape mismatch") pairs = intermediate // 32 - source_view = source.view( - experts, - reduction, - 2, - pairs, - 32, - ).permute(0, 1, 3, 2, 4) + source_view = source.view(experts, reduction, 2, pairs, 32).permute(0, 1, 3, 2, 4) target_view = target.as_strided( (experts, reduction, pairs, 2, 32), ( @@ -96,64 +119,8 @@ def _copy_gate_up_interleaved_reduction( if target.shape != source.shape or gate_up != 2 * intermediate: raise ValueError("backward W1-transpose training weight shape mismatch") pairs = intermediate // 32 - source_view = source.view( - experts, - 2, - pairs, - 32, - output, - ).permute(0, 2, 1, 3, 4) - target_view = target.as_strided( - (experts, pairs, 2, 32, output), - ( - target.stride(0), - 64 * target.stride(1), - 32 * target.stride(1), - target.stride(1), - target.stride(2), - ), - ) - target_view.copy_(source_view) - - -def _copy_blocked_scales_plain( - target: torch.Tensor, - source: torch.Tensor, - *, - raw_rows: int, - raw_columns: int, -) -> None: - """Pack public ``(E,Kblocks,N)`` scales for a non-interleaved weight.""" - - experts = source.shape[0] - if tuple(source.shape) != (experts, raw_columns, raw_rows): - raise ValueError("plain training scale shape mismatch: " f"{tuple(source.shape)} != " f"{(experts, raw_columns, raw_rows)}") - if raw_rows % 128 or raw_columns % 4: - raise ValueError("training scale pack requires rows divisible by 128 and " "columns divisible by 4") - row_blocks = raw_rows // 128 - column_blocks = raw_columns // 4 - source_view = ( - source.view( - torch.uint8, - ) - .view( - experts, - column_blocks, - 4, - row_blocks, - 4, - 32, - ) - .permute(0, 3, 1, 5, 4, 2) - ) - target.view(torch.uint8).view( - experts, - row_blocks, - column_blocks, - 32, - 4, - 4, - ).copy_(source_view) + source_view = source.view(experts, 2, pairs, 32, output).permute(0, 2, 1, 3, 4) + target.view(experts, pairs, 2, 32, output).copy_(source_view) def _copy_blocked_scales_gate_up_rows( @@ -167,34 +134,15 @@ def _copy_blocked_scales_gate_up_rows( experts = source.shape[0] raw_rows = 2 * intermediate - raw_columns = reduction_blocks - if tuple(source.shape) != (experts, raw_columns, raw_rows): - raise ValueError("forward FC1 training scale shape mismatch") - if intermediate % 64 or raw_columns % 4: - raise ValueError("forward FC1 scale pack requires intermediate divisible by 64 " "and reduction blocks divisible by 4") - row_blocks = raw_rows // 128 - column_blocks = raw_columns // 4 + if intermediate % 64 or reduction_blocks % 4: + raise ValueError("intermediate and reduction block alignment are invalid") source_view = ( source.view(torch.uint8) - .view( - experts, - column_blocks, - 4, - 2, - row_blocks, - 2, - 32, - ) + .view(experts, reduction_blocks // 4, 4, 2, raw_rows // 128, 2, 32) .permute(0, 4, 1, 6, 5, 3, 2) ) target.view(torch.uint8).view( - experts, - row_blocks, - column_blocks, - 32, - 2, - 2, - 4, + experts, raw_rows // 128, reduction_blocks // 4, 32, 2, 2, 4 ).copy_(source_view) @@ -209,37 +157,15 @@ def _copy_blocked_scales_gate_up_columns( experts = source.shape[0] reduction_blocks = intermediate // 32 - if tuple(source.shape) != ( - experts, - 2 * reduction_blocks, - output, - ): - raise ValueError("backward W1-transpose training scale shape mismatch") if output % 128 or reduction_blocks % 2: - raise ValueError("backward W1-transpose scale pack requires output divisible by " "128 and intermediate divisible by 64") - row_blocks = output // 128 - column_blocks = reduction_blocks // 2 + raise ValueError("output and reduction block alignment are invalid") source_view = ( source.view(torch.uint8) - .view( - experts, - 2, - column_blocks, - 2, - row_blocks, - 4, - 32, - ) + .view(experts, 2, reduction_blocks // 2, 2, output // 128, 4, 32) .permute(0, 4, 2, 6, 5, 3, 1) ) target.view(torch.uint8).view( - experts, - row_blocks, - column_blocks, - 32, - 4, - 2, - 2, + experts, output // 128, reduction_blocks // 2, 32, 4, 2, 2 ).copy_(source_view) @@ -254,7 +180,7 @@ class Mxfp8BackwardWeights: class Mxfp8TrainingWeightBindings: - """Stable staging tensors refreshed from four pre-quantized sources.""" + """Direct data bindings with persistent kernel-native scale staging.""" def __init__(self, weights: MoeEpTrainingWeights) -> None: self.weights = weights @@ -262,16 +188,30 @@ def __init__(self, weights: MoeEpTrainingWeights) -> None: fwd_fc2 = weights.forward_fc2 bwd_w2t = weights.backward_w2_transpose bwd_w1t = weights.backward_w1_transpose + self._uses_direct_weight_bindings = ( + fwd_fc1.data.stride(1) == 1 + and fwd_fc2.data.stride(1) == 1 + and bwd_w2t.data.is_contiguous() + and bwd_w1t.data.is_contiguous() + ) self.forward = Mxfp8Weights( - fc1_weight=_empty_k_major_like(fwd_fc1.data), + fc1_weight=( + fwd_fc1.data + if self._uses_direct_weight_bindings + else _empty_k_major_like(fwd_fc1.data) + ), fc1_weight_sf=_empty_blocked_scales( fwd_fc1, raw_rows=fwd_fc1.data.shape[2], raw_columns=fwd_fc1.data.shape[1] // 32, dtype=torch.uint8, ), - fc2_weight=_empty_k_major_like(fwd_fc2.data), + fc2_weight=( + fwd_fc2.data + if self._uses_direct_weight_bindings + else _empty_k_major_like(fwd_fc2.data) + ), fc2_weight_sf=_empty_blocked_scales( fwd_fc2, raw_rows=fwd_fc2.data.shape[2], @@ -280,14 +220,22 @@ def __init__(self, weights: MoeEpTrainingWeights) -> None: ), ) self.backward = Mxfp8BackwardWeights( - fc1_weight=_empty_k_major_like(bwd_w2t.data), + fc1_weight=( + bwd_w2t.data + if self._uses_direct_weight_bindings + else torch.empty_like(bwd_w2t.data) + ), fc1_weight_sf=_empty_blocked_scales( bwd_w2t, raw_rows=bwd_w2t.data.shape[2], raw_columns=bwd_w2t.data.shape[1] // 32, dtype=torch.float8_e8m0fnu, ), - fc2_weight=_empty_k_major_like(bwd_w1t.data), + fc2_weight=( + bwd_w1t.data + if self._uses_direct_weight_bindings + else torch.empty_like(bwd_w1t.data) + ), fc2_weight_sf=_empty_blocked_scales( bwd_w1t, raw_rows=bwd_w1t.data.shape[2], @@ -304,20 +252,28 @@ def refresh(self) -> None: fwd_fc2 = self.weights.forward_fc2 bwd_w2t = self.weights.backward_w2_transpose bwd_w1t = self.weights.backward_w1_transpose - intermediate = fwd_fc2.data.shape[1] - _copy_gate_up_interleaved_last( - self.forward.fc1_weight, - fwd_fc1.data, - intermediate, - ) - _copy_blocked_scales_gate_up_rows( - self.forward.fc1_weight_sf, - fwd_fc1.scale, - intermediate=intermediate, - reduction_blocks=fwd_fc1.data.shape[1] // 32, - ) - _copy_k_major(self.forward.fc2_weight, fwd_fc2.data) + intermediate = fwd_fc2.data.shape[1] + if self._uses_direct_weight_bindings: + _copy_blocked_scales_plain( + self.forward.fc1_weight_sf, + fwd_fc1.scale, + raw_rows=fwd_fc1.data.shape[2], + raw_columns=fwd_fc1.data.shape[1] // 32, + ) + else: + _copy_gate_up_interleaved_last( + self.forward.fc1_weight, + fwd_fc1.data, + intermediate, + ) + _copy_blocked_scales_gate_up_rows( + self.forward.fc1_weight_sf, + fwd_fc1.scale, + intermediate=intermediate, + reduction_blocks=fwd_fc1.data.shape[1] // 32, + ) + self.forward.fc2_weight.copy_(fwd_fc2.data) _copy_blocked_scales_plain( self.forward.fc2_weight_sf, fwd_fc2.scale, @@ -325,24 +281,33 @@ def refresh(self) -> None: raw_columns=fwd_fc2.data.shape[1] // 32, ) - _copy_k_major(self.backward.fc1_weight, bwd_w2t.data) + if not self._uses_direct_weight_bindings: + self.backward.fc1_weight.copy_(bwd_w2t.data) _copy_blocked_scales_plain( self.backward.fc1_weight_sf, bwd_w2t.scale, raw_rows=bwd_w2t.data.shape[2], raw_columns=bwd_w2t.data.shape[1] // 32, ) - _copy_gate_up_interleaved_reduction( - self.backward.fc2_weight, - bwd_w1t.data, - intermediate, - ) - _copy_blocked_scales_gate_up_columns( - self.backward.fc2_weight_sf, - bwd_w1t.scale, - intermediate=intermediate, - output=bwd_w1t.data.shape[2], - ) + if self._uses_direct_weight_bindings: + _copy_blocked_scales_plain( + self.backward.fc2_weight_sf, + bwd_w1t.scale, + raw_rows=bwd_w1t.data.shape[2], + raw_columns=bwd_w1t.data.shape[1] // 32, + ) + else: + _copy_gate_up_interleaved_reduction( + self.backward.fc2_weight, + bwd_w1t.data, + intermediate, + ) + _copy_blocked_scales_gate_up_columns( + self.backward.fc2_weight_sf, + bwd_w1t.scale, + intermediate=intermediate, + output=bwd_w1t.data.shape[2], + ) __all__ = [ diff --git a/python/cudnn/moe_ep/_types.py b/python/cudnn/moe_ep/_types.py index 4ae2287ec..625948038 100644 --- a/python/cudnn/moe_ep/_types.py +++ b/python/cudnn/moe_ep/_types.py @@ -223,7 +223,7 @@ class MoeEpExecutionLane: class MoeEpTrainingResources: - """TE-owned lease on fixed-capacity training slots and execution lanes.""" + """Caller-owned lease on fixed-capacity training slots and execution lanes.""" def __init__( self, @@ -264,7 +264,7 @@ def _check_binding( raise ValueError("execution lane does not belong to these resources") def refresh_weights(self) -> None: - """Enqueue fixed-address weight-layout refreshes on the current stream. + """Enqueue fixed-address scale-layout refreshes on the current stream. Call after every in-place data+scale update and before the first forward/backward that consumes that version. The caller must establish diff --git a/python/cudnn/moe_ep/_validation.py b/python/cudnn/moe_ep/_validation.py index e7ec4067f..b3338f34c 100644 --- a/python/cudnn/moe_ep/_validation.py +++ b/python/cudnn/moe_ep/_validation.py @@ -227,6 +227,14 @@ def validate_training_weights( ) -> torch.device: """Validate fixed MXFP8 weight bindings used by training resources.""" + def has_supported_layout(tensor: torch.Tensor) -> bool: + if tensor.is_contiguous(): + return True + if tensor.ndim != 3: + return False + experts, reduction, output = tensor.shape + return tensor.stride() == (reduction * output, 1, reduction) + if not isinstance(weights, MoeEpTrainingWeights): raise TypeError("weights must be a MoeEpTrainingWeights, " f"got {type(weights).__name__}") expected = ( @@ -273,8 +281,11 @@ def validate_training_weights( raise TypeError(f"{name} must be an MXFP8 BlockScaledTensor for " "fixed training resources") if tensor.format is not MoeFormat.MXFP8: raise NotImplementedError(f"{name} must use format='mxfp8', got {tensor.format.value!r}") - if not tensor.data.is_contiguous() or not tensor.scale.is_contiguous(): - raise ValueError(f"{name} data and scale must be contiguous for fixed " "training weight binding") + if not has_supported_layout(tensor.data) or not has_supported_layout(tensor.scale): + raise ValueError( + f"{name} data and scale must be contiguous or compact K-major " + "for fixed training weight binding" + ) device = weights.forward_fc1.device for name, tensor, _shape in expected[1:]: if tensor.device != device: diff --git a/test/python/moe_ep/test_moe_ep_backward.py b/test/python/moe_ep/test_moe_ep_backward.py index efcda1532..72b54227a 100644 --- a/test/python/moe_ep/test_moe_ep_backward.py +++ b/test/python/moe_ep/test_moe_ep_backward.py @@ -368,8 +368,6 @@ def test_training_sources_track_adapter_grad_y2_and_dfc2_contracts(): "plain_tensor", "axis", "format", - "data_noncontiguous", - "scale_noncontiguous", ) ], ) @@ -404,6 +402,19 @@ def test_validate_training_weights_accepts_complete_fixed_weight_set(): ) == torch.device("cpu") +@pytest.mark.L1 +@pytest.mark.parametrize("part", ["data_noncontiguous", "scale_noncontiguous"]) +def test_validate_training_weights_accepts_compact_k_major_views(part): + weights, _, _ = _training_weight_defect( + _training_weights(), + "forward_fc1", + part, + ) + assert validate_training_weights(_training_config(), weights) == torch.device("cpu") + bindings = Mxfp8TrainingWeightBindings(weights) + bindings.refresh() + + def _operator(**overrides) -> MoeEp: values = { "num_experts": 2, @@ -534,30 +545,48 @@ def test_distributed_error_mode_requires_nccl(monkeypatch): @pytest.mark.L0 -def test_training_weight_refresh_keeps_destination_addresses_stable(): +def test_training_weight_bindings_alias_data_and_stage_only_scales(): weights = _training_weights() + from cudnn.moe_ep import BlockScaledTensor, MoeEpTrainingWeights + + def compact_k_major(tensor): + return BlockScaledTensor( + data=tensor.data.transpose(1, 2).contiguous().transpose(1, 2), + scale=tensor.scale.transpose(1, 2).contiguous().transpose(1, 2), + format=tensor.format, + logical_shape=tensor.logical_shape, + axis=tensor.axis, + ) + + weights = MoeEpTrainingWeights( + forward_fc1=compact_k_major(weights.forward_fc1), + forward_fc2=compact_k_major(weights.forward_fc2), + backward_w2_transpose=weights.backward_w2_transpose, + backward_w1_transpose=weights.backward_w1_transpose, + ) bindings = Mxfp8TrainingWeightBindings(weights) bindings.refresh() - tensors = ( - bindings.forward.fc1_weight, + data_pairs = ( + (bindings.forward.fc1_weight, weights.forward_fc1.data), + (bindings.forward.fc2_weight, weights.forward_fc2.data), + (bindings.backward.fc1_weight, weights.backward_w2_transpose.data), + (bindings.backward.fc2_weight, weights.backward_w1_transpose.data), + ) + scales = ( bindings.forward.fc1_weight_sf, - bindings.forward.fc2_weight, bindings.forward.fc2_weight_sf, - bindings.backward.fc1_weight, bindings.backward.fc1_weight_sf, - bindings.backward.fc2_weight, bindings.backward.fc2_weight_sf, ) - pointers = tuple(tensor.data_ptr() for tensor in tensors) - snapshots = tuple(tensor.clone() for tensor in tensors) + scale_pointers = tuple(tensor.data_ptr() for tensor in scales) + scale_snapshots = tuple(tensor.clone() for tensor in scales) - weights.forward_fc1.data.view(torch.uint8).bitwise_xor_(1) + weights.forward_fc1.scale.view(torch.uint8).bitwise_xor_(1) bindings.refresh() - assert tuple(tensor.data_ptr() for tensor in tensors) == pointers - assert not torch.equal(bindings.forward.fc1_weight, snapshots[0]) - for tensor in tensors: - assert tensor.is_contiguous() or tensor.stride(1) == 1 + assert all(bound.data_ptr() == source.data_ptr() for bound, source in data_pairs) + assert tuple(tensor.data_ptr() for tensor in scales) == scale_pointers + assert not torch.equal(bindings.forward.fc1_weight_sf, scale_snapshots[0]) @pytest.mark.L0 From 1580ce27085ee6785a2389db9250563ec104fc4a Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 30 Aug 2026 14:13:10 -0700 Subject: [PATCH 23/31] wgrad workspace caller owned Signed-off-by: Varun Thumbe --- .../gemm_fusions/grouped_gemm_wgrad.md | 27 ++++ python/cudnn/__init__.py | 4 + python/cudnn/gemm/cutedsl/grouped/__init__.py | 2 + .../gemm/cutedsl/grouped/wgrad/__init__.py | 2 + .../cutedsl/grouped/wgrad/_blockscaled_api.py | 42 ++++++- .../cudnn/gemm/cutedsl/grouped/wgrad/api.py | 64 ++++++++-- .../grouped_gemm/test_grouped_gemm_wgrad.py | 118 ++++++++++++++++++ 7 files changed, 244 insertions(+), 15 deletions(-) diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md index 88514b2ba..d3e95aca5 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md @@ -92,6 +92,33 @@ operand pair. It preserves the pre-existing scale-factor contract: provide `global_scale_b` where the selected low-precision format requires them. BF16 does not reinterpret these controls; it rejects them instead. +Dense Torch callers that retain operations for CUDA Graph replay may provide a +caller-owned `descriptor_workspace`. Allocate its size with +`get_grouped_gemm_wgrad_workspace_size_sm100`, keep it alive for as long as the +captured call site may replay, and do not share it between call sites that may +overlap. This lets multiple same-signature calls share one compiled kernel +without sharing mutable runtime TMA descriptors. Callers that omit this +argument retain the compatibility behavior that isolates cached API instances +by explicit dense output address. + +```python +workspace = torch.empty( + cudnn.get_grouped_gemm_wgrad_workspace_size_sm100(num_experts), + dtype=torch.uint8, + device=a_tensor.device, +) +result = cudnn.grouped_gemm_wgrad_wrapper_sm100( + a_tensor=a_tensor, + b_tensor=b_tensor, + sfa_tensor=sfa_tensor, + sfb_tensor=sfb_tensor, + offsets_tensor=offsets_tensor, + wgrad_tensor=wgrad_tensor, + descriptor_workspace=workspace, + output_mode="dense", +) +``` + ## API usage ### BF16 diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 54277ba6c..de6eb29ec 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -399,6 +399,10 @@ def _dlopen_cudnn(): "GroupedGemmDgluSm100": (".gemm.cutedsl.grouped", "GroupedGemmDgluSm100"), "grouped_gemm_dglu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dglu_wrapper_sm100"), "GroupedGemmWgradSm100": (".gemm.cutedsl.grouped", "GroupedGemmWgradSm100"), + "get_grouped_gemm_wgrad_workspace_size_sm100": ( + ".gemm.cutedsl.grouped", + "get_grouped_gemm_wgrad_workspace_size_sm100", + ), "grouped_gemm_wgrad_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_wgrad_wrapper_sm100"), "discrete_grouped_gemm": (".gemm.cutedsl.discrete_grouped", None), "DiscreteGroupedGemmSwigluSm100": (".gemm.cutedsl.discrete_grouped", "DiscreteGroupedGemmSwigluSm100"), diff --git a/python/cudnn/gemm/cutedsl/grouped/__init__.py b/python/cudnn/gemm/cutedsl/grouped/__init__.py index b9b4ca937..e70a1d745 100644 --- a/python/cudnn/gemm/cutedsl/grouped/__init__.py +++ b/python/cudnn/gemm/cutedsl/grouped/__init__.py @@ -48,6 +48,7 @@ from .wgrad.api import ( GroupedGemmWgradSm100, + get_grouped_gemm_wgrad_workspace_size_sm100, grouped_gemm_wgrad_wrapper_sm100, ) @@ -76,6 +77,7 @@ "GroupedGemmDgluSm100", "grouped_gemm_dglu_wrapper_sm100", "GroupedGemmWgradSm100", + "get_grouped_gemm_wgrad_workspace_size_sm100", "grouped_gemm_wgrad_wrapper_sm100", "GroupedGemmSm100", "grouped_gemm_wrapper_sm100", diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/__init__.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/__init__.py index 85d717865..c81e20eec 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/__init__.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/__init__.py @@ -3,11 +3,13 @@ from .api import ( GroupedGemmWgradSm100, + get_grouped_gemm_wgrad_workspace_size_sm100, grouped_gemm_wgrad_wrapper_sm100, ) __all__ = [ "GroupedGemmWgradSm100", + "get_grouped_gemm_wgrad_workspace_size_sm100", "grouped_gemm_wgrad_wrapper_sm100", "grouped_gemm_wgrad_jax_sm100", ] diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py index 782fc3820..26c76b999 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py @@ -155,6 +155,8 @@ def __init__( self.accumulate_on_output = accumulate_on_output self._kernel = _get_rubin_kernel() if self._is_rubin_kernel else BlockScaledMoEGroupedGemmWgradKernel self._workspace = None + self._workspace_bytes = None + self._workspace_arg = None def _validate_offsets(self, offsets_tensor: torch.Tensor, tokens_sum: int, name: str) -> Tuple[int, ...]: self._value_error_if(offsets_tensor.ndim != 1, f"{name} must be rank-1, got shape {tuple(offsets_tensor.shape)}") @@ -335,7 +337,8 @@ def compile(self) -> None: hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - self._workspace = torch.empty(max(kernel.get_workspace_bytes(), 1), dtype=torch.uint8, device=self.a_desc.device) + self._workspace_bytes = max(kernel.get_workspace_bytes(), 1) + self._workspace = torch.empty(self._workspace_bytes, dtype=torch.uint8, device=self.a_desc.device) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) if self.weight_mode == MoEWeightMode.DENSE: @@ -419,8 +422,11 @@ def _compile_dense(self, kernel, max_active_clusters, fake_stream) -> None: None, options="--enable-tvm-ffi", ) - - cached_workspace = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) + self._workspace_arg = from_dlpack( + self._workspace, + assumed_align=128, + enable_tvm_ffi=True, + ) def tensor_api( a_tensor: torch.Tensor, @@ -429,6 +435,7 @@ def tensor_api( sfb_tensor: torch.Tensor, wgrad_tensor: torch.Tensor, offsets_tensor: torch.Tensor, + workspace, stream: cuda.CUstream, global_scale_a: Optional[torch.Tensor], global_scale_b: Optional[torch.Tensor], @@ -440,7 +447,7 @@ def tensor_api( sfb_tensor, wgrad_tensor, offsets_tensor, - cached_workspace, + workspace, stream, global_scale_a, global_scale_b, @@ -579,6 +586,7 @@ def execute( offsets_tensor: torch.Tensor, wgrad_tensor: Optional[torch.Tensor] = None, wgrad_ptrs: Optional[torch.Tensor] = None, + descriptor_workspace: Optional[torch.Tensor] = None, global_scale_a: Optional[torch.Tensor] = None, global_scale_b: Optional[torch.Tensor] = None, current_stream: Optional[cuda.CUstream] = None, @@ -590,6 +598,31 @@ def execute( if self.weight_mode == MoEWeightMode.DENSE: self._value_error_if(wgrad_tensor is None, "wgrad_tensor is required in dense mode") + if descriptor_workspace is None: + workspace_arg = self._workspace_arg + else: + self._value_error_if( + descriptor_workspace.dtype != torch.uint8, + f"descriptor_workspace must have dtype uint8, got {descriptor_workspace.dtype}", + ) + self._value_error_if( + descriptor_workspace.device != wgrad_tensor.device, + "descriptor_workspace and wgrad_tensor must be on the same device", + ) + self._value_error_if( + not descriptor_workspace.is_contiguous(), + "descriptor_workspace must be contiguous", + ) + self._value_error_if( + descriptor_workspace.numel() < self._workspace_bytes, + f"descriptor_workspace requires at least {self._workspace_bytes} bytes, " + f"got {descriptor_workspace.numel()}", + ) + workspace_arg = from_dlpack( + descriptor_workspace, + assumed_align=128, + enable_tvm_ffi=True, + ) self._compiled_kernel( a_tensor, b_tensor, @@ -597,6 +630,7 @@ def execute( sfb_tensor, wgrad_tensor, offsets_tensor, + workspace_arg, current_stream, global_scale_a, global_scale_b, diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py index 1aa90d76c..946660f57 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py @@ -30,7 +30,7 @@ backend_cache_key, select_grouped_gemm_backend, ) -from ..moe_utils import WGradInputOrder +from ..moe_utils import MoEWeightMode, WGradInputOrder, WgradSfTensormapConstructor def _block_scaled_dtype_pairs(): @@ -48,6 +48,33 @@ def _block_scaled_dtype_pairs(): _cache_of_GroupedGemmWgradSm100Objects = {} +def get_grouped_gemm_wgrad_workspace_size_sm100( + num_experts: int, + *, + output_mode: str = "dense", + input_order: WGradInputOrder | str = WGradInputOrder.Tensor2D, +) -> int: + """Return required runtime TMA-descriptor workspace bytes.""" + if num_experts <= 0: + raise ValueError(f"num_experts must be positive, got {num_experts}") + try: + weight_mode = MoEWeightMode(output_mode) + except ValueError as exc: + raise ValueError(f"unsupported output_mode {output_mode!r}") from exc + try: + normalized_input_order = WGradInputOrder(input_order) + except ValueError as exc: + raise ValueError(f"unsupported input_order {input_order!r}") from exc + return max( + WgradSfTensormapConstructor.get_workspace_size( + normalized_input_order, + weight_mode, + num_experts, + ), + 1, + ) + + from ._bf16_api import GroupedGemmWgradBf16API from ._blockscaled_api import ( GroupedGemmWgradBlockScaledAPI, @@ -173,6 +200,7 @@ def execute( offsets_tensor: torch.Tensor, wgrad_tensor: Optional[torch.Tensor] = None, wgrad_ptrs: Optional[torch.Tensor] = None, + descriptor_workspace: None = None, *, global_scale_a: None = None, global_scale_b: None = None, @@ -189,6 +217,7 @@ def execute( offsets_tensor: torch.Tensor, wgrad_tensor: Optional[torch.Tensor] = None, wgrad_ptrs: Optional[torch.Tensor] = None, + descriptor_workspace: Optional[torch.Tensor] = None, *, global_scale_a: Optional[torch.Tensor] = None, global_scale_b: Optional[torch.Tensor] = None, @@ -203,13 +232,19 @@ def execute( offsets_tensor: torch.Tensor, wgrad_tensor: Optional[torch.Tensor] = None, wgrad_ptrs: Optional[torch.Tensor] = None, + descriptor_workspace: Optional[torch.Tensor] = None, global_scale_a: Optional[torch.Tensor] = None, global_scale_b: Optional[torch.Tensor] = None, current_stream: Optional[cuda.CUstream] = None, ) -> None: if self._implementation is None: raise RuntimeError("Kernel not compiled; call compile() first") - self._implementation.execute( + if descriptor_workspace is not None and not isinstance( + self._implementation, + GroupedGemmWgradBlockScaledAPI, + ): + raise ValueError("descriptor_workspace requires the block-scaled WGrad backend") + execute_kwargs = dict( a_tensor=a_tensor, b_tensor=b_tensor, sfa_tensor=sfa_tensor, @@ -221,6 +256,9 @@ def execute( global_scale_b=global_scale_b, current_stream=current_stream, ) + if descriptor_workspace is not None: + execute_kwargs["descriptor_workspace"] = descriptor_workspace + self._implementation.execute(**execute_kwargs) def _wgrad_tensor_signature(tensor: Optional[torch.Tensor], *, dynamic_dims: tuple[int, ...] = (), exact_stride: bool): @@ -243,6 +281,7 @@ def grouped_gemm_wgrad_wrapper_sm100( output_mode: str = "dense", wgrad_tensor: Optional[torch.Tensor] = None, wgrad_ptrs: Optional[torch.Tensor] = None, + descriptor_workspace: Optional[torch.Tensor] = None, global_scale_a: Optional[torch.Tensor] = None, global_scale_b: Optional[torch.Tensor] = None, acc_dtype: Optional[torch.dtype] = None, @@ -294,22 +333,24 @@ def grouped_gemm_wgrad_wrapper_sm100( ) if framework == "jax" and backend is GroupedGemmBackend.BLOCK_SCALED: raise ValueError(_BLOCK_SCALED_JAX_ERROR) + if descriptor_workspace is not None and ( + backend is not GroupedGemmBackend.BLOCK_SCALED + or framework != "torch" + or output_mode != "dense" + ): + raise ValueError( + "descriptor_workspace is supported only for dense torch block-scaled WGrad" + ) explicit_dense_output_identity = None if ( backend is GroupedGemmBackend.BLOCK_SCALED and framework == "torch" and output_mode == "dense" and wgrad_tensor is not None + and descriptor_workspace is None ): - # Temporary workaround: - # 1. Problem behavior: Multiple same-signature launches with explicit - # outputs can corrupt later results when captured in one CUDA graph - # and backed by one cached API instance. - # 2. Possible root cause: The launches alias mutable TMA descriptor - # workspace owned by the cached API instance. - # 3. Possible long-term fix: Cache compiled kernels by shape, but own - # descriptor workspace independently per graph call site. Then - # remove output identity from the cache key. + # Compatibility path: callers that do not own descriptor workspace keep + # the validated one-API-instance-per-output isolation. explicit_dense_output_identity = int(wgrad_tensor.data_ptr()) if wgrad_tensor is None and wgrad_ptrs is None: wgrad_shape = (expert_cnt, hidden, intermediate) @@ -404,6 +445,7 @@ def _sample_wgrad_expert(): offsets_tensor=offsets_tensor, wgrad_tensor=wgrad_tensor, wgrad_ptrs=wgrad_ptrs, + descriptor_workspace=descriptor_workspace, global_scale_a=global_scale_a, global_scale_b=global_scale_b, current_stream=current_stream, diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py index 85fefa5ea..e239a8b80 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py @@ -761,6 +761,124 @@ def counted_compile(self): assert cache_entries == 1 +@pytest.mark.L0 +@pytest.mark.parametrize( + ("caller_owned_workspace", "expected_cache_entries"), + [(False, 2), (True, 1)], + ids=["compatibility-isolation", "caller-workspace"], +) +def test_grouped_gemm_wgrad_wrapper_explicit_dense_output_cache( + monkeypatch, + caller_owned_workspace, + expected_cache_entries, +): + from cudnn.gemm.cutedsl.grouped.wgrad import api as grouped_gemm_wgrad_api + + grouped_gemm_wgrad_api._cache_of_GroupedGemmWgradSm100Objects.clear() + compile_count = {"value": 0} + + def counted_compile(self): + compile_count["value"] += 1 + + monkeypatch.setattr(grouped_gemm_wgrad_api.GroupedGemmWgradSm100, "check_support", lambda self: True) + monkeypatch.setattr(grouped_gemm_wgrad_api.GroupedGemmWgradSm100, "compile", counted_compile) + monkeypatch.setattr(grouped_gemm_wgrad_api.GroupedGemmWgradSm100, "execute", lambda self, **kwargs: None) + monkeypatch.setattr( + grouped_gemm_wgrad_api, + "select_grouped_gemm_backend", + lambda **_: grouped_gemm_wgrad_api.GroupedGemmBackend.BLOCK_SCALED, + ) + + inputs = _make_wgrad_wrapper_cache_inputs([8, 12]) + outputs = [torch.empty((2, 32, 64), dtype=torch.bfloat16) for _ in range(2)] + workspaces = [torch.empty(512, dtype=torch.uint8) for _ in range(2)] + try: + for output, workspace in zip(outputs, workspaces): + workspace_kwargs = ( + {"descriptor_workspace": workspace} if caller_owned_workspace else {} + ) + cudnn.grouped_gemm_wgrad_wrapper_sm100( + **inputs, + **workspace_kwargs, + output_mode="dense", + wgrad_tensor=output, + acc_dtype=torch.float32, + wgrad_dtype=torch.bfloat16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=16, + ) + finally: + cache_entries = len(grouped_gemm_wgrad_api._cache_of_GroupedGemmWgradSm100Objects) + grouped_gemm_wgrad_api._cache_of_GroupedGemmWgradSm100Objects.clear() + + assert outputs[0].data_ptr() != outputs[1].data_ptr() + assert compile_count["value"] == expected_cache_entries + assert cache_entries == expected_cache_entries + + +@pytest.mark.L0 +def test_grouped_gemm_wgrad_workspace_size(): + assert cudnn.get_grouped_gemm_wgrad_workspace_size_sm100(2) == 512 + assert ( + cudnn.get_grouped_gemm_wgrad_workspace_size_sm100( + 2, + input_order="tensor_ragged", + ) + == 1024 + ) + + +@pytest.mark.L0 +def test_blockscaled_wgrad_execute_uses_caller_workspace(monkeypatch): + from cudnn.gemm.cutedsl.grouped.wgrad import _blockscaled_api + + api = object.__new__(_blockscaled_api.GroupedGemmWgradBlockScaledAPI) + api._workspace_bytes = 16 + api._workspace = torch.empty(16, dtype=torch.uint8) + api.a_desc = type("TensorDesc", (), {"device": torch.device("cpu")})() + api.weight_mode = _blockscaled_api.MoEWeightMode.DENSE + api._get_default_stream = lambda stream: stream + api._runtime_error_if = lambda condition, message: None + api._value_error_if = lambda condition, message: None + monkeypatch.setattr( + _blockscaled_api, + "from_dlpack", + lambda tensor, **kwargs: tensor, + ) + + launch_workspaces = [] + + def compiled_kernel(*args): + launch_workspaces.append(args[6]) + + api._compiled_kernel = compiled_kernel + operand = torch.empty((1, 1)) + offsets = torch.tensor([1], dtype=torch.int32) + outputs = [torch.empty((2, 32, 64), dtype=torch.bfloat16) for _ in range(2)] + workspaces = [torch.empty(16, dtype=torch.uint8) for _ in range(2)] + for output, workspace in ( + (outputs[0], workspaces[0]), + (outputs[0], workspaces[0]), + (outputs[1], workspaces[1]), + ): + api.execute( + operand, + operand, + operand, + operand, + offsets, + wgrad_tensor=output, + descriptor_workspace=workspace, + current_stream=object(), + ) + + assert launch_workspaces[0] is workspaces[0] + assert launch_workspaces[1] is workspaces[0] + assert launch_workspaces[2] is workspaces[1] + assert launch_workspaces[0].data_ptr() != launch_workspaces[2].data_ptr() + + @pytest.mark.L0 def test_grouped_gemm_wgrad_wrapper_input_order_cache_key(monkeypatch): from cudnn.gemm.cutedsl.grouped.wgrad import api as grouped_gemm_wgrad_api From e46be8ab983db1b637d30b00c1598b6ea1fff290 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 30 Aug 2026 16:39:28 -0700 Subject: [PATCH 24/31] allow for weight interleaving Signed-off-by: Varun Thumbe --- docs/fe-oss-apis/moe_ep.md | 12 ++- docs/operations/moe_ep.md | 6 ++ python/cudnn/moe_ep/_contracts.py | 1 + .../cudnn/moe_ep/_megamoe_backend/README.md | 21 +++-- .../moe_ep/_megamoe_backend/mxfp8/_adapter.py | 18 +++- .../moe_ep/_megamoe_backend/mxfp8/_config.py | 2 + .../mxfp8/_training_resources.py | 31 +------ .../mxfp8/_training_weights.py | 16 +++- .../_megamoe_backend/mxfp8/_training_wgrad.py | 42 +-------- .../mxfp8/_training_wgrad_kernel.py | 26 +----- python/cudnn/moe_ep/_types.py | 4 +- python/cudnn/moe_ep/_validation.py | 24 +++++ python/cudnn/moe_ep/api.py | 5 + test/python/moe_ep/moe_ep_reference.py | 48 ++++++++++ test/python/moe_ep/moe_ep_test_support.py | 14 +++ test/python/moe_ep/test_moe_ep_backward.py | 64 ++++++++++++- test/python/moe_ep/test_moe_ep_forward.py | 91 +++++++++++++++++++ 17 files changed, 315 insertions(+), 110 deletions(-) diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md index 551efb1cd..a3015286c 100644 --- a/docs/fe-oss-apis/moe_ep.md +++ b/docs/fe-oss-apis/moe_ep.md @@ -42,6 +42,7 @@ The `MoeEp` constructor accepts: | `output_format` | `"bf16"` only for current execution | | `combine_format` | `"bf16"` or `"mxfp8"` | | `apply_topk_in_fc1` | Must be `True` | +| `weight_interleave_size` | `None` for conventional gate-then-up weights or `32` for pre-interleaved MXFP8 weights | | `gate_up_clamp` | Optional finite clamp magnitude | | `token_padding_size` | Positive; training fixed resources use 128 internally | | `sf_padding_size` | Positive multiple of 128; training fixed resources use 128 internally | @@ -293,12 +294,15 @@ Callers must establish stream/event ordering for in-place weight updates. ### Explicit weight refresh contract -Rubin training recognizes compact K-major forward views plus -contiguous backward transposes. For that form, weight data is bound directly -and FC1 gate/up values must already use 32-element interleaving. +With `weight_interleave_size=32`, Rubin training recognizes compact K-major +forward views plus contiguous backward transposes as pre-interleaved. For that +form, weight data is bound directly and FC1 gate/up values must use alternating +32-element strips. Layout alone never selects this semantic convention. `resources.refresh_weights()` then swizzles only scales into fixed-address, kernel-native buffers and never copies the weight payload. Existing contiguous -public packs retain the compatible data-and-scale staging path. +public packs with the default `None` retain conventional gate-then-up semantics +and the compatible data-and-scale staging path. Plain BF16/FP16/FP32 inference +weights cannot use `weight_interleave_size=32`. The caller must obey all of the following: diff --git a/docs/operations/moe_ep.md b/docs/operations/moe_ep.md index 6f0ad1cf7..e5bbfad56 100644 --- a/docs/operations/moe_ep.md +++ b/docs/operations/moe_ep.md @@ -65,6 +65,7 @@ op = MoeEp( output_format="bf16", combine_format="bf16", # "bf16" or "mxfp8" apply_topk_in_fc1=True, + weight_interleave_size=None, # Or 32 for pre-interleaved MXFP8 W1 gate_up_clamp=None, ) ``` @@ -72,6 +73,11 @@ op = MoeEp( `topk_idx` contains global expert IDs. Each rank passes its local tokens and its contiguous shard of both expert-weight tensors: +`weight_interleave_size=32` declares that MXFP8 FC1 values already use +alternating 32-element gate/up strips. The default `None` uses conventional +gate-then-up order. Plain BF16/FP16/FP32 weights remain conventional and reject +the interleaved contract because they must be quantized and staged internally. + ```python output = op( activation, # (T, H) diff --git a/python/cudnn/moe_ep/_contracts.py b/python/cudnn/moe_ep/_contracts.py index 9acb907fb..1d3e741ef 100644 --- a/python/cudnn/moe_ep/_contracts.py +++ b/python/cudnn/moe_ep/_contracts.py @@ -39,6 +39,7 @@ class ForwardConfig: backward_wgrad_mode: Literal["none", "operands"] = "none" max_recv_size_per_rank: Optional[int] = None drop_on_overflow: bool = False + weight_interleave_size: Optional[int] = None @dataclass(frozen=True) diff --git a/python/cudnn/moe_ep/_megamoe_backend/README.md b/python/cudnn/moe_ep/_megamoe_backend/README.md index 8da780fc6..3a9837ad9 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/README.md +++ b/python/cudnn/moe_ep/_megamoe_backend/README.md @@ -73,12 +73,12 @@ forward, backward, and WGrad-export kernels are compiled before capture. `MoeEpTrainingWeights` contains four address-stable MXFP8 block-scaled tensors: forward W1/W2 and independently quantized backward W2-transpose/W1-transpose. -When forward weights use compact K-major storage and backward transpose weights -use standard contiguous storage, the kernels alias weight data directly. In -this layout, W1 gate/up values are interleaved in 32-element strips and only -scales require kernel-native staging. When forward weights use standard -contiguous storage, weight data and scales are copied and reordered into -persistent kernel buffers. After every in-place data+scale update, the caller +With `weight_interleave_size=32`, compact K-major forward weights and contiguous +backward transposes are interpreted as already using 32-element W1 gate/up +strips, and the kernels alias weight data directly; only scales require +kernel-native staging. With the default `None`, weights use conventional +gate-then-up order and are copied and interleaved into persistent kernel +buffers. After every in-place data+scale update, the caller must enqueue `resources.refresh_weights()` before the first consumer, with explicit stream/event ordering. A matching forward/backward pair must use one version; refresh cannot overlap any consumer on another slot/lane. Replacing @@ -93,8 +93,13 @@ contract and relaxed atomic accumulation order. `MoeEpTrainingWgradOperands` is a fixed-capacity producer ABI. Device `expert_offsets` and `valid_route_counts` describe the current valid K extent; -padding is zeroed. No specific downstream grouped-WGrad consumer is guaranteed -by this milestone. +padding is zeroed. Its data operands alias persistent forward or backward +outputs, using transpose views where required; exporting them performs no +full-tensor data copies. Scale operands are expanded into persistent +grouped-WGrad layouts. FC1 dY data and scales preserve the same 32-element +gate/up-interleaved row order as W1, so a downstream grouped-WGrad consumer +writes gradients in parameter storage order. No specific downstream +grouped-WGrad consumer is guaranteed by this milestone. ## Overflow policy diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py index 12057847a..05b0572bf 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py @@ -238,18 +238,25 @@ def _stack_blocked_scales(raw_scales: torch.Tensor) -> torch.Tensor: def _prepare_fc1( tensor: BlockScaledTensor, intermediate: int, + *, + already_interleaved: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - """Preserve logical bytes while building interleaved K-major FC1 tensors.""" + """Build kernel-native K-major FC1 tensors.""" payload_nkh = _as_bytes(tensor.data).permute(0, 2, 1).contiguous() - payload_interleaved = _interleave_gate_up_rows( - payload_nkh, - intermediate, + payload_interleaved = ( + payload_nkh + if already_interleaved + else _interleave_gate_up_rows(payload_nkh, intermediate) ) payload = payload_interleaved.view(_MXFP8_DATA_DTYPE).permute(0, 2, 1) scales_nk = _as_bytes(tensor.scale).permute(0, 2, 1).contiguous() - scales_interleaved = _interleave_gate_up_rows(scales_nk, intermediate) + scales_interleaved = ( + scales_nk + if already_interleaved + else _interleave_gate_up_rows(scales_nk, intermediate) + ) scale = _stack_blocked_scales(scales_interleaved) return payload, scale @@ -362,6 +369,7 @@ def _prepare_weights( fc1_weight, fc1_weight_sf = _prepare_fc1( fc1_source, config.intermediate, + already_interleaved=config.weight_interleave_size == 32, ) fc2_weight, fc2_weight_sf = _prepare_fc2(fc2_source) weights = Mxfp8Weights( diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py index 3f09f1d6f..0fcd8a669 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py @@ -25,6 +25,7 @@ class Mxfp8KernelConfig: top_k: int max_tokens_per_rank: int apply_topk_in_fc1: bool + weight_interleave_size: int | None gate_up_clamp: float | None generate_c: bool max_recv_size_per_rank: int | None = None @@ -78,6 +79,7 @@ def from_forward_config(cls, config: ForwardConfig) -> "Mxfp8KernelConfig": top_k=config.top_k, max_tokens_per_rank=config.max_tokens_per_rank, apply_topk_in_fc1=config.apply_topk_in_fc1, + weight_interleave_size=config.weight_interleave_size, gate_up_clamp=config.gate_up_clamp, generate_c=config.generate_c, max_recv_size_per_rank=max_recv_size_per_rank, diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py index 4a5e1427a..d48480414 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py @@ -243,10 +243,6 @@ def build_training_workspace_requirements( } scale_columns = _align_scale_columns(forward.pool_token_capacity) wgrad_shapes = { - "wgrad_fc1_b": ( - forward.pool_token_capacity, - 2 * config.intermediate_size, - ), "wgrad_fc1_sfa": ( _round_up(config.hidden_size, 128), scale_columns, @@ -255,10 +251,6 @@ def build_training_workspace_requirements( _round_up(2 * config.intermediate_size, 128), scale_columns, ), - "wgrad_fc2_a": ( - config.intermediate_size, - forward.pool_token_capacity, - ), "wgrad_fc2_sfa": ( _round_up(config.intermediate_size, 128), scale_columns, @@ -612,10 +604,8 @@ class Mxfp8TrainingSlotViews: fc1_col_output_sf: torch.Tensor grad_y2: torch.Tensor grad_y2_sf: torch.Tensor - wgrad_fc1_b: torch.Tensor wgrad_fc1_sfa: torch.Tensor wgrad_fc1_sfb: torch.Tensor - wgrad_fc2_a: torch.Tensor wgrad_fc2_sfa: torch.Tensor wgrad_fc2_sfb: torch.Tensor @@ -651,7 +641,10 @@ def __init__( self.device = torch.device(device) self.forward_prepared = forward self.backward_prepared = backward - self.weight_bindings = Mxfp8TrainingWeightBindings(weights) + self.weight_bindings = Mxfp8TrainingWeightBindings( + weights, + weight_interleave_size=config.weight_interleave_size, + ) self.stager = Mxfp8TrainingStager(config.hidden_size, config.top_k) self.wgrad_exporter = Mxfp8TrainingWgradExporter( experts=config.experts_per_rank, @@ -1016,14 +1009,6 @@ def local_bytes(name: str) -> torch.Tensor: torch.uint8, bwd_shapes["grad_y2_sf"], ), - wgrad_fc1_b=_typed_k_major_view( - local_bytes("wgrad_fc1_b"), - _DATA_DTYPE, - ( - self.forward_prepared.pool_token_capacity, - 2 * config.intermediate_size, - ), - ), wgrad_fc1_sfa=_typed_view( local_bytes("wgrad_fc1_sfa"), _SCALE_DTYPE, @@ -1037,14 +1022,6 @@ def local_bytes(name: str) -> torch.Tensor: scale_columns, ), ), - wgrad_fc2_a=_typed_view( - local_bytes("wgrad_fc2_a"), - _DATA_DTYPE, - ( - config.intermediate_size, - self.forward_prepared.pool_token_capacity, - ), - ), wgrad_fc2_sfa=_typed_view( local_bytes("wgrad_fc2_sfa"), _SCALE_DTYPE, diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py index 8d900b3e5..cdb61faa4 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py @@ -182,18 +182,30 @@ class Mxfp8BackwardWeights: class Mxfp8TrainingWeightBindings: """Direct data bindings with persistent kernel-native scale staging.""" - def __init__(self, weights: MoeEpTrainingWeights) -> None: + def __init__( + self, + weights: MoeEpTrainingWeights, + *, + weight_interleave_size: int | None = None, + ) -> None: self.weights = weights + self.weight_interleave_size = weight_interleave_size fwd_fc1 = weights.forward_fc1 fwd_fc2 = weights.forward_fc2 bwd_w2t = weights.backward_w2_transpose bwd_w1t = weights.backward_w1_transpose self._uses_direct_weight_bindings = ( - fwd_fc1.data.stride(1) == 1 + weight_interleave_size == 32 + and fwd_fc1.data.stride(1) == 1 and fwd_fc2.data.stride(1) == 1 and bwd_w2t.data.is_contiguous() and bwd_w1t.data.is_contiguous() ) + if weight_interleave_size == 32 and not self._uses_direct_weight_bindings: + raise ValueError( + "weight_interleave_size=32 requires compact K-major forward " + "weights and contiguous backward transpose weights" + ) self.forward = Mxfp8Weights( fc1_weight=( diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py index 9f6687bb9..6e1017037 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py @@ -32,34 +32,9 @@ def __init__( self.hidden = int(hidden) self.intermediate = int(intermediate) self.sf_padding = int(sf_padding) - self._compiled: dict[tuple[int, int, int | None], object] = {} + self._compiled: dict[tuple[int, int], object] = {} self._lock = threading.RLock() - @staticmethod - def _copy_gate_up_data( - target: torch.Tensor, - source: torch.Tensor, - intermediate: int, - ) -> None: - pool_rows = source.shape[0] - pairs = intermediate // 32 - source_view = source.view(pool_rows, pairs, 2, 32).permute( - 0, - 2, - 1, - 3, - ) - target_view = target.as_strided( - (pool_rows, 2, pairs, 32), - ( - target.stride(0), - intermediate * target.stride(1), - 32 * target.stride(1), - target.stride(1), - ), - ) - target_view.copy_(source_view) - def _expand_scales( self, source: torch.Tensor, @@ -68,7 +43,6 @@ def _expand_scales( output: torch.Tensor, *, non_k_size: int, - deinterleave_gate_up: int | None = None, ) -> None: if source.dtype not in (torch.uint8, torch.float8_e8m0fnu): raise TypeError("WGrad source scales must use Uint8 or E8M0") @@ -76,7 +50,7 @@ def _expand_scales( raise TypeError("WGrad output scales must use E8M0") source_bytes = source.view(torch.uint8).reshape(-1) output_bytes = output.view(torch.uint8).reshape(-1) - key = (int(non_k_size), self.sf_padding, deinterleave_gate_up) + key = (int(non_k_size), self.sf_padding) import cuda.bindings.driver as cuda stream = torch.cuda.current_stream(output.device) @@ -102,7 +76,6 @@ def _expand_scales( non_k_size=non_k_size, expert_count=self.experts, source_sf_padding=self.sf_padding, - deinterleave_gate_up=deinterleave_gate_up, ) compiled = cute.compile(kernel, *args) self._compiled[key] = compiled @@ -120,12 +93,6 @@ def export( if slot.col_quant_data.shape[0] != pool_rows: raise RuntimeError("forward/backward WGrad pool capacities differ") - self._copy_gate_up_data( - slot.wgrad_fc1_b, - slot.fc1_col_output, - self.intermediate, - ) - slot.wgrad_fc2_a.copy_(slot.fc1_recompute.transpose(0, 1)) self._expand_scales( slot.col_quant_sf, slot.valid_route_counts, @@ -139,7 +106,6 @@ def export( slot.expert_offsets, slot.wgrad_fc1_sfb, non_k_size=2 * self.intermediate, - deinterleave_gate_up=self.intermediate, ) self._expand_scales( slot.fc1_recompute_sf, @@ -159,9 +125,9 @@ def export( return MoeEpTrainingWgradOperands( fc1_a=slot.col_quant_data.transpose(0, 1), fc1_sfa=slot.wgrad_fc1_sfa, - fc1_b=slot.wgrad_fc1_b, + fc1_b=slot.fc1_col_output, fc1_sfb=slot.wgrad_fc1_sfb, - fc2_a=slot.wgrad_fc2_a, + fc2_a=slot.fc1_recompute.transpose(0, 1), fc2_sfa=slot.wgrad_fc2_sfa, fc2_b=slot.grad_y2, fc2_sfb=slot.wgrad_fc2_sfb, diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py index 3e2909b31..244e809ea 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py @@ -25,20 +25,16 @@ def __init__( non_k_size: int, expert_count: int, source_sf_padding: int, - deinterleave_gate_up: int | None = None, ) -> None: self.non_k_size = int(non_k_size) self.expert_count = int(expert_count) self.source_sf_padding = int(source_sf_padding) - self.deinterleave_gate_up = None if deinterleave_gate_up is None else int(deinterleave_gate_up) if self.non_k_size <= 0 or self.non_k_size % 128: raise ValueError("WGrad scale non-K size must be divisible by 128") if self.expert_count <= 0: raise ValueError("WGrad scale expansion requires experts") if self.source_sf_padding <= 0 or self.source_sf_padding % 128: raise ValueError("WGrad source SF padding must be a positive multiple of 128") - if self.deinterleave_gate_up is not None and self.non_k_size != 2 * self.deinterleave_gate_up: - raise ValueError("gate/up scale deinterleave size mismatch") @cute.jit def __call__( @@ -98,26 +94,8 @@ def _kernel( hidden_atom = relative_atom // target_token_atoms token_atom = relative_atom % target_token_atoms if token_atom < source_token_atoms: - source_hidden_atom = hidden_atom - source_byte = byte_in_atom - if cutlass.const_expr(self.deinterleave_gate_up is not None): - lane = byte_in_atom // Int32(16) - byte_tail = byte_in_atom % Int32(16) - group = byte_tail // Int32(4) - column_lane = byte_tail % Int32(4) - feature = hidden_atom * Int32(128) + group * Int32(32) + lane - intermediate = Int32(self.deinterleave_gate_up) - source_feature = Int32(0) - if feature < intermediate: - source_feature = (feature // Int32(32)) * Int32(64) + feature % Int32(32) - else: - up_feature = feature - intermediate - source_feature = (up_feature // Int32(32)) * Int32(64) + Int32(32) + up_feature % Int32(32) - source_hidden_atom = source_feature // Int32(128) - source_feature_in_atom = source_feature % Int32(128) - source_byte = (source_feature_in_atom % Int32(32)) * Int32(16) + (source_feature_in_atom // Int32(32)) * Int32(4) + column_lane - source_atom = source_atom_base + source_hidden_atom * source_token_atoms + token_atom - value = source[source_atom * Int32(atom_bytes) + source_byte] + source_atom = source_atom_base + hidden_atom * source_token_atoms + token_atom + value = source[source_atom * Int32(atom_bytes) + byte_in_atom] target_atom_base += target_atom_count source_atom_base += Int32(non_k_atoms) * source_token_atoms previous_end = end diff --git a/python/cudnn/moe_ep/_types.py b/python/cudnn/moe_ep/_types.py index 625948038..11927eab9 100644 --- a/python/cudnn/moe_ep/_types.py +++ b/python/cudnn/moe_ep/_types.py @@ -181,7 +181,9 @@ class MoeEpTrainingWeights: quantized transposes: ``backward_w2_transpose=(E,H,I)`` for ``dH=dY@W2.T`` and ``backward_w1_transpose=(E,2I,H)`` for ``dX=dC@W1.T``. Every tensor is block-scaled along logical axis 1, the - reduction axis of its corresponding GEMM. + reduction axis of its corresponding GEMM. When the owning ``MoeEp`` has + ``weight_interleave_size=32``, forward W1's output axis and backward W1T's + reduction axis must contain alternating 32-element gate/up strips. """ forward_fc1: BlockScaledTensor diff --git a/python/cudnn/moe_ep/_validation.py b/python/cudnn/moe_ep/_validation.py index b3338f34c..d5338c335 100644 --- a/python/cudnn/moe_ep/_validation.py +++ b/python/cudnn/moe_ep/_validation.py @@ -185,6 +185,14 @@ def validate_forward( config.hidden_size, ), ) + if config.weight_interleave_size is not None and ( + not isinstance(fc1_weight, BlockScaledTensor) + or fc1_weight.format is not MoeFormat.MXFP8 + ): + raise ValueError( + "weight_interleave_size=32 requires an MXFP8 BlockScaledTensor " + "for fc1_weight" + ) _validate_routes( config, token_count, @@ -235,6 +243,12 @@ def has_supported_layout(tensor: torch.Tensor) -> bool: experts, reduction, output = tensor.shape return tensor.stride() == (reduction * output, 1, reduction) + def is_compact_k_major(tensor: torch.Tensor) -> bool: + if tensor.ndim != 3: + return False + experts, reduction, output = tensor.shape + return tensor.stride() == (reduction * output, 1, reduction) + if not isinstance(weights, MoeEpTrainingWeights): raise TypeError("weights must be a MoeEpTrainingWeights, " f"got {type(weights).__name__}") expected = ( @@ -286,6 +300,16 @@ def has_supported_layout(tensor: torch.Tensor) -> bool: f"{name} data and scale must be contiguous or compact K-major " "for fixed training weight binding" ) + if config.weight_interleave_size == 32 and not ( + is_compact_k_major(weights.forward_fc1.data) + and is_compact_k_major(weights.forward_fc2.data) + and weights.backward_w2_transpose.data.is_contiguous() + and weights.backward_w1_transpose.data.is_contiguous() + ): + raise ValueError( + "weight_interleave_size=32 requires compact K-major forward " + "weights and contiguous backward transpose weights" + ) device = weights.forward_fc1.device for name, tensor, _shape in expected[1:]: if tensor.device != device: diff --git a/python/cudnn/moe_ep/api.py b/python/cudnn/moe_ep/api.py index 317870433..b71dd439b 100644 --- a/python/cudnn/moe_ep/api.py +++ b/python/cudnn/moe_ep/api.py @@ -113,6 +113,7 @@ def __init__( output_format: Union[MoeFormat, str] = MoeFormat.BF16, combine_format: Union[MoeFormat, str] = MoeFormat.BF16, apply_topk_in_fc1: bool = True, + weight_interleave_size: Optional[int] = None, gate_up_clamp: Optional[float] = None, token_padding_size: int = 128, sf_padding_size: int = 128, @@ -139,6 +140,8 @@ def __init__( raise ValueError("drop_on_overflow must be a bool") if not isinstance(apply_topk_in_fc1, bool): raise ValueError("apply_topk_in_fc1 must be a bool") + if weight_interleave_size not in (None, 32): + raise ValueError("weight_interleave_size must be None or 32") for name, value in ( ("token_padding_size", token_padding_size), ("sf_padding_size", sf_padding_size), @@ -177,6 +180,7 @@ def __init__( self.output_format = _parse_format(output_format) self.combine_format = _parse_format(combine_format) self.apply_topk_in_fc1 = apply_topk_in_fc1 + self.weight_interleave_size = weight_interleave_size self.gate_up_clamp = None if gate_up_clamp is None else abs(gate_up_clamp) self.token_padding_size = token_padding_size self.sf_padding_size = sf_padding_size @@ -210,6 +214,7 @@ def __init__( output_format=self.output_format.value, combine_format=self.combine_format.value, apply_topk_in_fc1=self.apply_topk_in_fc1, + weight_interleave_size=self.weight_interleave_size, gate_up_clamp=self.gate_up_clamp, generate_c=False, token_padding_size=self.token_padding_size, diff --git a/test/python/moe_ep/moe_ep_reference.py b/test/python/moe_ep/moe_ep_reference.py index 0358edd66..ac4ed6721 100644 --- a/test/python/moe_ep/moe_ep_reference.py +++ b/test/python/moe_ep/moe_ep_reference.py @@ -433,6 +433,36 @@ def _padded_expert_rows( return padded +def _deinterleave_glu(tensor: torch.Tensor, interleave_size: int) -> torch.Tensor: + """Convert fixed-width gate/up strips to contiguous gate and up halves.""" + shape = tensor.shape + return ( + tensor.reshape( + *shape[:-1], + shape[-1] // (2 * interleave_size), + 2, + interleave_size, + ) + .transpose(-3, -2) + .reshape(shape) + ) + + +def _interleave_glu(tensor: torch.Tensor, interleave_size: int) -> torch.Tensor: + """Convert contiguous gate and up halves to fixed-width strips.""" + shape = tensor.shape + return ( + tensor.reshape( + *shape[:-1], + 2, + shape[-1] // (2 * interleave_size), + interleave_size, + ) + .transpose(-3, -2) + .reshape(shape) + ) + + class MoeEpReference: """Reference implementation of routed SwiGLU experts plus EP dispatch. @@ -462,6 +492,7 @@ def __init__( intermediate_format: Optional[Union[MoeFormat, str]] = None, backward_operand_format: Optional[Union[MoeFormat, str]] = None, apply_topk_in_fc1: bool = True, + weight_interleave_size: Optional[int] = None, gate_up_clamp: Optional[float] = None, generate_c: bool = False, backward_wgrad_mode: str = "none", @@ -487,6 +518,8 @@ def __init__( raise ValueError("token_padding_size must be a positive integer") if backward_wgrad_mode == "operands" and token_padding_size != 256: raise ValueError("backward_wgrad_mode='operands' requires " "token_padding_size=256") + if weight_interleave_size not in (None, 32): + raise ValueError("weight_interleave_size must be None or 32") if ep_group is None: ep_size, ep_rank = 1, 0 @@ -512,6 +545,7 @@ def __init__( self.intermediate_format = None if intermediate_format is None else _parse_format(intermediate_format) self.backward_operand_format = None if backward_operand_format is None else _parse_format(backward_operand_format) self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) + self.weight_interleave_size = weight_interleave_size self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) self.generate_c = bool(generate_c) self.backward_wgrad_mode = backward_wgrad_mode @@ -629,6 +663,8 @@ def _run_local_experts( 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)) + if self.weight_interleave_size is not None: + gate_up = _deinterleave_glu(gate_up, self.weight_interleave_size) 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) @@ -697,6 +733,14 @@ def __call__( re-dispatch. """ + if self.weight_interleave_size == 32 and ( + not isinstance(fc1_weight, BlockScaledTensor) + or fc1_weight.format is not MoeFormat.MXFP8 + ): + raise ValueError( + "weight_interleave_size=32 requires an MXFP8 BlockScaledTensor " + "for fc1_weight" + ) 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] @@ -1035,6 +1079,8 @@ def backward( d_y = dy_rows.index_select(0, positions) semantic_d_y = semantic_dy_rows.index_select(0, positions) + if self.weight_interleave_size is not None: + c = _deinterleave_glu(c, self.weight_interleave_size) 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) @@ -1068,6 +1114,8 @@ def backward( else: d_gate, d_up = d_g, d_u d_c = torch.cat((d_gate, d_up), dim=-1) + if self.weight_interleave_size is not None: + d_c = _interleave_glu(d_c, self.weight_interleave_size) dc_rows.index_copy_(0, positions, d_c) if self.intermediate_format is not None: d_c = _format_round_trip(d_c, self.intermediate_format) diff --git a/test/python/moe_ep/moe_ep_test_support.py b/test/python/moe_ep/moe_ep_test_support.py index efb90718e..72ee0c7f1 100644 --- a/test/python/moe_ep/moe_ep_test_support.py +++ b/test/python/moe_ep/moe_ep_test_support.py @@ -1268,6 +1268,20 @@ def _assert_wgrads_match_reference( actual_dense = _dense_wgrads_from_operands(actual) if expected_dense is None: expected_dense = expected.dense_wgrads() + expected_fc1, expected_fc2 = expected_dense + interleave_size = 32 + fc1_out_features = expected_fc1.shape[-1] + expected_fc1 = ( + expected_fc1.view( + *expected_fc1.shape[:-1], + 2, + fc1_out_features // (2 * interleave_size), + interleave_size, + ) + .transpose(-3, -2) + .reshape(expected_fc1.shape) + ) + expected_dense = (expected_fc1, expected_fc2) for name, actual_dw, expected_dw in zip( ("grad_fc1_weight", "grad_fc2_weight"), actual_dense, diff --git a/test/python/moe_ep/test_moe_ep_backward.py b/test/python/moe_ep/test_moe_ep_backward.py index 72b54227a..d457cbe11 100644 --- a/test/python/moe_ep/test_moe_ep_backward.py +++ b/test/python/moe_ep/test_moe_ep_backward.py @@ -44,6 +44,9 @@ from cudnn.moe_ep._megamoe_backend.mxfp8._training_stage import ( Mxfp8TrainingStager, ) +from cudnn.moe_ep._megamoe_backend.mxfp8._training_wgrad import ( + Mxfp8TrainingWgradExporter, +) from cudnn.moe_ep._megamoe_backend.mxfp8._fingerprint import ( canonical_json_sha256, ) @@ -402,6 +405,15 @@ def test_validate_training_weights_accepts_complete_fixed_weight_set(): ) == torch.device("cpu") +@pytest.mark.L1 +def test_interleaved_training_weights_require_direct_layouts(): + with pytest.raises(ValueError, match="requires compact K-major forward weights"): + validate_training_weights( + _training_config(weight_interleave_size=32), + _training_weights(), + ) + + @pytest.mark.L1 @pytest.mark.parametrize("part", ["data_noncontiguous", "scale_noncontiguous"]) def test_validate_training_weights_accepts_compact_k_major_views(part): @@ -475,6 +487,42 @@ def test_k_major_workspace_view_matches_upstream_token_major_abi(): assert torch.equal(view, storage.reshape(4, 3).transpose(0, 1)) +@pytest.mark.L0 +def test_training_wgrad_data_operands_alias_backward_outputs(): + pool_rows, hidden, intermediate = 8, 4, 6 + slot = SimpleNamespace( + col_quant_data=torch.empty((pool_rows, hidden), dtype=torch.uint8), + col_quant_sf=torch.empty(1, dtype=torch.uint8), + valid_route_counts=torch.zeros(1, dtype=torch.int32), + expert_offsets=torch.zeros(1, dtype=torch.int32), + fc1_recompute=torch.empty((pool_rows, intermediate), dtype=torch.uint8), + fc1_recompute_sf=torch.empty(1, dtype=torch.uint8), + fc1_col_output=torch.empty((pool_rows, 2 * intermediate), dtype=torch.uint8), + fc1_col_output_sf=torch.empty(1, dtype=torch.uint8), + grad_y2=torch.empty((pool_rows, hidden), dtype=torch.uint8), + grad_y2_sf=torch.empty(1, dtype=torch.uint8), + wgrad_fc1_sfa=torch.empty(1, dtype=torch.uint8), + wgrad_fc1_sfb=torch.empty(1, dtype=torch.uint8), + wgrad_fc2_sfa=torch.empty(1, dtype=torch.uint8), + wgrad_fc2_sfb=torch.empty(1, dtype=torch.uint8), + ) + exporter = Mxfp8TrainingWgradExporter( + experts=1, + hidden=hidden, + intermediate=intermediate, + ) + exporter._expand_scales = Mock() + + operands = exporter.export(slot) + + assert operands.fc1_b is slot.fc1_col_output + assert operands.fc1_b.data_ptr() == slot.fc1_col_output.data_ptr() + assert operands.fc1_b.stride() == slot.fc1_col_output.stride() + assert operands.fc2_a.data_ptr() == slot.fc1_recompute.data_ptr() + assert operands.fc2_a.shape == (intermediate, pool_rows) + assert operands.fc2_a.stride() == slot.fc1_recompute.transpose(0, 1).stride() + + @pytest.mark.L0 def test_only_fixed_training_wgrad_types_are_public(): expected = [f"fc{layer}_{part}" for layer in (1, 2) for part in ("a", "sfa", "b", "sfb")] @@ -564,7 +612,21 @@ def compact_k_major(tensor): backward_w2_transpose=weights.backward_w2_transpose, backward_w1_transpose=weights.backward_w1_transpose, ) - bindings = Mxfp8TrainingWeightBindings(weights) + assert validate_training_weights( + _training_config(weight_interleave_size=32), + weights, + ) == torch.device("cpu") + compatibility_bindings = Mxfp8TrainingWeightBindings(weights) + assert not compatibility_bindings._uses_direct_weight_bindings + assert ( + compatibility_bindings.forward.fc1_weight.data_ptr() + != weights.forward_fc1.data.data_ptr() + ) + + bindings = Mxfp8TrainingWeightBindings( + weights, + weight_interleave_size=32, + ) bindings.refresh() data_pairs = ( (bindings.forward.fc1_weight, weights.forward_fc1.data), diff --git a/test/python/moe_ep/test_moe_ep_forward.py b/test/python/moe_ep/test_moe_ep_forward.py index 9eebbf5ad..353e5e248 100644 --- a/test/python/moe_ep/test_moe_ep_forward.py +++ b/test/python/moe_ep/test_moe_ep_forward.py @@ -424,6 +424,33 @@ def test_moe_ep_rejects_invalid_padding(kwargs): MoeEp(**_forward_config(), **kwargs) +@pytest.mark.L0 +def test_moe_ep_rejects_unsupported_weight_interleave_size(): + from cudnn import MoeEp + + with pytest.raises(ValueError, match="weight_interleave_size must be None or 32"): + MoeEp(**_forward_config(), weight_interleave_size=16) + + +@pytest.mark.L0 +def test_moe_ep_rejects_interleaved_plain_fc1_weight(): + from cudnn import MoeEp + + config = _forward_config() + activation = torch.zeros((1, config["hidden_size"])) + fc1_weight = torch.zeros( + (config["num_experts"], config["hidden_size"], 2 * config["intermediate_size"]) + ) + fc2_weight = torch.zeros( + (config["num_experts"], config["intermediate_size"], config["hidden_size"]) + ) + topk_idx = torch.zeros((1, config["top_k"]), dtype=torch.int32) + topk_weights = torch.ones((1, config["top_k"])) + with MoeEp(**config, weight_interleave_size=32) as op: + with pytest.raises(ValueError, match="requires an MXFP8"): + op(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) + + @pytest.mark.L0 def test_moe_ep_rejects_untyped_tuning(): from cudnn import MoeEp @@ -1091,6 +1118,70 @@ def test_reference_mxfp8_inputs_bf16_combine_matches_naive( torch.testing.assert_close(actual, expected, atol=0, rtol=0) +@pytest.mark.L0 +def test_reference_interleaved_fc1_matches_logical_fc1(): + from cudnn import BlockScaledTensor + + torch.manual_seed(29) + experts, tokens, hidden, intermediate = 2, 3, 128, 128 + activation = quantize_blockwise( + torch.randn(tokens, hidden), + MoeFormat.MXFP8, + axis=1, + ) + logical_fc1 = quantize_blockwise( + torch.randn(experts, hidden, 2 * intermediate) / 8, + MoeFormat.MXFP8, + axis=1, + ) + fc2 = quantize_blockwise( + torch.randn(experts, intermediate, hidden) / 8, + MoeFormat.MXFP8, + axis=1, + ) + + def interleave_last(tensor): + shape = tensor.shape + return ( + tensor.view(*shape[:-1], 2, intermediate // 32, 32) + .transpose(-3, -2) + .reshape(shape) + ) + + interleaved_fc1 = BlockScaledTensor( + data=interleave_last(logical_fc1.data), + scale=interleave_last(logical_fc1.scale), + format=logical_fc1.format, + logical_shape=logical_fc1.logical_shape, + axis=logical_fc1.axis, + ) + topk_idx = torch.tensor([[0], [1], [0]], dtype=torch.int64) + topk_weights = torch.ones(tokens, 1) + kwargs = dict( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + top_k=1, + ) + + logical = MoeEpReference(**kwargs)( + activation, + logical_fc1, + fc2, + topk_idx, + topk_weights, + ) + interleaved = MoeEpReference(**kwargs, weight_interleave_size=32)( + activation, + interleaved_fc1, + fc2, + topk_idx, + topk_weights, + ) + + torch.testing.assert_close(interleaved, logical, atol=0, rtol=0) + + # Host-side EP topology and runtime bootstrap. From bbbe1e2db808825d1e762e1ca6fe57a38b4e55ec Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 30 Aug 2026 20:52:14 -0700 Subject: [PATCH 25/31] other changes needed for cudnn weight interleave option Signed-off-by: Varun Thumbe --- .../cudnn/moe_ep/_megamoe_backend/README.md | 13 +++--- .../mxfp8/_training_resources.py | 19 ++++++++ .../_megamoe_backend/mxfp8/_training_wgrad.py | 45 +++++++++++++++++-- .../mxfp8/_training_wgrad_kernel.py | 43 +++++++++++++++++- test/python/moe_ep/test_moe_ep_backward.py | 1 + 5 files changed, 110 insertions(+), 11 deletions(-) diff --git a/python/cudnn/moe_ep/_megamoe_backend/README.md b/python/cudnn/moe_ep/_megamoe_backend/README.md index 3a9837ad9..f3beda3bb 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/README.md +++ b/python/cudnn/moe_ep/_megamoe_backend/README.md @@ -94,12 +94,13 @@ contract and relaxed atomic accumulation order. `MoeEpTrainingWgradOperands` is a fixed-capacity producer ABI. Device `expert_offsets` and `valid_route_counts` describe the current valid K extent; padding is zeroed. Its data operands alias persistent forward or backward -outputs, using transpose views where required; exporting them performs no -full-tensor data copies. Scale operands are expanded into persistent -grouped-WGrad layouts. FC1 dY data and scales preserve the same 32-element -gate/up-interleaved row order as W1, so a downstream grouped-WGrad consumer -writes gradients in parameter storage order. No specific downstream -grouped-WGrad consumer is guaranteed by this milestone. +outputs, using transpose views where required. With +`weight_interleave_size=32`, exporting them performs no full-tensor data +copies, and FC1 dY data and scales preserve the same 32-element gate/up order +as W1. With `None`, FC1 dY is copied and deinterleaved to match conventional +gate-then-up W1 storage. Scale operands are expanded into persistent +grouped-WGrad layouts. No specific downstream grouped-WGrad consumer is +guaranteed by this milestone. ## Overflow policy diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py index d48480414..2868cdd71 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py @@ -260,6 +260,11 @@ def build_training_workspace_requirements( scale_columns, ), } + if config.weight_interleave_size is None: + wgrad_shapes["wgrad_fc1_b"] = ( + forward.pool_token_capacity, + 2 * config.intermediate_size, + ) for slot in range(slot_count): for name in sorted(_FORWARD_SLOT_SYMMETRIC): @@ -604,6 +609,7 @@ class Mxfp8TrainingSlotViews: fc1_col_output_sf: torch.Tensor grad_y2: torch.Tensor grad_y2_sf: torch.Tensor + wgrad_fc1_b: torch.Tensor | None wgrad_fc1_sfa: torch.Tensor wgrad_fc1_sfb: torch.Tensor wgrad_fc2_sfa: torch.Tensor @@ -651,6 +657,7 @@ def __init__( hidden=config.hidden_size, intermediate=config.intermediate_size, sf_padding=backward.config.sf_padding_block, + weight_interleave_size=config.weight_interleave_size, ) self.beta = torch.ones( (config.experts_per_rank,), @@ -1009,6 +1016,18 @@ def local_bytes(name: str) -> torch.Tensor: torch.uint8, bwd_shapes["grad_y2_sf"], ), + wgrad_fc1_b=( + None + if config.weight_interleave_size == 32 + else _typed_k_major_view( + local_bytes("wgrad_fc1_b"), + _DATA_DTYPE, + ( + self.forward_prepared.pool_token_capacity, + 2 * config.intermediate_size, + ), + ) + ), wgrad_fc1_sfa=_typed_view( local_bytes("wgrad_fc1_sfa"), _SCALE_DTYPE, diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py index 6e1017037..d185fde02 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py @@ -27,14 +27,37 @@ def __init__( hidden: int, intermediate: int, sf_padding: int = 128, + weight_interleave_size: int | None = None, ) -> None: self.experts = int(experts) self.hidden = int(hidden) self.intermediate = int(intermediate) self.sf_padding = int(sf_padding) - self._compiled: dict[tuple[int, int], object] = {} + self.weight_interleave_size = weight_interleave_size + self._compiled: dict[tuple[int, int, int | None], object] = {} self._lock = threading.RLock() + @staticmethod + def _copy_gate_up_data( + target: torch.Tensor, + source: torch.Tensor, + intermediate: int, + ) -> None: + """Deinterleave FC1 dY while copying into K-major staging.""" + pool_rows = source.shape[0] + pairs = intermediate // 32 + source_view = source.view(pool_rows, pairs, 2, 32).permute(0, 2, 1, 3) + target_view = target.as_strided( + (pool_rows, 2, pairs, 32), + ( + target.stride(0), + intermediate * target.stride(1), + 32 * target.stride(1), + target.stride(1), + ), + ) + target_view.copy_(source_view) + def _expand_scales( self, source: torch.Tensor, @@ -43,6 +66,7 @@ def _expand_scales( output: torch.Tensor, *, non_k_size: int, + deinterleave_gate_up: int | None = None, ) -> None: if source.dtype not in (torch.uint8, torch.float8_e8m0fnu): raise TypeError("WGrad source scales must use Uint8 or E8M0") @@ -50,7 +74,7 @@ def _expand_scales( raise TypeError("WGrad output scales must use E8M0") source_bytes = source.view(torch.uint8).reshape(-1) output_bytes = output.view(torch.uint8).reshape(-1) - key = (int(non_k_size), self.sf_padding) + key = (int(non_k_size), self.sf_padding, deinterleave_gate_up) import cuda.bindings.driver as cuda stream = torch.cuda.current_stream(output.device) @@ -76,6 +100,7 @@ def _expand_scales( non_k_size=non_k_size, expert_count=self.experts, source_sf_padding=self.sf_padding, + deinterleave_gate_up=deinterleave_gate_up, ) compiled = cute.compile(kernel, *args) self._compiled[key] = compiled @@ -93,6 +118,17 @@ def export( if slot.col_quant_data.shape[0] != pool_rows: raise RuntimeError("forward/backward WGrad pool capacities differ") + fc1_b = slot.fc1_col_output + if self.weight_interleave_size is None: + if slot.wgrad_fc1_b is None: + raise RuntimeError("conventional W1 requires FC1 WGrad staging") + self._copy_gate_up_data( + slot.wgrad_fc1_b, + slot.fc1_col_output, + self.intermediate, + ) + fc1_b = slot.wgrad_fc1_b + self._expand_scales( slot.col_quant_sf, slot.valid_route_counts, @@ -106,6 +142,9 @@ def export( slot.expert_offsets, slot.wgrad_fc1_sfb, non_k_size=2 * self.intermediate, + deinterleave_gate_up=( + self.intermediate if self.weight_interleave_size is None else None + ), ) self._expand_scales( slot.fc1_recompute_sf, @@ -125,7 +164,7 @@ def export( return MoeEpTrainingWgradOperands( fc1_a=slot.col_quant_data.transpose(0, 1), fc1_sfa=slot.wgrad_fc1_sfa, - fc1_b=slot.fc1_col_output, + fc1_b=fc1_b, fc1_sfb=slot.wgrad_fc1_sfb, fc2_a=slot.fc1_recompute.transpose(0, 1), fc2_sfa=slot.wgrad_fc2_sfa, diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py index 244e809ea..16f63243e 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py @@ -25,16 +25,25 @@ def __init__( non_k_size: int, expert_count: int, source_sf_padding: int, + deinterleave_gate_up: int | None = None, ) -> None: self.non_k_size = int(non_k_size) self.expert_count = int(expert_count) self.source_sf_padding = int(source_sf_padding) + self.deinterleave_gate_up = ( + None if deinterleave_gate_up is None else int(deinterleave_gate_up) + ) if self.non_k_size <= 0 or self.non_k_size % 128: raise ValueError("WGrad scale non-K size must be divisible by 128") if self.expert_count <= 0: raise ValueError("WGrad scale expansion requires experts") if self.source_sf_padding <= 0 or self.source_sf_padding % 128: raise ValueError("WGrad source SF padding must be a positive multiple of 128") + if ( + self.deinterleave_gate_up is not None + and self.non_k_size != 2 * self.deinterleave_gate_up + ): + raise ValueError("gate/up scale deinterleave size mismatch") @cute.jit def __call__( @@ -94,8 +103,38 @@ def _kernel( hidden_atom = relative_atom // target_token_atoms token_atom = relative_atom % target_token_atoms if token_atom < source_token_atoms: - source_atom = source_atom_base + hidden_atom * source_token_atoms + token_atom - value = source[source_atom * Int32(atom_bytes) + byte_in_atom] + source_hidden_atom = hidden_atom + source_byte = byte_in_atom + if cutlass.const_expr(self.deinterleave_gate_up is not None): + lane = byte_in_atom // Int32(16) + byte_tail = byte_in_atom % Int32(16) + group = byte_tail // Int32(4) + column_lane = byte_tail % Int32(4) + feature = hidden_atom * Int32(128) + group * Int32(32) + lane + intermediate = Int32(self.deinterleave_gate_up) + source_feature = Int32(0) + if feature < intermediate: + source_feature = ( + feature // Int32(32) + ) * Int32(64) + feature % Int32(32) + else: + up_feature = feature - intermediate + source_feature = ( + up_feature // Int32(32) + ) * Int32(64) + Int32(32) + up_feature % Int32(32) + source_hidden_atom = source_feature // Int32(128) + source_feature_in_atom = source_feature % Int32(128) + source_byte = ( + (source_feature_in_atom % Int32(32)) * Int32(16) + + (source_feature_in_atom // Int32(32)) * Int32(4) + + column_lane + ) + source_atom = ( + source_atom_base + + source_hidden_atom * source_token_atoms + + token_atom + ) + value = source[source_atom * Int32(atom_bytes) + source_byte] target_atom_base += target_atom_count source_atom_base += Int32(non_k_atoms) * source_token_atoms previous_end = end diff --git a/test/python/moe_ep/test_moe_ep_backward.py b/test/python/moe_ep/test_moe_ep_backward.py index d457cbe11..da811d2df 100644 --- a/test/python/moe_ep/test_moe_ep_backward.py +++ b/test/python/moe_ep/test_moe_ep_backward.py @@ -510,6 +510,7 @@ def test_training_wgrad_data_operands_alias_backward_outputs(): experts=1, hidden=hidden, intermediate=intermediate, + weight_interleave_size=32, ) exporter._expand_scales = Mock() From 5e7114b449697ce5702eb829a10af08138a23e59 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 31 Aug 2026 00:45:21 -0700 Subject: [PATCH 26/31] address review comments Signed-off-by: Varun Thumbe --- python/cudnn/moe_ep/_contracts.py | 22 ++++- .../cudnn/moe_ep/_megamoe_backend/README.md | 3 + .../moe_ep/_megamoe_backend/mxfp8/_adapter.py | 6 +- .../moe_ep/_megamoe_backend/mxfp8/_config.py | 7 +- .../mxfp8/_training_resources.py | 11 ++- .../mxfp8/_training_weights.py | 12 ++- .../_megamoe_backend/mxfp8/_training_wgrad.py | 11 ++- python/cudnn/moe_ep/_validation.py | 6 +- python/cudnn/moe_ep/api.py | 8 +- test/python/moe_ep/moe_ep_test_support.py | 31 ++++--- test/python/moe_ep/test_moe_ep_backward.py | 92 ++++++++++++++++++- test/python/moe_ep/test_moe_ep_forward.py | 21 +++++ 12 files changed, 188 insertions(+), 42 deletions(-) diff --git a/python/cudnn/moe_ep/_contracts.py b/python/cudnn/moe_ep/_contracts.py index 1d3e741ef..cfafd3350 100644 --- a/python/cudnn/moe_ep/_contracts.py +++ b/python/cudnn/moe_ep/_contracts.py @@ -6,6 +6,7 @@ from __future__ import annotations from dataclasses import dataclass +from enum import Enum from typing import Any, Literal, Optional import torch @@ -14,6 +15,23 @@ from ._types import MoeTensor +class Fc1WeightLayout(str, Enum): + """Logical gate/up ordering used by FC1 weights and their gradients.""" + + GATE_THEN_UP = "gate_then_up" + GATE_UP_INTERLEAVED_32 = "gate_up_interleaved_32" + + +def normalize_fc1_weight_layout(weight_interleave_size: Optional[int]) -> Fc1WeightLayout: + """Normalize the public compatibility flag into the internal layout ABI.""" + + if weight_interleave_size is None: + return Fc1WeightLayout.GATE_THEN_UP + if weight_interleave_size == 32: + return Fc1WeightLayout.GATE_UP_INTERLEAVED_32 + raise ValueError("weight_interleave_size must be None or 32") + + @dataclass(frozen=True) class ForwardConfig: """Static configuration snapshot for one ``MoeEp`` instance.""" @@ -39,7 +57,7 @@ class ForwardConfig: backward_wgrad_mode: Literal["none", "operands"] = "none" max_recv_size_per_rank: Optional[int] = None drop_on_overflow: bool = False - weight_interleave_size: Optional[int] = None + fc1_weight_layout: Fc1WeightLayout = Fc1WeightLayout.GATE_THEN_UP @dataclass(frozen=True) @@ -57,6 +75,8 @@ class ValidatedForwardRequest: __all__ = [ + "Fc1WeightLayout", "ForwardConfig", "ValidatedForwardRequest", + "normalize_fc1_weight_layout", ] diff --git a/python/cudnn/moe_ep/_megamoe_backend/README.md b/python/cudnn/moe_ep/_megamoe_backend/README.md index f3beda3bb..0b2c5e6d9 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/README.md +++ b/python/cudnn/moe_ep/_megamoe_backend/README.md @@ -63,6 +63,9 @@ forward, backward, and WGrad-export kernels are compiled before capture. - An execution lane owns mutable router, barrier, and kernel scratch. - Every symmetric region is built in deterministic order and its size is normalized by name across EP ranks before allocation. +- The training ABI fingerprint includes the normalized FC1 layout policy + (`gate_then_up` or `gate_up_interleaved_32`), so ranks with different + gate/up semantics fail the collective handshake before allocation. - Multiple streams require distinct lanes. Distributed MegaMoE kernels must be ordered consistently on every rank with captured CUDA events; independent lane storage does not permit unordered communication overlap. diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py index 05b0572bf..4e59fc0ff 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_adapter.py @@ -9,7 +9,7 @@ import torch -from ..._contracts import ValidatedForwardRequest +from ..._contracts import Fc1WeightLayout, ValidatedForwardRequest from ..._types import BlockScaledTensor, MoeFormat from .._plan import PreparedResources from .._workspace import padded_mxfp8_scale_columns @@ -369,7 +369,9 @@ def _prepare_weights( fc1_weight, fc1_weight_sf = _prepare_fc1( fc1_source, config.intermediate, - already_interleaved=config.weight_interleave_size == 32, + already_interleaved=( + config.fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32 + ), ) fc2_weight, fc2_weight_sf = _prepare_fc2(fc2_source) weights = Mxfp8Weights( diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py index 0fcd8a669..8243093ec 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_config.py @@ -9,7 +9,7 @@ import torch -from ..._contracts import ForwardConfig +from ..._contracts import Fc1WeightLayout, ForwardConfig from ._formats import combine_wire_format @@ -25,7 +25,7 @@ class Mxfp8KernelConfig: top_k: int max_tokens_per_rank: int apply_topk_in_fc1: bool - weight_interleave_size: int | None + fc1_weight_layout: Fc1WeightLayout gate_up_clamp: float | None generate_c: bool max_recv_size_per_rank: int | None = None @@ -79,7 +79,7 @@ def from_forward_config(cls, config: ForwardConfig) -> "Mxfp8KernelConfig": top_k=config.top_k, max_tokens_per_rank=config.max_tokens_per_rank, apply_topk_in_fc1=config.apply_topk_in_fc1, - weight_interleave_size=config.weight_interleave_size, + fc1_weight_layout=config.fc1_weight_layout, gate_up_clamp=config.gate_up_clamp, generate_c=config.generate_c, max_recv_size_per_rank=max_recv_size_per_rank, @@ -134,6 +134,7 @@ def effective_config(self, launch_cluster_count: int) -> dict[str, object]: "max_recv_size_per_rank": self.max_recv_size_per_rank, "drop_on_overflow": self.drop_on_overflow, "apply_topk_in_fc1": self.apply_topk_in_fc1, + "fc1_weight_layout": self.fc1_weight_layout.value, "gate_up_clamp": self.gate_up_clamp, "generate_c": self.generate_c, "enable_col_quant": self.enable_col_quant, diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py index 2868cdd71..1f097a583 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py @@ -16,7 +16,7 @@ import torch import torch.distributed as dist -from ..._contracts import ForwardConfig +from ..._contracts import Fc1WeightLayout, ForwardConfig from ..._types import MoeEpTrainingWeights from .._comm import SymmetricMemoryProvider from .._plan import PreparedResources @@ -260,7 +260,7 @@ def build_training_workspace_requirements( scale_columns, ), } - if config.weight_interleave_size is None: + if config.fc1_weight_layout is Fc1WeightLayout.GATE_THEN_UP: wgrad_shapes["wgrad_fc1_b"] = ( forward.pool_token_capacity, 2 * config.intermediate_size, @@ -547,6 +547,7 @@ def _build_training_abi_facts( "combine_format": config.combine_format, "output_format": config.output_format, "apply_topk_in_fc1": bool(config.apply_topk_in_fc1), + "fc1_weight_layout": config.fc1_weight_layout.value, "gate_up_clamp": config.gate_up_clamp, }, "resources": { @@ -649,7 +650,7 @@ def __init__( self.backward_prepared = backward self.weight_bindings = Mxfp8TrainingWeightBindings( weights, - weight_interleave_size=config.weight_interleave_size, + fc1_weight_layout=config.fc1_weight_layout, ) self.stager = Mxfp8TrainingStager(config.hidden_size, config.top_k) self.wgrad_exporter = Mxfp8TrainingWgradExporter( @@ -657,7 +658,7 @@ def __init__( hidden=config.hidden_size, intermediate=config.intermediate_size, sf_padding=backward.config.sf_padding_block, - weight_interleave_size=config.weight_interleave_size, + fc1_weight_layout=config.fc1_weight_layout, ) self.beta = torch.ones( (config.experts_per_rank,), @@ -1018,7 +1019,7 @@ def local_bytes(name: str) -> torch.Tensor: ), wgrad_fc1_b=( None - if config.weight_interleave_size == 32 + if config.fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32 else _typed_k_major_view( local_bytes("wgrad_fc1_b"), _DATA_DTYPE, diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py index cdb61faa4..02f946f83 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py @@ -9,6 +9,7 @@ import torch +from ..._contracts import Fc1WeightLayout from ..._types import BlockScaledTensor, MoeEpTrainingWeights from ._adapter import Mxfp8Weights @@ -186,22 +187,25 @@ def __init__( self, weights: MoeEpTrainingWeights, *, - weight_interleave_size: int | None = None, + fc1_weight_layout: Fc1WeightLayout = Fc1WeightLayout.GATE_THEN_UP, ) -> None: self.weights = weights - self.weight_interleave_size = weight_interleave_size + self.fc1_weight_layout = fc1_weight_layout fwd_fc1 = weights.forward_fc1 fwd_fc2 = weights.forward_fc2 bwd_w2t = weights.backward_w2_transpose bwd_w1t = weights.backward_w1_transpose self._uses_direct_weight_bindings = ( - weight_interleave_size == 32 + fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32 and fwd_fc1.data.stride(1) == 1 and fwd_fc2.data.stride(1) == 1 and bwd_w2t.data.is_contiguous() and bwd_w1t.data.is_contiguous() ) - if weight_interleave_size == 32 and not self._uses_direct_weight_bindings: + if ( + fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32 + and not self._uses_direct_weight_bindings + ): raise ValueError( "weight_interleave_size=32 requires compact K-major forward " "weights and contiguous backward transpose weights" diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py index d185fde02..4a413e349 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py @@ -10,6 +10,7 @@ import torch +from ..._contracts import Fc1WeightLayout from ..._types import MoeEpTrainingWgradOperands from ._launch import _to_cute @@ -27,13 +28,13 @@ def __init__( hidden: int, intermediate: int, sf_padding: int = 128, - weight_interleave_size: int | None = None, + fc1_weight_layout: Fc1WeightLayout = Fc1WeightLayout.GATE_THEN_UP, ) -> None: self.experts = int(experts) self.hidden = int(hidden) self.intermediate = int(intermediate) self.sf_padding = int(sf_padding) - self.weight_interleave_size = weight_interleave_size + self.fc1_weight_layout = fc1_weight_layout self._compiled: dict[tuple[int, int, int | None], object] = {} self._lock = threading.RLock() @@ -119,7 +120,7 @@ def export( raise RuntimeError("forward/backward WGrad pool capacities differ") fc1_b = slot.fc1_col_output - if self.weight_interleave_size is None: + if self.fc1_weight_layout is Fc1WeightLayout.GATE_THEN_UP: if slot.wgrad_fc1_b is None: raise RuntimeError("conventional W1 requires FC1 WGrad staging") self._copy_gate_up_data( @@ -143,7 +144,9 @@ def export( slot.wgrad_fc1_sfb, non_k_size=2 * self.intermediate, deinterleave_gate_up=( - self.intermediate if self.weight_interleave_size is None else None + self.intermediate + if self.fc1_weight_layout is Fc1WeightLayout.GATE_THEN_UP + else None ), ) self._expand_scales( diff --git a/python/cudnn/moe_ep/_validation.py b/python/cudnn/moe_ep/_validation.py index d5338c335..1370b55e6 100644 --- a/python/cudnn/moe_ep/_validation.py +++ b/python/cudnn/moe_ep/_validation.py @@ -9,7 +9,7 @@ import torch -from ._contracts import ForwardConfig, ValidatedForwardRequest +from ._contracts import Fc1WeightLayout, ForwardConfig, ValidatedForwardRequest from ._types import ( BlockScaledTensor, MoeEpTrainingWeights, @@ -185,7 +185,7 @@ def validate_forward( config.hidden_size, ), ) - if config.weight_interleave_size is not None and ( + if config.fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32 and ( not isinstance(fc1_weight, BlockScaledTensor) or fc1_weight.format is not MoeFormat.MXFP8 ): @@ -300,7 +300,7 @@ def is_compact_k_major(tensor: torch.Tensor) -> bool: f"{name} data and scale must be contiguous or compact K-major " "for fixed training weight binding" ) - if config.weight_interleave_size == 32 and not ( + if config.fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32 and not ( is_compact_k_major(weights.forward_fc1.data) and is_compact_k_major(weights.forward_fc2.data) and weights.backward_w2_transpose.data.is_contiguous() diff --git a/python/cudnn/moe_ep/api.py b/python/cudnn/moe_ep/api.py index b71dd439b..10c182d4a 100644 --- a/python/cudnn/moe_ep/api.py +++ b/python/cudnn/moe_ep/api.py @@ -20,7 +20,7 @@ import torch import torch.distributed as dist -from ._contracts import ForwardConfig +from ._contracts import ForwardConfig, normalize_fc1_weight_layout from ._tuning import MoeEpTuningConfig from ._types import ( BlockScaledTensor, @@ -140,8 +140,7 @@ def __init__( raise ValueError("drop_on_overflow must be a bool") if not isinstance(apply_topk_in_fc1, bool): raise ValueError("apply_topk_in_fc1 must be a bool") - if weight_interleave_size not in (None, 32): - raise ValueError("weight_interleave_size must be None or 32") + fc1_weight_layout = normalize_fc1_weight_layout(weight_interleave_size) for name, value in ( ("token_padding_size", token_padding_size), ("sf_padding_size", sf_padding_size), @@ -181,6 +180,7 @@ def __init__( self.combine_format = _parse_format(combine_format) self.apply_topk_in_fc1 = apply_topk_in_fc1 self.weight_interleave_size = weight_interleave_size + self._fc1_weight_layout = fc1_weight_layout self.gate_up_clamp = None if gate_up_clamp is None else abs(gate_up_clamp) self.token_padding_size = token_padding_size self.sf_padding_size = sf_padding_size @@ -214,7 +214,7 @@ def __init__( output_format=self.output_format.value, combine_format=self.combine_format.value, apply_topk_in_fc1=self.apply_topk_in_fc1, - weight_interleave_size=self.weight_interleave_size, + fc1_weight_layout=self._fc1_weight_layout, gate_up_clamp=self.gate_up_clamp, generate_c=False, token_padding_size=self.token_padding_size, diff --git a/test/python/moe_ep/moe_ep_test_support.py b/test/python/moe_ep/moe_ep_test_support.py index 72ee0c7f1..304d7979e 100644 --- a/test/python/moe_ep/moe_ep_test_support.py +++ b/test/python/moe_ep/moe_ep_test_support.py @@ -260,9 +260,10 @@ def quantize_mxfp8(tensor: torch.Tensor, *, axis: int = -1): def _training_config(**overrides): - from cudnn.moe_ep._contracts import ForwardConfig + from cudnn.moe_ep._contracts import ForwardConfig, normalize_fc1_weight_layout from cudnn.moe_ep._tuning import MoeEpTuningConfig + weight_interleave_size = overrides.pop("weight_interleave_size", None) values = { "num_experts": 2, "hidden_size": 128, @@ -285,6 +286,7 @@ def _training_config(**overrides): "sf_padding_size": 128, "tuning": MoeEpTuningConfig(), "backward_wgrad_mode": "operands", + "fc1_weight_layout": normalize_fc1_weight_layout(weight_interleave_size), } values.update(overrides) return ForwardConfig(**values) @@ -1255,6 +1257,7 @@ def _assert_wgrads_match_reference( expected, *, expected_dense=None, + weight_interleave_size=None, ) -> None: """Compare fixed-capacity production operands with standalone dense dW.""" @@ -1268,20 +1271,20 @@ def _assert_wgrads_match_reference( actual_dense = _dense_wgrads_from_operands(actual) if expected_dense is None: expected_dense = expected.dense_wgrads() - expected_fc1, expected_fc2 = expected_dense - interleave_size = 32 - fc1_out_features = expected_fc1.shape[-1] - expected_fc1 = ( - expected_fc1.view( - *expected_fc1.shape[:-1], - 2, - fc1_out_features // (2 * interleave_size), - interleave_size, + if weight_interleave_size is not None: + expected_fc1, expected_fc2 = expected_dense + fc1_out_features = expected_fc1.shape[-1] + expected_fc1 = ( + expected_fc1.view( + *expected_fc1.shape[:-1], + 2, + fc1_out_features // (2 * weight_interleave_size), + weight_interleave_size, + ) + .transpose(-3, -2) + .reshape(expected_fc1.shape) ) - .transpose(-3, -2) - .reshape(expected_fc1.shape) - ) - expected_dense = (expected_fc1, expected_fc2) + expected_dense = (expected_fc1, expected_fc2) for name, actual_dw, expected_dw in zip( ("grad_fc1_weight", "grad_fc2_weight"), actual_dense, diff --git a/test/python/moe_ep/test_moe_ep_backward.py b/test/python/moe_ep/test_moe_ep_backward.py index da811d2df..5d8d78d33 100644 --- a/test/python/moe_ep/test_moe_ep_backward.py +++ b/test/python/moe_ep/test_moe_ep_backward.py @@ -26,6 +26,7 @@ MoeEpTrainingSlot, MoeEpTrainingWgradOperands, ) +from cudnn.moe_ep._contracts import Fc1WeightLayout from cudnn.moe_ep._validation import validate_training_weights from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( _typed_k_major_view, @@ -233,9 +234,26 @@ def test_training_abi_fingerprint_is_stable_and_structural(): lane_count=2, source_tree_digest="source", ) + changed_layout = _build_training_abi_facts( + _training_config( + ep_size=2, + ep_global_ranks=(0, 1), + weight_interleave_size=32, + ), + forward, + backward, + weights, + requirements, + slot_count=2, + lane_count=1, + source_tree_digest="source", + ) + assert first["policy"]["fc1_weight_layout"] == "gate_then_up" + assert changed_layout["policy"]["fc1_weight_layout"] == "gate_up_interleaved_32" assert canonical_json_sha256(first) == canonical_json_sha256(second) assert canonical_json_sha256(first) != canonical_json_sha256(changed) + assert canonical_json_sha256(first) != canonical_json_sha256(changed_layout) @pytest.mark.L0 @@ -510,7 +528,7 @@ def test_training_wgrad_data_operands_alias_backward_outputs(): experts=1, hidden=hidden, intermediate=intermediate, - weight_interleave_size=32, + fc1_weight_layout=Fc1WeightLayout.GATE_UP_INTERLEAVED_32, ) exporter._expand_scales = Mock() @@ -524,6 +542,76 @@ def test_training_wgrad_data_operands_alias_backward_outputs(): assert operands.fc2_a.stride() == slot.fc1_recompute.transpose(0, 1).stride() +@pytest.mark.L0 +@pytest.mark.parametrize( + "layout", + [ + Fc1WeightLayout.GATE_THEN_UP, + Fc1WeightLayout.GATE_UP_INTERLEAVED_32, + ], +) +def test_training_wgrad_fc1_layout_matches_reference(layout): + pool_rows, hidden, intermediate = 3, 4, 64 + semantic_dc = ( + torch.arange(pool_rows * 2 * intermediate, dtype=torch.int64) + .remainder(251) + .to(torch.uint8) + .reshape(pool_rows, 2 * intermediate) + ) + interleaved_dc = ( + semantic_dc.view(pool_rows, 2, intermediate // 32, 32) + .transpose(1, 2) + .reshape_as(semantic_dc) + ) + x = torch.arange(pool_rows * hidden, dtype=torch.uint8).reshape(pool_rows, hidden) + slot = SimpleNamespace( + col_quant_data=x, + col_quant_sf=torch.empty(1, dtype=torch.uint8), + valid_route_counts=torch.zeros(1, dtype=torch.int32), + expert_offsets=torch.zeros(1, dtype=torch.int32), + fc1_recompute=torch.empty((pool_rows, intermediate), dtype=torch.uint8), + fc1_recompute_sf=torch.empty(1, dtype=torch.uint8), + fc1_col_output=interleaved_dc, + fc1_col_output_sf=torch.empty(1, dtype=torch.uint8), + grad_y2=torch.empty((pool_rows, hidden), dtype=torch.uint8), + grad_y2_sf=torch.empty(1, dtype=torch.uint8), + wgrad_fc1_b=torch.empty_like(interleaved_dc), + wgrad_fc1_sfa=torch.empty(1, dtype=torch.uint8), + wgrad_fc1_sfb=torch.empty(1, dtype=torch.uint8), + wgrad_fc2_sfa=torch.empty(1, dtype=torch.uint8), + wgrad_fc2_sfb=torch.empty(1, dtype=torch.uint8), + ) + exporter = Mxfp8TrainingWgradExporter( + experts=1, + hidden=hidden, + intermediate=intermediate, + fc1_weight_layout=layout, + ) + exporter._expand_scales = Mock() + + operands = exporter.export(slot) + expected_dc = ( + semantic_dc + if layout is Fc1WeightLayout.GATE_THEN_UP + else interleaved_dc + ) + expected_wgrad = x.transpose(0, 1).float() @ expected_dc.float() + actual_wgrad = operands.fc1_a.float() @ operands.fc1_b.float() + + torch.testing.assert_close(actual_wgrad, expected_wgrad, atol=0, rtol=0) + if layout is Fc1WeightLayout.GATE_THEN_UP: + assert operands.fc1_b is slot.wgrad_fc1_b + assert operands.fc1_b.data_ptr() != slot.fc1_col_output.data_ptr() + assert exporter._expand_scales.call_args_list[1].kwargs[ + "deinterleave_gate_up" + ] == intermediate + else: + assert operands.fc1_b is slot.fc1_col_output + assert exporter._expand_scales.call_args_list[1].kwargs[ + "deinterleave_gate_up" + ] is None + + @pytest.mark.L0 def test_only_fixed_training_wgrad_types_are_public(): expected = [f"fc{layer}_{part}" for layer in (1, 2) for part in ("a", "sfa", "b", "sfb")] @@ -626,7 +714,7 @@ def compact_k_major(tensor): bindings = Mxfp8TrainingWeightBindings( weights, - weight_interleave_size=32, + fc1_weight_layout=Fc1WeightLayout.GATE_UP_INTERLEAVED_32, ) bindings.refresh() data_pairs = ( diff --git a/test/python/moe_ep/test_moe_ep_forward.py b/test/python/moe_ep/test_moe_ep_forward.py index 353e5e248..f03eacf1c 100644 --- a/test/python/moe_ep/test_moe_ep_forward.py +++ b/test/python/moe_ep/test_moe_ep_forward.py @@ -432,6 +432,27 @@ def test_moe_ep_rejects_unsupported_weight_interleave_size(): MoeEp(**_forward_config(), weight_interleave_size=16) +@pytest.mark.L0 +@pytest.mark.parametrize( + ("weight_interleave_size", "expected_layout"), + [ + (None, "gate_then_up"), + (32, "gate_up_interleaved_32"), + ], +) +def test_moe_ep_normalizes_fc1_weight_layout( + weight_interleave_size, + expected_layout, +): + from cudnn import MoeEp + + with MoeEp( + **_forward_config(), + weight_interleave_size=weight_interleave_size, + ) as op: + assert op._forward_config.fc1_weight_layout.value == expected_layout + + @pytest.mark.L0 def test_moe_ep_rejects_interleaved_plain_fc1_weight(): from cudnn import MoeEp From 6cc24e35c64999fe733262b3936e85c5309ffc34 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 31 Aug 2026 15:34:14 -0700 Subject: [PATCH 27/31] have discrete wgrad also route through external workspace Signed-off-by: Varun Thumbe --- .../gemm_fusions/grouped_gemm_wgrad.md | 8 +- .../cutedsl/grouped/wgrad/_blockscaled_api.py | 35 +++++- .../cudnn/gemm/cutedsl/grouped/wgrad/api.py | 3 +- .../grouped_gemm/test_grouped_gemm_wgrad.py | 109 ++++++++++++++++++ 4 files changed, 148 insertions(+), 7 deletions(-) diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md index d3e95aca5..4842244f8 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md @@ -92,14 +92,16 @@ operand pair. It preserves the pre-existing scale-factor contract: provide `global_scale_b` where the selected low-precision format requires them. BF16 does not reinterpret these controls; it rejects them instead. -Dense Torch callers that retain operations for CUDA Graph replay may provide a -caller-owned `descriptor_workspace`. Allocate its size with +Torch block-scaled callers in dense or discrete output mode that retain +operations for CUDA Graph replay may provide a caller-owned +`descriptor_workspace`. Allocate its size with `get_grouped_gemm_wgrad_workspace_size_sm100`, keep it alive for as long as the captured call site may replay, and do not share it between call sites that may overlap. This lets multiple same-signature calls share one compiled kernel without sharing mutable runtime TMA descriptors. Callers that omit this argument retain the compatibility behavior that isolates cached API instances -by explicit dense output address. +by explicit dense output address; discrete callers retain the compiled +operation's internal workspace. ```python workspace = torch.empty( diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py index 26c76b999..c546768b4 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py @@ -537,7 +537,11 @@ def _compile_discrete(self, kernel, max_active_clusters, fake_stream) -> None: options="--enable-tvm-ffi", ) - cached_workspace = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) + self._workspace_arg = from_dlpack( + self._workspace, + assumed_align=128, + enable_tvm_ffi=True, + ) single_expert_placeholder = torch.empty_strided( self.single_expert_wgrad_desc.shape, self.single_expert_wgrad_desc.stride, @@ -557,6 +561,7 @@ def tensor_api( sfb_tensor: torch.Tensor, wgrad_ptrs: torch.Tensor, offsets_tensor: torch.Tensor, + workspace, stream: cuda.CUstream, global_scale_a: Optional[torch.Tensor], global_scale_b: Optional[torch.Tensor], @@ -568,7 +573,7 @@ def tensor_api( sfb_tensor, wgrad_ptrs.data_ptr(), offsets_tensor, - cached_workspace, + workspace, stream, global_scale_a, global_scale_b, @@ -648,6 +653,31 @@ def execute( ptrs = [wgrad_tensor.data_ptr() + i * expert_stride_bytes for i in range(wgrad_tensor.shape[0])] wgrad_ptrs = torch.tensor(ptrs, dtype=torch.int64, device=wgrad_tensor.device) _validate_pointer_tensor(wgrad_ptrs, "wgrad_ptrs", self.expert_cnt) + if descriptor_workspace is None: + workspace_arg = self._workspace_arg + else: + self._value_error_if( + descriptor_workspace.dtype != torch.uint8, + f"descriptor_workspace must have dtype uint8, got {descriptor_workspace.dtype}", + ) + self._value_error_if( + descriptor_workspace.device != a_tensor.device, + "descriptor_workspace and WGrad operands must be on the same device", + ) + self._value_error_if( + not descriptor_workspace.is_contiguous(), + "descriptor_workspace must be contiguous", + ) + self._value_error_if( + descriptor_workspace.numel() < self._workspace_bytes, + f"descriptor_workspace requires at least {self._workspace_bytes} bytes, " + f"got {descriptor_workspace.numel()}", + ) + workspace_arg = from_dlpack( + descriptor_workspace, + assumed_align=128, + enable_tvm_ffi=True, + ) self._compiled_kernel( a_tensor, b_tensor, @@ -655,6 +685,7 @@ def execute( sfb_tensor, wgrad_ptrs, offsets_tensor, + workspace_arg, current_stream, global_scale_a, global_scale_b, diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py index 946660f57..44597bd96 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py @@ -336,10 +336,9 @@ def grouped_gemm_wgrad_wrapper_sm100( if descriptor_workspace is not None and ( backend is not GroupedGemmBackend.BLOCK_SCALED or framework != "torch" - or output_mode != "dense" ): raise ValueError( - "descriptor_workspace is supported only for dense torch block-scaled WGrad" + "descriptor_workspace is supported only for torch block-scaled WGrad" ) explicit_dense_output_identity = None if ( diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py index e239a8b80..1dc022b35 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py @@ -817,6 +817,63 @@ def counted_compile(self): assert cache_entries == expected_cache_entries +@pytest.mark.L0 +def test_grouped_gemm_wgrad_wrapper_discrete_accepts_caller_workspace(monkeypatch): + from cudnn.gemm.cutedsl.grouped.wgrad import api as grouped_gemm_wgrad_api + + grouped_gemm_wgrad_api._cache_of_GroupedGemmWgradSm100Objects.clear() + compile_count = {"value": 0} + + def counted_compile(self): + compile_count["value"] += 1 + + monkeypatch.setattr( + grouped_gemm_wgrad_api.GroupedGemmWgradSm100, + "check_support", + lambda self: True, + ) + monkeypatch.setattr( + grouped_gemm_wgrad_api.GroupedGemmWgradSm100, + "compile", + counted_compile, + ) + monkeypatch.setattr( + grouped_gemm_wgrad_api.GroupedGemmWgradSm100, + "execute", + lambda self, **kwargs: None, + ) + monkeypatch.setattr( + grouped_gemm_wgrad_api, + "select_grouped_gemm_backend", + lambda **_: grouped_gemm_wgrad_api.GroupedGemmBackend.BLOCK_SCALED, + ) + + inputs = _make_wgrad_wrapper_cache_inputs([8, 12]) + outputs = [torch.empty((2, 32, 64), dtype=torch.bfloat16) for _ in range(2)] + workspaces = [torch.empty(512, dtype=torch.uint8) for _ in range(2)] + try: + for output, workspace in zip(outputs, workspaces): + cudnn.grouped_gemm_wgrad_wrapper_sm100( + **inputs, + output_mode="discrete", + wgrad_tensor=output, + descriptor_workspace=workspace, + acc_dtype=torch.float32, + wgrad_dtype=torch.bfloat16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=16, + ) + finally: + cache_entries = len( + grouped_gemm_wgrad_api._cache_of_GroupedGemmWgradSm100Objects + ) + grouped_gemm_wgrad_api._cache_of_GroupedGemmWgradSm100Objects.clear() + + assert compile_count["value"] == 1 + assert cache_entries == 1 + + @pytest.mark.L0 def test_grouped_gemm_wgrad_workspace_size(): assert cudnn.get_grouped_gemm_wgrad_workspace_size_sm100(2) == 512 @@ -879,6 +936,58 @@ def compiled_kernel(*args): assert launch_workspaces[0].data_ptr() != launch_workspaces[2].data_ptr() +@pytest.mark.L0 +def test_blockscaled_discrete_wgrad_execute_uses_caller_workspace(monkeypatch): + from cudnn.gemm.cutedsl.grouped.wgrad import _blockscaled_api + + api = object.__new__(_blockscaled_api.GroupedGemmWgradBlockScaledAPI) + api._workspace_bytes = 16 + api._workspace_arg = torch.empty(16, dtype=torch.uint8) + api.a_desc = type("TensorDesc", (), {"device": torch.device("cpu")})() + api.weight_mode = _blockscaled_api.MoEWeightMode.DISCRETE + api.expert_cnt = 2 + api._get_default_stream = lambda stream: stream + api._runtime_error_if = lambda condition, message: None + api._value_error_if = lambda condition, message: None + monkeypatch.setattr( + _blockscaled_api, + "from_dlpack", + lambda tensor, **kwargs: tensor, + ) + monkeypatch.setattr( + _blockscaled_api, + "_validate_pointer_tensor", + lambda tensor, name, count: None, + ) + + launch_workspaces = [] + + def compiled_kernel(*args): + launch_workspaces.append(args[6]) + + api._compiled_kernel = compiled_kernel + operand = torch.empty((1, 1)) + offsets = torch.tensor([1, 2], dtype=torch.int32) + wgrad_ptrs = torch.empty(2, dtype=torch.int64) + workspaces = [torch.empty(16, dtype=torch.uint8) for _ in range(2)] + for workspace in (workspaces[0], workspaces[0], workspaces[1]): + api.execute( + operand, + operand, + operand, + operand, + offsets, + wgrad_ptrs=wgrad_ptrs, + descriptor_workspace=workspace, + current_stream=object(), + ) + + assert launch_workspaces[0] is workspaces[0] + assert launch_workspaces[1] is workspaces[0] + assert launch_workspaces[2] is workspaces[1] + assert launch_workspaces[0].data_ptr() != launch_workspaces[2].data_ptr() + + @pytest.mark.L0 def test_grouped_gemm_wgrad_wrapper_input_order_cache_key(monkeypatch): from cudnn.gemm.cutedsl.grouped.wgrad import api as grouped_gemm_wgrad_api From 26954605e8803d8cc4b6a5632107049d85cd685c Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 31 Aug 2026 19:59:11 -0700 Subject: [PATCH 28/31] make contiguos restriction for backward colwise data Signed-off-by: Varun Thumbe --- docs/fe-oss-apis/moe_ep.md | 12 +++++---- .../cudnn/moe_ep/_megamoe_backend/README.md | 2 ++ .../mxfp8/_training_weights.py | 10 +++++-- python/cudnn/moe_ep/_validation.py | 7 +++++ test/python/moe_ep/test_moe_ep_backward.py | 27 +++++++++++++++++++ 5 files changed, 51 insertions(+), 7 deletions(-) diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md index a3015286c..47ea7d658 100644 --- a/docs/fe-oss-apis/moe_ep.md +++ b/docs/fe-oss-apis/moe_ep.md @@ -283,11 +283,13 @@ they are not revalidated by host code after graph capture. - `backward_w2_transpose`: `(E_local, H, I)` - `backward_w1_transpose`: `(E_local, 2I, H)` -Each data and scale tensor must be contiguous or use compact K-major strides -`(K*N, 1, K)`, reside on one device, and use logical block axis 1. The K-major -form lets framework integrations bind transposed weight views without an extra -copy. Plain FP16 operands are accepted by inference staging, but fixed-resource -training accepts only BF16 or FP32 `activation` and `grad_output`. +Forward data and all scale tensors must be contiguous or use compact K-major +strides `(K*N, 1, K)`. Backward-transpose data must be contiguous to match the +training AOT kernel signature. All tensors reside on one device and use logical +block axis 1. The K-major forward form lets framework integrations bind +transposed weight views without an extra copy. Plain FP16 operands are accepted +by inference staging, but fixed-resource training accepts only BF16 or FP32 +`activation` and `grad_output`. Replacing weight storage requires preparing resources and capturing again. Callers must establish stream/event ordering for in-place weight updates. diff --git a/python/cudnn/moe_ep/_megamoe_backend/README.md b/python/cudnn/moe_ep/_megamoe_backend/README.md index 0b2c5e6d9..31391001d 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/README.md +++ b/python/cudnn/moe_ep/_megamoe_backend/README.md @@ -76,6 +76,8 @@ forward, backward, and WGrad-export kernels are compiled before capture. `MoeEpTrainingWeights` contains four address-stable MXFP8 block-scaled tensors: forward W1/W2 and independently quantized backward W2-transpose/W1-transpose. +Backward-transpose data is C-contiguous under both FC1 layout policies, matching +the single training AOT signature. With `weight_interleave_size=32`, compact K-major forward weights and contiguous backward transposes are interpreted as already using 32-element W1 gate/up strips, and the kernels alias weight data directly; only scales require diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py index 02f946f83..d1e3fe882 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py @@ -239,7 +239,10 @@ def __init__( fc1_weight=( bwd_w2t.data if self._uses_direct_weight_bindings - else torch.empty_like(bwd_w2t.data) + else torch.empty_like( + bwd_w2t.data, + memory_format=torch.contiguous_format, + ) ), fc1_weight_sf=_empty_blocked_scales( bwd_w2t, @@ -250,7 +253,10 @@ def __init__( fc2_weight=( bwd_w1t.data if self._uses_direct_weight_bindings - else torch.empty_like(bwd_w1t.data) + else torch.empty_like( + bwd_w1t.data, + memory_format=torch.contiguous_format, + ) ), fc2_weight_sf=_empty_blocked_scales( bwd_w1t, diff --git a/python/cudnn/moe_ep/_validation.py b/python/cudnn/moe_ep/_validation.py index 1370b55e6..31ea364e6 100644 --- a/python/cudnn/moe_ep/_validation.py +++ b/python/cudnn/moe_ep/_validation.py @@ -295,6 +295,13 @@ def is_compact_k_major(tensor: torch.Tensor) -> bool: raise TypeError(f"{name} must be an MXFP8 BlockScaledTensor for " "fixed training resources") if tensor.format is not MoeFormat.MXFP8: raise NotImplementedError(f"{name} must use format='mxfp8', got {tensor.format.value!r}") + if name in ( + "weights.backward_w2_transpose", + "weights.backward_w1_transpose", + ) and not tensor.data.is_contiguous(): + raise ValueError( + f"{name} data must be contiguous for fixed training weight binding" + ) if not has_supported_layout(tensor.data) or not has_supported_layout(tensor.scale): raise ValueError( f"{name} data and scale must be contiguous or compact K-major " diff --git a/test/python/moe_ep/test_moe_ep_backward.py b/test/python/moe_ep/test_moe_ep_backward.py index 5d8d78d33..eaf30b891 100644 --- a/test/python/moe_ep/test_moe_ep_backward.py +++ b/test/python/moe_ep/test_moe_ep_backward.py @@ -445,6 +445,31 @@ def test_validate_training_weights_accepts_compact_k_major_views(part): bindings.refresh() +@pytest.mark.L1 +@pytest.mark.parametrize( + "field", + ["backward_w2_transpose", "backward_w1_transpose"], +) +@pytest.mark.parametrize("weight_interleave_size", [None, 32]) +def test_validate_training_weights_rejects_k_major_backward_data( + field, + weight_interleave_size, +): + weights, _, _ = _training_weight_defect( + _training_weights(), + field, + "data_noncontiguous", + ) + with pytest.raises( + ValueError, + match=rf"weights\.{field} data must be contiguous", + ): + validate_training_weights( + _training_config(weight_interleave_size=weight_interleave_size), + weights, + ) + + def _operator(**overrides) -> MoeEp: values = { "num_experts": 2, @@ -711,6 +736,8 @@ def compact_k_major(tensor): compatibility_bindings.forward.fc1_weight.data_ptr() != weights.forward_fc1.data.data_ptr() ) + assert compatibility_bindings.backward.fc1_weight.is_contiguous() + assert compatibility_bindings.backward.fc2_weight.is_contiguous() bindings = Mxfp8TrainingWeightBindings( weights, From 873270d873907a24bd5139531c15cba8d21cc642 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Wed, 2 Sep 2026 07:23:41 -0700 Subject: [PATCH 29/31] add stateless MoeEP training API Use caller-owned native weights and output buffers to decouple forward and backward execution and support MXFP8 CUDA graph replay. --- docs/fe-oss-apis/moe_ep.md | 654 ++--- docs/fe-oss-apis/overview.md | 5 +- docs/operations/moe_ep.md | 73 +- python/cudnn/__init__.py | 30 +- python/cudnn/moe_ep/__init__.py | 30 +- python/cudnn/moe_ep/_math.py | 17 + .../cudnn/moe_ep/_megamoe_backend/README.md | 203 +- .../moe_ep/_megamoe_backend/_workspace.py | 23 +- .../moe_ep/_megamoe_backend/mxfp8/_backend.py | 35 +- .../mxfp8/_backward_compile.py | 28 +- .../mxfp8/_backward_launch.py | 57 - .../moe_ep/_megamoe_backend/mxfp8/_compile.py | 41 +- .../_megamoe_backend/mxfp8/_compile_common.py | 62 + .../mxfp8/_training_execute.py | 316 ++- .../mxfp8/_training_resources.py | 700 ++--- .../_megamoe_backend/mxfp8/_training_stage.py | 2 +- .../mxfp8/_training_stage_kernel.py | 2 +- .../mxfp8/_training_weights.py | 493 ++-- .../_megamoe_backend/mxfp8/_training_wgrad.py | 215 +- .../mxfp8/_training_wgrad_kernel.py | 145 - python/cudnn/moe_ep/_types.py | 362 ++- python/cudnn/moe_ep/_validation.py | 525 +++- python/cudnn/moe_ep/api.py | 472 +++- .../moe_ep/moe_ep_distributed_workers.py | 162 +- test/python/moe_ep/moe_ep_test_support.py | 705 +---- .../moe_ep/probe_moe_ep_training_graph.py | 1561 ++-------- test/python/moe_ep/test_moe_ep_backward.py | 2516 ++++++----------- test/python/moe_ep/test_moe_ep_cutedsl.py | 61 +- test/python/moe_ep/test_moe_ep_forward.py | 42 +- test/python/moe_ep/test_moe_ep_multinode.py | 23 +- 30 files changed, 3507 insertions(+), 6053 deletions(-) create mode 100644 python/cudnn/moe_ep/_math.py delete mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py create mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile_common.py delete mode 100644 python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md index 47ea7d658..1e40d278e 100644 --- a/docs/fe-oss-apis/moe_ep.md +++ b/docs/fe-oss-apis/moe_ep.md @@ -1,110 +1,28 @@ # Mixture of Experts with Expert Parallelism -`cudnn.moe_ep` provides a fused SwiGLU MoE implementation for Rubin SM107. -Experts are sharded contiguously across an optional expert-parallel process -group. This page documents the Python API and lifecycle. See the -[MoeEP operation reference](../operations/moe_ep.md) for supported -architectures, data formats, tensor contracts, and expert-parallel topology. +`cudnn.moe_ep` provides Rubin SM107 fused SwiGLU MoE execution with optional +expert parallelism. Inference and training share one `MoeEp` object but use +separate call surfaces: -## Installation +- `MoeEp.__call__` and `warmup` for inference; +- `prepare_training`, `training_forward`, and `training_backward` for training. + +Training is stateless with respect to caller tensors. The operator retains +compiled kernels, runtime state, and private per-lane NVSHMEM scratch, but it +does not retain weights, forward state, output buffers, or fallback weight +staging. -Install the reusable CuTeDSL and communication extras, then the PyTorch -integration dependencies: +## Installation ```bash pip install "nvidia-cudnn-frontend[cutedsl,comm]" torch torch-c-dlpack-ext ``` -The package keeps the general CuTeDSL installation floor at 4.5.0. Rubin -MegaMoE execution checks for `nvidia-cutlass-dsl>=4.8.0` when preparing its -kernels. EP2+ additionally requires an initialized NCCL process group and a -usable NVSHMEM peer topology. - -## Public API and constructor - -The public surface exports `MoeEp`, `MoeEpTrainingWeights`, -`MoeEpTrainingResources`, `MoeEpTrainingSlot`, `MoeEpExecutionLane`, -`MoeEpTrainingWgradOperands`, `BlockScaledTensor`, `MoeFormat`, and -`MoeEpTuningConfig`. - -The `MoeEp` constructor accepts: - -| Parameter | Current contract | -| --- | --- | -| `num_experts` | Positive global expert count; divisible by EP size | -| `hidden_size` | Positive and divisible by 128 | -| `intermediate_size` | Positive and divisible by 256 | -| `top_k` | Positive and no larger than 32 or `num_experts` | -| `ep_group` | Optional initialized `torch.distributed.ProcessGroup`; `None` selects EP1 | -| `max_tokens_per_rank` | Required by the executable backend and must be positive | -| `max_recv_size_per_rank` | Optional positive receive-pool capacity | -| `drop_on_overflow` | `False` by default; selects fatal-assert versus reporting/drop policy | -| `output_format` | `"bf16"` only for current execution | -| `combine_format` | `"bf16"` or `"mxfp8"` | -| `apply_topk_in_fc1` | Must be `True` | -| `weight_interleave_size` | `None` for conventional gate-then-up weights or `32` for pre-interleaved MXFP8 weights | -| `gate_up_clamp` | Optional finite clamp magnitude | -| `token_padding_size` | Positive; training fixed resources use 128 internally | -| `sf_padding_size` | Positive multiple of 128; training fixed resources use 128 internally | -| `tuning` | Optional `MoeEpTuningConfig`; must match on every EP rank | - -When `max_recv_size_per_rank` is omitted, the backend allocates for the -worst-case receive count: - -```text -ep_size * max_tokens_per_rank * top_k -``` - -An explicit value is capped at that same worst-case count. +Rubin MegaMoE requires `nvidia-cutlass-dsl>=4.8.0`. EP2+ also requires an +initialized NCCL process group and an NVSHMEM topology in which all +participating ranks are directly peer-addressable. -## Breaking training API migration - -This release removes the legacy dynamic compact training API: - -- `MoeEp.backward(...)` -- constructor arguments `generate_c` and `backward_wgrad_mode` -- forward returns containing compact `fc1_c` and `route_metadata` -- `MoeEpWgradForwardStash` and `MoeEpWgradOperands` - -Old: - -```python -output, fc1_c, route_metadata = op( - activation, w1, w2, topk_idx, topk_weights -) -dx, dprob = op.backward( - grad_output, w1, w2, topk_idx, topk_weights, fc1_c, route_metadata -) -``` - -New: - -```python -resources = op.prepare_training_resources( - training_weights, - slot_count=2, - lane_count=1, -) -slot = resources.slots[0] -lane = resources.lanes[0] -resources.refresh_weights() -output = resources.forward( - slot, lane, activation, topk_idx, topk_weights -) -dx, dprob, operands = resources.backward(slot, lane, grad_output) -overflow = resources.finalize_overflow((slot,), lane) -``` - -`dprob` now follows the MXFP8-staged kernel numerical contract and relaxed -atomic accumulation order. Dynamic inputs use a trusted-caller contract. -Distributed graph support requires one direct-P2P MNNVL domain, and -distributed lanes must be ordered consistently across ranks with captured -events. The fixed-capacity WGrad result is a producer ABI; no specific grouped -WGrad consumer is guaranteed in this release. - -## Inference forward - -`MoeEp.__call__` is the inference-forward surface: +## Constructing the operator ```python from cudnn import MoeEp @@ -117,390 +35,240 @@ op = MoeEp( ep_group=ep_group, max_tokens_per_rank=max_tokens, max_recv_size_per_rank=recv_capacity, -) - -output = op( - activation, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights, + drop_on_overflow=False, + output_format="bf16", + combine_format="bf16", + apply_topk_in_fc1=True, + weight_interleave_size=32, ) ``` -`activation`, `fc1_weight`, and `fc2_weight` may independently be plain BF16, -FP16, or FP32 tensors, or MXFP8 `BlockScaledTensor` values. The logical shapes -are: +Native training requires `weight_interleave_size=32`. FC1 payloads then use +alternating 32-element gate/up strips. -- `activation`: `(T, H)` -- `fc1_weight`: `(E_local, H, 2I)` -- `fc2_weight`: `(E_local, I, H)` -- `topk_idx` and `topk_weights`: `(T, K)` +## Stateless training preparation -MXFP8 operands are block-scaled along logical axis 1. The output always has -shape `(T, H)` and dtype `torch.bfloat16`. This surface does not return compact -FC1 or route metadata stashes and does not provide backward. - -### Inference CUDA Graph capture - -Call `warmup` with the exact tensors that will be captured: +Preparation is collective over `ep_group` and must run outside CUDA Graph +capture: ```python -import torch +requirements = op.prepare_training( + lane_count=1, + device=None, # current CUDA device; pass an explicit device for multi-GPU hosts +) +lane = op.training_lanes[0] +``` -op.warmup(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) -if ep_group is not None: - torch.distributed.barrier(group=ep_group) +`prepare_training` does not accept or bind weights. It returns a plain mapping +whose values are: -graph = torch.cuda.CUDAGraph() -with torch.cuda.graph(graph): - graph_output = op( - activation, fc1_weight, fc2_weight, topk_idx, topk_weights - ) +```text +(shape, stride, dtype, alignment_bytes) ``` -`warmup` completes runtime bootstrap, symmetric allocation, weight staging, -JIT compilation, and one real launch. It is collective by contract for EP2+ -but intentionally does not issue a process-group barrier. Captured inference -weights must expose usable PyTorch version counters and must match the warmed -weight cache. Replay may update captured tensor contents in place but may not -replace their storage. +The mapping contains `output`, `fc1_preact`, `fc1_a`, `fc1_sfa`, +`valid_route_counts`, `expert_offsets`, `grad_activation`, `dprob`, `fc1_b`, +`fc1_sfb`, `fc2_a`, `fc2_sfa`, `fc2_b`, and `fc2_sfb`. TE allocates these +buffers and passes them to each invocation. cuDNN validates exact shape, +stride, dtype, alignment, device, and non-aliasing before launch. -`MoeEp` supports explicit `close()` and context-manager use. One instance is -bound to one CUDA device. Do not close an operator while a stream is capturing -or while graph work using its resources remains outstanding. +`device=None` binds the current CUDA device. An explicit CUDA device takes +precedence. Every later training tensor must use the bound device. -## Fixed-resource training +## Native weight ABI -Training uses `prepare_training_resources`: +Forward and backward receive independent packs: ```python -from cudnn import MoeEpTrainingWeights - -weights = MoeEpTrainingWeights( - forward_fc1=forward_fc1_mxfp8, - forward_fc2=forward_fc2_mxfp8, - backward_w2_transpose=backward_w2t_mxfp8, - backward_w1_transpose=backward_w1t_mxfp8, -) - -resources = op.prepare_training_resources( - weights, - slot_count=2, - lane_count=1, +from cudnn import ( + MoeEpNativeForwardWeights, + MoeEpNativeBackwardWeights, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, ) - -slot0, slot1 = resources.slots -lane0 = resources.lanes[0] - -# Required after each in-place source-weight update and before the first -# forward/backward that consumes that version. -resources.refresh_weights() - -y0 = resources.forward( - slot0, - lane0, - activation0, - topk_idx0, - topk_weights0, -) -dx0, dprob0, operands0 = resources.backward(slot0, lane0, grad_output0) -overflow = resources.finalize_overflow((slot0,), lane0) ``` -`prepare_training_resources` is collective across `ep_group` and must execute -outside CUDA Graph capture. All ranks must use matching static configuration, -slot/lane counts, and tuning. The training backend internally enables FC1 -preactivation generation and fixed-capacity WGrad operands, and fixes token and -scale-factor padding to 128. - -`forward`, `backward`, and `finalize_overflow` enqueue the same device -operations in ordinary execution and in a caller-owned CUDA Graph. The caller -owns capture, replay, stream/event dependencies, slot reuse, and resource -lifetime. - -A `MoeEp` instance can own only one training-resource set. Closing that set is -terminal for the operator: replacing source-weight storage requires a new -operator, new resources, and new graph captures. - -### Slots and lanes - -A persistent slot owns state that survives from matching forward to backward: - -- routing indices and weights -- pool-native FC1 preactivation -- expert counts and padded offsets -- kernel dprob -- backward auxiliaries and outputs -- fixed-capacity WGrad operands -- per-slot overflow flags - -An execution lane owns mutable router, protocol, and kernel scratch. Multiple -active streams require distinct lanes. Distributed MegaMoE communication -kernels must be ordered identically on every rank with captured CUDA events; -the kernels cannot be launched in unordered concurrent lanes. - -All peer-visible regions are constructed in deterministic order. Their sizes -are validated and normalized across ranks before symmetric allocation so every -named region has the same peer offset. - -## Tensor contracts - -For inference through `MoeEp.__call__`, `activation` has: - -- shape `(T, H)` -- BF16, FP16, FP32, or MXFP8 block-scaled input - -Inference `fc1_weight` and `fc2_weight` accept the same plain or MXFP8 operand -families. MXFP8 activation and weights must be represented by -`BlockScaledTensor` with logical block axis 1. - -The inference `topk_idx` contract is: +Each `MoeEpNativeWeight` contains: -- shape `(T, K)` -- Int32 or Int64 -- each element is `-1` or a valid global expert ID +- `payload`: kernel-native E4M3 data; +- `scale`: contiguous Rubin-blocked E8M0 scales; +- `layout_id`: the exact versioned payload-and-scale layout. -The inference `topk_weights` contract is: +Execution validates the `layout_id` and passes payload and scale pointers to +the kernel without transformation or retention. Eager calls may use different +weight addresses. CUDA Graph capture pins every referenced address until the +graph executable is destroyed. -- shape `(T, K)` -- floating point - -Fixed-resource training uses a narrower, graph-stable staging ABI: - -- contiguous BF16 or FP32 `activation` and `grad_output`, each shaped `(T, H)`; -- contiguous Int32 `topk_idx` shaped `(T, K)`; -- contiguous FP32 `topk_weights` shaped `(T, K)`; -- all tensors on one device and `T <= max_tokens_per_rank`. - -Expert IDs and finite dynamic values remain a trusted-caller replay contract; -they are not revalidated by host code after graph capture. - -`MoeEpTrainingWeights` contains four MXFP8 block-scaled tensors: - -- `forward_fc1`: `(E_local, H, 2I)` -- `forward_fc2`: `(E_local, I, H)` -- `backward_w2_transpose`: `(E_local, H, I)` -- `backward_w1_transpose`: `(E_local, 2I, H)` - -Forward data and all scale tensors must be contiguous or use compact K-major -strides `(K*N, 1, K)`. Backward-transpose data must be contiguous to match the -training AOT kernel signature. All tensors reside on one device and use logical -block axis 1. The K-major forward form lets framework integrations bind -transposed weight views without an extra copy. Plain FP16 operands are accepted -by inference staging, but fixed-resource training accepts only BF16 or FP32 -`activation` and `grad_output`. - -Replacing weight storage requires preparing resources and capturing again. -Callers must establish stream/event ordering for in-place weight updates. - -### Explicit weight refresh contract - -With `weight_interleave_size=32`, Rubin training recognizes compact K-major -forward views plus contiguous backward transposes as pre-interleaved. For that -form, weight data is bound directly and FC1 gate/up values must use alternating -32-element strips. Layout alone never selects this semantic convention. -`resources.refresh_weights()` then swizzles only scales into fixed-address, -kernel-native buffers and never copies the weight payload. Existing contiguous -public packs with the default `None` retain conventional gate-then-up semantics -and the compatible data-and-scale staging path. Plain BF16/FP16/FP32 inference -weights cannot use `weight_interleave_size=32`. +Let `B(R, C) = round_up(R, 128) * round_up(C, 4)`. The native V1 contracts +are: -The caller must obey all of the following: +- forward FC1: payload `(E_local, H, 2I)`, stride `(2HI, 1, H)`, with + gate/up 32-column strips; scale `(E_local, B(2I, H/32))`; +- forward FC2: payload `(E_local, I, H)`, stride `(IH, 1, I)`; scale + `(E_local, B(H, I/32))`; +- backward W2-transpose: contiguous payload `(E_local, H, I)`; scale + `(E_local, B(I, H/32))`; +- backward W1-transpose: contiguous payload `(E_local, 2I, H)`, with + gate/up 32-row strips; scale `(E_local, B(H, 2I/32))`. -- update both the data and scale tensors in place; their storage addresses, - shape, stride, dtype, device, and capacity must remain unchanged; -- call `resources.refresh_weights()` after every source-weight update and - before any forward or backward that consumes the new version; -- establish stream ordering from the weight update to the refresh and from the - refresh to the first consumer, using the same stream or CUDA events; -- do not refresh between a matching forward and backward; both operations must - observe the same four-tensor weight version; -- do not overlap a refresh with any forward/backward that reads the shared - internal weight bindings, including operations using another slot or lane; -- replace any source storage only by closing the existing operator, creating a - new `MoeEp` instance and resources, and recapturing every graph that - references them. Closed resources cannot be reopened or replaced on the same - operator. +Every native scale tensor is contiguous E8M0. The corresponding +`MoeEpNativeWeightLayout` enum value is required; a compact or differently +swizzled scale tensor is rejected even when its element count matches. -For CUDA Graph execution, capture the refresh at the appropriate update -boundary: +When upstream does not already produce native weights, use caller-owned +staging: ```python -with torch.cuda.graph(graph, stream=stream): - # An optimizer or external producer must complete its in-place updates - # before this node. - resources.refresh_weights() - y = resources.forward(slot, lane, x, topk_idx, topk_weights) - dx, dprob, operands = resources.backward(slot, lane, grad_output) - overflow = resources.finalize_overflow((slot,), lane) +native_fw = op.pack_forward_weights(source_fw, out=forward_staging) +native_bw = op.pack_backward_weights(source_bw, out=backward_staging) ``` -Replay then executes the captured device refresh; Python is not called during -replay. Activation `x` has the same storage rule under CUDA Graph capture: -its contents may change in place, but replacing its captured storage requires -recapture. - -## Backward outputs - -Fixed-resource backward returns: - -- `grad_activation`: fixed-slot `(T, H)` FP32 view -- `dprob`: source-order `(T, K)` kernel dprob -- `MoeEpTrainingWgradOperands` - -Kernel dprob follows the MXFP8-staged backward numerical contract and relaxed -atomic accumulation order. Bitwise determinism is not guaranteed. - -`MoeEpTrainingWgradOperands` contains: - -- FC1 operands: `fc1_a`, `fc1_sfa`, `fc1_b`, `fc1_sfb` -- FC2 operands: `fc2_a`, `fc2_sfa`, `fc2_b`, `fc2_sfb` -- `expert_offsets` and `valid_route_counts` - -These tensors have fixed addresses and fixed capacity. Device -`expert_offsets`/`valid_route_counts` describe the live expert segments and -padding rows are zero. This release guarantees the producer ABI only; a -specific grouped-WGrad consumer is future integration work. The result is not -a pair of dense gradients that can be passed directly to an optimizer. - -## Tuning - -`MoeEpTuningConfig` exposes semantic-preserving performance controls: - -- `token_back_mode`: `"epi_warps"`, `"standalone_warps"`, or - `"reuse_dispatch_warps"` -- `epi_flag_batch`: one of the validated `(M, N)` flag-batch pairs -- `token_in_flag_batch`: `1`, `2`, `4`, `8`, or `16` -- `group_hint`: `None`, `64`, `128`, `256`, `512`, `768`, or `1024` -- `reduce_topk_in_kernel`: Boolean - -Every rank in an EP group must use the same tuning configuration. -`reduce_topk_in_kernel=True` requires BF16 combine/output, -`apply_topk_in_fc1=True`, and `token_back_mode="epi_warps"`. - -## Capacity and overflow - -`max_recv_size_per_rank` defines bounded receive capacity. Resources cannot -grow during capture. A capacity change requires preparation and recapture. - -Inference checks its per-call overflow result after the fused launch. The -fixed-resource training transport deterministically truncates overflow so all -ranks complete the communication protocol. `finalize_overflow` combines the -forward and backward flags for all selected slots and performs one captured -scalar MAX all-reduce for EP2+. - -With `drop_on_overflow=True`, `finalize_overflow` returns a one-element Int32 -CUDA tensor: zero means no overflow and nonzero means truncation occurred. -Dropped routes contribute zero. With `drop_on_overflow=False`, overflow is a -fatal device assertion; this mode requires `torch._assert_async`, and EP2+ -requires an NCCL process group. The training transport still truncates first -to let every rank finish the protocol before the public policy is applied at -graph tail. - -## CUDA Graph execution - -For inference: +The equivalent standalone `pack_forward_weights` and `pack_backward_weights` +functions are also exported. Packing allocates nothing: every transformed +payload or scale is written to the supplied staging bundle. These fallback +packers consume logical gate-then-up `MoeEpForwardWeights` / +`MoeEpBackwardWeights` with compact axis-1 scales; already interleaved, +blocked producers should construct the native packs directly instead of +packing them again. -1. all ranks call `MoeEp.warmup` with the exact capture bindings; -2. the caller aligns ranks after warmup; -3. each rank captures its forward graph; -4. graph execs are replayed in the same cross-rank order. +## Forward -For fixed-resource training: - -1. all ranks collectively call `prepare_training_resources`; -2. all ranks perform an ordinary - `refresh_weights -> forward -> backward -> finalize_overflow` warmup so - every scale-staging, MegaMoE, and WGrad-export kernel is compiled; -3. each rank captures its outer graph; -4. ranks align after capture; -5. graph execs are submitted in lockstep without host synchronization inside - a replay burst; -6. all stream work completes before resources are closed. - -For EP2+: - -Distributed MegaMoE launches must have the same order on every rank. -Independent lane storage does not permit unordered collective-kernel overlap. -Use distinct lanes for simultaneously active streams and captured CUDA events -to impose the same cross-stream order on every rank. - -Each graph binds fixed tensor shapes, addresses, slots, lanes, and token -extent. Dynamic routing values may change in place; dynamic shapes may not. -Different token extents may share prepared resources, but each extent must be -warmed and captured as its own graph specialization. - -## NVSHMEM topology - -MegaMoE kernels use direct symmetric peer pointers. With -`NVSHMEM_REMOTE_TRANSPORT=none`, every EP rank must appear in the P2P connected -list and `NVSHMEM_TEAM_SHARED` must span the complete EP world. - -Selecting `ibrc` does not by itself make a non-P2P peer directly addressable. -Cross-MNNVL execution is not part of the current support matrix. - -Global expert `e` belongs to group-relative EP rank -`e // experts_per_rank`. Noncontiguous global-rank process groups are accepted -when their group-relative topology and direct peer access are valid. When -creating multiple subgroups, every world rank must create them in the same -order. - -## Validation coverage - -The following configurations are exercised by the current test and probe -suite. This list describes validation coverage and does not broaden the -topology contract beyond one direct-P2P MNNVL domain: - -- EP1 inference, fixed-resource training, overflow, and CUDA Graph replay -- single-node EP2/EP3/EP4 inference -- single-node EP2/EP4 training -- noncontiguous EP2 inference and training subgroups -- multi-node forward acceptance for EP4/EP6/EP12/EP16 -- multi-node backward acceptance for EP8/EP16/EP32 -- fixed-resource CUDA Graph launchers for EP8/EP16/EP32 - -## Validation - -Run host-side and local tests: +```python +from cudnn import MoeEpTrainingForwardOutputs -```bash -python -m pytest \ - test/python/moe_ep/test_moe_ep_cutedsl.py \ - test/python/moe_ep/test_moe_ep_forward.py \ - test/python/moe_ep/test_moe_ep_backward.py \ - -m L0 +y = op.training_forward( + lane, + activation, + topk_idx, + topk_weights, + weights=native_fw, + out=MoeEpTrainingForwardOutputs( + output=y_out, + fc1_preact=fc1_preact, + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + valid_route_counts=valid_route_counts, + expert_offsets=expert_offsets, + ), +) ``` -Run SM107 single-node distributed tests from an exclusive GPU allocation: +`activation` may be contiguous BF16/FP32 or an axis-1 MXFP8 +`BlockScaledTensor`. MXFP8 input bypasses the BF16-to-MXFP8 quantization +stager. Routing still copies data into private symmetric memory because remote +ranks address that memory directly. -```bash -python -m pytest \ - test/python/moe_ep/test_moe_ep_forward.py \ - test/python/moe_ep/test_moe_ep_backward.py \ - -m L1 -``` +`fc1_preact` is required because the training forward kernel always runs with +`generate_c=True`; TE must provide its destination and retain it through the +matching backward. `fc1_a`, `fc1_sfa`, `valid_route_counts`, and +`expert_offsets` are also required caller-owned destinations after +`prepare_training()`. -`test/python/moe_ep/test_moe_ep_multinode.py` is torchrun-native and requires -the standard `LOCAL_RANK`, `LOCAL_WORLD_SIZE`, `RANK`, and `WORLD_SIZE` -environment. The multi-node fixture initializes NCCL and defaults -`NVIDIA_IMEX_CHANNELS=0`. +`output` is required. The return is a logical `(T, H)` view of that +caller-owned capacity buffer. -Run distributed fixed-resource probes from an existing Slurm allocation: +## Backward and WGrad -```bash -NVSHMEM_REMOTE_TRANSPORT=none \ -data/script/run_moe_ep_forward_multinode_slurm.sh backward-ep8 +```python +from cudnn import MoeEpTrainingBackwardOutputs -NVSHMEM_REMOTE_TRANSPORT=none \ -data/script/run_moe_ep_forward_multinode_slurm.sh graph-ep8 +dx, dprob, operands = op.training_backward( + lane, + grad_output, + topk_idx, + topk_weights, + weights=native_bw, + fc1_preact=fc1_preact, + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + valid_route_counts=valid_route_counts, + expert_offsets=expert_offsets, + out=MoeEpTrainingBackwardOutputs( + grad_activation=dx_out, + dprob=dprob_out, + fc1_b=fc1_b, + fc1_sfb=fc1_sfb, + fc2_a=fc2_a, + fc2_sfa=fc2_sfa, + fc2_b=fc2_b, + fc2_sfb=fc2_sfb, + ), +) ``` -The launcher provides `backward-ep16`, `backward-ep32`, `graph-ep16`, and -`graph-ep32` for the larger EP configurations. The fatal captured-overflow -assertion is a separate `graph-ep8-error` expected-failure task. - -The graph tasks invoke `test/python/moe_ep/probe_moe_ep_training_graph.py`, -which covers collective warmup, capture alignment, lockstep replay bursts, -dynamic routing/overflow recovery, ordered multi-lane execution, and -collective teardown. +`grad_output` has the same BF16/FP32/MXFP8 input choices as forward. +`fc1_preact` and the four forward WGrad values are required and passed +explicitly because cuDNN does not retain the forward output bundle. + +`grad_activation` and `dprob` are required caller-owned destinations. + +All six backward WGrad fields are required. `operands` is always a +`MoeEpTrainingWgradOperands` containing non-owning views of the exact caller +buffers. + +The producer-native ABI is directly consumable by the separately invoked +grouped WGrad kernel. Here `K_pool` is the fixed routed-token pool capacity, +not the model's top-k value: + +- `fc1_b` remains gate/up-interleaved with shape `(K_pool, 2I)` and stride + `(2I, 1)`; +- `fc1_a` and `fc2_a` use the advertised transpose-view layouts; +- all four scale tensors are written in the final grouped-WGrad 128x4 + interleaved layout; +- no public compact scale, deinterleave copy, physical transpose, slot export, + or scale-expansion kernel is used. + +## Ownership and lifetime + +- TE owns all native weights, output bundles, saved forward state, WGrad + operands, and optional pack staging. +- cuDNN borrows these tensors for one call and does not cache their Python + objects or pointers. +- TE must provide `fc1_preact` to forward and keep it live through the matching + backward; cuDNN has no private preactivation fallback or workspace alias. +- Forward WGrad outputs, segment metadata, and backward WGrad outputs remain + live until the independent grouped WGrad consumer completes. +- cuDNN owns private per-lane local and NVSHMEM symmetric scratch. +- One lane may be active on only one stream at a time. +- All EP ranks must submit distributed forward/backward calls in identical + order. +- The caller owns forward/backward weight-version consistency. +- `MoeEp.close()` releases only private runtime resources and never clears or + frees caller memory. + +## Overflow + +Overflow is private per-launch state. Each forward and backward applies the +configured policy before returning; there is no public overflow tensor or +`finalize_overflow` method. EP2+ retains the scalar MAX reduction required to +make the policy rank-consistent. + +## CUDA Graph capture + +1. Collectively call `prepare_training`. +2. Allocate every capture binding from the returned requirements. +3. Materialize or provide native weights at stable addresses. +4. Run ordinary forward/backward warmups for every captured specialization. +5. Capture calls using every caller-owned destination returned by + `prepare_training()`, including primary outputs, saved forward state, and + forward/backward WGrad tensors. +6. Keep all captured input, output, saved-state, staging, and native-pack + addresses stable until every referencing graph executable is destroyed. + +Dynamic contents may change at fixed addresses. Eager invocations may replace +addresses between calls. + +## Breaking migration + +Removed: + +- `MoeEpTrainingResources` +- `MoeEpTrainingSlot` +- `MoeEpTrainingWeights` +- `prepare_training_resources` +- `refresh_weights` +- `finalize_overflow` + +The old resource-owned forward/backward state is replaced by explicit +per-invocation native weight packs and caller-owned output buffers. No +compatibility shim is retained. diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 0243082a2..6583aab4d 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -37,7 +37,7 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d - [SDPA Backward (SM120)](attention/sdpa_bwd_sm120.md) - [RMSNorm + SiLU](rmsnorm_silu.md) - [MoE + Expert Parallel API](moe_ep.md) — Rubin SM107 fused SwiGLU with - fixed-resource training and CUDA Graph support; see the + stateless caller-owned training buffers and CUDA Graph support; see the [MoeEP operation reference](../operations/moe_ep.md) for support details ## Installation and setup @@ -72,7 +72,8 @@ Each operation exposes two APIs: MoeEP is an exception to the generic wrapper/kernel pattern below. It exposes an object API: `MoeEp.__call__` for inference and -`MoeEp.prepare_training_resources` for fixed-resource training. See +`prepare_training` plus stateless `training_forward`/`training_backward` calls +for training. See [MoE + Expert Parallel API](moe_ep.md) for its lifecycle and CUDA Graph contract. diff --git a/docs/operations/moe_ep.md b/docs/operations/moe_ep.md index e5bbfad56..d2ae0ba87 100644 --- a/docs/operations/moe_ep.md +++ b/docs/operations/moe_ep.md @@ -91,34 +91,35 @@ output = op( For inference CUDA Graph capture, call `op.warmup(...)` with the exact bindings before capture. `MoeEp` supports `close()` and context-manager use. -Fixed-resource training uses the same operator and binds graph-stable weights, -slots, and execution lanes: +Stateless training prepares only private execution lanes. Every invocation +receives independent native weights and caller-owned outputs: ```python -from cudnn import MoeEpTrainingWeights +requirements = op.prepare_training(lane_count=1, device=device) +lane = op.training_lanes[0] -weights = MoeEpTrainingWeights( - forward_fc1=forward_fc1_mxfp8, - forward_fc2=forward_fc2_mxfp8, - backward_w2_transpose=backward_w2t_mxfp8, - backward_w1_transpose=backward_w1t_mxfp8, +output = op.training_forward( + lane, activation, topk_idx, topk_weights, + weights=native_forward_weights, + out=forward_outputs, ) -resources = op.prepare_training_resources(weights, slot_count=2, lane_count=1) -slot, lane = resources.slots[0], resources.lanes[0] - -resources.refresh_weights() -output = resources.forward(slot, lane, activation, topk_idx, topk_weights) -grad_activation, dprob, wgrad_operands = resources.backward( - slot, lane, grad_output +grad_activation, dprob, wgrad_operands = op.training_backward( + lane, grad_output, topk_idx, topk_weights, + weights=native_backward_weights, + fc1_preact=forward_outputs.fc1_preact, + fc1_a=forward_outputs.fc1_a, + fc1_sfa=forward_outputs.fc1_sfa, + valid_route_counts=forward_outputs.valid_route_counts, + expert_offsets=forward_outputs.expert_offsets, + out=backward_outputs, ) -overflow = resources.finalize_overflow((slot,), lane) ``` The WGrad result is a fixed-capacity grouped-GEMM operand bundle, not dense optimizer-ready weight gradients. See the detailed [MoE + Expert Parallel API](../fe-oss-apis/moe_ep.md) reference for -installation, all constructor arguments, tensor formats, training resource -lifecycle, tuning, overflow handling, and CUDA Graph requirements. MoeEP is +installation, all constructor arguments, native layouts, buffer ownership, +overflow handling, and CUDA Graph requirements. MoeEP is distinct from the cuDNN graph [MoE Grouped Matmul](MoeGroupedMatmul.md) operation. @@ -137,7 +138,7 @@ operation. - `num_experts` divisible by the expert-parallel group size. - An explicit positive `max_tokens_per_rank`. -The fixed-resource CUDA Graph path has hardware acceptance through EP32 when +The stateless training CUDA Graph path has hardware acceptance through EP32 when all ranks are in one direct-P2P MNNVL peer-access domain. The Python capability layer does not impose an EP-size ceiling; cross-MNNVL execution is not part of the validated support surface. @@ -153,9 +154,9 @@ The current executable output format is BF16. The expert-combine path accepts BF16 or MXFP8. NVFP4 types are represented by the public API but native NVFP4 operands, combine, and output are not executable by this backend. -Fixed-resource training narrows dynamic activation and gradient inputs to -contiguous BF16 or FP32 tensors. Training weights are contiguous MXFP8 -block-scaled tensors. +Training accepts contiguous BF16/FP32 or MXFP8 block-scaled activation and +gradient inputs. Execution weights use versioned kernel-native E4M3 payload +and Rubin-blocked E8M0 scale layouts. ## Tensor contracts @@ -180,19 +181,21 @@ Inference uses: All inference tensors must reside on one device, and the local token count must satisfy `T <= max_tokens_per_rank`. -Fixed-resource training uses a narrower graph-stable contract: +Stateless training uses: -- `activation` and `grad_output`: contiguous `(T, H)`, BF16 or FP32; +- `activation` and `grad_output`: contiguous `(T, H)`, BF16, FP32, or MXFP8; - `topk_idx`: contiguous `(T, K)`, Int32; - `topk_weights`: contiguous `(T, K)`, FP32; -- forward FC1 and FC2 weights: contiguous MXFP8 block-scaled tensors with - shapes `(E_local, H, 2I)` and `(E_local, I, H)`; -- transposed backward weights: contiguous MXFP8 block-scaled tensors with - shapes `(E_local, H, I)` and `(E_local, 2I, H)`; -- forward output: `(T, H)`, BF16; -- `grad_activation`: fixed-slot `(T, H)`, FP32; -- `dprob`: source-order `(T, K)`, FP32; -- `wgrad_operands`: a fixed-capacity `MoeEpTrainingWgradOperands` bundle. +- independent forward and backward native weight packs with exact versioned + `layout_id` values; +- required caller-owned forward output: `(T, H)`, BF16; +- required caller-owned `fc1_preact`, produced by training forward with + `generate_c=True` and retained through matching backward; +- required caller-owned `grad_activation`: `(T, H)` view of a capacity buffer, + FP32; +- required caller-owned `dprob`: source-order `(T, K)`, FP32; +- required caller-owned WGrad saved state and a fixed-capacity + `MoeEpTrainingWgradOperands` bundle. All dynamic training tensors must reside on one device and satisfy `T <= max_tokens_per_rank`. @@ -204,7 +207,7 @@ EP2+ execution requires: - an initialized NCCL process group; - `nvshmem4py` and usable NVSHMEM libraries; - direct peer access among every pair of participating ranks; and -- consistent rank ordering, resource sizes, tuning, slot selection, and lane +- consistent rank ordering, buffer schemas, tuning, lane selection, and launch ordering across the group. `max_recv_size_per_rank` bounds receive capacity. When omitted, it defaults to @@ -214,5 +217,5 @@ the worst-case route count: ep_size * max_tokens_per_rank * top_k ``` -Resources cannot grow during CUDA Graph replay. Capacity or storage changes -require resource preparation and graph capture again. +Private lane resources cannot grow during CUDA Graph replay. Capacity changes +require a new operator preparation; caller-address changes require recapture. diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index de6eb29ec..cddc99535 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -311,14 +311,23 @@ def _dlopen_cudnn(): "moe_ep", "BlockScaledTensor", "MoeEp", + "MoeEpBackwardWeightStaging", + "MoeEpBackwardWeights", "MoeEpExecutionLane", - "MoeEpTrainingResources", - "MoeEpTrainingSlot", - "MoeEpTrainingWeights", + "MoeEpForwardWeightStaging", + "MoeEpForwardWeights", + "MoeEpNativeBackwardWeights", + "MoeEpNativeForwardWeights", + "MoeEpNativeWeight", + "MoeEpNativeWeightLayout", + "MoeEpTrainingBackwardOutputs", + "MoeEpTrainingForwardOutputs", "MoeEpTrainingWgradOperands", "MoeEpTuningConfig", "MoeFormat", "MoeTensor", + "pack_backward_weights", + "pack_forward_weights", } _LAZY_OPTIONAL_IMPORTS = { @@ -326,10 +335,17 @@ def _dlopen_cudnn(): "moe_ep": (".moe_ep", None), "BlockScaledTensor": (".moe_ep", "BlockScaledTensor"), "MoeEp": (".moe_ep", "MoeEp"), + "MoeEpBackwardWeightStaging": (".moe_ep", "MoeEpBackwardWeightStaging"), + "MoeEpBackwardWeights": (".moe_ep", "MoeEpBackwardWeights"), "MoeEpExecutionLane": (".moe_ep", "MoeEpExecutionLane"), - "MoeEpTrainingResources": (".moe_ep", "MoeEpTrainingResources"), - "MoeEpTrainingSlot": (".moe_ep", "MoeEpTrainingSlot"), - "MoeEpTrainingWeights": (".moe_ep", "MoeEpTrainingWeights"), + "MoeEpForwardWeightStaging": (".moe_ep", "MoeEpForwardWeightStaging"), + "MoeEpForwardWeights": (".moe_ep", "MoeEpForwardWeights"), + "MoeEpNativeBackwardWeights": (".moe_ep", "MoeEpNativeBackwardWeights"), + "MoeEpNativeForwardWeights": (".moe_ep", "MoeEpNativeForwardWeights"), + "MoeEpNativeWeight": (".moe_ep", "MoeEpNativeWeight"), + "MoeEpNativeWeightLayout": (".moe_ep", "MoeEpNativeWeightLayout"), + "MoeEpTrainingBackwardOutputs": (".moe_ep", "MoeEpTrainingBackwardOutputs"), + "MoeEpTrainingForwardOutputs": (".moe_ep", "MoeEpTrainingForwardOutputs"), "MoeEpTrainingWgradOperands": ( ".moe_ep", "MoeEpTrainingWgradOperands", @@ -337,6 +353,8 @@ def _dlopen_cudnn(): "MoeEpTuningConfig": (".moe_ep", "MoeEpTuningConfig"), "MoeFormat": (".moe_ep", "MoeFormat"), "MoeTensor": (".moe_ep", "MoeTensor"), + "pack_backward_weights": (".moe_ep", "pack_backward_weights"), + "pack_forward_weights": (".moe_ep", "pack_forward_weights"), "BSA": (".block_sparse_attention", "BSA"), "block_sparse_attention_forward": (".block_sparse_attention", "block_sparse_attention_forward"), "block_sparse_attention_fp8_forward": (".block_sparse_attention", "block_sparse_attention_fp8_forward"), diff --git a/python/cudnn/moe_ep/__init__.py b/python/cudnn/moe_ep/__init__.py index 8557cff2d..e918a2b37 100644 --- a/python/cudnn/moe_ep/__init__.py +++ b/python/cudnn/moe_ep/__init__.py @@ -4,25 +4,41 @@ from ._tuning import MoeEpTuningConfig from ._types import ( BlockScaledTensor, + MoeEpBackwardWeightStaging, + MoeEpBackwardWeights, MoeEpExecutionLane, - MoeEpTrainingResources, - MoeEpTrainingSlot, - MoeEpTrainingWeights, + MoeEpForwardWeightStaging, + MoeEpForwardWeights, + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, + MoeEpTrainingBackwardOutputs, + MoeEpTrainingForwardOutputs, MoeEpTrainingWgradOperands, MoeFormat, MoeTensor, ) -from .api import MoeEp +from .api import MoeEp, pack_backward_weights, pack_forward_weights __all__ = [ "BlockScaledTensor", "MoeEp", + "MoeEpBackwardWeightStaging", + "MoeEpBackwardWeights", "MoeEpExecutionLane", - "MoeEpTrainingResources", - "MoeEpTrainingSlot", - "MoeEpTrainingWeights", + "MoeEpForwardWeightStaging", + "MoeEpForwardWeights", + "MoeEpNativeBackwardWeights", + "MoeEpNativeForwardWeights", + "MoeEpNativeWeight", + "MoeEpNativeWeightLayout", + "MoeEpTrainingBackwardOutputs", + "MoeEpTrainingForwardOutputs", "MoeEpTrainingWgradOperands", "MoeEpTuningConfig", "MoeFormat", "MoeTensor", + "pack_backward_weights", + "pack_forward_weights", ] diff --git a/python/cudnn/moe_ep/_math.py b/python/cudnn/moe_ep/_math.py new file mode 100644 index 000000000..466649d19 --- /dev/null +++ b/python/cudnn/moe_ep/_math.py @@ -0,0 +1,17 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Small integer helpers shared by public contracts and private backends.""" + +from __future__ import annotations + + +def ceil_div(value: int, divisor: int) -> int: + return (value + divisor - 1) // divisor + + +def round_up(value: int, multiple: int) -> int: + return ceil_div(value, multiple) * multiple + + +__all__ = ["ceil_div", "round_up"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/README.md b/python/cudnn/moe_ep/_megamoe_backend/README.md index 31391001d..f1bd9c9e8 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/README.md +++ b/python/cudnn/moe_ep/_megamoe_backend/README.md @@ -1,139 +1,84 @@ # MegaMoE backend -The private MegaMoE backend provides Rubin SM107 MXFP8 execution for -`cudnn.moe_ep`. +The private backend provides Rubin SM107 MXFP8 execution for `cudnn.moe_ep`. -## Executable capability +## Capability -- CUDA Rubin SM107 (compute capability 10.7) +- CUDA Rubin SM107 - BF16 output with BF16 or MXFP8 combine - `hidden_size % 128 == 0` - `intermediate_size % 256 == 0` - `top_k <= min(32, num_experts)` -- explicit positive `max_tokens_per_rank` +- positive `max_tokens_per_rank` - `apply_topk_in_fc1=True` -Inference accepts plain BF16/FP16/FP32 or MXFP8 operands and stages plain -operands to MXFP8. Fixed-resource training accepts contiguous BF16/FP32 -activation and grad-output tensors, contiguous Int32 routing indices, -contiguous FP32 routing weights, and four contiguous MXFP8 training-weight -packs. NVFP4 operands, non-BF16 output, and `apply_topk_in_fc1=False` are not -executable. - -## Public execution paths - -`MoeEp.__call__` is the inference-forward surface. It returns only the fused -BF16 `(T, H)` output and does not expose a compact training stash or backward. -Inference CUDA Graph capture requires `MoeEp.warmup` with the exact capture -bindings before capture. EP ranks must align after warmup and replay in the -same cross-rank order. - -Training uses fixed resources: - -```python -resources = op.prepare_training_resources( - weights, - slot_count=2, - lane_count=1, -) -slot0, slot1 = resources.slots -lane0 = resources.lanes[0] - -resources.refresh_weights() -y0 = resources.forward(slot0, lane0, x0, topk_idx0, topk_weights0) -dx0, dprob0, wgrad0 = resources.backward(slot0, lane0, grad0) -overflow = resources.finalize_overflow((slot0,), lane0) -``` - -`prepare_training_resources` is collective over the EP group, executes outside -capture, and fixes the training kernel to FC1-preactivation generation, -fixed-capacity WGrad operands, and token/scale-factor padding 128. - -The same methods execute ordinarily during warmup and enqueue identical nodes -inside a caller-owned outer CUDA Graph. MoeEP does not own or wrap graph -replay. The ordinary warmup must cover -`refresh_weights -> forward -> backward -> finalize_overflow` so staging, -forward, backward, and WGrad-export kernels are compiled before capture. - -## Fixed resource model - -- A persistent slot owns one microbatch's routing snapshot, pool-native FC1 - preactivation, kernel dprob, outputs, backward auxiliaries, overflow flags, - and fixed-capacity WGrad operands. -- An execution lane owns mutable router, barrier, and kernel scratch. -- Every symmetric region is built in deterministic order and its size is - normalized by name across EP ranks before allocation. -- The training ABI fingerprint includes the normalized FC1 layout policy - (`gate_then_up` or `gate_up_interleaved_32`), so ranks with different - gate/up semantics fail the collective handshake before allocation. -- Multiple streams require distinct lanes. Distributed MegaMoE kernels must be - ordered consistently on every rank with captured CUDA events; independent - lane storage does not permit unordered communication overlap. -- `max_recv_size_per_rank` bounds allocation. Capacity never grows during - capture; changing it requires new resources and graph capture. - -## Weights and WGrad outputs - -`MoeEpTrainingWeights` contains four address-stable MXFP8 block-scaled tensors: -forward W1/W2 and independently quantized backward W2-transpose/W1-transpose. -Backward-transpose data is C-contiguous under both FC1 layout policies, matching -the single training AOT signature. -With `weight_interleave_size=32`, compact K-major forward weights and contiguous -backward transposes are interpreted as already using 32-element W1 gate/up -strips, and the kernels alias weight data directly; only scales require -kernel-native staging. With the default `None`, weights use conventional -gate-then-up order and are copied and interleaved into persistent kernel -buffers. After every in-place data+scale update, the caller -must enqueue `resources.refresh_weights()` before the first consumer, -with explicit stream/event ordering. A matching forward/backward pair must use -one version; refresh cannot overlap any consumer on another slot/lane. Replacing -source storage requires closing the old operator, creating a new `MoeEp` -instance and resources, and capturing a new graph. Closed resources are -terminal and cannot be replaced on the same operator. Capturing the refresh -turns these transforms into fixed-address graph nodes, so replay does not call -Python. - -Backward returns kernel dprob directly. It follows the MXFP8-staged numerical -contract and relaxed atomic accumulation order. - -`MoeEpTrainingWgradOperands` is a fixed-capacity producer ABI. Device -`expert_offsets` and `valid_route_counts` describe the current valid K extent; -padding is zeroed. Its data operands alias persistent forward or backward -outputs, using transpose views where required. With -`weight_interleave_size=32`, exporting them performs no full-tensor data -copies, and FC1 dY data and scales preserve the same 32-element gate/up order -as W1. With `None`, FC1 dY is copied and deinterleaved to match conventional -gate-then-up W1 storage. Scale operands are expanded into persistent -grouped-WGrad layouts. No specific downstream grouped-WGrad consumer is -guaranteed by this milestone. - -## Overflow policy - -`max_recv_size_per_rank` bounds the fixed receive pool. When omitted, it uses -the worst-case `ep_size * max_tokens_per_rank * top_k`; an explicit value is -capped at that count. - -The fixed-resource transport truncates deterministically so every rank -completes its communication protocol. `finalize_overflow` aggregates the -selected slots and performs a scalar MAX all-reduce for EP2+. With -`drop_on_overflow=True`, it returns a one-element Int32 status tensor and -dropped routes contribute zero. With `drop_on_overflow=False`, the graph tail -uses `torch._assert_async`; EP2+ error mode requires NCCL. - -## Distributed support - -Hardware acceptance covers EP1, EP2/4, EP8, EP16, and EP32 on one MNNVL -peer-access domain. The Python capability layer has no hard EP-size ceiling; -the listed sizes are validated scope rather than cross-MNNVL support. - -Current tests additionally cover single-node EP3 inference, noncontiguous EP2 -subgroups, multi-node EP4/6/12/16 forward, multi-node EP8/16/32 backward, and -EP8/16/32 fixed-resource graph launchers. EP2+ probes perform collective -warmup, independent capture, capture alignment, diagnostic replay, lockstep -production-like replay bursts, overflow/recovery, ordered multi-lane -execution, and collective teardown. - -The kernels use direct peer pointers obtained from NVSHMEM symmetric tensors. -`NVSHMEM_REMOTE_TRANSPORT=none` is valid only when every EP rank is directly -P2P-accessible (`NVSHMEM_TEAM_SHARED` spans the EP world). IBRC initialization -alone does not make non-P2P peers directly addressable by these kernels. +Inference accepts BF16/FP16/FP32 or MXFP8 operands. Training accepts +BF16/FP32 or MXFP8 activation and grad-output, contiguous Int32 routing +indices, contiguous FP32 routing weights, and independent native forward and +backward weight packs. + +## Execution state + +`MoeEp.prepare_training` creates one private `Mxfp8TrainingState`, which owns +only: + +- prepared forward/backward kernels and compile caches; +- NVSHMEM/runtime handles; +- one local and symmetric scratch slab per execution lane; +- private fixed-capacity transport and routing scratch used only during a call. + +It does not own or retain caller weights, output bundles, saved forward state, +WGrad operands, or weight-staging bundles. No slot is exposed by the public +API. + +## Native weights + +Training execution accepts only `MoeEpNativeForwardWeights` or +`MoeEpNativeBackwardWeights`. Validation checks the exact versioned +`layout_id`, shape, stride, dtype, alignment, and device. The launch adapter +creates aliases to payload and blocked E8M0 scale tensors without allocation, +copy, refresh, or persistent binding. + +`materialize_forward` and `materialize_backward` are allocation-free fallback +transforms. They write only caller-provided staging bundles and return native +packs that alias those destinations. + +## Inputs and outputs + +Plain training inputs use `Mxfp8TrainingStager`. MXFP8 +`BlockScaledTensor` inputs bypass quantization and copy their payload/scales +only into the symmetric transport plane required for peer addressing. + +Caller outputs are borrowed for one launch: + +- required FC1 preactivation is passed directly to forward and backward + kernels; +- all forward and backward WGrad payloads, scales, and route metadata are + required after `prepare_training()` and passed directly to the kernels; +- combine output and dprob first land in private symmetric buffers, then copy + to caller buffers because remote ranks address the symmetric plane; +- primary forward/backward outputs are required caller-owned destinations. + +The producing kernels already expose the final grouped-WGrad scale carriers +when token and scale-factor padding are both 128. Caller E8M0 matrices are +viewed through the producer's flat or matrix signature, so no scale expansion +kernel is launched. FC1-B remains gate/up-interleaved, and FC1-A/FC2-A use +legal transpose views without physical transpose copies. + +## Overflow and distributed ordering + +Each phase keeps overflow state private and applies the configured policy +before returning. EP2+ performs the scalar MAX needed for a rank-consistent +decision. There is no public `finalize_overflow`. + +One lane is exclusive to one active stream. Every EP rank must submit +distributed forward/backward launches in identical order. Distinct lanes do +not make unordered collective-kernel overlap valid. + +## CUDA Graph + +Preparation and first-time compilation happen before capture. Training calls +require every destination advertised by `prepare_training()` to be +caller-owned. Every input, output, saved-state, native weight, and staging +address referenced by a graph remains stable until that graph executable is +destroyed. Eager calls may change addresses between invocations. diff --git a/python/cudnn/moe_ep/_megamoe_backend/_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/_workspace.py index 24a24d71d..59d546f1f 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/_workspace.py +++ b/python/cudnn/moe_ep/_megamoe_backend/_workspace.py @@ -13,6 +13,7 @@ import torch from .._contracts import ForwardConfig +from .._math import round_up from ._comm import ( PeerMapping, SymmetricMemoryProvider, @@ -22,15 +23,11 @@ from ._runtime import RuntimeHandle, _runtime_debug -def _align_up(value: int, alignment: int) -> int: - return (value + alignment - 1) // alignment * alignment - - def padded_mxfp8_scale_columns(hidden: int) -> int: """Return the E8M0 row width required by Rubin's 16-byte token-in copy.""" logical_columns = (hidden + 31) // 32 - return _align_up(logical_columns, 16) + return round_up(logical_columns, 16) @dataclass(frozen=True) @@ -76,7 +73,7 @@ def build(cls, regions: Sequence[BufferRegion]) -> "BufferLayout": if region.name in names: raise ValueError(f"duplicate workspace region {region.name!r}") names.add(region.name) - offset = _align_up(offset, region.alignment) + offset = round_up(offset, region.alignment) placements.append( BufferPlacement( name=region.name, @@ -88,7 +85,7 @@ def build(cls, regions: Sequence[BufferRegion]) -> "BufferLayout": max_alignment = max(max_alignment, region.alignment) return cls( placements=tuple(placements), - total_bytes=_align_up(offset, max_alignment), + total_bytes=round_up(offset, max_alignment), ) def placement(self, name: str) -> BufferPlacement: @@ -129,7 +126,6 @@ def for_mxfp8( kernel_shared_workspace_bytes: int, col_quant_data_bytes: int = 0, col_quant_sf_bytes: int = 0, - backward_fc1_preact_bytes: int = 0, backward_dprob_bytes: int = 0, backward_aux_data_bytes: int = 0, backward_aux_scale_bytes: int = 0, @@ -141,7 +137,6 @@ def for_mxfp8( ("kernel_shared_workspace_bytes", kernel_shared_workspace_bytes), ("col_quant_data_bytes", col_quant_data_bytes), ("col_quant_sf_bytes", col_quant_sf_bytes), - ("backward_fc1_preact_bytes", backward_fc1_preact_bytes), ("backward_dprob_bytes", backward_dprob_bytes), ("backward_aux_data_bytes", backward_aux_data_bytes), ("backward_aux_scale_bytes", backward_aux_scale_bytes), @@ -151,13 +146,12 @@ def for_mxfp8( if bool(col_quant_data_bytes) != bool(col_quant_sf_bytes): raise ValueError("column requant data and scale workspace must be enabled together") backward_sizes = ( - backward_fc1_preact_bytes, backward_dprob_bytes, backward_aux_data_bytes, backward_aux_scale_bytes, ) if any(backward_sizes) and not all(backward_sizes): - raise ValueError("backward preactivation, dprob, data, and scale workspace " "must be enabled together") + raise ValueError("backward dprob, data, and scale workspace must be enabled together") tokens = config.max_tokens_per_rank hidden = config.hidden_size @@ -189,11 +183,6 @@ def for_mxfp8( ) backward_local_regions = ( ( - BufferRegion( - "backward_fc1_preact", - backward_fc1_preact_bytes, - alignment=128, - ), BufferRegion( "backward_aux_data", backward_aux_data_bytes, @@ -205,7 +194,7 @@ def for_mxfp8( alignment=128, ), ) - if backward_fc1_preact_bytes + if backward_aux_data_bytes else () ) local_regions = ( diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py index f32978734..a76b97886 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backend.py @@ -13,7 +13,6 @@ from ..._backend import BackendUnavailableError from ..._contracts import ForwardConfig, ValidatedForwardRequest -from ..._types import MoeEpTrainingWeights from .._plan import ExecutionPlanOwner from ._adapter import Mxfp8InputAdapter from ._backward_compile import prepare_backward_kernel @@ -44,7 +43,7 @@ def __init__(self, config: ForwardConfig, device: torch.device) -> None: self._completion_recorded = False self._device_work_may_be_pending = False self._ep_launch_ready = config.ep_size == 1 - self._training_resource_owner = None + self._training_state = None self._lock = threading.RLock() @property @@ -180,20 +179,18 @@ def forward(self, request: ValidatedForwardRequest): self._warmed_up = True return output - def prepare_training_resources( + def prepare_training( self, - weights: MoeEpTrainingWeights, *, - slot_count: int, lane_count: int, ): - """Allocate the fixed slot/lane roots used by the training graph path.""" + """Allocate private per-lane state for stateless training calls.""" with self._lock: if self._closed: raise RuntimeError("MoeEp MXFP8 backend is closed") - if self._training_resource_owner is not None: - raise RuntimeError("MoeEp training resources already exist") + if self._training_state is not None: + raise RuntimeError("MoeEp training is already prepared") training_config = replace( self.config, generate_c=True, @@ -224,24 +221,22 @@ def prepare_training_resources( graph_kernel_config, self.device, ) - from ._training_resources import Mxfp8TrainingResourceOwner + from ._training_resources import Mxfp8TrainingState - owner = Mxfp8TrainingResourceOwner( + state = Mxfp8TrainingState( training_config, self.device, forward, backward, - weights, - slot_count=slot_count, lane_count=lane_count, ) try: - owner.prepare() + state.prepare() except Exception: - owner.close() + state.close() raise - self._training_resource_owner = owner - return owner + self._training_state = state + return state def close(self) -> None: with self._lock: @@ -250,12 +245,12 @@ def close(self) -> None: with torch.cuda.device(self.device): if torch.cuda.is_current_stream_capturing(): raise RuntimeError("MoeEp MXFP8 backend cannot be closed during " "CUDA graph capture") - if self._plan is not None or self._training_resource_owner is not None: + if self._plan is not None or self._training_state is not None: torch.cuda.synchronize(self.device) self._adapter.close() - if self._training_resource_owner is not None: - self._training_resource_owner.close() - self._training_resource_owner = None + if self._training_state is not None: + self._training_state.close() + self._training_state = None if self._plan is not None: self._plan.close() self._plan = None diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py index 190647e9d..5e69c76ea 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_compile.py @@ -6,7 +6,6 @@ from __future__ import annotations import math -import os import threading from dataclasses import dataclass from typing import Any @@ -20,8 +19,8 @@ _pre_reduced_sf_workspace_metadata, _pre_reduced_workspace_metadata, ) +from ._compile_common import _compile_kernel, _prepare_rubin_environment from ._config import Mxfp8KernelConfig -from ._cutedsl import require_rubin_cutedsl from ._formats import combine_wire_format from ._launch import _to_cute, _to_cute_ptr @@ -89,27 +88,18 @@ def prepare_backward_kernel( ) -> PreparedMxfp8BackwardKernel: """Instantiate the fixed Rubin dGLU specialization.""" - require_rubin_cutedsl() - torch.cuda.set_device(device) - architecture = torch.cuda.get_device_capability(device) - if architecture != (10, 7): - raise RuntimeError("Rubin MXFP8 backward requires compute capability (10, 7), " f"got {architecture}") - configured_architecture = os.environ.get("CUTE_DSL_ARCH") - if configured_architecture is None: - os.environ["CUTE_DSL_ARCH"] = "sm_107a" - elif configured_architecture not in ("sm_107", "sm_107a"): - raise RuntimeError("CUTE_DSL_ARCH must target SM107 for the Rubin MXFP8 backward") + architecture, launch_cluster_count = _prepare_rubin_environment( + device, + config, + context="backward", + ) import cutlass - import cutlass.utils as utils from ..cutedsl_src.kernel_src.rubin.training.mega.bwd_dglu import ( Sm107MegaMoEMxfp8DgluKernel, ) from ..cutedsl_src.quant_def import CombineFormat - launch_cluster_count = int(utils.HardwareInfo().get_max_active_clusters(config.cluster_size)) - if launch_cluster_count <= 0: - raise RuntimeError("hardware occupancy query returned no launchable Rubin clusters") group_hint = launch_cluster_count if config.group_hint is None else config.group_hint operands_mode = forward_config.backward_wgrad_mode == "operands" dfc2_recompute = operands_mode @@ -165,7 +155,6 @@ def prepare_backward_kernel( if fc1_preact_shape != expected_preact_shape: raise RuntimeError("Rubin dGLU fc1_preact shape mismatch: " f"{fc1_preact_shape} != {expected_preact_shape}") aux_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in kernel.get_aux_output_shapes().items()} - fc1_preact_bytes = math.prod(fc1_preact_shape) * torch.bfloat16.itemsize dprob_bytes = math.prod(aux_shapes["dprob"]) * torch.float32.itemsize aux_data_bytes = ( max( @@ -187,7 +176,6 @@ def prepare_backward_kernel( forward_config, kernel_local_workspace_bytes=local_bytes, kernel_shared_workspace_bytes=shared_bytes, - backward_fc1_preact_bytes=fc1_preact_bytes, backward_dprob_bytes=dprob_bytes, backward_aux_data_bytes=aux_data_bytes, backward_aux_scale_bytes=aux_scale_bytes, @@ -316,12 +304,10 @@ def compile_backward_or_get( return cached if torch.cuda.is_current_stream_capturing(): raise RuntimeError("MXFP8 backward kernel must be compiled before capture") - import cutlass.cute as cute - runtime_kwargs = build_backward_runtime_kwargs(inputs, resources) compiled = CompiledMxfp8BackwardKernel( key=key, - callable=cute.compile(prepared.kernel, **runtime_kwargs), + callable=_compile_kernel(prepared.kernel, runtime_kwargs), ) _COMPILE_CACHE[key] = compiled return compiled diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py deleted file mode 100644 index 9e4fe27ed..000000000 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_backward_launch.py +++ /dev/null @@ -1,57 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Current-stream launch of the Rubin MXFP8 dGLU backward product.""" - -from __future__ import annotations - -from dataclasses import dataclass - -import torch - -from .._plan import PreparedResources -from ._backward_compile import ( - CompiledMxfp8BackwardKernel, - Mxfp8BackwardLaunchInputs, - build_backward_runtime_kwargs, -) -from ._launch import _check_overflow - - -@dataclass(frozen=True) -class Mxfp8DgluResult: - grad_activation: torch.Tensor - grad_topk_weights: torch.Tensor - fc1_recompute: torch.Tensor - fc1_recompute_sf: torch.Tensor - fc1_col_output: torch.Tensor - fc1_col_output_sf: torch.Tensor - grad_y2: torch.Tensor - grad_y2_sf: torch.Tensor - - -def launch_backward_dglu( - compiled: CompiledMxfp8BackwardKernel, - inputs: Mxfp8BackwardLaunchInputs, - resources: PreparedResources, -) -> Mxfp8DgluResult: - runtime_kwargs = build_backward_runtime_kwargs(inputs, resources) - compiled.callable(**runtime_kwargs) - _check_overflow(inputs.overflow_flag) - - return Mxfp8DgluResult( - grad_activation=inputs.output_activation[: inputs.token_count].float(), - # The dGLU epilogue has already returned source-order dprob through - # the symmetric token-communication plane. Own the public result so a - # later launch cannot overwrite it. - grad_topk_weights=inputs.dprob[: inputs.token_count].clone(), - fc1_recompute=inputs.fc1_recompute, - fc1_recompute_sf=inputs.fc1_recompute_sf, - fc1_col_output=inputs.fc1_col_output, - fc1_col_output_sf=inputs.fc1_col_output_sf, - grad_y2=inputs.grad_y2, - grad_y2_sf=inputs.grad_y2_sf, - ) - - -__all__ = ["Mxfp8DgluResult", "launch_backward_dglu"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py index 26769c201..9af47003b 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile.py @@ -5,7 +5,6 @@ from __future__ import annotations -import os import threading from dataclasses import dataclass from typing import Any @@ -16,8 +15,8 @@ from .._plan import PreparedResources from .._workspace import WorkspaceRequirements from ._adapter import Mxfp8LaunchInputs +from ._compile_common import _compile_kernel, _prepare_rubin_environment from ._config import Mxfp8KernelConfig -from ._cutedsl import require_rubin_cutedsl from ._fingerprint import build_kernel_fingerprint from ._launch import build_runtime_kwargs, layout_signature @@ -60,15 +59,6 @@ class CompiledMxfp8Kernel: _PRE_REDUCED_ACTIVATION_SF_REGION = "nvlink.token_comm.pre_reduced_activation_sf" -def _compile_kernel(kernel: Any, compile_kwargs: dict[str, Any]) -> Any: - """Import CuTeDSL only on a cache miss and compile one callable.""" - - require_rubin_cutedsl() - import cutlass.cute as cute - - return cute.compile(kernel, **compile_kwargs) - - def _pre_reduced_workspace_metadata( device_workspace: Any, config: Mxfp8KernelConfig, @@ -133,28 +123,18 @@ def prepare_kernel( ) -> PreparedMxfp8Kernel: """Instantiate the kernel and derive exact allocation requirements.""" - require_rubin_cutedsl() - torch.cuda.set_device(device) - architecture = torch.cuda.get_device_capability(device) - if architecture != (10, 7): - raise RuntimeError("Rubin MXFP8 kernel preparation requires compute capability " f"(10, 7), got {architecture}") - configured_architecture = os.environ.get("CUTE_DSL_ARCH") - if configured_architecture is None: - os.environ["CUTE_DSL_ARCH"] = "sm_107a" - elif configured_architecture not in ("sm_107", "sm_107a"): - raise RuntimeError("CUTE_DSL_ARCH must target SM107 for the Rubin MXFP8 backend, " f"got {configured_architecture!r}") - + architecture, launch_cluster_count = _prepare_rubin_environment( + device, + config, + context="forward", + ) import cutlass - import cutlass.utils as utils from ..cutedsl_src.kernel_src.rubin.training.mega.fwd_glu import ( Sm107MegaMoEMxfp8GluKernel, ) from ..cutedsl_src.quant_def import CombineFormat - launch_cluster_count = int(utils.HardwareInfo().get_max_active_clusters(config.cluster_size)) - if launch_cluster_count <= 0: - raise RuntimeError("hardware occupancy query returned no launchable Rubin clusters") group_hint = launch_cluster_count if config.group_hint is None else config.group_hint kernel_kwargs = dict( mma_tiler_mnk=config.mma_tiler_mnk, @@ -316,18 +296,9 @@ def compile_or_get( _COMPILE_CACHE[key] = compiled return compiled - -def clear_compile_cache() -> None: - """Drop process-local compiled callable references.""" - - with _COMPILE_LOCK: - _COMPILE_CACHE.clear() - - __all__ = [ "CompiledMxfp8Kernel", "PreparedMxfp8Kernel", - "clear_compile_cache", "compile_or_get", "prepare_kernel", ] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile_common.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile_common.py new file mode 100644 index 000000000..ac9257f0a --- /dev/null +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_compile_common.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared Rubin environment setup for direction-specific MXFP8 compilers.""" + +from __future__ import annotations + +import os +from typing import Any + +import torch + +from ._config import Mxfp8KernelConfig +from ._cutedsl import require_rubin_cutedsl + + +def _prepare_rubin_environment( + device: torch.device, + config: Mxfp8KernelConfig, + *, + context: str, +) -> tuple[tuple[int, int], int]: + require_rubin_cutedsl() + torch.cuda.set_device(device) + architecture = torch.cuda.get_device_capability(device) + if architecture != (10, 7): + raise RuntimeError( + f"Rubin MXFP8 {context} preparation requires compute capability " + f"(10, 7), got {architecture}" + ) + + configured_architecture = os.environ.get("CUTE_DSL_ARCH") + if configured_architecture is None: + os.environ["CUTE_DSL_ARCH"] = "sm_107a" + elif configured_architecture not in ("sm_107", "sm_107a"): + raise RuntimeError( + "CUTE_DSL_ARCH must target SM107 for Rubin MXFP8 " + f"{context}, got {configured_architecture!r}" + ) + + import cutlass.utils as utils + + launch_cluster_count = int( + utils.HardwareInfo().get_max_active_clusters(config.cluster_size) + ) + if launch_cluster_count <= 0: + raise RuntimeError( + "hardware occupancy query returned no launchable Rubin clusters" + ) + return architecture, launch_cluster_count + + +def _compile_kernel(kernel: Any, compile_kwargs: dict[str, Any]) -> Any: + """Import CuTeDSL only on a cache miss and compile one callable.""" + + require_rubin_cutedsl() + import cutlass.cute as cute + + return cute.compile(kernel, **compile_kwargs) + + +__all__ = ["_compile_kernel", "_prepare_rubin_environment"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py index 0a48f81a9..5750f9166 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_execute.py @@ -1,13 +1,21 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Ordinary/capturable launch path over fixed MXFP8 training resources.""" +"""Ordinary/capturable stateless launch path over private lane resources.""" from __future__ import annotations import torch -from ..._types import MoeEpTrainingWgradOperands +from ..._types import ( + BlockScaledTensor, + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpTrainingBackwardOutputs, + MoeEpTrainingForwardOutputs, + MoeEpTrainingWgradOperands, + MoeTensor, +) from .._runtime import _runtime_debug from .._workspace import padded_mxfp8_scale_columns from ._adapter import ( @@ -23,8 +31,13 @@ from ._launch import build_runtime_kwargs from ._training_resources import ( Mxfp8TrainingExecutionViews, - Mxfp8TrainingResourceOwner, + Mxfp8TrainingState, ) +from ._training_weights import ( + backward_native_to_kernel, + forward_native_to_kernel, +) +from ._training_wgrad import assemble_training_wgrad_operands def _zero_pre_reduced(inputs, prepared) -> None: @@ -69,15 +82,51 @@ def _activation_views( ) +def _stage_input( + state: Mxfp8TrainingState, + value: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + activation_data: torch.Tensor, + activation_sf: torch.Tensor, + routing_topk_idx: torch.Tensor, + routing_topk_weights: torch.Tensor, +) -> None: + """Stage into private symmetric memory, bypassing quantization for MXFP8.""" + + if not isinstance(value, BlockScaledTensor): + state.stager.stage( + value, + topk_idx, + topk_weights, + activation_data, + activation_sf, + routing_topk_idx, + routing_topk_weights, + ) + return + + token_count = int(value.logical_shape[0]) + scale_columns = int(value.scale.shape[1]) + activation_data.zero_() + activation_sf.zero_() + routing_topk_idx.fill_(-1) + routing_topk_weights.zero_() + if token_count == 0: + return + activation_data[:token_count].copy_(value.data) + activation_sf[:token_count, :scale_columns].copy_(value.scale) + routing_topk_idx[:token_count].copy_(topk_idx) + routing_topk_weights[:token_count].copy_(topk_weights) + + def _write_expert_offsets( execution: Mxfp8TrainingExecutionViews, padding: int, + counts: torch.Tensor, + offsets: torch.Tensor, ) -> None: snapshot = execution.forward_expert_size_snapshot - if snapshot is None: - raise RuntimeError("training forward requires the persistent expert-size snapshot") - counts = execution.slot.valid_route_counts - offsets = execution.slot.expert_offsets counts.copy_(snapshot) torch.add(counts, padding - 1, out=offsets) torch.div(offsets, padding, rounding_mode="floor", out=offsets) @@ -86,22 +135,26 @@ def _write_expert_offsets( def launch_training_forward( - owner: Mxfp8TrainingResourceOwner, + state: Mxfp8TrainingState, execution: Mxfp8TrainingExecutionViews, - activation: torch.Tensor, + activation: MoeTensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor, + *, + weights: MoeEpNativeForwardWeights, + out: MoeEpTrainingForwardOutputs, ) -> torch.Tensor: - """Stage and launch one fixed-slot forward without host-visible routing.""" + """Launch one stateless forward over caller-owned outputs.""" - prepared = owner.forward_prepared + prepared = state.forward_prepared config = prepared.config capacity = config.max_tokens_per_rank - slot = execution.slot + scratch = execution.scratch + token_count = int(activation.logical_shape[0] if isinstance(activation, BlockScaledTensor) else activation.shape[0]) _runtime_debug( "training-forward.begin", - slot=execution.slot.index, - token_count=int(activation.shape[0]), + lane=scratch.index, + token_count=token_count, ) activation_data, activation_sf = _activation_views( execution, @@ -109,78 +162,113 @@ def launch_training_forward( capacity=capacity, hidden=config.hidden, ) - _runtime_debug("training-forward.stage.begin", slot=execution.slot.index) - owner.stager.stage( + _runtime_debug("training-forward.stage.begin", lane=scratch.index) + _stage_input( + state, activation, topk_idx, topk_weights, activation_data, activation_sf, - slot.routing_topk_idx, - slot.routing_topk_weights, + scratch.routing_topk_idx, + scratch.routing_topk_weights, ) - _runtime_debug("training-forward.stage.end", slot=execution.slot.index) - slot.forward_output.zero_() - slot.forward_overflow.zero_() - if slot.col_quant_data is not None: - slot.col_quant_data.zero_() - if slot.col_quant_sf is not None: - slot.col_quant_sf.zero_() - _runtime_debug("training-forward.reset.end", slot=execution.slot.index) + _runtime_debug("training-forward.stage.end", lane=scratch.index) + + assert out.output is not None + assert out.fc1_a is not None + assert out.fc1_sfa is not None + assert out.valid_route_counts is not None + assert out.expert_offsets is not None + fc1_preact = out.fc1_preact + col_quant_data = out.fc1_a.transpose(0, 1) + col_quant_sf = out.fc1_sfa.view(torch.uint8).reshape(-1) + expected_elements = int(prepared.col_quant_sf_elements) + if col_quant_sf.numel() != expected_elements: + raise ValueError("out.fc1_sfa storage does not match the forward producer ABI: " f"{col_quant_sf.numel()} != {expected_elements}") + valid_route_counts = out.valid_route_counts + expert_offsets = out.expert_offsets + + scratch.forward_output.zero_() + scratch.forward_overflow.zero_() + col_quant_data.zero_() + # E8M0 byte 127 encodes scale 1.0. The producer only overwrites active + # expert segments, so the unused grouped-WGrad capacity must stay neutral. + col_quant_sf.fill_(127) + _runtime_debug("training-forward.reset.end", lane=scratch.index) workspace = execution.forward.workspace inputs = Mxfp8LaunchInputs( activation=activation_data, activation_sf=activation_sf, - topk_indices=slot.routing_topk_idx, - topk_scores=slot.routing_topk_weights, - weights=owner.weight_bindings.forward, - fc1_c=slot.fc1_preact, - output_data=slot.forward_output, - col_quant_data=slot.col_quant_data, - col_quant_sf=slot.col_quant_sf, - overflow_flag=slot.forward_overflow, + topk_indices=scratch.routing_topk_idx, + topk_scores=scratch.routing_topk_weights, + weights=forward_native_to_kernel(weights), + fc1_c=fc1_preact, + output_data=scratch.forward_output, + col_quant_data=col_quant_data, + col_quant_sf=col_quant_sf, + overflow_flag=scratch.forward_overflow, local_workspace=workspace.local["kernel_local_workspace"], shared_workspace=workspace.symmetric["kernel_shared_workspace"], - token_count=int(activation.shape[0]), + token_count=token_count, ) _zero_pre_reduced(inputs, prepared) - _runtime_debug("training-forward.compile.begin", slot=execution.slot.index) + _runtime_debug("training-forward.compile.begin", lane=scratch.index) compiled = compile_or_get( prepared, inputs, execution.forward, ) - _runtime_debug("training-forward.compile.end", slot=execution.slot.index) - _runtime_debug("training-forward.launch.begin", slot=execution.slot.index) + _runtime_debug("training-forward.compile.end", lane=scratch.index) + _runtime_debug("training-forward.launch.begin", lane=scratch.index) compiled.callable(**build_runtime_kwargs(inputs, execution.forward)) - _runtime_debug("training-forward.launch.end", slot=execution.slot.index) - _runtime_debug("training-forward.offsets.begin", slot=execution.slot.index) - _write_expert_offsets(execution, config.token_padding_block) - _runtime_debug("training-forward.offsets.end", slot=execution.slot.index) - _runtime_debug("training-forward.end", slot=execution.slot.index) - return slot.forward_output[: inputs.token_count] + _runtime_debug("training-forward.launch.end", lane=scratch.index) + _runtime_debug("training-forward.offsets.begin", lane=scratch.index) + _write_expert_offsets( + execution, + config.token_padding_block, + valid_route_counts, + expert_offsets, + ) + _runtime_debug("training-forward.offsets.end", lane=scratch.index) + state.apply_overflow(lane=scratch.index, phase="forward") + + output = out.output[:token_count] + output.copy_(scratch.forward_output[:token_count]) + _runtime_debug("training-forward.end", lane=scratch.index) + return output def launch_training_backward( - owner: Mxfp8TrainingResourceOwner, + state: Mxfp8TrainingState, execution: Mxfp8TrainingExecutionViews, - grad_output: torch.Tensor, + grad_output: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + weights: MoeEpNativeBackwardWeights, + fc1_preact: torch.Tensor, + fc1_a: torch.Tensor | None, + fc1_sfa: torch.Tensor | None, + valid_route_counts: torch.Tensor | None, + expert_offsets: torch.Tensor | None, + out: MoeEpTrainingBackwardOutputs, ) -> tuple[ torch.Tensor, torch.Tensor, MoeEpTrainingWgradOperands, ]: - """Stage and launch one fixed-slot backward using forward's raw pool.""" + """Launch one stateless backward using explicit caller-owned saved state.""" - prepared = owner.backward_prepared + prepared = state.backward_prepared config = prepared.config capacity = config.max_tokens_per_rank - slot = execution.slot - token_count = int(grad_output.shape[0]) + scratch = execution.scratch + token_count = int(grad_output.logical_shape[0] if isinstance(grad_output, BlockScaledTensor) else grad_output.shape[0]) _runtime_debug( "training-backward.begin", - slot=execution.slot.index, + lane=scratch.index, token_count=token_count, ) activation_data, activation_sf = _activation_views( @@ -189,77 +277,103 @@ def launch_training_backward( capacity=capacity, hidden=config.hidden, ) - _runtime_debug("training-backward.stage.begin", slot=execution.slot.index) - owner.stager.stage( + _runtime_debug("training-backward.stage.begin", lane=scratch.index) + _stage_input( + state, grad_output, - slot.routing_topk_idx[:token_count], - slot.routing_topk_weights[:token_count], + topk_idx, + topk_weights, activation_data, activation_sf, - slot.routing_topk_idx, - slot.routing_topk_weights, + scratch.routing_topk_idx, + scratch.routing_topk_weights, ) - _runtime_debug("training-backward.stage.end", slot=execution.slot.index) - - slot.backward_output.zero_() - slot.grad_activation.zero_() - slot.backward_overflow.zero_() - slot.dprob.zero_() - slot.fc1_recompute.zero_() - slot.fc1_recompute_sf.view(torch.uint8).fill_(127) - slot.fc1_col_output.zero_() - slot.fc1_col_output_sf.view(torch.uint8).fill_(127) - slot.grad_y2.zero_() - slot.grad_y2_sf.fill_(127) - _runtime_debug("training-backward.reset.end", slot=execution.slot.index) + _runtime_debug("training-backward.stage.end", lane=scratch.index) + + assert fc1_a is not None + assert fc1_sfa is not None + assert valid_route_counts is not None + assert expert_offsets is not None + assert out.grad_activation is not None + assert out.dprob is not None + assert out.fc1_b is not None + assert out.fc1_sfb is not None + assert out.fc2_a is not None + assert out.fc2_sfa is not None + assert out.fc2_b is not None + assert out.fc2_sfb is not None + fc1_recompute = out.fc2_a.transpose(0, 1) + fc1_recompute_sf = out.fc2_sfa + fc1_col_output = out.fc1_b + fc1_col_output_sf = out.fc1_sfb + grad_y2 = out.fc2_b + grad_y2_sf = out.fc2_sfb.view(torch.uint8).reshape(-1) + + scratch.backward_output.zero_() + scratch.backward_overflow.zero_() + scratch.dprob.zero_() + fc1_recompute.zero_() + fc1_recompute_sf.view(torch.uint8).fill_(127) + fc1_col_output.zero_() + fc1_col_output_sf.view(torch.uint8).fill_(127) + grad_y2.zero_() + grad_y2_sf.fill_(127) + _runtime_debug("training-backward.reset.end", lane=scratch.index) workspace = execution.backward.workspace - weights = owner.weight_bindings.backward + kernel_weights = backward_native_to_kernel(weights) inputs = Mxfp8BackwardLaunchInputs( grad_out=activation_data, grad_out_sf=activation_sf, - topk_idx=slot.routing_topk_idx, - topk_weights=slot.routing_topk_weights, - fc1_weight=weights.fc1_weight, - fc1_weight_sf=weights.fc1_weight_sf, - fc2_weight=weights.fc2_weight, - fc2_weight_sf=weights.fc2_weight_sf, - beta=owner.beta, - fc1_preact=slot.fc1_preact, - output_activation=slot.backward_output, - overflow_flag=slot.backward_overflow, - dprob=slot.dprob, - fc1_recompute=slot.fc1_recompute, - fc1_recompute_sf=slot.fc1_recompute_sf, - fc1_col_output=slot.fc1_col_output, - fc1_col_output_sf=slot.fc1_col_output_sf, - grad_y2=slot.grad_y2, - grad_y2_sf=slot.grad_y2_sf, + topk_idx=scratch.routing_topk_idx, + topk_weights=scratch.routing_topk_weights, + fc1_weight=kernel_weights.fc1_weight, + fc1_weight_sf=kernel_weights.fc1_weight_sf, + fc2_weight=kernel_weights.fc2_weight, + fc2_weight_sf=kernel_weights.fc2_weight_sf, + beta=state.beta, + fc1_preact=fc1_preact, + output_activation=scratch.backward_output, + overflow_flag=scratch.backward_overflow, + dprob=scratch.dprob, + fc1_recompute=fc1_recompute, + fc1_recompute_sf=fc1_recompute_sf, + fc1_col_output=fc1_col_output, + fc1_col_output_sf=fc1_col_output_sf, + grad_y2=grad_y2, + grad_y2_sf=grad_y2_sf, local_workspace=workspace.local["kernel_local_workspace"], shared_workspace=workspace.symmetric["kernel_shared_workspace"], token_count=token_count, ) _zero_pre_reduced(inputs, prepared) - _runtime_debug("training-backward.compile.begin", slot=execution.slot.index) + _runtime_debug("training-backward.compile.begin", lane=scratch.index) compiled = compile_backward_or_get( prepared, inputs, execution.backward, ) - _runtime_debug("training-backward.compile.end", slot=execution.slot.index) - _runtime_debug("training-backward.launch.begin", slot=execution.slot.index) + _runtime_debug("training-backward.compile.end", lane=scratch.index) + _runtime_debug("training-backward.launch.begin", lane=scratch.index) compiled.callable(**build_backward_runtime_kwargs(inputs, execution.backward)) - _runtime_debug("training-backward.launch.end", slot=execution.slot.index) - slot.grad_activation.copy_(slot.backward_output) - _runtime_debug("training-backward.wgrad-export.begin", slot=execution.slot.index) - operands = owner.wgrad_exporter.export(slot) - _runtime_debug("training-backward.wgrad-export.end", slot=execution.slot.index) - _runtime_debug("training-backward.end", slot=execution.slot.index) - return ( - slot.grad_activation[:token_count], - slot.dprob[:token_count], - operands, + _runtime_debug("training-backward.launch.end", lane=scratch.index) + state.apply_overflow(lane=scratch.index, phase="backward") + + grad_activation = out.grad_activation[:token_count] + grad_activation.copy_(scratch.backward_output[:token_count]) + + dprob = out.dprob[:token_count] + dprob.copy_(scratch.dprob[:token_count]) + + operands = assemble_training_wgrad_operands( + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + expert_offsets=expert_offsets, + valid_route_counts=valid_route_counts, + backward=out, ) + _runtime_debug("training-backward.end", lane=scratch.index) + return grad_activation, dprob, operands __all__ = [ diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py index 1f097a583..8c7bd936c 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_resources.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Fixed-capacity slot/lane resources for graph-capable MXFP8 training.""" +"""Private per-lane state for graph-capable stateless MXFP8 training.""" from __future__ import annotations @@ -16,8 +16,9 @@ import torch import torch.distributed as dist -from ..._contracts import Fc1WeightLayout, ForwardConfig -from ..._types import MoeEpTrainingWeights +from ..._contracts import ForwardConfig +from ..._math import round_up +from ..._types import MoeEpNativeWeightLayout from .._comm import SymmetricMemoryProvider from .._plan import PreparedResources from .._runtime import ( @@ -34,31 +35,22 @@ WorkspaceRequirements, WorkspaceViews, ) -from ._adapter import _typed_k_major_view, _typed_view +from ._adapter import _typed_view from ._backward_compile import PreparedMxfp8BackwardKernel from ._compile import PreparedMxfp8Kernel from ._fingerprint import canonical_json_sha256, source_tree_sha256 from ._training_stage import Mxfp8TrainingStager -from ._training_weights import Mxfp8TrainingWeightBindings -from ._training_wgrad import Mxfp8TrainingWgradExporter _DATA_DTYPE = torch.float8_e4m3fn _SCALE_DTYPE = torch.float8_e8m0fnu _ROUTING_SYMMETRIC = frozenset({"topk_weights"}) _ROUTING_LOCAL = frozenset({"topk_idx"}) -_FORWARD_SLOT_SYMMETRIC = frozenset({"output_data", *_ROUTING_SYMMETRIC}) -_FORWARD_SLOT_LOCAL = frozenset({"overflow_flag", "col_quant_data", "col_quant_sf", *_ROUTING_LOCAL}) -_BACKWARD_SLOT_SYMMETRIC = frozenset({"output_data", "backward_dprob", *_ROUTING_SYMMETRIC}) -_BACKWARD_SLOT_LOCAL = frozenset({"overflow_flag", "backward_aux_data", "backward_aux_scale", *_ROUTING_LOCAL}) - - -def _round_up(value: int, multiple: int) -> int: - return (value + multiple - 1) // multiple * multiple - - -def _align_scale_columns(token_capacity: int) -> int: - return _round_up((token_capacity + 31) // 32, 4) +_CALLER_OWNED_FORWARD_LOCAL = frozenset({"col_quant_data", "col_quant_sf"}) +_FORWARD_PRIVATE_SYMMETRIC = frozenset({"output_data", *_ROUTING_SYMMETRIC}) +_FORWARD_PRIVATE_LOCAL = frozenset({"overflow_flag", *_CALLER_OWNED_FORWARD_LOCAL, *_ROUTING_LOCAL}) +_BACKWARD_PRIVATE_SYMMETRIC = frozenset({"output_data", "backward_dprob", *_ROUTING_SYMMETRIC}) +_BACKWARD_PRIVATE_LOCAL = frozenset({"overflow_flag", "backward_aux_data", "backward_aux_scale", *_ROUTING_LOCAL}) def _lane_name( @@ -70,21 +62,12 @@ def _lane_name( return f"lane.{lane}.{phase}.{space}.{name}" -def _slot_name( - slot: int, - phase: str, +def _lane_fallback_name( + lane: int, space: str, name: str, ) -> str: - return f"slot.{slot}.{phase}.{space}.{name}" - - -def _custom_slot_name(slot: int, name: str) -> str: - return f"slot.{slot}.persistent.local.{name}" - - -def _custom_slot_symmetric_name(slot: int, name: str) -> str: - return f"slot.{slot}.persistent.symmetric.{name}" + return f"lane.{lane}.fallback.{space}.{name}" def _clone_region(name: str, region: BufferRegion) -> BufferRegion: @@ -103,17 +86,6 @@ def _region_map( return {region.name: region for region in regions} -def _required_region( - requirements: WorkspaceRequirements, - space: str, - name: str, -) -> BufferRegion: - try: - return _region_map(requirements, space)[name] - except KeyError as exc: - raise ValueError(f"{space} workspace requirements do not contain {name!r}") from exc - - def _add_lane_regions( output: list[BufferRegion], requirements: WorkspaceRequirements, @@ -121,14 +93,11 @@ def _add_lane_regions( lane: int, phase: str, space: str, - slot_names: frozenset[str], + excluded_names: frozenset[str], ) -> None: regions = requirements.symmetric_regions if space == "symmetric" else requirements.local_regions for region in regions: - if region.name in slot_names: - continue - if phase == "backward" and space == "local" and region.name == ("backward_fc1_preact"): - # The graph path aliases forward's raw receiver pool directly. + if region.name in excluded_names: continue output.append( _clone_region( @@ -143,16 +112,14 @@ def build_training_workspace_requirements( forward: PreparedMxfp8Kernel, backward: PreparedMxfp8BackwardKernel, *, - slot_count: int, lane_count: int, ) -> WorkspaceRequirements: - """Build one deterministic root layout for N slots and M lanes.""" + """Build one deterministic root layout for private execution lanes.""" - for name, value in (("slot_count", slot_count), ("lane_count", lane_count)): - if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise ValueError(f"{name} must be a positive integer, got {value!r}") + if isinstance(lane_count, bool) or not isinstance(lane_count, int) or lane_count <= 0: + raise ValueError(f"lane_count must be a positive integer, got {lane_count!r}") if not config.generate_c: - raise ValueError("training resources require generate_c=True") + raise ValueError("training preparation requires generate_c=True") if forward.pool_token_capacity != backward.pool_token_capacity: raise ValueError("forward/backward pool capacities must match, got " f"{forward.pool_token_capacity} and " f"{backward.pool_token_capacity}") @@ -192,7 +159,7 @@ def build_training_workspace_requirements( lane=lane, phase="forward", space="symmetric", - slot_names=_FORWARD_SLOT_SYMMETRIC, + excluded_names=_FORWARD_PRIVATE_SYMMETRIC, ) _add_lane_regions( local_regions, @@ -200,7 +167,7 @@ def build_training_workspace_requirements( lane=lane, phase="forward", space="local", - slot_names=_FORWARD_SLOT_LOCAL, + excluded_names=_FORWARD_PRIVATE_LOCAL, ) _add_lane_regions( symmetric_regions, @@ -208,7 +175,7 @@ def build_training_workspace_requirements( lane=lane, phase="backward", space="symmetric", - slot_names=_BACKWARD_SLOT_SYMMETRIC, + excluded_names=_BACKWARD_PRIVATE_SYMMETRIC, ) _add_lane_regions( local_regions, @@ -216,7 +183,7 @@ def build_training_workspace_requirements( lane=lane, phase="backward", space="local", - slot_names=_BACKWARD_SLOT_LOCAL, + excluded_names=_BACKWARD_PRIVATE_LOCAL, ) forward_symmetric = _region_map(forward_requirements, "symmetric") @@ -224,140 +191,62 @@ def build_training_workspace_requirements( backward_symmetric = _region_map(backward_requirements, "symmetric") backward_local = _region_map(backward_requirements, "local") fc1_c_shape = tuple(int(extent) for extent in forward.kernel.get_aux_output_shapes()["fc1_c"]) - fc1_c_bytes = math.prod(fc1_c_shape) * torch.bfloat16.itemsize - backward_preact = _required_region( - backward_requirements, - "local", - "backward_fc1_preact", - ) - if fc1_c_bytes != backward_preact.nbytes: - raise ValueError("forward fc1_c and backward preactivation byte sizes differ: " f"{fc1_c_bytes} != {backward_preact.nbytes}") - aux_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in backward.kernel.get_aux_output_shapes().items()} - aux_dtypes = { - "fc1_recompute": _DATA_DTYPE, - "fc1_recompute_sf": _SCALE_DTYPE, - "fc1_col_output": _DATA_DTYPE, - "fc1_col_output_sf": _SCALE_DTYPE, - "grad_y2": _DATA_DTYPE, - "grad_y2_sf": torch.uint8, - } - scale_columns = _align_scale_columns(forward.pool_token_capacity) - wgrad_shapes = { - "wgrad_fc1_sfa": ( - _round_up(config.hidden_size, 128), - scale_columns, - ), - "wgrad_fc1_sfb": ( - _round_up(2 * config.intermediate_size, 128), - scale_columns, - ), - "wgrad_fc2_sfa": ( - _round_up(config.intermediate_size, 128), - scale_columns, - ), - "wgrad_fc2_sfb": ( - _round_up(config.hidden_size, 128), - scale_columns, - ), - } - if config.fc1_weight_layout is Fc1WeightLayout.GATE_THEN_UP: - wgrad_shapes["wgrad_fc1_b"] = ( - forward.pool_token_capacity, - 2 * config.intermediate_size, - ) - - for slot in range(slot_count): - for name in sorted(_FORWARD_SLOT_SYMMETRIC): + backward_fc1_preact_shape = tuple(int(extent) for extent in backward.kernel.get_fc1_preact_shape()) + if fc1_c_shape != backward_fc1_preact_shape: + raise ValueError("forward fc1_c and backward fc1_preact shapes differ: " f"{fc1_c_shape} != {backward_fc1_preact_shape}") + for lane in range(lane_count): + for name in sorted(_FORWARD_PRIVATE_SYMMETRIC): if name in _ROUTING_SYMMETRIC: continue symmetric_regions.append( _clone_region( - _slot_name(slot, "forward", "symmetric", name), + _lane_name(lane, "forward", "symmetric", name), forward_symmetric[name], ) ) - for name in sorted(_BACKWARD_SLOT_SYMMETRIC): + for name in sorted(_BACKWARD_PRIVATE_SYMMETRIC): if name in _ROUTING_SYMMETRIC: continue symmetric_regions.append( _clone_region( - _slot_name(slot, "backward", "symmetric", name), + _lane_name(lane, "backward", "symmetric", name), backward_symmetric[name], ) ) - for name in sorted(_FORWARD_SLOT_LOCAL): - if name in _ROUTING_LOCAL: + for name in sorted(_FORWARD_PRIVATE_LOCAL): + if name in _ROUTING_LOCAL or name in _CALLER_OWNED_FORWARD_LOCAL: continue region = forward_local.get(name) if region is not None: local_regions.append( _clone_region( - _slot_name(slot, "forward", "local", name), + _lane_name(lane, "forward", "local", name), region, ) ) - for name in sorted(_BACKWARD_SLOT_LOCAL): + for name in sorted(_BACKWARD_PRIVATE_LOCAL): if name in _ROUTING_LOCAL: continue local_regions.append( _clone_region( - _slot_name(slot, "backward", "local", name), + _lane_name(lane, "backward", "local", name), backward_local[name], ) ) - local_regions.extend( - ( - BufferRegion( - _custom_slot_name(slot, "fc1_preact"), - fc1_c_bytes, - alignment=128, - ), - BufferRegion( - _custom_slot_name(slot, "routing_topk_idx"), - int(config.max_tokens_per_rank) * config.top_k * torch.int32.itemsize, - alignment=16, - ), - BufferRegion( - _custom_slot_name(slot, "valid_route_counts"), - config.experts_per_rank * torch.int32.itemsize, - alignment=16, - ), - BufferRegion( - _custom_slot_name(slot, "expert_offsets"), - config.experts_per_rank * torch.int32.itemsize, - alignment=16, - ), - BufferRegion( - _custom_slot_name(slot, "grad_activation"), - int(config.max_tokens_per_rank) * config.hidden_size * torch.float32.itemsize, - alignment=16, - ), + local_regions.append( + BufferRegion( + _lane_fallback_name(lane, "local", "routing_topk_idx"), + int(config.max_tokens_per_rank) * config.top_k * torch.int32.itemsize, + alignment=16, ) ) symmetric_regions.append( BufferRegion( - _custom_slot_symmetric_name(slot, "routing_topk_weights"), + _lane_fallback_name(lane, "symmetric", "routing_topk_weights"), int(config.max_tokens_per_rank) * config.top_k * torch.float32.itemsize, alignment=16, ) ) - for name, dtype in aux_dtypes.items(): - local_regions.append( - BufferRegion( - _custom_slot_name(slot, name), - math.prod(aux_shapes[name]) * dtype.itemsize, - alignment=128 if name != "grad_y2_sf" else 16, - ) - ) - for name, shape in wgrad_shapes.items(): - local_regions.append( - BufferRegion( - _custom_slot_name(slot, name), - math.prod(shape), - alignment=128, - ) - ) - return WorkspaceRequirements( max_tokens_per_rank=int(config.max_tokens_per_rank), symmetric_regions=tuple(symmetric_regions), @@ -427,7 +316,7 @@ def _harmonize_symmetric_regions( f"{region.name}:{region.nbytes}->{harmonized_size}" for region, harmonized_size in zip(regions, harmonized_sizes) if region.nbytes != harmonized_size ) _runtime_debug( - "training-resources.symmetric-layout-harmonized", + "training-state.symmetric-layout-harmonized", region_count=len(regions), changed_regions=changes, ) @@ -448,24 +337,6 @@ def _harmonize_symmetric_regions( ) -def _block_scaled_tensor_abi(tensor) -> dict[str, object]: - return { - "format": tensor.format.value, - "axis": int(tensor.axis), - "logical_shape": list(tensor.logical_shape), - "data": { - "shape": list(tensor.data.shape), - "stride": list(tensor.data.stride()), - "dtype": str(tensor.data.dtype), - }, - "scale": { - "shape": list(tensor.scale.shape), - "stride": list(tensor.scale.stride()), - "dtype": str(tensor.scale.dtype), - }, - } - - def _workspace_abi(requirements: WorkspaceRequirements) -> dict[str, object]: def regions(values) -> list[dict[str, object]]: return [ @@ -493,8 +364,8 @@ def _prepared_kernel_abi(prepared) -> dict[str, object]: "launch": { "cluster_count": int(prepared.launch_cluster_count), "threads_per_cta": int(kernel.threads_per_cta), - "occupancy": int(getattr(kernel, "occupancy", 1)), - "smem_capacity": int(getattr(kernel, "smem_capacity", 0)), + "occupancy": int(kernel.occupancy), + "smem_capacity": int(kernel.smem_capacity), }, "workspace": _workspace_abi(prepared.workspace_requirements), "pool_token_capacity": int(prepared.pool_token_capacity), @@ -505,29 +376,18 @@ def _build_training_abi_facts( config: ForwardConfig, forward: PreparedMxfp8Kernel, backward: PreparedMxfp8BackwardKernel, - weights: MoeEpTrainingWeights, requirements: WorkspaceRequirements, *, - slot_count: int, lane_count: int, source_tree_digest: str | None = None, ) -> dict[str, object]: - """Return rank-independent JSON-safe facts for one training resource ABI.""" + """Return rank-independent JSON-safe facts for the stateless training ABI.""" if source_tree_digest is None: source_root = Path(__file__).resolve().parents[1] / "cutedsl_src" source_tree_digest = source_tree_sha256(source_root) - weight_facts = { - name: _block_scaled_tensor_abi(getattr(weights, name)) - for name in ( - "forward_fc1", - "forward_fc2", - "backward_w2_transpose", - "backward_w1_transpose", - ) - } return { - "schema_version": 1, + "schema_version": 2, "source_tree_sha256": source_tree_digest, "ep": { "size": int(config.ep_size), @@ -551,11 +411,10 @@ def _build_training_abi_facts( "gate_up_clamp": config.gate_up_clamp, }, "resources": { - "slot_count": int(slot_count), "lane_count": int(lane_count), "workspace": _workspace_abi(requirements), }, - "weights": weight_facts, + "native_weight_layouts": [layout.value for layout in MoeEpNativeWeightLayout], "forward_kernel": _prepared_kernel_abi(forward), "backward_kernel": _prepared_kernel_abi(backward), } @@ -587,48 +446,31 @@ def _verify_training_abi_across_ranks( @dataclass(frozen=True) -class Mxfp8TrainingSlotViews: - """Persistent tensors that survive from forward through wgrad consumption.""" +class Mxfp8TrainingLaneScratch: + """Private fixed-capacity transport and routing tensors for one lane.""" index: int routing_topk_idx: torch.Tensor routing_topk_weights: torch.Tensor - fc1_preact: torch.Tensor - col_quant_data: torch.Tensor | None - col_quant_sf: torch.Tensor | None - valid_route_counts: torch.Tensor - expert_offsets: torch.Tensor forward_output: torch.Tensor backward_output: torch.Tensor - grad_activation: torch.Tensor dprob: torch.Tensor forward_overflow: torch.Tensor backward_overflow: torch.Tensor - fc1_recompute: torch.Tensor - fc1_recompute_sf: torch.Tensor - fc1_col_output: torch.Tensor - fc1_col_output_sf: torch.Tensor - grad_y2: torch.Tensor - grad_y2_sf: torch.Tensor - wgrad_fc1_b: torch.Tensor | None - wgrad_fc1_sfa: torch.Tensor - wgrad_fc1_sfb: torch.Tensor - wgrad_fc2_sfa: torch.Tensor - wgrad_fc2_sfb: torch.Tensor @dataclass(frozen=True) class Mxfp8TrainingExecutionViews: - """One slot bound to one mutable execution lane.""" + """Prepared workspaces and private scratch for one execution lane.""" - slot: Mxfp8TrainingSlotViews + scratch: Mxfp8TrainingLaneScratch forward: PreparedResources backward: PreparedResources - forward_expert_size_snapshot: torch.Tensor | None + forward_expert_size_snapshot: torch.Tensor -class Mxfp8TrainingResourceOwner: - """Own one combined symmetric/local root for N slots and M lanes.""" +class Mxfp8TrainingState: + """Own only private runtime and per-lane training scratch.""" def __init__( self, @@ -636,9 +478,7 @@ def __init__( device: torch.device, forward: PreparedMxfp8Kernel, backward: PreparedMxfp8BackwardKernel, - weights: MoeEpTrainingWeights, *, - slot_count: int, lane_count: int, runtime_manager: Optional[RuntimeManager] = None, symmetric_provider: Optional[SymmetricMemoryProvider] = None, @@ -648,30 +488,17 @@ def __init__( self.device = torch.device(device) self.forward_prepared = forward self.backward_prepared = backward - self.weight_bindings = Mxfp8TrainingWeightBindings( - weights, - fc1_weight_layout=config.fc1_weight_layout, - ) self.stager = Mxfp8TrainingStager(config.hidden_size, config.top_k) - self.wgrad_exporter = Mxfp8TrainingWgradExporter( - experts=config.experts_per_rank, - hidden=config.hidden_size, - intermediate=config.intermediate_size, - sf_padding=backward.config.sf_padding_block, - fc1_weight_layout=config.fc1_weight_layout, - ) self.beta = torch.ones( (config.experts_per_rank,), dtype=torch.float32, device=self.device, ) - self.slot_count = slot_count self.lane_count = lane_count self.requirements = build_training_workspace_requirements( config, forward, backward, - slot_count=slot_count, lane_count=lane_count, ) self._runtime_manager = runtime_manager or get_runtime_manager() @@ -679,48 +506,34 @@ def __init__( self._local_provider = local_provider self._runtime: RuntimeHandle | None = None self._workspace: WorkspaceOwner | None = None - self._abi_fingerprint: str | None = None self._closed = False self._lock = threading.RLock() - @property - def prepared(self) -> bool: - return not self._closed and self._runtime is not None and self._workspace is not None and self._workspace.allocated - def prepare(self) -> None: with self._lock: if self._closed: - raise RuntimeError("training resources are closed") - if self.prepared: + raise RuntimeError("private training state is closed") + if self._runtime is not None and self._workspace is not None and self._workspace.allocated: return if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("training resources must be prepared before CUDA graph capture") + raise RuntimeError("private training state must be prepared before CUDA graph capture") _runtime_debug( - "training-resources.prepare.begin", - slot_count=self.slot_count, + "training-state.prepare.begin", lane_count=self.lane_count, - local_bytes=( - self.requirements.local_layout.total_bytes - if hasattr(self.requirements, "local_layout") - else sum(region.nbytes for region in self.requirements.local_regions) - ), + local_bytes=sum(region.nbytes for region in self.requirements.local_regions), symmetric_bytes=sum(region.nbytes for region in self.requirements.symmetric_regions), ) - _runtime_debug("training-resources.runtime-acquire.begin") + _runtime_debug("training-state.runtime-acquire.begin") runtime = self._runtime_manager.acquire(self.config, self.device) _runtime_debug( - "training-resources.runtime-acquire.end", - runtime_ref_count=getattr( - self._runtime_manager, - "ref_count", - "?", - ), + "training-state.runtime-acquire.end", + runtime_ref_count=self._runtime_manager.ref_count, ) self._runtime = runtime try: - layout_watchdog = _RuntimeWatchdog("training-resources.symmetric-layout-harmonize") + layout_watchdog = _RuntimeWatchdog("training-state.symmetric-layout-harmonize") layout_watchdog.start() - _runtime_debug("training-resources.symmetric-layout-harmonize.begin") + _runtime_debug("training-state.symmetric-layout-harmonize.begin") try: self.requirements = _harmonize_symmetric_regions( self.requirements, @@ -729,22 +542,20 @@ def prepare(self) -> None: ) finally: layout_watchdog.close() - _runtime_debug("training-resources.symmetric-layout-harmonize.end") + _runtime_debug("training-state.symmetric-layout-harmonize.end") if runtime.world_size > 1: - abi_watchdog = _RuntimeWatchdog("training-resources.abi-handshake") + abi_watchdog = _RuntimeWatchdog("training-state.abi-handshake") abi_watchdog.start() - _runtime_debug("training-resources.abi-handshake.begin") + _runtime_debug("training-state.abi-handshake.begin") try: abi_facts = _build_training_abi_facts( self.config, self.forward_prepared, self.backward_prepared, - self.weight_bindings.weights, self.requirements, - slot_count=self.slot_count, lane_count=self.lane_count, ) - self._abi_fingerprint = _verify_training_abi_across_ranks( + abi_fingerprint = _verify_training_abi_across_ranks( abi_facts, runtime, self.device, @@ -752,10 +563,10 @@ def prepare(self) -> None: finally: abi_watchdog.close() _runtime_debug( - "training-resources.abi-handshake.end", - fingerprint=self._abi_fingerprint, + "training-state.abi-handshake.end", + fingerprint=abi_fingerprint, ) - _runtime_debug("training-resources.workspace-create.begin") + _runtime_debug("training-state.workspace-create.begin") workspace = WorkspaceOwner( self.requirements, runtime, @@ -763,39 +574,39 @@ def prepare(self) -> None: local_provider=self._local_provider, ) _runtime_debug( - "training-resources.workspace-create.end", + "training-state.workspace-create.end", local_bytes=workspace.local_layout.total_bytes, symmetric_bytes=workspace.symmetric_layout.total_bytes, ) self._workspace = workspace - allocation_watchdog = _RuntimeWatchdog("training-resources.workspace-allocate") + allocation_watchdog = _RuntimeWatchdog("training-state.workspace-allocate") allocation_watchdog.start() try: workspace.ensure_allocated() finally: allocation_watchdog.close() - _runtime_debug("training-resources.workspace-allocate.end") + _runtime_debug("training-state.workspace-allocate.end") if runtime.world_size > 1: # Symmetric-root zeroing is asynchronous. No rank may # enter the first device barrier until every peer has # completed allocation and root initialization. - stream_watchdog = _RuntimeWatchdog("training-resources.stream-synchronize") + stream_watchdog = _RuntimeWatchdog("training-state.stream-synchronize") stream_watchdog.start() - _runtime_debug("training-resources.stream-synchronize.begin") + _runtime_debug("training-state.stream-synchronize.begin") try: torch.cuda.current_stream(self.device).synchronize() finally: stream_watchdog.close() - _runtime_debug("training-resources.stream-synchronize.end") + _runtime_debug("training-state.stream-synchronize.end") - barrier_watchdog = _RuntimeWatchdog("training-resources.rank-barrier") + barrier_watchdog = _RuntimeWatchdog("training-state.rank-barrier") barrier_watchdog.start() - _runtime_debug("training-resources.rank-barrier.begin") + _runtime_debug("training-state.rank-barrier.begin") try: dist.barrier(group=runtime.group) finally: barrier_watchdog.close() - _runtime_debug("training-resources.rank-barrier.end") + _runtime_debug("training-state.rank-barrier.end") except Exception: if self._workspace is not None: self._workspace.close() @@ -803,7 +614,7 @@ def prepare(self) -> None: runtime.close() self._runtime = None raise - _runtime_debug("training-resources.prepare.end") + _runtime_debug("training-state.prepare.end") def _flat_views(self, token_count: int) -> WorkspaceViews: self.prepare() @@ -815,31 +626,23 @@ def _phase_workspace( flat: WorkspaceViews, requirements: WorkspaceRequirements, *, - slot: int, lane: int, phase: str, ) -> WorkspaceViews: symmetric = {} local = {} - slot_symmetric = _FORWARD_SLOT_SYMMETRIC if phase == "forward" else _BACKWARD_SLOT_SYMMETRIC - slot_local = _FORWARD_SLOT_LOCAL if phase == "forward" else _BACKWARD_SLOT_LOCAL for region in requirements.symmetric_regions: if region.name in _ROUTING_SYMMETRIC: - symmetric[region.name] = flat.symmetric[_custom_slot_symmetric_name(slot, "routing_topk_weights")] + symmetric[region.name] = flat.symmetric[_lane_fallback_name(lane, "symmetric", "routing_topk_weights")] continue - scope_name = ( - _slot_name(slot, phase, "symmetric", region.name) if region.name in slot_symmetric else _lane_name(lane, phase, "symmetric", region.name) - ) - symmetric[region.name] = flat.symmetric[scope_name] + symmetric[region.name] = flat.symmetric[_lane_name(lane, phase, "symmetric", region.name)] for region in requirements.local_regions: - if region.name in _ROUTING_LOCAL: - local[region.name] = flat.local[_custom_slot_name(slot, "routing_topk_idx")] + if phase == "forward" and region.name in _CALLER_OWNED_FORWARD_LOCAL: continue - if phase == "backward" and region.name == "backward_fc1_preact": - local[region.name] = flat.local[_custom_slot_name(slot, "fc1_preact")] + if region.name in _ROUTING_LOCAL: + local[region.name] = flat.local[_lane_fallback_name(lane, "local", "routing_topk_idx")] continue - scope_name = _slot_name(slot, phase, "local", region.name) if region.name in slot_local else _lane_name(lane, phase, "local", region.name) - local[region.name] = flat.local[scope_name] + local[region.name] = flat.local[_lane_name(lane, phase, "local", region.name)] return WorkspaceViews( token_count=flat.token_count, symmetric=MappingProxyType(symmetric), @@ -847,49 +650,20 @@ def _phase_workspace( peer_mapping=flat.peer_mapping, ) - def _slot_views( + def _lane_scratch_views( self, flat: WorkspaceViews, - slot: int, - ) -> Mxfp8TrainingSlotViews: + lane: int, + ) -> Mxfp8TrainingLaneScratch: config = self.config capacity = int(config.max_tokens_per_rank) - fwd_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in self.forward_prepared.kernel.get_aux_output_shapes().items()} bwd_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in self.backward_prepared.kernel.get_aux_output_shapes().items()} - scale_columns = _align_scale_columns(self.forward_prepared.pool_token_capacity) def local_bytes(name: str) -> torch.Tensor: - return flat.local[_custom_slot_name(slot, name)] - - col_quant_data = None - col_quant_sf = None - col_data_name = _slot_name( - slot, - "forward", - "local", - "col_quant_data", - ) - if col_data_name in flat.local: - col_quant_data = _typed_k_major_view( - flat.local[col_data_name], - _DATA_DTYPE, - fwd_shapes["col_quant_data"], - ) - col_quant_sf = _typed_view( - flat.local[ - _slot_name( - slot, - "forward", - "local", - "col_quant_sf", - ) - ], - torch.uint8, - fwd_shapes["col_quant_sf"], - ) + return flat.local[_lane_fallback_name(lane, "local", name)] - return Mxfp8TrainingSlotViews( - index=slot, + return Mxfp8TrainingLaneScratch( + index=lane, routing_topk_idx=_typed_view( local_bytes("routing_topk_idx"), torch.int32, @@ -897,35 +671,19 @@ def local_bytes(name: str) -> torch.Tensor: ), routing_topk_weights=_typed_view( flat.symmetric[ - _custom_slot_symmetric_name( - slot, + _lane_fallback_name( + lane, + "symmetric", "routing_topk_weights", ) ], torch.float32, (capacity, config.top_k), ), - fc1_preact=_typed_view( - local_bytes("fc1_preact"), - torch.bfloat16, - fwd_shapes["fc1_c"], - ), - col_quant_data=col_quant_data, - col_quant_sf=col_quant_sf, - valid_route_counts=_typed_view( - local_bytes("valid_route_counts"), - torch.int32, - (config.experts_per_rank,), - ), - expert_offsets=_typed_view( - local_bytes("expert_offsets"), - torch.int32, - (config.experts_per_rank,), - ), forward_output=_typed_view( flat.symmetric[ - _slot_name( - slot, + _lane_name( + lane, "forward", "symmetric", "output_data", @@ -936,8 +694,8 @@ def local_bytes(name: str) -> torch.Tensor: ), backward_output=_typed_view( flat.symmetric[ - _slot_name( - slot, + _lane_name( + lane, "backward", "symmetric", "output_data", @@ -946,15 +704,10 @@ def local_bytes(name: str) -> torch.Tensor: torch.bfloat16, (capacity, config.hidden_size), ), - grad_activation=_typed_view( - local_bytes("grad_activation"), - torch.float32, - (capacity, config.hidden_size), - ), dprob=_typed_view( flat.symmetric[ - _slot_name( - slot, + _lane_name( + lane, "backward", "symmetric", "backward_dprob", @@ -965,8 +718,8 @@ def local_bytes(name: str) -> torch.Tensor: ), forward_overflow=_typed_view( flat.local[ - _slot_name( - slot, + _lane_name( + lane, "forward", "local", "overflow_flag", @@ -977,8 +730,8 @@ def local_bytes(name: str) -> torch.Tensor: ), backward_overflow=_typed_view( flat.local[ - _slot_name( - slot, + _lane_name( + lane, "backward", "local", "overflow_flag", @@ -987,118 +740,155 @@ def local_bytes(name: str) -> torch.Tensor: torch.int32, (1,), ), - fc1_recompute=_typed_view( - local_bytes("fc1_recompute"), - _DATA_DTYPE, - bwd_shapes["fc1_recompute"], + ) + + def public_requirements( + self, + ) -> Mapping[ + str, + tuple[tuple[int, ...], tuple[int, ...], torch.dtype, int], + ]: + """Return exact caller-owned output contracts without buffer objects.""" + + config = self.config + capacity = int(config.max_tokens_per_rank) + pool_rows = int(self.forward_prepared.pool_token_capacity) + forward_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in self.forward_prepared.kernel.get_aux_output_shapes().items()} + backward_shapes = {name: tuple(int(extent) for extent in shape) for name, shape in self.backward_prepared.kernel.get_aux_output_shapes().items()} + fc1_sfa_rows = round_up(config.hidden_size, 128) + fc1_sfa_elements = math.prod(forward_shapes["col_quant_sf"]) + if fc1_sfa_elements % fc1_sfa_rows: + raise ValueError("forward fc1_sfa producer size is not atom aligned") + fc2_sfb_rows = round_up(config.hidden_size, 128) + fc2_sfb_elements = math.prod(backward_shapes["grad_y2_sf"]) + if fc2_sfb_elements % fc2_sfb_rows: + raise ValueError("backward fc2_sfb producer size is not atom aligned") + requirements = { + "output": ( + (capacity, config.hidden_size), + (config.hidden_size, 1), + torch.bfloat16, + 16, ), - fc1_recompute_sf=_typed_view( - local_bytes("fc1_recompute_sf"), - _SCALE_DTYPE, - bwd_shapes["fc1_recompute_sf"], + "fc1_preact": ( + forward_shapes["fc1_c"], + (forward_shapes["fc1_c"][1], 1), + torch.bfloat16, + 128, ), - fc1_col_output=_typed_view( - local_bytes("fc1_col_output"), + "fc1_a": ( + (config.hidden_size, pool_rows), + (pool_rows, 1), _DATA_DTYPE, - bwd_shapes["fc1_col_output"], + 128, ), - fc1_col_output_sf=_typed_view( - local_bytes("fc1_col_output_sf"), + "fc1_sfa": ( + (fc1_sfa_rows, fc1_sfa_elements // fc1_sfa_rows), + (fc1_sfa_elements // fc1_sfa_rows, 1), _SCALE_DTYPE, - bwd_shapes["fc1_col_output_sf"], + 128, ), - grad_y2=_typed_k_major_view( - local_bytes("grad_y2"), - _DATA_DTYPE, - bwd_shapes["grad_y2"], + "valid_route_counts": ( + (config.experts_per_rank,), + (1,), + torch.int32, + 16, ), - grad_y2_sf=_typed_view( - local_bytes("grad_y2_sf"), - torch.uint8, - bwd_shapes["grad_y2_sf"], + "expert_offsets": ( + (config.experts_per_rank,), + (1,), + torch.int32, + 16, ), - wgrad_fc1_b=( - None - if config.fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32 - else _typed_k_major_view( - local_bytes("wgrad_fc1_b"), - _DATA_DTYPE, - ( - self.forward_prepared.pool_token_capacity, - 2 * config.intermediate_size, - ), - ) + "grad_activation": ( + (capacity, config.hidden_size), + (config.hidden_size, 1), + torch.float32, + 16, ), - wgrad_fc1_sfa=_typed_view( - local_bytes("wgrad_fc1_sfa"), - _SCALE_DTYPE, - (_round_up(config.hidden_size, 128), scale_columns), + "dprob": ( + backward_shapes["dprob"], + (backward_shapes["dprob"][1], 1), + torch.float32, + 16, ), - wgrad_fc1_sfb=_typed_view( - local_bytes("wgrad_fc1_sfb"), + "fc1_b": ( + backward_shapes["fc1_col_output"], + (2 * config.intermediate_size, 1), + _DATA_DTYPE, + 128, + ), + "fc1_sfb": ( + backward_shapes["fc1_col_output_sf"], + (backward_shapes["fc1_col_output_sf"][1], 1), _SCALE_DTYPE, - ( - _round_up(2 * config.intermediate_size, 128), - scale_columns, - ), + 128, ), - wgrad_fc2_sfa=_typed_view( - local_bytes("wgrad_fc2_sfa"), + "fc2_a": ( + (config.intermediate_size, pool_rows), + (1, config.intermediate_size), + _DATA_DTYPE, + 128, + ), + "fc2_sfa": ( + backward_shapes["fc1_recompute_sf"], + (backward_shapes["fc1_recompute_sf"][1], 1), _SCALE_DTYPE, - ( - _round_up(config.intermediate_size, 128), - scale_columns, - ), + 128, ), - wgrad_fc2_sfb=_typed_view( - local_bytes("wgrad_fc2_sfb"), + "fc2_b": ( + backward_shapes["grad_y2"], + (1, pool_rows), + _DATA_DTYPE, + 128, + ), + "fc2_sfb": ( + (fc2_sfb_rows, fc2_sfb_elements // fc2_sfb_rows), + (fc2_sfb_elements // fc2_sfb_rows, 1), _SCALE_DTYPE, - (_round_up(config.hidden_size, 128), scale_columns), + 128, ), - ) + } + return MappingProxyType(requirements) def views( self, *, - slot: int, lane: int, token_count: int, ) -> Mxfp8TrainingExecutionViews: with self._lock: - if slot < 0 or slot >= self.slot_count: - raise ValueError(f"slot {slot} is outside [0, {self.slot_count})") if lane < 0 or lane >= self.lane_count: raise ValueError(f"lane {lane} is outside [0, {self.lane_count})") + col_quant_sizes_offset = self.forward_prepared.col_quant_sizes_offset + if col_quant_sizes_offset is None: + raise RuntimeError("training preparation requires a persistent col-quant expert-size snapshot") flat = self._flat_views(token_count) forward_workspace = self._phase_workspace( flat, self.forward_prepared.workspace_requirements, - slot=slot, lane=lane, phase="forward", ) backward_workspace = self._phase_workspace( flat, self.backward_prepared.workspace_requirements, - slot=slot, lane=lane, phase="backward", ) - snapshot = None - if self.forward_prepared.col_quant_sizes_offset is not None: - snapshot_bytes = forward_workspace.local["kernel_local_workspace"].narrow( - 0, - self.forward_prepared.col_quant_sizes_offset, - self.forward_prepared.col_quant_sizes_bytes, - ) - snapshot = _typed_view( - snapshot_bytes, - torch.int32, - (self.config.experts_per_rank,), - ) + snapshot_bytes = forward_workspace.local["kernel_local_workspace"].narrow( + 0, + col_quant_sizes_offset, + self.forward_prepared.col_quant_sizes_bytes, + ) + snapshot = _typed_view( + snapshot_bytes, + torch.int32, + (self.config.experts_per_rank,), + ) assert self._runtime is not None return Mxfp8TrainingExecutionViews( - slot=self._slot_views(flat, slot), + scratch=self._lane_scratch_views(flat, lane), forward=PreparedResources( runtime=self._runtime, workspace=forward_workspace, @@ -1110,14 +900,6 @@ def views( forward_expert_size_snapshot=snapshot, ) - def refresh_weights(self) -> None: - """Enqueue fixed-address layout refreshes from the bound MXFP8 pack.""" - - with self._lock: - if self._closed: - raise RuntimeError("training resources are closed") - self.weight_bindings.refresh() - def close(self) -> None: with self._lock: if self._closed: @@ -1130,23 +912,18 @@ def close(self) -> None: self._runtime = None self._closed = True - def finalize_overflow( + def apply_overflow( self, - slots: tuple[int, ...], *, lane: int, + phase: str, ) -> torch.Tensor: - """Aggregate slot flags and apply the public error/drop policy.""" - - if not slots: - raise ValueError("finalize_overflow requires at least one slot") - if len(set(slots)) != len(slots): - raise ValueError("finalize_overflow slots must be unique") - for slot in slots: - if slot < 0 or slot >= self.slot_count: - raise ValueError(f"slot {slot} is outside [0, {self.slot_count})") + """Apply the configured policy to one phase's private overflow flag.""" + if lane < 0 or lane >= self.lane_count: raise ValueError(f"lane {lane} is outside [0, {self.lane_count})") + if phase not in ("forward", "backward"): + raise ValueError(f"phase must be 'forward' or 'backward', got {phase!r}") flat = self._flat_views(0) global_overflow = _typed_view( flat.local[ @@ -1160,22 +937,19 @@ def finalize_overflow( torch.int32, (1,), ) - global_overflow.zero_() - for slot in slots: - for phase in ("forward", "backward"): - flag = _typed_view( - flat.local[ - _slot_name( - slot, - phase, - "local", - "overflow_flag", - ) - ], - torch.int32, - (1,), + flag = _typed_view( + flat.local[ + _lane_name( + lane, + phase, + "local", + "overflow_flag", ) - torch.maximum(global_overflow, flag, out=global_overflow) + ], + torch.int32, + (1,), + ) + global_overflow.copy_(flag) assert self._runtime is not None if self._runtime.world_size > 1: dist.all_reduce( @@ -1186,7 +960,7 @@ def finalize_overflow( if not self.config.drop_on_overflow: assert_async = getattr(torch, "_assert_async", None) if assert_async is None: - raise RuntimeError("drop_on_overflow=False training resources require " "torch._assert_async") + raise RuntimeError("drop_on_overflow=False training requires torch._assert_async") overflow_ok = _typed_view( flat.local[ _lane_name( @@ -1202,14 +976,12 @@ def finalize_overflow( torch.eq(global_overflow, 0, out=overflow_ok) assert_async( overflow_ok, - "Rubin MegaMoE receive route-pool overflow; " "the fixed-slot outputs are invalid", + f"Rubin MegaMoE receive route-pool overflow; the {phase} " "outputs are invalid", ) return global_overflow __all__ = [ "Mxfp8TrainingExecutionViews", - "Mxfp8TrainingResourceOwner", - "Mxfp8TrainingSlotViews", - "build_training_workspace_requirements", + "Mxfp8TrainingState", ] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py index a410e1740..7c9f63549 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Cached fused MXFP8 staging for fixed-address training resources.""" +"""Cached BF16/FP32 staging into private symmetric training scratch.""" from __future__ import annotations diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py index f570b9b3d..c9f5946d1 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_stage_kernel.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""One-launch BF16/FP32 to MXFP8 staging for fixed training resources.""" +"""One-launch BF16/FP32 to MXFP8 private symmetric staging.""" from __future__ import annotations diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py index d1e3fe882..c74b2552f 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_weights.py @@ -10,40 +10,22 @@ import torch from ..._contracts import Fc1WeightLayout -from ..._types import BlockScaledTensor, MoeEpTrainingWeights +from ..._math import round_up +from ..._types import ( + BlockScaledTensor, + MoeEpBackwardWeightStaging, + MoeEpBackwardWeights, + MoeEpForwardWeightStaging, + MoeEpForwardWeights, + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, +) +from ..._validation import validate_training_non_aliasing from ._adapter import Mxfp8Weights -def _round_up(value: int, multiple: int) -> int: - return (value + multiple - 1) // multiple * multiple - - -def _empty_k_major_like(tensor: torch.Tensor) -> torch.Tensor: - experts, reduction, output = tensor.shape - return torch.empty_strided( - tensor.shape, - (reduction * output, 1, reduction), - dtype=tensor.dtype, - device=tensor.device, - ) - - -def _empty_blocked_scales( - source: BlockScaledTensor, - *, - raw_rows: int, - raw_columns: int, - dtype: torch.dtype, -) -> torch.Tensor: - experts = source.data.shape[0] - packed_bytes = _round_up(raw_rows, 128) * _round_up(raw_columns, 4) - return torch.empty( - (experts, packed_bytes), - dtype=torch.uint8, - device=source.device, - ).view(dtype) - - def _copy_blocked_scales_plain( target: torch.Tensor, source: torch.Tensor, @@ -137,14 +119,8 @@ def _copy_blocked_scales_gate_up_rows( raw_rows = 2 * intermediate if intermediate % 64 or reduction_blocks % 4: raise ValueError("intermediate and reduction block alignment are invalid") - source_view = ( - source.view(torch.uint8) - .view(experts, reduction_blocks // 4, 4, 2, raw_rows // 128, 2, 32) - .permute(0, 4, 1, 6, 5, 3, 2) - ) - target.view(torch.uint8).view( - experts, raw_rows // 128, reduction_blocks // 4, 32, 2, 2, 4 - ).copy_(source_view) + source_view = source.view(torch.uint8).view(experts, reduction_blocks // 4, 4, 2, raw_rows // 128, 2, 32).permute(0, 4, 1, 6, 5, 3, 2) + target.view(torch.uint8).view(experts, raw_rows // 128, reduction_blocks // 4, 32, 2, 2, 4).copy_(source_view) def _copy_blocked_scales_gate_up_columns( @@ -160,14 +136,8 @@ def _copy_blocked_scales_gate_up_columns( reduction_blocks = intermediate // 32 if output % 128 or reduction_blocks % 2: raise ValueError("output and reduction block alignment are invalid") - source_view = ( - source.view(torch.uint8) - .view(experts, 2, reduction_blocks // 2, 2, output // 128, 4, 32) - .permute(0, 4, 2, 6, 5, 3, 1) - ) - target.view(torch.uint8).view( - experts, output // 128, reduction_blocks // 2, 32, 4, 2, 2 - ).copy_(source_view) + source_view = source.view(torch.uint8).view(experts, 2, reduction_blocks // 2, 2, output // 128, 4, 32).permute(0, 4, 2, 6, 5, 3, 1) + target.view(torch.uint8).view(experts, output // 128, reduction_blocks // 2, 32, 4, 2, 2).copy_(source_view) @dataclass(frozen=True) @@ -180,159 +150,288 @@ class Mxfp8BackwardWeights: fc2_weight_sf: torch.Tensor -class Mxfp8TrainingWeightBindings: - """Direct data bindings with persistent kernel-native scale staging.""" - - def __init__( - self, - weights: MoeEpTrainingWeights, - *, - fc1_weight_layout: Fc1WeightLayout = Fc1WeightLayout.GATE_THEN_UP, - ) -> None: - self.weights = weights - self.fc1_weight_layout = fc1_weight_layout - fwd_fc1 = weights.forward_fc1 - fwd_fc2 = weights.forward_fc2 - bwd_w2t = weights.backward_w2_transpose - bwd_w1t = weights.backward_w1_transpose - self._uses_direct_weight_bindings = ( - fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32 - and fwd_fc1.data.stride(1) == 1 - and fwd_fc2.data.stride(1) == 1 - and bwd_w2t.data.is_contiguous() - and bwd_w1t.data.is_contiguous() - ) - if ( - fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32 - and not self._uses_direct_weight_bindings - ): - raise ValueError( - "weight_interleave_size=32 requires compact K-major forward " - "weights and contiguous backward transpose weights" - ) - - self.forward = Mxfp8Weights( - fc1_weight=( - fwd_fc1.data - if self._uses_direct_weight_bindings - else _empty_k_major_like(fwd_fc1.data) - ), - fc1_weight_sf=_empty_blocked_scales( - fwd_fc1, - raw_rows=fwd_fc1.data.shape[2], - raw_columns=fwd_fc1.data.shape[1] // 32, - dtype=torch.uint8, - ), - fc2_weight=( - fwd_fc2.data - if self._uses_direct_weight_bindings - else _empty_k_major_like(fwd_fc2.data) - ), - fc2_weight_sf=_empty_blocked_scales( - fwd_fc2, - raw_rows=fwd_fc2.data.shape[2], - raw_columns=fwd_fc2.data.shape[1] // 32, - dtype=torch.uint8, - ), - ) - self.backward = Mxfp8BackwardWeights( - fc1_weight=( - bwd_w2t.data - if self._uses_direct_weight_bindings - else torch.empty_like( - bwd_w2t.data, - memory_format=torch.contiguous_format, - ) - ), - fc1_weight_sf=_empty_blocked_scales( - bwd_w2t, - raw_rows=bwd_w2t.data.shape[2], - raw_columns=bwd_w2t.data.shape[1] // 32, - dtype=torch.float8_e8m0fnu, - ), - fc2_weight=( - bwd_w1t.data - if self._uses_direct_weight_bindings - else torch.empty_like( - bwd_w1t.data, - memory_format=torch.contiguous_format, - ) - ), - fc2_weight_sf=_empty_blocked_scales( - bwd_w1t, - raw_rows=bwd_w1t.data.shape[2], - raw_columns=bwd_w1t.data.shape[1] // 32, - dtype=torch.float8_e8m0fnu, - ), - ) - self.refresh() - - def refresh(self) -> None: - """Enqueue fixed-address layout copies; safe to record in a graph.""" - - fwd_fc1 = self.weights.forward_fc1 - fwd_fc2 = self.weights.forward_fc2 - bwd_w2t = self.weights.backward_w2_transpose - bwd_w1t = self.weights.backward_w1_transpose - - intermediate = fwd_fc2.data.shape[1] - if self._uses_direct_weight_bindings: - _copy_blocked_scales_plain( - self.forward.fc1_weight_sf, - fwd_fc1.scale, - raw_rows=fwd_fc1.data.shape[2], - raw_columns=fwd_fc1.data.shape[1] // 32, - ) - else: - _copy_gate_up_interleaved_last( - self.forward.fc1_weight, - fwd_fc1.data, - intermediate, - ) - _copy_blocked_scales_gate_up_rows( - self.forward.fc1_weight_sf, - fwd_fc1.scale, - intermediate=intermediate, - reduction_blocks=fwd_fc1.data.shape[1] // 32, - ) - self.forward.fc2_weight.copy_(fwd_fc2.data) - _copy_blocked_scales_plain( - self.forward.fc2_weight_sf, - fwd_fc2.scale, - raw_rows=fwd_fc2.data.shape[2], - raw_columns=fwd_fc2.data.shape[1] // 32, - ) +def _expect_staging_tensor( + name: str, + tensor: torch.Tensor, + *, + shape: tuple[int, ...], + stride: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> None: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.layout is not torch.strided: + raise ValueError(f"{name} must use torch.strided layout, got {tensor.layout}") + if tuple(tensor.shape) != shape or tuple(tensor.stride()) != stride: + raise ValueError(f"{name} must have shape={shape}, stride={stride}; got " f"shape={tuple(tensor.shape)}, stride={tuple(tensor.stride())}") + if tensor.dtype is not dtype: + raise ValueError(f"{name} must have dtype {dtype}, got {tensor.dtype}") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}") + if tensor.data_ptr() % 16: + raise ValueError(f"{name} must be at least 16-byte aligned") + + +def _expect_scale_staging( + name: str, + tensor: torch.Tensor, + *, + experts: int, + raw_rows: int, + raw_columns: int, + dtype: torch.dtype, + device: torch.device, +) -> None: + elements = round_up(raw_rows, 128) * round_up(raw_columns, 4) + _expect_staging_tensor( + name, + tensor, + shape=(experts, elements), + stride=(elements, 1), + dtype=dtype, + device=device, + ) - if not self._uses_direct_weight_bindings: - self.backward.fc1_weight.copy_(bwd_w2t.data) - _copy_blocked_scales_plain( - self.backward.fc1_weight_sf, - bwd_w2t.scale, - raw_rows=bwd_w2t.data.shape[2], - raw_columns=bwd_w2t.data.shape[1] // 32, - ) - if self._uses_direct_weight_bindings: - _copy_blocked_scales_plain( - self.backward.fc2_weight_sf, - bwd_w1t.scale, - raw_rows=bwd_w1t.data.shape[2], - raw_columns=bwd_w1t.data.shape[1] // 32, - ) - else: - _copy_gate_up_interleaved_reduction( - self.backward.fc2_weight, - bwd_w1t.data, - intermediate, - ) - _copy_blocked_scales_gate_up_columns( - self.backward.fc2_weight_sf, - bwd_w1t.scale, - intermediate=intermediate, - output=bwd_w1t.data.shape[2], - ) + +def forward_native_to_kernel(weights: MoeEpNativeForwardWeights) -> Mxfp8Weights: + """Create kernel views without allocating, copying, or retaining inputs.""" + + return Mxfp8Weights( + fc1_weight=weights.fc1.payload, + fc1_weight_sf=weights.fc1.scale, + fc2_weight=weights.fc2.payload, + fc2_weight_sf=weights.fc2.scale, + ) + + +def backward_native_to_kernel( + weights: MoeEpNativeBackwardWeights, +) -> Mxfp8BackwardWeights: + """Create kernel views without allocating, copying, or retaining inputs.""" + + return Mxfp8BackwardWeights( + fc1_weight=weights.w2_transpose.payload, + fc1_weight_sf=weights.w2_transpose.scale, + fc2_weight=weights.w1_transpose.payload, + fc2_weight_sf=weights.w1_transpose.scale, + ) + + +def materialize_forward( + weights: MoeEpForwardWeights, + *, + out: MoeEpForwardWeightStaging, + fc1_weight_layout: Fc1WeightLayout, +) -> MoeEpNativeForwardWeights: + """Materialize source forward weights into caller-owned native storage.""" + + if fc1_weight_layout is not Fc1WeightLayout.GATE_UP_INTERLEAVED_32: + raise ValueError("native training materialization requires weight_interleave_size=32") + if not isinstance(weights, MoeEpForwardWeights): + raise TypeError("weights must be a MoeEpForwardWeights") + if not isinstance(out, MoeEpForwardWeightStaging): + raise TypeError("out must be a MoeEpForwardWeightStaging") + fc1 = weights.fc1 + fc2 = weights.fc2 + for name, value in (("weights.fc1", fc1), ("weights.fc2", fc2)): + if not isinstance(value, BlockScaledTensor) or value.format.value != "mxfp8" or value.axis != 1: + raise TypeError(f"{name} must be an axis-1 MXFP8 BlockScaledTensor") + if not value.data.is_contiguous() or not value.scale.is_contiguous(): + raise ValueError(f"{name} data and scale must be contiguous") + experts, hidden, gate_up = fc1.data.shape + intermediate = fc2.data.shape[1] + if tuple(fc2.data.shape) != (experts, intermediate, hidden) or gate_up != 2 * intermediate: + raise ValueError("forward source weight shapes are inconsistent") + if fc2.device != fc1.device: + raise ValueError("forward source weights must share one device") + sf_dtype = torch.float8_e8m0fnu + _expect_staging_tensor( + "out.fc1_payload", + out.fc1_payload, + shape=tuple(fc1.data.shape), + stride=(hidden * gate_up, 1, hidden), + dtype=fc1.data.dtype, + device=fc1.device, + ) + _expect_staging_tensor( + "out.fc2_payload", + out.fc2_payload, + shape=tuple(fc2.data.shape), + stride=(intermediate * hidden, 1, intermediate), + dtype=fc2.data.dtype, + device=fc1.device, + ) + _expect_scale_staging( + "out.fc1_scale", + out.fc1_scale, + experts=experts, + raw_rows=gate_up, + raw_columns=hidden // 32, + dtype=sf_dtype, + device=fc1.device, + ) + _expect_scale_staging( + "out.fc2_scale", + out.fc2_scale, + experts=experts, + raw_rows=hidden, + raw_columns=intermediate // 32, + dtype=sf_dtype, + device=fc1.device, + ) + validate_training_non_aliasing( + { + "weights.fc1.data": fc1.data, + "weights.fc1.scale": fc1.scale, + "weights.fc2.data": fc2.data, + "weights.fc2.scale": fc2.scale, + "out.fc1_payload": out.fc1_payload, + "out.fc1_scale": out.fc1_scale, + "out.fc2_payload": out.fc2_payload, + "out.fc2_scale": out.fc2_scale, + } + ) + _copy_gate_up_interleaved_last(out.fc1_payload, fc1.data, intermediate) + _copy_blocked_scales_gate_up_rows( + out.fc1_scale, + fc1.scale, + intermediate=intermediate, + reduction_blocks=hidden // 32, + ) + out.fc2_payload.copy_(fc2.data) + _copy_blocked_scales_plain( + out.fc2_scale, + fc2.scale, + raw_rows=hidden, + raw_columns=intermediate // 32, + ) + return MoeEpNativeForwardWeights( + fc1=MoeEpNativeWeight( + out.fc1_payload, + out.fc1_scale, + MoeEpNativeWeightLayout.FORWARD_FC1_GATE_UP_INTERLEAVED_32_V1, + ), + fc2=MoeEpNativeWeight( + out.fc2_payload, + out.fc2_scale, + MoeEpNativeWeightLayout.FORWARD_FC2_K_MAJOR_V1, + ), + ) + + +def materialize_backward( + weights: MoeEpBackwardWeights, + *, + out: MoeEpBackwardWeightStaging, + fc1_weight_layout: Fc1WeightLayout, +) -> MoeEpNativeBackwardWeights: + """Materialize source backward weights into caller-owned native storage.""" + + if fc1_weight_layout is not Fc1WeightLayout.GATE_UP_INTERLEAVED_32: + raise ValueError("native training materialization requires weight_interleave_size=32") + if not isinstance(weights, MoeEpBackwardWeights): + raise TypeError("weights must be a MoeEpBackwardWeights") + if not isinstance(out, MoeEpBackwardWeightStaging): + raise TypeError("out must be a MoeEpBackwardWeightStaging") + w2t = weights.w2_transpose + w1t = weights.w1_transpose + for name, value in ( + ("weights.w2_transpose", w2t), + ("weights.w1_transpose", w1t), + ): + if not isinstance(value, BlockScaledTensor) or value.format.value != "mxfp8" or value.axis != 1: + raise TypeError(f"{name} must be an axis-1 MXFP8 BlockScaledTensor") + if not value.data.is_contiguous() or not value.scale.is_contiguous(): + raise ValueError(f"{name} data and scale must be contiguous") + experts, hidden, intermediate = w2t.data.shape + if tuple(w1t.data.shape) != (experts, 2 * intermediate, hidden): + raise ValueError("backward source weight shapes are inconsistent") + if w1t.device != w2t.device: + raise ValueError("backward source weights must share one device") + _expect_staging_tensor( + "out.w2_transpose_payload", + out.w2_transpose_payload, + shape=tuple(w2t.data.shape), + stride=(hidden * intermediate, intermediate, 1), + dtype=w2t.data.dtype, + device=w2t.device, + ) + _expect_staging_tensor( + "out.w1_transpose_payload", + out.w1_transpose_payload, + shape=tuple(w1t.data.shape), + stride=(2 * intermediate * hidden, hidden, 1), + dtype=w1t.data.dtype, + device=w2t.device, + ) + sf_dtype = torch.float8_e8m0fnu + _expect_scale_staging( + "out.w2_transpose_scale", + out.w2_transpose_scale, + experts=experts, + raw_rows=intermediate, + raw_columns=hidden // 32, + dtype=sf_dtype, + device=w2t.device, + ) + _expect_scale_staging( + "out.w1_transpose_scale", + out.w1_transpose_scale, + experts=experts, + raw_rows=hidden, + raw_columns=2 * intermediate // 32, + dtype=sf_dtype, + device=w2t.device, + ) + validate_training_non_aliasing( + { + "weights.w2_transpose.data": w2t.data, + "weights.w2_transpose.scale": w2t.scale, + "weights.w1_transpose.data": w1t.data, + "weights.w1_transpose.scale": w1t.scale, + "out.w2_transpose_payload": out.w2_transpose_payload, + "out.w2_transpose_scale": out.w2_transpose_scale, + "out.w1_transpose_payload": out.w1_transpose_payload, + "out.w1_transpose_scale": out.w1_transpose_scale, + } + ) + out.w2_transpose_payload.copy_(w2t.data) + _copy_blocked_scales_plain( + out.w2_transpose_scale, + w2t.scale, + raw_rows=intermediate, + raw_columns=hidden // 32, + ) + _copy_gate_up_interleaved_reduction( + out.w1_transpose_payload, + w1t.data, + intermediate, + ) + _copy_blocked_scales_gate_up_columns( + out.w1_transpose_scale, + w1t.scale, + intermediate=intermediate, + output=hidden, + ) + return MoeEpNativeBackwardWeights( + w2_transpose=MoeEpNativeWeight( + out.w2_transpose_payload, + out.w2_transpose_scale, + MoeEpNativeWeightLayout.BACKWARD_W2_TRANSPOSE_V1, + ), + w1_transpose=MoeEpNativeWeight( + out.w1_transpose_payload, + out.w1_transpose_scale, + MoeEpNativeWeightLayout.BACKWARD_W1_TRANSPOSE_GATE_UP_INTERLEAVED_32_V1, + ), + ) __all__ = [ "Mxfp8BackwardWeights", - "Mxfp8TrainingWeightBindings", + "backward_native_to_kernel", + "forward_native_to_kernel", + "materialize_backward", + "materialize_forward", ] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py index 4a413e349..a0eea6167 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py +++ b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad.py @@ -1,181 +1,50 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Fixed-address WGrad operand materialization for the training graph path.""" +"""Pure, zero-copy assembly of caller-owned WGrad operand views.""" from __future__ import annotations -import threading -from typing import TYPE_CHECKING - import torch -from ..._contracts import Fc1WeightLayout -from ..._types import MoeEpTrainingWgradOperands -from ._launch import _to_cute - -if TYPE_CHECKING: - from ._training_resources import Mxfp8TrainingSlotViews - - -class Mxfp8TrainingWgradExporter: - """Own scale-expansion compiles; every export writes existing buffers.""" - - def __init__( - self, - *, - experts: int, - hidden: int, - intermediate: int, - sf_padding: int = 128, - fc1_weight_layout: Fc1WeightLayout = Fc1WeightLayout.GATE_THEN_UP, - ) -> None: - self.experts = int(experts) - self.hidden = int(hidden) - self.intermediate = int(intermediate) - self.sf_padding = int(sf_padding) - self.fc1_weight_layout = fc1_weight_layout - self._compiled: dict[tuple[int, int, int | None], object] = {} - self._lock = threading.RLock() - - @staticmethod - def _copy_gate_up_data( - target: torch.Tensor, - source: torch.Tensor, - intermediate: int, - ) -> None: - """Deinterleave FC1 dY while copying into K-major staging.""" - pool_rows = source.shape[0] - pairs = intermediate // 32 - source_view = source.view(pool_rows, pairs, 2, 32).permute(0, 2, 1, 3) - target_view = target.as_strided( - (pool_rows, 2, pairs, 32), - ( - target.stride(0), - intermediate * target.stride(1), - 32 * target.stride(1), - target.stride(1), - ), - ) - target_view.copy_(source_view) - - def _expand_scales( - self, - source: torch.Tensor, - counts: torch.Tensor, - offsets: torch.Tensor, - output: torch.Tensor, - *, - non_k_size: int, - deinterleave_gate_up: int | None = None, - ) -> None: - if source.dtype not in (torch.uint8, torch.float8_e8m0fnu): - raise TypeError("WGrad source scales must use Uint8 or E8M0") - if output.dtype is not torch.float8_e8m0fnu: - raise TypeError("WGrad output scales must use E8M0") - source_bytes = source.view(torch.uint8).reshape(-1) - output_bytes = output.view(torch.uint8).reshape(-1) - key = (int(non_k_size), self.sf_padding, deinterleave_gate_up) - import cuda.bindings.driver as cuda - - stream = torch.cuda.current_stream(output.device) - args = ( - _to_cute(source_bytes, dynamic_layout=False), - _to_cute(counts, assumed_align=4, dynamic_layout=False), - _to_cute(offsets, assumed_align=4, dynamic_layout=False), - _to_cute(output_bytes, dynamic_layout=False), - cuda.CUstream(stream.cuda_stream), - ) - with self._lock: - compiled = self._compiled.get(key) - if compiled is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("WGrad scale expansion must be compiled before " "CUDA graph capture") - import cutlass.cute as cute - - from ._training_wgrad_kernel import ( - Mxfp8TrainingScaleExpandKernel, - ) - - kernel = Mxfp8TrainingScaleExpandKernel( - non_k_size=non_k_size, - expert_count=self.experts, - source_sf_padding=self.sf_padding, - deinterleave_gate_up=deinterleave_gate_up, - ) - compiled = cute.compile(kernel, *args) - self._compiled[key] = compiled - compiled(*args) - - def export( - self, - slot: "Mxfp8TrainingSlotViews", - ) -> MoeEpTrainingWgradOperands: - """Write and return the fixed-capacity grouped-WGrad operands.""" - - if slot.col_quant_data is None or slot.col_quant_sf is None: - raise RuntimeError("training WGrad export requires forward col-quant") - pool_rows = slot.fc1_recompute.shape[0] - if slot.col_quant_data.shape[0] != pool_rows: - raise RuntimeError("forward/backward WGrad pool capacities differ") - - fc1_b = slot.fc1_col_output - if self.fc1_weight_layout is Fc1WeightLayout.GATE_THEN_UP: - if slot.wgrad_fc1_b is None: - raise RuntimeError("conventional W1 requires FC1 WGrad staging") - self._copy_gate_up_data( - slot.wgrad_fc1_b, - slot.fc1_col_output, - self.intermediate, - ) - fc1_b = slot.wgrad_fc1_b - - self._expand_scales( - slot.col_quant_sf, - slot.valid_route_counts, - slot.expert_offsets, - slot.wgrad_fc1_sfa, - non_k_size=self.hidden, - ) - self._expand_scales( - slot.fc1_col_output_sf, - slot.valid_route_counts, - slot.expert_offsets, - slot.wgrad_fc1_sfb, - non_k_size=2 * self.intermediate, - deinterleave_gate_up=( - self.intermediate - if self.fc1_weight_layout is Fc1WeightLayout.GATE_THEN_UP - else None - ), - ) - self._expand_scales( - slot.fc1_recompute_sf, - slot.valid_route_counts, - slot.expert_offsets, - slot.wgrad_fc2_sfa, - non_k_size=self.intermediate, - ) - self._expand_scales( - slot.grad_y2_sf, - slot.valid_route_counts, - slot.expert_offsets, - slot.wgrad_fc2_sfb, - non_k_size=self.hidden, - ) - - return MoeEpTrainingWgradOperands( - fc1_a=slot.col_quant_data.transpose(0, 1), - fc1_sfa=slot.wgrad_fc1_sfa, - fc1_b=fc1_b, - fc1_sfb=slot.wgrad_fc1_sfb, - fc2_a=slot.fc1_recompute.transpose(0, 1), - fc2_sfa=slot.wgrad_fc2_sfa, - fc2_b=slot.grad_y2, - fc2_sfb=slot.wgrad_fc2_sfb, - expert_offsets=slot.expert_offsets, - valid_route_counts=slot.valid_route_counts, - ) - - -__all__ = ["Mxfp8TrainingWgradExporter"] +from ..._types import ( + MoeEpTrainingBackwardOutputs, + MoeEpTrainingWgradOperands, +) + + +def assemble_training_wgrad_operands( + *, + fc1_a: torch.Tensor, + fc1_sfa: torch.Tensor, + valid_route_counts: torch.Tensor, + expert_offsets: torch.Tensor, + backward: MoeEpTrainingBackwardOutputs, +) -> MoeEpTrainingWgradOperands: + """Return non-owning views after producers wrote their final layouts.""" + + required = ( + backward.fc1_b, + backward.fc1_sfb, + backward.fc2_a, + backward.fc2_sfa, + backward.fc2_b, + backward.fc2_sfb, + ) + if any(value is None for value in required): + raise ValueError("all backward WGrad outputs are required for assembly") + return MoeEpTrainingWgradOperands( + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + fc1_b=backward.fc1_b, + fc1_sfb=backward.fc1_sfb, + fc2_a=backward.fc2_a, + fc2_sfa=backward.fc2_sfa, + fc2_b=backward.fc2_b, + fc2_sfb=backward.fc2_sfb, + expert_offsets=expert_offsets, + valid_route_counts=valid_route_counts, + ) + + +__all__ = ["assemble_training_wgrad_operands"] diff --git a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py deleted file mode 100644 index 16f63243e..000000000 --- a/python/cudnn/moe_ep/_megamoe_backend/mxfp8/_training_wgrad_kernel.py +++ /dev/null @@ -1,145 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Device expansion of per-expert SF atoms into fixed-capacity WGrad ABI.""" - -from __future__ import annotations - -import cuda.bindings.driver as cuda - -import cutlass -import cutlass.cute as cute -from cutlass.cutlass_dsl import Int32 - - -class Mxfp8TrainingScaleExpandKernel: - """Expand 128-row SF extents into 256-row fixed-capacity segments.""" - - _threads = 256 - _atom_bytes = 512 - _neutral_e8m0 = 127 - - def __init__( - self, - *, - non_k_size: int, - expert_count: int, - source_sf_padding: int, - deinterleave_gate_up: int | None = None, - ) -> None: - self.non_k_size = int(non_k_size) - self.expert_count = int(expert_count) - self.source_sf_padding = int(source_sf_padding) - self.deinterleave_gate_up = ( - None if deinterleave_gate_up is None else int(deinterleave_gate_up) - ) - if self.non_k_size <= 0 or self.non_k_size % 128: - raise ValueError("WGrad scale non-K size must be divisible by 128") - if self.expert_count <= 0: - raise ValueError("WGrad scale expansion requires experts") - if self.source_sf_padding <= 0 or self.source_sf_padding % 128: - raise ValueError("WGrad source SF padding must be a positive multiple of 128") - if ( - self.deinterleave_gate_up is not None - and self.non_k_size != 2 * self.deinterleave_gate_up - ): - raise ValueError("gate/up scale deinterleave size mismatch") - - @cute.jit - def __call__( - self, - source: cute.Tensor, - valid_counts: cute.Tensor, - expert_offsets: cute.Tensor, - output: cute.Tensor, - stream: cuda.CUstream, - ) -> None: - output_bytes = cute.size(output) - self._kernel( - source, - valid_counts, - expert_offsets, - output, - ).launch( - grid=[ - output_bytes // self._threads, - 1, - 1, - ], - block=[self._threads, 1, 1], - stream=stream, - min_blocks_per_mp=1, - ) - - @cute.kernel - def _kernel( - self, - source: cute.Tensor, - valid_counts: cute.Tensor, - expert_offsets: cute.Tensor, - output: cute.Tensor, - ) -> None: - linear = cute.arch.block_idx()[0] * Int32(self._threads) + cute.arch.thread_idx()[0] - - atom_bytes: cutlass.Constexpr[int] = self._atom_bytes - non_k_atoms: cutlass.Constexpr[int] = self.non_k_size // 128 - atom = linear // Int32(atom_bytes) - byte_in_atom = linear % Int32(atom_bytes) - value = cutlass.Uint8(self._neutral_e8m0) - target_atom_base = Int32(0) - source_atom_base = Int32(0) - previous_end = Int32(0) - - for expert in cutlass.range_constexpr(self.expert_count): - end = Int32(expert_offsets[expert]) - target_token_atoms = (end - previous_end) // Int32(128) - source_token_atoms = ((Int32(valid_counts[expert]) + Int32(self.source_sf_padding - 1)) // Int32(self.source_sf_padding)) * Int32( - self.source_sf_padding // 128 - ) - target_atom_count = Int32(non_k_atoms) * target_token_atoms - in_expert = (atom >= target_atom_base) & (atom < target_atom_base + target_atom_count) - if in_expert & (target_token_atoms > Int32(0)): - relative_atom = atom - target_atom_base - hidden_atom = relative_atom // target_token_atoms - token_atom = relative_atom % target_token_atoms - if token_atom < source_token_atoms: - source_hidden_atom = hidden_atom - source_byte = byte_in_atom - if cutlass.const_expr(self.deinterleave_gate_up is not None): - lane = byte_in_atom // Int32(16) - byte_tail = byte_in_atom % Int32(16) - group = byte_tail // Int32(4) - column_lane = byte_tail % Int32(4) - feature = hidden_atom * Int32(128) + group * Int32(32) + lane - intermediate = Int32(self.deinterleave_gate_up) - source_feature = Int32(0) - if feature < intermediate: - source_feature = ( - feature // Int32(32) - ) * Int32(64) + feature % Int32(32) - else: - up_feature = feature - intermediate - source_feature = ( - up_feature // Int32(32) - ) * Int32(64) + Int32(32) + up_feature % Int32(32) - source_hidden_atom = source_feature // Int32(128) - source_feature_in_atom = source_feature % Int32(128) - source_byte = ( - (source_feature_in_atom % Int32(32)) * Int32(16) - + (source_feature_in_atom // Int32(32)) * Int32(4) - + column_lane - ) - source_atom = ( - source_atom_base - + source_hidden_atom * source_token_atoms - + token_atom - ) - value = source[source_atom * Int32(atom_bytes) + source_byte] - target_atom_base += target_atom_count - source_atom_base += Int32(non_k_atoms) * source_token_atoms - previous_end = end - - output[linear] = value - - -__all__ = ["Mxfp8TrainingScaleExpandKernel"] diff --git a/python/cudnn/moe_ep/_types.py b/python/cudnn/moe_ep/_types.py index 11927eab9..eb575a869 100644 --- a/python/cudnn/moe_ep/_types.py +++ b/python/cudnn/moe_ep/_types.py @@ -8,10 +8,12 @@ import operator from dataclasses import dataclass from enum import Enum -from typing import Any, Tuple, Union +from typing import Tuple, Union import torch +from ._math import ceil_div + class MoeFormat(str, Enum): """Data formats supported by the MoE+EP interface.""" @@ -46,6 +48,41 @@ def _normalize_axis(axis: int, ndim: int) -> int: return normalized +def _block_scaled_representation( + fmt: MoeFormat, + logical_shape: Tuple[int, ...], + axis: int, +) -> tuple[Tuple[int, ...], Tuple[int, ...], torch.dtype, torch.dtype]: + if fmt is MoeFormat.BF16: + raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") + logical_extent = logical_shape[axis] + block_size = 32 if fmt is MoeFormat.MXFP8 else 16 + payload_extent = ( + logical_extent + if fmt is MoeFormat.MXFP8 + else ceil_div(logical_extent, 2) + ) + data_shape = list(logical_shape) + data_shape[axis] = payload_extent + scale_shape = list(logical_shape) + scale_shape[axis] = ceil_div(logical_extent, block_size) + + e4m3_dtype = getattr(torch, "float8_e4m3fn", None) + if e4m3_dtype is None: + raise RuntimeError("this PyTorch build does not provide torch.float8_e4m3fn") + if fmt is MoeFormat.MXFP8: + scale_dtype = getattr(torch, "float8_e8m0fnu", None) + if scale_dtype is None: + raise RuntimeError( + "this PyTorch build does not provide torch.float8_e8m0fnu" + ) + data_dtype = e4m3_dtype + else: + data_dtype = torch.uint8 + scale_dtype = e4m3_dtype + return tuple(data_shape), tuple(scale_shape), data_dtype, scale_dtype + + @dataclass(frozen=True) class BlockScaledTensor: """Data-plus-scale result returned for MXFP8 and NVFP4 outputs.""" @@ -83,31 +120,16 @@ def __post_init__(self) -> None: logical_shape.append(dim) normalized_shape = tuple(logical_shape) axis = _normalize_axis(self.axis, len(normalized_shape)) - logical_extent = normalized_shape[axis] - block_size = 32 if fmt is MoeFormat.MXFP8 else 16 - payload_extent = logical_extent if fmt is MoeFormat.MXFP8 else (logical_extent + 1) // 2 - scale_extent = (logical_extent + block_size - 1) // block_size - expected_data_shape = list(normalized_shape) - expected_data_shape[axis] = payload_extent - expected_scale_shape = list(normalized_shape) - expected_scale_shape[axis] = scale_extent - expected_data_shape = tuple(expected_data_shape) - expected_scale_shape = tuple(expected_scale_shape) + ( + expected_data_shape, + expected_scale_shape, + expected_data_dtype, + expected_scale_dtype, + ) = _block_scaled_representation(fmt, normalized_shape, axis) if tuple(self.data.shape) != expected_data_shape: raise ValueError(f"{fmt.value} data shape must be {expected_data_shape}, " f"got {tuple(self.data.shape)}") if tuple(self.scale.shape) != expected_scale_shape: raise ValueError(f"{fmt.value} scale shape must be {expected_scale_shape}, " f"got {tuple(self.scale.shape)}") - e4m3_dtype = getattr(torch, "float8_e4m3fn", None) - if e4m3_dtype is None: - raise RuntimeError("this PyTorch build does not provide torch.float8_e4m3fn") - if fmt is MoeFormat.MXFP8: - expected_data_dtype = e4m3_dtype - expected_scale_dtype = getattr(torch, "float8_e8m0fnu", None) - if expected_scale_dtype is None: - raise RuntimeError("this PyTorch build does not provide torch.float8_e8m0fnu") - else: - expected_data_dtype = torch.uint8 - expected_scale_dtype = e4m3_dtype if self.data.dtype is not expected_data_dtype: raise ValueError(f"{fmt.value} data must have dtype {expected_data_dtype}, " f"got {self.data.dtype}") if self.scale.dtype is not expected_scale_dtype: @@ -172,29 +194,123 @@ def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: return (values * expanded_scale).movedim(-1, self.axis).to(dtype) +class MoeEpNativeWeightLayout(str, Enum): + """Versioned kernel-native MXFP8 weight layouts.""" + + FORWARD_FC1_GATE_UP_INTERLEAVED_32_V1 = "mxfp8.forward_fc1.gate_up_interleaved_32.blocked_sf.v1" + FORWARD_FC2_K_MAJOR_V1 = "mxfp8.forward_fc2.k_major.blocked_sf.v1" + BACKWARD_W2_TRANSPOSE_V1 = "mxfp8.backward_w2_transpose.contiguous.blocked_sf.v1" + BACKWARD_W1_TRANSPOSE_GATE_UP_INTERLEAVED_32_V1 = "mxfp8.backward_w1_transpose.gate_up_interleaved_32.blocked_sf.v1" + + +@dataclass(frozen=True) +class MoeEpForwardWeights: + """Logical gate-then-up MXFP8 sources accepted by the fallback packer.""" + + fc1: BlockScaledTensor + fc2: BlockScaledTensor + + +@dataclass(frozen=True) +class MoeEpBackwardWeights: + """Logical gate-then-up MXFP8 transpose sources for the fallback packer.""" + + w2_transpose: BlockScaledTensor + w1_transpose: BlockScaledTensor + + +@dataclass(frozen=True) +class MoeEpNativeWeight: + """One kernel-executable payload plus Rubin blocked/interleaved scales.""" + + payload: torch.Tensor + scale: torch.Tensor + layout_id: Union[MoeEpNativeWeightLayout, str] + + def __post_init__(self) -> None: + if not isinstance(self.payload, torch.Tensor): + raise TypeError("payload must be a torch.Tensor, " f"got {type(self.payload).__name__}") + if not isinstance(self.scale, torch.Tensor): + raise TypeError("scale must be a torch.Tensor, " f"got {type(self.scale).__name__}") + if self.payload.device != self.scale.device: + raise ValueError(f"payload device {self.payload.device} does not match " f"scale device {self.scale.device}") + try: + layout_id = MoeEpNativeWeightLayout(self.layout_id) + except (TypeError, ValueError) as exc: + choices = ", ".join(layout.value for layout in MoeEpNativeWeightLayout) + raise ValueError(f"unsupported native weight layout_id {self.layout_id!r}; " f"expected one of: {choices}") from exc + object.__setattr__(self, "layout_id", layout_id) + + @property + def device(self) -> torch.device: + return self.payload.device + + +@dataclass(frozen=True) +class MoeEpNativeForwardWeights: + """Independent kernel-native weights consumed by one forward call.""" + + fc1: MoeEpNativeWeight + fc2: MoeEpNativeWeight + + +@dataclass(frozen=True) +class MoeEpNativeBackwardWeights: + """Independent kernel-native transpose weights consumed by one backward.""" + + w2_transpose: MoeEpNativeWeight + w1_transpose: MoeEpNativeWeight + + @dataclass(frozen=True) -class MoeEpTrainingWeights: - """Stable MXFP8 bindings for forward and dgrad GEMMs. +class MoeEpForwardWeightStaging: + """Caller-owned destinations used by forward weight materialization.""" - Forward consumes ``forward_fc1`` with logical shape ``(E,H,2I)`` and - ``forward_fc2`` with ``(E,I,H)``. Backward consumes independently - quantized transposes: ``backward_w2_transpose=(E,H,I)`` for - ``dH=dY@W2.T`` and ``backward_w1_transpose=(E,2I,H)`` for - ``dX=dC@W1.T``. Every tensor is block-scaled along logical axis 1, the - reduction axis of its corresponding GEMM. When the owning ``MoeEp`` has - ``weight_interleave_size=32``, forward W1's output axis and backward W1T's - reduction axis must contain alternating 32-element gate/up strips. - """ + fc1_payload: torch.Tensor + fc1_scale: torch.Tensor + fc2_payload: torch.Tensor + fc2_scale: torch.Tensor - forward_fc1: BlockScaledTensor - forward_fc2: BlockScaledTensor - backward_w2_transpose: BlockScaledTensor - backward_w1_transpose: BlockScaledTensor + +@dataclass(frozen=True) +class MoeEpBackwardWeightStaging: + """Caller-owned destinations used by backward weight materialization.""" + + w2_transpose_payload: torch.Tensor + w2_transpose_scale: torch.Tensor + w1_transpose_payload: torch.Tensor + w1_transpose_scale: torch.Tensor + + +@dataclass(frozen=True) +class MoeEpTrainingForwardOutputs: + """Caller-owned forward destinations.""" + + fc1_preact: torch.Tensor + output: torch.Tensor | None = None + fc1_a: torch.Tensor | None = None + fc1_sfa: torch.Tensor | None = None + valid_route_counts: torch.Tensor | None = None + expert_offsets: torch.Tensor | None = None + + +@dataclass(frozen=True) +class MoeEpTrainingBackwardOutputs: + """Caller-owned backward and final grouped-WGrad destinations.""" + + grad_activation: torch.Tensor | None = None + dprob: torch.Tensor | None = None + fc1_b: torch.Tensor | None = None + fc1_sfb: torch.Tensor | None = None + fc2_a: torch.Tensor | None = None + fc2_sfa: torch.Tensor | None = None + fc2_b: torch.Tensor | None = None + fc2_sfb: torch.Tensor | None = None @dataclass(frozen=True) class MoeEpTrainingWgradOperands: - """Fixed-capacity MXFP8 operands produced by the training resource path.""" + """Non-owning views over caller-owned grouped-WGrad operand buffers.""" fc1_a: torch.Tensor fc1_sfa: torch.Tensor @@ -208,165 +324,12 @@ class MoeEpTrainingWgradOperands: valid_route_counts: torch.Tensor -@dataclass(frozen=True) -class MoeEpTrainingSlot: - """Opaque index of one persistent forward/backward training slot.""" - - index: int - _resource_token: object - - @dataclass(frozen=True) class MoeEpExecutionLane: - """Opaque index of one mutable per-stream execution lane.""" + """Operator-bound index of one mutable per-stream execution lane.""" index: int - _resource_token: object - - -class MoeEpTrainingResources: - """Caller-owned lease on fixed-capacity training slots and execution lanes.""" - - def __init__( - self, - *, - owner: Any, - operator_token: object, - weights: MoeEpTrainingWeights, - slot_count: int, - lane_count: int, - device: torch.device, - ) -> None: - self._owner = owner - self._operator_token = operator_token - self._resource_token = object() - self.weights = weights - self.device = torch.device(device) - self.slots = tuple(MoeEpTrainingSlot(index, self._resource_token) for index in range(slot_count)) - self.lanes = tuple(MoeEpExecutionLane(index, self._resource_token) for index in range(lane_count)) - self._closed = False - - @property - def closed(self) -> bool: - return self._closed - - def _check_binding( - self, - operator_token: object, - slot: MoeEpTrainingSlot, - lane: MoeEpExecutionLane, - ) -> None: - if self._closed: - raise RuntimeError("MoeEp training resources are closed") - if self._operator_token is not operator_token: - raise ValueError("training resources belong to another MoeEp instance") - if slot._resource_token is not self._resource_token or slot not in self.slots: - raise ValueError("training slot does not belong to these resources") - if lane._resource_token is not self._resource_token or lane not in self.lanes: - raise ValueError("execution lane does not belong to these resources") - - def refresh_weights(self) -> None: - """Enqueue fixed-address scale-layout refreshes on the current stream. - - Call after every in-place data+scale update and before the first - forward/backward that consumes that version. The caller must establish - stream/event ordering, must not refresh between a matching forward and - backward, and must not overlap refresh with any consumer of these - resources. Replacing source storage requires closing the old operator, - creating a new ``MoeEp`` instance and resources, and capturing a new - graph. This method may itself be captured, in which case replay executes - only the recorded device transforms. - """ - - if self._closed: - raise RuntimeError("MoeEp training resources are closed") - self._owner.refresh_weights() - - def forward( - self, - slot: MoeEpTrainingSlot, - lane: MoeEpExecutionLane, - activation: torch.Tensor, - topk_idx: torch.Tensor, - topk_weights: torch.Tensor, - ) -> torch.Tensor: - """Run the fixed-slot forward in ordinary or capture mode.""" - - self._check_binding(self._operator_token, slot, lane) - execution = self._owner.views( - slot=slot.index, - lane=lane.index, - token_count=int(activation.shape[0]), - ) - from ._megamoe_backend.mxfp8._training_execute import ( - launch_training_forward, - ) - - return launch_training_forward( - self._owner, - execution, - activation, - topk_idx, - topk_weights, - ) - - def backward( - self, - slot: MoeEpTrainingSlot, - lane: MoeEpExecutionLane, - grad_output: torch.Tensor, - ) -> tuple[ - torch.Tensor, - torch.Tensor, - MoeEpTrainingWgradOperands, - ]: - """Run fixed-slot dgrad/dprob in ordinary or capture mode.""" - - self._check_binding(self._operator_token, slot, lane) - execution = self._owner.views( - slot=slot.index, - lane=lane.index, - token_count=int(grad_output.shape[0]), - ) - from ._megamoe_backend.mxfp8._training_execute import ( - launch_training_backward, - ) - - grad_activation, grad_topk_weights, operands = launch_training_backward( - self._owner, - execution, - grad_output, - ) - return grad_activation, grad_topk_weights, operands - - def finalize_overflow( - self, - slots: Tuple[MoeEpTrainingSlot, ...], - lane: MoeEpExecutionLane | None = None, - ) -> torch.Tensor: - """Aggregate one computation group's flags and apply its policy.""" - - if self._closed: - raise RuntimeError("MoeEp training resources are closed") - if lane is None: - lane = self.lanes[0] - if not isinstance(lane, MoeEpExecutionLane) or lane._resource_token is not self._resource_token or lane not in self.lanes: - raise ValueError("overflow execution lane does not belong to these resources") - slot_indices = [] - for slot in slots: - if not isinstance(slot, MoeEpTrainingSlot) or slot._resource_token is not self._resource_token or slot not in self.slots: - raise ValueError("overflow slot does not belong to these resources") - slot_indices.append(slot.index) - return self._owner.finalize_overflow( - tuple(slot_indices), - lane=lane.index, - ) - - def close(self) -> None: - if self._closed: - return - self._owner.close() - self._closed = True + _operator_token: object MoeTensor = Union[torch.Tensor, BlockScaledTensor] @@ -375,9 +338,16 @@ def close(self) -> None: __all__ = [ "BlockScaledTensor", "MoeEpExecutionLane", - "MoeEpTrainingResources", - "MoeEpTrainingSlot", - "MoeEpTrainingWeights", + "MoeEpBackwardWeightStaging", + "MoeEpBackwardWeights", + "MoeEpForwardWeightStaging", + "MoeEpForwardWeights", + "MoeEpNativeBackwardWeights", + "MoeEpNativeForwardWeights", + "MoeEpNativeWeight", + "MoeEpNativeWeightLayout", + "MoeEpTrainingBackwardOutputs", + "MoeEpTrainingForwardOutputs", "MoeEpTrainingWgradOperands", "MoeFormat", "MoeTensor", diff --git a/python/cudnn/moe_ep/_validation.py b/python/cudnn/moe_ep/_validation.py index 31ea364e6..1fc31c4eb 100644 --- a/python/cudnn/moe_ep/_validation.py +++ b/python/cudnn/moe_ep/_validation.py @@ -5,31 +5,36 @@ from __future__ import annotations -from typing import Tuple +from typing import Mapping, Tuple import torch from ._contracts import Fc1WeightLayout, ForwardConfig, ValidatedForwardRequest +from ._math import round_up from ._types import ( BlockScaledTensor, - MoeEpTrainingWeights, + MoeEpBackwardWeights, + MoeEpForwardWeights, + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, + MoeEpTrainingBackwardOutputs, + MoeEpTrainingForwardOutputs, MoeFormat, MoeTensor, + _block_scaled_representation, ) - -def _replace_axis( - shape: Tuple[int, ...], - axis: int, - extent: int, -) -> Tuple[int, ...]: - result = list(shape) - result[axis] = extent - return tuple(result) - - -def _ceil_div(value: int, divisor: int) -> int: - return (value + divisor - 1) // divisor +_SourceWeightSpec = tuple[str, BlockScaledTensor, Tuple[int, ...]] +_NativeWeightSpec = tuple[ + str, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, + Tuple[int, ...], + Tuple[int, ...], + int, +] def _require_torch_dtype(name: str) -> torch.dtype: @@ -80,26 +85,15 @@ def _validate_tensor_representation( _validate_strided(f"{name}.data", tensor.data) _validate_strided(f"{name}.scale", tensor.scale) - logical_extent = expected_logical_shape[tensor.axis] - if tensor.format is MoeFormat.MXFP8: - payload_extent = logical_extent - block_size = 32 - expected_data_dtype = _require_torch_dtype("float8_e4m3fn") - expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") - else: - payload_extent = _ceil_div(logical_extent, 2) - block_size = 16 - expected_data_dtype = torch.uint8 - expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") - expected_data_shape = _replace_axis( - expected_logical_shape, - tensor.axis, - payload_extent, - ) - expected_scale_shape = _replace_axis( + ( + expected_data_shape, + expected_scale_shape, + expected_data_dtype, + expected_scale_dtype, + ) = _block_scaled_representation( + tensor.format, expected_logical_shape, tensor.axis, - _ceil_div(logical_extent, block_size), ) if tuple(tensor.data.shape) != expected_data_shape: raise ValueError(f"{name}.data shape must be {expected_data_shape}, " f"got {tuple(tensor.data.shape)}") @@ -186,13 +180,9 @@ def validate_forward( ), ) if config.fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32 and ( - not isinstance(fc1_weight, BlockScaledTensor) - or fc1_weight.format is not MoeFormat.MXFP8 + not isinstance(fc1_weight, BlockScaledTensor) or fc1_weight.format is not MoeFormat.MXFP8 ): - raise ValueError( - "weight_interleave_size=32 requires an MXFP8 BlockScaledTensor " - "for fc1_weight" - ) + raise ValueError("weight_interleave_size=32 requires an MXFP8 BlockScaledTensor " "for fc1_weight") _validate_routes( config, token_count, @@ -229,99 +219,394 @@ def validate_forward( ) -def validate_training_weights( +def _validate_source_weight( + name: str, + tensor: BlockScaledTensor, + shape: Tuple[int, ...], +) -> None: + _validate_tensor_representation(name, tensor, shape) + if not isinstance(tensor, BlockScaledTensor): + raise TypeError(f"{name} must be an MXFP8 BlockScaledTensor") + if tensor.format is not MoeFormat.MXFP8: + raise NotImplementedError(f"{name} must use format='mxfp8', got {tensor.format.value!r}") + if not tensor.data.is_contiguous() or not tensor.scale.is_contiguous(): + raise ValueError(f"{name} data and scale must be contiguous") + + +def _validate_source_weight_pair( + expected: tuple[_SourceWeightSpec, _SourceWeightSpec], +) -> torch.device: + for name, tensor, shape in expected: + _validate_source_weight(name, tensor, shape) + device = expected[0][1].device + second_name, second_tensor, _ = expected[1] + if second_tensor.device != device: + raise ValueError(f"{second_name} must be on {device}, got {second_tensor.device}") + return device + + +def validate_forward_source_weights( + config: ForwardConfig, + weights: MoeEpForwardWeights, +) -> torch.device: + """Validate source weights used by allocation-free forward packing.""" + + if not isinstance(weights, MoeEpForwardWeights): + raise TypeError("weights must be a MoeEpForwardWeights, " f"got {type(weights).__name__}") + expected = ( + ("weights.fc1", weights.fc1, (config.experts_per_rank, config.hidden_size, 2 * config.intermediate_size)), + ("weights.fc2", weights.fc2, (config.experts_per_rank, config.intermediate_size, config.hidden_size)), + ) + return _validate_source_weight_pair(expected) + + +def validate_backward_source_weights( config: ForwardConfig, - weights: MoeEpTrainingWeights, + weights: MoeEpBackwardWeights, ) -> torch.device: - """Validate fixed MXFP8 weight bindings used by training resources.""" - - def has_supported_layout(tensor: torch.Tensor) -> bool: - if tensor.is_contiguous(): - return True - if tensor.ndim != 3: - return False - experts, reduction, output = tensor.shape - return tensor.stride() == (reduction * output, 1, reduction) - - def is_compact_k_major(tensor: torch.Tensor) -> bool: - if tensor.ndim != 3: - return False - experts, reduction, output = tensor.shape - return tensor.stride() == (reduction * output, 1, reduction) - - if not isinstance(weights, MoeEpTrainingWeights): - raise TypeError("weights must be a MoeEpTrainingWeights, " f"got {type(weights).__name__}") + """Validate source weights used by allocation-free backward packing.""" + + if not isinstance(weights, MoeEpBackwardWeights): + raise TypeError("weights must be a MoeEpBackwardWeights, " f"got {type(weights).__name__}") + expected = ( + ("weights.w2_transpose", weights.w2_transpose, (config.experts_per_rank, config.hidden_size, config.intermediate_size)), + ("weights.w1_transpose", weights.w1_transpose, (config.experts_per_rank, 2 * config.intermediate_size, config.hidden_size)), + ) + return _validate_source_weight_pair(expected) + + +def _blocked_scale_elements(raw_rows: int, raw_columns: int) -> int: + return round_up(raw_rows, 128) * round_up(raw_columns, 4) + + +def _validate_native_weight( + name: str, + weight: MoeEpNativeWeight, + *, + layout_id: MoeEpNativeWeightLayout, + payload_shape: Tuple[int, ...], + payload_stride: Tuple[int, ...], + scale_elements: int, + scale_dtype: torch.dtype, + device: torch.device | None, +) -> torch.device: + if not isinstance(weight, MoeEpNativeWeight): + raise TypeError(f"{name} must be a MoeEpNativeWeight, got {type(weight).__name__}") + if weight.layout_id is not layout_id: + raise ValueError(f"{name}.layout_id must be {layout_id.value!r}, " f"got {weight.layout_id.value!r}") + _validate_strided(f"{name}.payload", weight.payload) + _validate_strided(f"{name}.scale", weight.scale) + if tuple(weight.payload.shape) != payload_shape: + raise ValueError(f"{name}.payload shape must be {payload_shape}, " f"got {tuple(weight.payload.shape)}") + if tuple(weight.payload.stride()) != payload_stride: + raise ValueError(f"{name}.payload stride must be {payload_stride}, " f"got {tuple(weight.payload.stride())}") + expected_payload_dtype = _require_torch_dtype("float8_e4m3fn") + if weight.payload.dtype is not expected_payload_dtype: + raise ValueError(f"{name}.payload must have dtype {expected_payload_dtype}, " f"got {weight.payload.dtype}") + expected_scale_shape = (payload_shape[0], scale_elements) + if tuple(weight.scale.shape) != expected_scale_shape: + raise ValueError(f"{name}.scale shape must be {expected_scale_shape}, " f"got {tuple(weight.scale.shape)}") + if not weight.scale.is_contiguous(): + raise ValueError(f"{name}.scale must be contiguous") + if weight.scale.dtype is not scale_dtype: + raise ValueError(f"{name}.scale must have dtype {scale_dtype}, " f"got {weight.scale.dtype}") + for field_name, tensor in ( + ("payload", weight.payload), + ("scale", weight.scale), + ): + if tensor.data_ptr() % 16: + raise ValueError(f"{name}.{field_name} must be at least 16-byte aligned") + if device is not None and weight.device != device: + raise ValueError(f"{name} must be on {device}, got {weight.device}") + return weight.device + + +def _validate_native_weight_pair( + expected: tuple[_NativeWeightSpec, _NativeWeightSpec], + *, + scale_dtype: torch.dtype, + device: torch.device | None, +) -> torch.device: + resolved = device + for name, weight, layout_id, payload_shape, payload_stride, scale_elements in expected: + resolved = _validate_native_weight( + name, + weight, + layout_id=layout_id, + payload_shape=payload_shape, + payload_stride=payload_stride, + scale_elements=scale_elements, + scale_dtype=scale_dtype, + device=resolved, + ) + assert resolved is not None + return resolved + + +def validate_native_forward_weights( + config: ForwardConfig, + weights: MoeEpNativeForwardWeights, + *, + device: torch.device | None = None, +) -> torch.device: + if not isinstance(weights, MoeEpNativeForwardWeights): + raise TypeError("weights must be a MoeEpNativeForwardWeights, " f"got {type(weights).__name__}") + if config.fc1_weight_layout is not Fc1WeightLayout.GATE_UP_INTERLEAVED_32: + raise ValueError("native training weights require weight_interleave_size=32") + experts = config.experts_per_rank + hidden = config.hidden_size + intermediate = config.intermediate_size + sf_dtype = _require_torch_dtype("float8_e8m0fnu") expected = ( ( - "weights.forward_fc1", - weights.forward_fc1, - ( - config.experts_per_rank, - config.hidden_size, - 2 * config.intermediate_size, - ), + "weights.fc1", + weights.fc1, + MoeEpNativeWeightLayout.FORWARD_FC1_GATE_UP_INTERLEAVED_32_V1, + (experts, hidden, 2 * intermediate), + (hidden * 2 * intermediate, 1, hidden), + _blocked_scale_elements(2 * intermediate, hidden // 32), ), ( - "weights.forward_fc2", - weights.forward_fc2, - ( - config.experts_per_rank, - config.intermediate_size, - config.hidden_size, - ), + "weights.fc2", + weights.fc2, + MoeEpNativeWeightLayout.FORWARD_FC2_K_MAJOR_V1, + (experts, intermediate, hidden), + (intermediate * hidden, 1, intermediate), + _blocked_scale_elements(hidden, intermediate // 32), ), + ) + return _validate_native_weight_pair( + expected, + scale_dtype=sf_dtype, + device=device, + ) + + +def validate_native_backward_weights( + config: ForwardConfig, + weights: MoeEpNativeBackwardWeights, + *, + device: torch.device | None = None, +) -> torch.device: + if not isinstance(weights, MoeEpNativeBackwardWeights): + raise TypeError("weights must be a MoeEpNativeBackwardWeights, " f"got {type(weights).__name__}") + if config.fc1_weight_layout is not Fc1WeightLayout.GATE_UP_INTERLEAVED_32: + raise ValueError("native training weights require weight_interleave_size=32") + experts = config.experts_per_rank + hidden = config.hidden_size + intermediate = config.intermediate_size + sf_dtype = _require_torch_dtype("float8_e8m0fnu") + expected = ( ( - "weights.backward_w2_transpose", - weights.backward_w2_transpose, - ( - config.experts_per_rank, - config.hidden_size, - config.intermediate_size, - ), + "weights.w2_transpose", + weights.w2_transpose, + MoeEpNativeWeightLayout.BACKWARD_W2_TRANSPOSE_V1, + (experts, hidden, intermediate), + (hidden * intermediate, intermediate, 1), + _blocked_scale_elements(intermediate, hidden // 32), ), ( - "weights.backward_w1_transpose", - weights.backward_w1_transpose, - ( - config.experts_per_rank, - 2 * config.intermediate_size, - config.hidden_size, - ), + "weights.w1_transpose", + weights.w1_transpose, + MoeEpNativeWeightLayout.BACKWARD_W1_TRANSPOSE_GATE_UP_INTERLEAVED_32_V1, + (experts, 2 * intermediate, hidden), + (2 * intermediate * hidden, hidden, 1), + _blocked_scale_elements(hidden, 2 * intermediate // 32), ), ) - for name, tensor, shape in expected: - _validate_tensor_representation(name, tensor, shape) - if not isinstance(tensor, BlockScaledTensor): - raise TypeError(f"{name} must be an MXFP8 BlockScaledTensor for " "fixed training resources") - if tensor.format is not MoeFormat.MXFP8: - raise NotImplementedError(f"{name} must use format='mxfp8', got {tensor.format.value!r}") - if name in ( - "weights.backward_w2_transpose", - "weights.backward_w1_transpose", - ) and not tensor.data.is_contiguous(): - raise ValueError( - f"{name} data must be contiguous for fixed training weight binding" - ) - if not has_supported_layout(tensor.data) or not has_supported_layout(tensor.scale): - raise ValueError( - f"{name} data and scale must be contiguous or compact K-major " - "for fixed training weight binding" - ) - if config.fc1_weight_layout is Fc1WeightLayout.GATE_UP_INTERLEAVED_32 and not ( - is_compact_k_major(weights.forward_fc1.data) - and is_compact_k_major(weights.forward_fc2.data) - and weights.backward_w2_transpose.data.is_contiguous() - and weights.backward_w1_transpose.data.is_contiguous() - ): - raise ValueError( - "weight_interleave_size=32 requires compact K-major forward " - "weights and contiguous backward transpose weights" - ) - device = weights.forward_fc1.device - for name, tensor, _shape in expected[1:]: + return _validate_native_weight_pair( + expected, + scale_dtype=sf_dtype, + device=device, + ) + + +def validate_training_input( + config: ForwardConfig, + name: str, + value: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + device: torch.device, +) -> int: + logical_shape = _logical_shape(value) + if len(logical_shape) != 2 or logical_shape[1] != config.hidden_size: + raise ValueError(f"{name} logical shape must be (T, {config.hidden_size}), " f"got {logical_shape}") + _validate_tensor_representation(name, value, logical_shape) + if isinstance(value, BlockScaledTensor): + if value.format is not MoeFormat.MXFP8: + raise NotImplementedError(f"{name} only supports MXFP8 block scaling") + if not value.data.is_contiguous() or not value.scale.is_contiguous(): + raise ValueError(f"{name} MXFP8 data and scale must be contiguous") + elif value.dtype not in (torch.bfloat16, torch.float32): + raise TypeError(f"{name} must be BF16, FP32, or an MXFP8 BlockScaledTensor") + elif not value.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + token_count = logical_shape[0] + if device.type == "cuda": + with torch.cuda.device(device): + capturing = torch.cuda.is_current_stream_capturing() + else: + capturing = False + _validate_routes( + config, + token_count, + topk_idx, + topk_weights, + validate_expert_ids=not capturing, + ) + if topk_idx.dtype is not torch.int32 or not topk_idx.is_contiguous(): + raise TypeError("training topk_idx must be contiguous torch.int32") + if topk_weights.dtype is not torch.float32 or not topk_weights.is_contiguous(): + raise TypeError("training topk_weights must be contiguous torch.float32") + tensors = ( + (name, value), + ("topk_idx", topk_idx), + ("topk_weights", topk_weights), + ) + for tensor_name, tensor in tensors: + tensor_device = _tensor_device(tensor) + if tensor_device != device: + raise ValueError(f"{tensor_name} must be on {device}, got {tensor_device}") + return token_count + + +def _tensor_byte_range(tensor: torch.Tensor) -> tuple[int, int]: + byte_start = tensor.data_ptr() + max_element_offset = sum( + (int(extent) - 1) * int(step) + for extent, step in zip(tensor.shape, tensor.stride()) + if int(extent) > 0 + ) + byte_end = byte_start + ( + 0 + if tensor.numel() == 0 + else (max_element_offset + 1) * tensor.element_size() + ) + return byte_start, byte_end + + +def _assert_no_overlap( + name: str, + byte_range: tuple[int, int], + ranges: list[tuple[int, int, str]], +) -> None: + byte_start, byte_end = byte_range + for other_start, other_end, other_name in ranges: + if byte_start < other_end and other_start < byte_end: + raise ValueError(f"{name} must not alias {other_name}") + + +def _validate_named_buffers( + tensors: Mapping[str, object], + requirements: Mapping[str, tuple[Tuple[int, ...], Tuple[int, ...], torch.dtype, int]], + *, + device: torch.device, +) -> None: + ranges: list[tuple[int, int, str]] = [] + for name, requirement in requirements.items(): + tensor = tensors[name] + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"out.{name} must be a torch.Tensor") + shape, stride, dtype, alignment = requirement + if tuple(tensor.shape) != tuple(shape): + raise ValueError(f"out.{name} shape must be {tuple(shape)}, got {tuple(tensor.shape)}") + if tuple(tensor.stride()) != tuple(stride): + raise ValueError(f"out.{name} stride must be {tuple(stride)}, " f"got {tuple(tensor.stride())}") + if tensor.dtype is not dtype: + raise ValueError(f"out.{name} dtype must be {dtype}, got {tensor.dtype}") if tensor.device != device: - raise ValueError(f"{name} must be on {device}, got {tensor.device}") - return device + raise ValueError(f"out.{name} must be on {device}, got {tensor.device}") + if tensor.data_ptr() % alignment: + raise ValueError(f"out.{name} must be {alignment}-byte aligned") + qualified_name = f"out.{name}" + byte_start, byte_end = _tensor_byte_range(tensor) + _assert_no_overlap(qualified_name, (byte_start, byte_end), ranges) + ranges.append((byte_start, byte_end, qualified_name)) + + +def _validate_output_buffers( + output: object, + requirements: Mapping[str, tuple[Tuple[int, ...], Tuple[int, ...], torch.dtype, int]], + *, + device: torch.device, +) -> None: + _validate_named_buffers( + {name: getattr(output, name) for name in requirements}, + requirements, + device=device, + ) + +def validate_training_non_aliasing( + tensors: Mapping[str, torch.Tensor | None], +) -> None: + """Reject overlapping caller inputs, saved state, weights, and outputs.""" -__all__ = ["validate_forward", "validate_training_weights"] + ranges: list[tuple[int, int, str]] = [] + for name, tensor in tensors.items(): + if tensor is None or tensor.numel() == 0: + continue + byte_start, byte_end = _tensor_byte_range(tensor) + _assert_no_overlap(name, (byte_start, byte_end), ranges) + ranges.append((byte_start, byte_end, name)) + + +def validate_training_forward_outputs( + output: MoeEpTrainingForwardOutputs, + requirements: Mapping[str, tuple[Tuple[int, ...], Tuple[int, ...], torch.dtype, int]], + *, + device: torch.device, +) -> None: + if not isinstance(output, MoeEpTrainingForwardOutputs): + raise TypeError("out must be a MoeEpTrainingForwardOutputs, " f"got {type(output).__name__}") + _validate_output_buffers(output, requirements, device=device) + + +def validate_training_forward_state( + *, + fc1_preact: torch.Tensor, + fc1_a: torch.Tensor | None, + fc1_sfa: torch.Tensor | None, + valid_route_counts: torch.Tensor | None, + expert_offsets: torch.Tensor | None, + requirements: Mapping[str, tuple[Tuple[int, ...], Tuple[int, ...], torch.dtype, int]], + device: torch.device, +) -> None: + _validate_named_buffers( + { + "fc1_preact": fc1_preact, + "fc1_a": fc1_a, + "fc1_sfa": fc1_sfa, + "valid_route_counts": valid_route_counts, + "expert_offsets": expert_offsets, + }, + requirements, + device=device, + ) + + +def validate_training_backward_outputs( + output: MoeEpTrainingBackwardOutputs, + requirements: Mapping[str, tuple[Tuple[int, ...], Tuple[int, ...], torch.dtype, int]], + *, + device: torch.device, +) -> None: + if not isinstance(output, MoeEpTrainingBackwardOutputs): + raise TypeError("out must be a MoeEpTrainingBackwardOutputs, " f"got {type(output).__name__}") + _validate_output_buffers(output, requirements, device=device) + + +__all__ = [ + "validate_backward_source_weights", + "validate_forward", + "validate_forward_source_weights", + "validate_native_backward_weights", + "validate_native_forward_weights", + "validate_training_backward_outputs", + "validate_training_forward_outputs", + "validate_training_forward_state", + "validate_training_input", + "validate_training_non_aliasing", +] diff --git a/python/cudnn/moe_ep/api.py b/python/cudnn/moe_ep/api.py index 10c182d4a..19cb47a91 100644 --- a/python/cudnn/moe_ep/api.py +++ b/python/cudnn/moe_ep/api.py @@ -15,24 +15,41 @@ import threading import warnings from numbers import Real -from typing import Optional, Union +from typing import Mapping, Optional, Union import torch import torch.distributed as dist -from ._contracts import ForwardConfig, normalize_fc1_weight_layout +from ._contracts import Fc1WeightLayout, ForwardConfig, normalize_fc1_weight_layout from ._tuning import MoeEpTuningConfig from ._types import ( BlockScaledTensor, + MoeEpBackwardWeightStaging, + MoeEpBackwardWeights, MoeEpExecutionLane, - MoeEpTrainingResources, - MoeEpTrainingSlot, - MoeEpTrainingWeights, + MoeEpForwardWeightStaging, + MoeEpForwardWeights, + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpTrainingBackwardOutputs, + MoeEpTrainingForwardOutputs, + MoeEpTrainingWgradOperands, MoeFormat, MoeTensor, parse_format as _parse_format, ) -from ._validation import validate_forward, validate_training_weights +from ._validation import ( + validate_backward_source_weights, + validate_forward, + validate_forward_source_weights, + validate_native_backward_weights, + validate_native_forward_weights, + validate_training_backward_outputs, + validate_training_forward_outputs, + validate_training_forward_state, + validate_training_input, + validate_training_non_aliasing, +) def _resolve_ep_topology( @@ -64,12 +81,78 @@ def _validate_training_assert_capability(config: ForwardConfig) -> None: if config.drop_on_overflow: return if not callable(getattr(torch, "_assert_async", None)): - raise RuntimeError("drop_on_overflow=False training resources require callable " "torch._assert_async before CUDA Graph capture") + raise RuntimeError("drop_on_overflow=False training requires callable " "torch._assert_async before CUDA Graph capture") if config.ep_size <= 1: return backend = dist.get_backend(config.ep_group) if backend != dist.Backend.NCCL and str(backend).lower() != "nccl": - raise NotImplementedError("drop_on_overflow=False EP2+ training resources require an NCCL " "process group for the captured scalar global overflow OR") + raise NotImplementedError("drop_on_overflow=False EP2+ training requires an NCCL " "process group for the captured scalar global overflow OR") + + +def _resolve_training_device( + device: torch.device | str | int | None, +) -> torch.device: + if device is None: + if not torch.cuda.is_available(): + raise RuntimeError("prepare_training requires an available CUDA device") + return torch.device("cuda", torch.cuda.current_device()) + if isinstance(device, bool): + raise TypeError("device must be a CUDA device, ordinal, or None") + if isinstance(device, int): + resolved = torch.device("cuda", device) + else: + resolved = torch.device(device) + if resolved.type == "cuda" and resolved.index is None: + resolved = torch.device("cuda", torch.cuda.current_device()) + if resolved.type != "cuda": + raise ValueError(f"training device must be CUDA, got {resolved}") + if resolved.index is None or resolved.index < 0 or resolved.index >= torch.cuda.device_count(): + raise ValueError(f"CUDA device {resolved} is not available") + return resolved + + +def _named_moe_tensors( + name: str, + value: MoeTensor, +) -> dict[str, torch.Tensor]: + if isinstance(value, BlockScaledTensor): + return { + f"{name}.data": value.data, + f"{name}.scale": value.scale, + } + return {name: value} + + +def pack_forward_weights( + weights: MoeEpForwardWeights, + *, + out: MoeEpForwardWeightStaging, +) -> MoeEpNativeForwardWeights: + """Standalone allocation-free forward weight materialization.""" + + from ._megamoe_backend.mxfp8._training_weights import materialize_forward + + return materialize_forward( + weights, + out=out, + fc1_weight_layout=Fc1WeightLayout.GATE_UP_INTERLEAVED_32, + ) + + +def pack_backward_weights( + weights: MoeEpBackwardWeights, + *, + out: MoeEpBackwardWeightStaging, +) -> MoeEpNativeBackwardWeights: + """Standalone allocation-free backward weight materialization.""" + + from ._megamoe_backend.mxfp8._training_weights import materialize_backward + + return materialize_backward( + weights, + out=out, + fc1_weight_layout=Fc1WeightLayout.GATE_UP_INTERLEAVED_32, + ) class MoeEp: @@ -79,7 +162,7 @@ class MoeEp: ``e // experts_per_rank``. The constructor captures static configuration; calling the instance accepts runtime tensors for this rank. - The Rubin training-Mega backend accepts plain BF16/FP16/FP32 operands + The Rubin training-Mega backend accepts plain BF16/FP32 operands (staged to MXFP8 E4M3) or MXFP8 ``BlockScaledTensor`` operands. Final output is BF16. ``combine_format`` may be BF16 or MXFP8; forward MXFP8 combine quantizes each FP32 route accumulator directly before top-k @@ -88,9 +171,10 @@ class MoeEp: Native NVFP4 operands and NVFP4 combine/output are not executable. ``__call__`` is the inference-only forward surface. Training uses - :meth:`prepare_training_resources`; the returned fixed-slot resource handle - provides ordinary/capturable ``forward`` and ``backward`` methods without - compact host-visible stashes. + :meth:`prepare_training`, :meth:`training_forward`, and + :meth:`training_backward`. Caller-owned output bundles carry all explicit + cross-phase state; the operator retains only private runtime and lane + scratch. The backend is created lazily on the first supported forward call. Valid combinations outside the current backend capability matrix fail explicitly @@ -227,7 +311,15 @@ def __init__( self._validated_topk_idx = None self._validated_topk_version = None self._operator_token = object() - self._training_resources: MoeEpTrainingResources | None = None + self._training_state = None + self._training_lanes: tuple[MoeEpExecutionLane, ...] = () + self._training_requirements: ( + Mapping[ + str, + tuple[tuple[int, ...], tuple[int, ...], torch.dtype, int], + ] + | None + ) = None self._closed = False @staticmethod @@ -275,8 +367,8 @@ def __call__( ``fc1_weight=(E_local,H,2I)``, ``fc2_weight=(E_local,I,H)``, and ``topk_idx=topk_weights=(T,K)``. - Training callers must use :meth:`prepare_training_resources` and the - returned fixed-slot resource handle. + Training callers must use :meth:`prepare_training` followed by the + stateless training methods. """ with self._lifecycle_lock: @@ -337,67 +429,313 @@ def warmup( if device.type == "cuda": torch.cuda.synchronize(device) - def prepare_training_resources( + @property + def training_lanes(self) -> tuple[MoeEpExecutionLane, ...]: + """Operator-bound lanes created by :meth:`prepare_training`.""" + + return self._training_lanes + + def prepare_training( self, - weights: MoeEpTrainingWeights, *, - slot_count: int = 2, lane_count: int = 1, - ) -> MoeEpTrainingResources: - """Bind MXFP8 weights and allocate fixed-capacity training resources. - - This collective preparation must run on every EP rank before CUDA - Graph capture. The returned handle owns persistent microbatch slots - and mutable per-stream execution lanes; closing the operator also - closes the handle. A closed handle cannot be replaced on this - operator; create a new ``MoeEp`` instance to bind new weight storage. + device: torch.device | str | int | None = None, + ) -> Mapping[ + str, + tuple[tuple[int, ...], tuple[int, ...], torch.dtype, int], + ]: + """Collectively prepare private training runtime and return contracts. + + ``device`` defaults to the current CUDA device. No weights or caller + output buffers are retained by the operator. """ with self._lifecycle_lock: if self._closed: raise RuntimeError("MoeEp is closed") - for name, value in ( - ("slot_count", slot_count), - ("lane_count", lane_count), - ): - if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise ValueError(f"{name} must be a positive integer, got {value!r}") - if self._training_resources is not None: - if not self._training_resources.closed: - raise RuntimeError("MoeEp training resources already exist") - raise RuntimeError("MoeEp training resources were closed; create a new " "MoeEp instance before preparing replacement weights") - - device = validate_training_weights( - self._forward_config, - weights, - ) + if isinstance(lane_count, bool) or not isinstance(lane_count, int) or lane_count <= 0: + raise ValueError(f"lane_count must be a positive integer, got {lane_count!r}") + if self._training_state is not None: + raise RuntimeError("MoeEp training is already prepared") + if self._fc1_weight_layout is not Fc1WeightLayout.GATE_UP_INTERLEAVED_32: + raise ValueError("prepare_training requires weight_interleave_size=32") + resolved_device = _resolve_training_device(device) _validate_training_assert_capability(self._forward_config) from . import _backend _backend.validate_config(self._forward_config) - if self._forward_backend is not None and device != self._forward_backend_device: - raise ValueError(f"MoeEp backend is bound to " f"{self._forward_backend_device}; got {device}") + if self._forward_backend is not None and resolved_device != self._forward_backend_device: + raise ValueError(f"MoeEp backend is bound to {self._forward_backend_device}; " f"got {resolved_device}") if self._forward_backend is None: self._forward_backend = _backend.create_backend( self._forward_config, - device, + resolved_device, ) - self._forward_backend_device = device - owner = self._forward_backend.prepare_training_resources( + self._forward_backend_device = resolved_device + with torch.cuda.device(resolved_device): + state = self._forward_backend.prepare_training( + lane_count=lane_count, + ) + self._training_state = state + self._training_lanes = tuple(MoeEpExecutionLane(index, self._operator_token) for index in range(lane_count)) + self._training_requirements = state.public_requirements() + return self._training_requirements + + def _require_training_lane( + self, + lane: MoeEpExecutionLane, + ) -> None: + if self._closed: + raise RuntimeError("MoeEp is closed") + if self._training_state is None or self._training_requirements is None: + raise RuntimeError("prepare_training() must be called first") + if not isinstance(lane, MoeEpExecutionLane) or lane._operator_token is not self._operator_token or lane not in self._training_lanes: + raise ValueError("execution lane does not belong to this MoeEp") + + def _training_requirement_subset( + self, + names: tuple[str, ...], + ) -> dict[str, tuple[tuple[int, ...], tuple[int, ...], torch.dtype, int]]: + assert self._training_requirements is not None + return {name: self._training_requirements[name] for name in names} + + def pack_forward_weights( + self, + weights: MoeEpForwardWeights, + *, + out: MoeEpForwardWeightStaging, + ) -> MoeEpNativeForwardWeights: + """Materialize source weights into caller-owned native storage.""" + + validate_forward_source_weights(self._forward_config, weights) + from ._megamoe_backend.mxfp8._training_weights import materialize_forward + + return materialize_forward( + weights, + out=out, + fc1_weight_layout=self._forward_config.fc1_weight_layout, + ) + + def pack_backward_weights( + self, + weights: MoeEpBackwardWeights, + *, + out: MoeEpBackwardWeightStaging, + ) -> MoeEpNativeBackwardWeights: + """Materialize source transpose weights into caller-owned storage.""" + + validate_backward_source_weights(self._forward_config, weights) + from ._megamoe_backend.mxfp8._training_weights import materialize_backward + + return materialize_backward( + weights, + out=out, + fc1_weight_layout=self._forward_config.fc1_weight_layout, + ) + + def training_forward( + self, + lane: MoeEpExecutionLane, + activation: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + weights: MoeEpNativeForwardWeights, + out: MoeEpTrainingForwardOutputs, + ) -> torch.Tensor: + """Run forward into caller-owned prepared-training outputs.""" + + with self._lifecycle_lock: + self._require_training_lane(lane) + assert self._training_state is not None + assert self._training_requirements is not None + assert self._forward_backend_device is not None + token_count = validate_training_input( + self._forward_config, + "activation", + activation, + topk_idx, + topk_weights, + device=self._forward_backend_device, + ) + validate_native_forward_weights( + self._forward_config, weights, - slot_count=slot_count, - lane_count=lane_count, + device=self._forward_backend_device, + ) + validate_training_forward_outputs( + out, + self._training_requirement_subset( + ( + "output", + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + ) + ), + device=self._forward_backend_device, ) - resources = MoeEpTrainingResources( - owner=owner, - operator_token=self._operator_token, - weights=weights, - slot_count=slot_count, - lane_count=lane_count, - device=device, + validate_training_non_aliasing( + { + **_named_moe_tensors("activation", activation), + "topk_idx": topk_idx, + "topk_weights": topk_weights, + "weights.fc1.payload": weights.fc1.payload, + "weights.fc1.scale": weights.fc1.scale, + "weights.fc2.payload": weights.fc2.payload, + "weights.fc2.scale": weights.fc2.scale, + "out.output": out.output, + "out.fc1_preact": out.fc1_preact, + "out.fc1_a": out.fc1_a, + "out.fc1_sfa": out.fc1_sfa, + "out.valid_route_counts": out.valid_route_counts, + "out.expert_offsets": out.expert_offsets, + } ) - self._training_resources = resources - return resources + from ._megamoe_backend.mxfp8._training_execute import ( + launch_training_forward, + ) + + with torch.cuda.device(self._forward_backend_device): + execution = self._training_state.views( + lane=lane.index, + token_count=token_count, + ) + return launch_training_forward( + self._training_state, + execution, + activation, + topk_idx, + topk_weights, + weights=weights, + out=out, + ) + + def training_backward( + self, + lane: MoeEpExecutionLane, + grad_output: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + weights: MoeEpNativeBackwardWeights, + fc1_preact: torch.Tensor, + fc1_a: torch.Tensor | None = None, + fc1_sfa: torch.Tensor | None = None, + valid_route_counts: torch.Tensor | None = None, + expert_offsets: torch.Tensor | None = None, + out: MoeEpTrainingBackwardOutputs | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + MoeEpTrainingWgradOperands, + ]: + """Run backward into caller-owned outputs using explicit forward state.""" + + with self._lifecycle_lock: + self._require_training_lane(lane) + assert self._training_state is not None + assert self._forward_backend_device is not None + if out is None: + raise TypeError("out must be a MoeEpTrainingBackwardOutputs") + token_count = validate_training_input( + self._forward_config, + "grad_output", + grad_output, + topk_idx, + topk_weights, + device=self._forward_backend_device, + ) + validate_native_backward_weights( + self._forward_config, + weights, + device=self._forward_backend_device, + ) + if fc1_preact is None: + raise ValueError("fc1_preact from the matching forward is required") + backward_output = out + validate_training_backward_outputs( + backward_output, + self._training_requirement_subset( + ( + "grad_activation", + "dprob", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + ) + ), + device=self._forward_backend_device, + ) + validate_training_forward_state( + fc1_preact=fc1_preact, + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + valid_route_counts=valid_route_counts, + expert_offsets=expert_offsets, + requirements=self._training_requirement_subset( + ( + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + ) + ), + device=self._forward_backend_device, + ) + validate_training_non_aliasing( + { + **_named_moe_tensors("grad_output", grad_output), + "topk_idx": topk_idx, + "topk_weights": topk_weights, + "weights.w2_transpose.payload": weights.w2_transpose.payload, + "weights.w2_transpose.scale": weights.w2_transpose.scale, + "weights.w1_transpose.payload": weights.w1_transpose.payload, + "weights.w1_transpose.scale": weights.w1_transpose.scale, + "fc1_preact": fc1_preact, + "fc1_a": fc1_a, + "fc1_sfa": fc1_sfa, + "valid_route_counts": valid_route_counts, + "expert_offsets": expert_offsets, + "out.grad_activation": backward_output.grad_activation, + "out.dprob": backward_output.dprob, + "out.fc1_b": backward_output.fc1_b, + "out.fc1_sfb": backward_output.fc1_sfb, + "out.fc2_a": backward_output.fc2_a, + "out.fc2_sfa": backward_output.fc2_sfa, + "out.fc2_b": backward_output.fc2_b, + "out.fc2_sfb": backward_output.fc2_sfb, + } + ) + from ._megamoe_backend.mxfp8._training_execute import ( + launch_training_backward, + ) + + with torch.cuda.device(self._forward_backend_device): + execution = self._training_state.views( + lane=lane.index, + token_count=token_count, + ) + return launch_training_backward( + self._training_state, + execution, + grad_output, + topk_idx, + topk_weights, + weights=weights, + fc1_preact=fc1_preact, + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + valid_route_counts=valid_route_counts, + expert_offsets=expert_offsets, + out=backward_output, + ) def close(self) -> None: """Release compiled-backend instance resources; idempotent.""" @@ -413,9 +751,9 @@ def close(self) -> None: self._forward_backend_device = None self._validated_topk_idx = None self._validated_topk_version = None - if self._training_resources is not None: - self._training_resources.close() - self._training_resources = None + self._training_state = None + self._training_lanes = () + self._training_requirements = None self._closed = True def __enter__(self) -> "MoeEp": @@ -449,10 +787,18 @@ def __del__(self) -> None: __all__ = [ "BlockScaledTensor", "MoeEp", + "MoeEpBackwardWeightStaging", + "MoeEpBackwardWeights", "MoeEpExecutionLane", - "MoeEpTrainingResources", - "MoeEpTrainingSlot", - "MoeEpTrainingWeights", + "MoeEpForwardWeightStaging", + "MoeEpForwardWeights", + "MoeEpNativeBackwardWeights", + "MoeEpNativeForwardWeights", + "MoeEpTrainingBackwardOutputs", + "MoeEpTrainingForwardOutputs", + "MoeEpTrainingWgradOperands", "MoeFormat", "MoeTensor", + "pack_backward_weights", + "pack_forward_weights", ] diff --git a/test/python/moe_ep/moe_ep_distributed_workers.py b/test/python/moe_ep/moe_ep_distributed_workers.py index 99cca10a6..0acf2ef2e 100644 --- a/test/python/moe_ep/moe_ep_distributed_workers.py +++ b/test/python/moe_ep/moe_ep_distributed_workers.py @@ -11,6 +11,8 @@ import torch.distributed as dist from moe_ep.moe_ep_test_support import ( + _allocate_stateless_training_outputs, + _allocate_training_weight_staging, _assert_backward_matches, _assert_grouped_wgrads_match_reference, _assert_matches_reference, @@ -21,6 +23,7 @@ _fixed_training_weights, _forward_config, _grad_output, + _interleave_fc1_wgrad, _output_as_float, _reference_forward, make_distributed_forward_inputs, @@ -28,9 +31,7 @@ ) __all__ = [ - "_distributed_backward_reference_worker", "_distributed_output_worker", - "_distributed_subgroup_backward_reference_worker", "_distributed_subgroup_output_worker", "_run_backward_reference_case", "_run_forward_output_case", @@ -237,7 +238,7 @@ def _run_backward_reference_case( gate_up_clamp: float | None = None, expected_global_ranks: tuple[int, ...] | None = None, ) -> None: - """Run fixed-resource training after the independent distributed oracle.""" + """Run stateless training after the independent distributed oracle.""" from cudnn import MoeEp @@ -262,7 +263,11 @@ def _run_backward_reference_case( drop_on_overflow=True, ) expected_y, expected_dx, expected_dprob, expected_wgrads = expected - expected_dense_wgrads = expected_wgrads.dense_wgrads() + expected_fc1_wgrad, expected_fc2_wgrad = expected_wgrads.dense_wgrads() + expected_dense_wgrads = ( + _interleave_fc1_wgrad(expected_fc1_wgrad), + expected_fc2_wgrad, + ) weights = _fixed_training_weights(args) op = MoeEp( @@ -276,29 +281,48 @@ def _run_backward_reference_case( drop_on_overflow=True, combine_format=combine_format, gate_up_clamp=gate_up_clamp, + weight_interleave_size=32, ) try: - resources = op.prepare_training_resources( - weights, - slot_count=1, + requirements = op.prepare_training( lane_count=1, + device=device, + ) + lane = op.training_lanes[0] + forward_staging, backward_staging = _allocate_training_weight_staging(weights) + native_forward = op.pack_forward_weights( + weights[0], + out=forward_staging, + ) + native_backward = op.pack_backward_weights( + weights[1], + out=backward_staging, + ) + forward_out, backward_out = _allocate_stateless_training_outputs( + requirements, + device, ) - slot = resources.slots[0] - lane = resources.lanes[0] - resources.refresh_weights() - actual_y = resources.forward( - slot, + actual_y = op.training_forward( lane, args[0], args[3], args[4], + weights=native_forward, + out=forward_out, ) - actual_dx, actual_dprob, actual_wgrads = resources.backward( - slot, + actual_dx, actual_dprob, actual_wgrads = op.training_backward( lane, grad_output, + args[3], + args[4], + weights=native_backward, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, ) - overflow = resources.finalize_overflow((slot,), lane) grouped_wgrads = _dense_wgrads_from_grouped_kernel(actual_wgrads) torch.cuda.synchronize(device) @@ -311,7 +335,6 @@ def _run_backward_reference_case( assert op.ep_size == ep_size if expected_global_ranks is not None: assert op.ep_global_ranks == expected_global_ranks - assert overflow.eq(0).all() assert args[3][0, 0] // 2 == ep_rank assert args[3][0, 1] // 2 == (ep_rank + 1) % ep_size assert args[3].eq(-1).any() @@ -364,110 +387,3 @@ def _run_backward_reference_case( raise assertion_error finally: op.close() - - -def _distributed_subgroup_backward_reference_worker( - global_rank: int, - global_world_size: int, - init_file: str, -) -> None: - """Run fixed-resource backward in two non-contiguous EP2 groups.""" - - device = torch.device("cuda", global_rank) - torch.cuda.set_device(device) - dist.init_process_group( - backend="nccl", - init_method=f"file://{init_file}", - rank=global_rank, - world_size=global_world_size, - device_id=device, - timeout=timedelta(minutes=10), - ) - ep_group = None - try: - subgroup_memberships = ((0, 2), (1, 3)) - # Every WORLD rank must create every subgroup in the same order. - subgroups = [ - dist.new_group( - list(members), - backend="nccl", - timeout=timedelta(minutes=10), - ) - for members in subgroup_memberships - ] - subgroup_index = global_rank % 2 - expected_global_ranks = subgroup_memberships[subgroup_index] - ep_group = subgroups[subgroup_index] - ep_rank = dist.get_rank(ep_group) - ep_size = dist.get_world_size(ep_group) - actual_global_ranks = tuple(dist.get_global_rank(ep_group, group_rank) for group_rank in range(ep_size)) - assert ep_size == len(expected_global_ranks) - assert ep_rank == expected_global_ranks.index(global_rank) - assert actual_global_ranks == expected_global_ranks - - _run_backward_reference_case( - device=device, - ep_group=ep_group, - ep_rank=ep_rank, - ep_size=ep_size, - combine_format="bf16", - expected_global_ranks=expected_global_ranks, - ) - finally: - if dist.is_initialized(): - try: - # Keep both independent groups alive until all work is done, - # then collectively finalize the process-local runtime. - dist.barrier() - from cudnn.moe_ep._megamoe_backend._runtime import ( - get_runtime_manager, - ) - - get_runtime_manager().shutdown() - dist.barrier() - finally: - if ep_group is not None: - dist.destroy_process_group(ep_group) - dist.destroy_process_group() - - -def _distributed_backward_reference_worker( - rank: int, - world_size: int, - init_file: str, - combine_format: str, - gate_up_clamp: float | None = None, -) -> None: - """Initialize one local rank and run distributed training parity.""" - - device = torch.device("cuda", rank) - torch.cuda.set_device(device) - dist.init_process_group( - backend="nccl", - init_method=f"file://{init_file}", - rank=rank, - world_size=world_size, - device_id=device, - timeout=timedelta(minutes=10), - ) - try: - _run_backward_reference_case( - device=device, - ep_group=dist.group.WORLD, - ep_rank=rank, - ep_size=world_size, - combine_format=combine_format, - gate_up_clamp=gate_up_clamp, - ) - finally: - if dist.is_initialized(): - try: - dist.barrier() - from cudnn.moe_ep._megamoe_backend._runtime import ( - get_runtime_manager, - ) - - get_runtime_manager().shutdown() - dist.barrier() - finally: - dist.destroy_process_group() diff --git a/test/python/moe_ep/moe_ep_test_support.py b/test/python/moe_ep/moe_ep_test_support.py index 304d7979e..11a3825c4 100644 --- a/test/python/moe_ep/moe_ep_test_support.py +++ b/test/python/moe_ep/moe_ep_test_support.py @@ -24,84 +24,37 @@ ) __all__ = [ - "_allocate_dense_grouped_wgrad_outputs", "_assert_backward_matches", - "_assert_fixed_training_drop_overflow_result", - "_assert_fixed_training_matches_reference", "_assert_grouped_wgrads_match_reference", "_assert_matches_reference", - "_assert_training_graph_tails_are_reset", - "_assert_training_weight_sources_changed", "_assert_wgrads_match_reference", - "_capture_fixed_training_batch", - "_copy_training_weight_sources_", "_dense_wgrads_from_operands", "_dense_wgrads_from_grouped_kernel", - "_expected_backward", - "_fixed_training_case", - "_fixed_training_drop_overflow_case", - "_fixed_training_drop_overflow_reference", "_fixed_training_reference", "_fixed_training_weights", + "_allocate_stateless_training_outputs", + "_allocate_training_weight_staging", "_forward_config", "_grad_output", "_make_forward_case", "_naive_reference", "_output_as_float", - "_prefill_training_graph_sentinels", "_reference_backward", "_reference_forward", "_replay_cuda_graph", "_require_distributed_sm107", - "_run_fixed_training_batch", "_run_grouped_wgrad_kernel", "_sm107_device", "_stress_backend_reuse", - "_training_public_pointers", - "_training_source_pointers", - "_training_weight_source_pointers", - "_training_weight_source_values", - "_TrainingResourceContractOwner", "_training_abi_prepared", "_training_config", - "_training_contract_resources", - "_training_inputs", "_training_prepared_pair", - "_training_staging_tensors", - "_training_weight_defect", - "_training_weights", "make_distributed_forward_inputs", "make_forward_inputs", "quantize_mxfp8", ] -def _allocate_dense_grouped_wgrad_outputs( - operands, - *, - fill_value=None, -): - """Allocate fixed-address dense BF16 outputs for FC1 and FC2 WGrad.""" - - expert_count = operands.expert_offsets.numel() - outputs = tuple( - torch.empty( - ( - expert_count, - getattr(operands, f"{prefix}_a").shape[0], - getattr(operands, f"{prefix}_b").shape[1], - ), - dtype=torch.bfloat16, - device=operands.expert_offsets.device, - ) - for prefix in ("fc1", "fc2") - ) - if fill_value is not None: - for output in outputs: - output.fill_(fill_value) - return outputs - - # Data @@ -292,16 +245,6 @@ def _training_config(**overrides): return ForwardConfig(**values) -def _training_inputs(): - return ( - torch.randn(2, 128, dtype=torch.bfloat16), - torch.randn(2, 128, 512, dtype=torch.bfloat16), - torch.randn(2, 256, 128, dtype=torch.bfloat16), - torch.tensor([[0, -1], [1, 0]], dtype=torch.int32), - torch.randn(2, 2, dtype=torch.float32), - ) - - def _training_prepared_pair(config, pool_rows: int = 512): from cudnn.moe_ep._megamoe_backend._workspace import WorkspaceRequirements @@ -339,12 +282,14 @@ def _training_prepared_pair(config, pool_rows: int = 512): config, kernel_local_workspace_bytes=3072, kernel_shared_workspace_bytes=4096, - backward_fc1_preact_bytes=pool_rows * 512 * 2, backward_dprob_bytes=4 * 2 * 4, backward_aux_data_bytes=pool_rows * 512, backward_aux_scale_bytes=512 * 8, ), - kernel=SimpleNamespace(get_aux_output_shapes=lambda: backward_shapes), + kernel=SimpleNamespace( + get_aux_output_shapes=lambda: backward_shapes, + get_fc1_preact_shape=lambda: forward_shapes["fc1_c"], + ), ) return forward, backward @@ -383,208 +328,6 @@ def _training_abi_prepared(name: str, max_recv_size: int = 4): ) -def _training_weights(args=None): - from cudnn.moe_ep import MoeEpTrainingWeights - from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( - _quantize_plain_mxfp8, - ) - - if args is None: - args = _training_inputs() - return MoeEpTrainingWeights( - forward_fc1=_quantize_plain_mxfp8(args[1], axis=1), - forward_fc2=_quantize_plain_mxfp8(args[2], axis=1), - backward_w2_transpose=_quantize_plain_mxfp8( - args[2].transpose(1, 2).contiguous(), - axis=1, - ), - backward_w1_transpose=_quantize_plain_mxfp8( - args[1].transpose(1, 2).contiguous(), - axis=1, - ), - ) - - -def _training_empty_block_scaled_like(tensor, *, axis: int, format: str): - import cudnn - - logical_shape = tensor.logical_shape - data_shape = list(logical_shape) - scale_shape = list(logical_shape) - if format == "mxfp8": - data_dtype = tensor.data.dtype - scale_dtype = tensor.scale.dtype - scale_shape[axis] = (logical_shape[axis] + 31) // 32 - else: - data_dtype = torch.uint8 - scale_dtype = tensor.data.dtype - data_shape[axis] = (logical_shape[axis] + 1) // 2 - scale_shape[axis] = (logical_shape[axis] + 15) // 16 - return cudnn.BlockScaledTensor( - data=torch.empty(tuple(data_shape), dtype=data_dtype, device=tensor.device), - scale=torch.empty( - tuple(scale_shape), - dtype=scale_dtype, - device=tensor.device, - ), - format=format, - logical_shape=logical_shape, - axis=axis, - ) - - -def _training_same_shape_noncontiguous(tensor: torch.Tensor) -> torch.Tensor: - result = tensor.transpose(-2, -1).contiguous().transpose(-2, -1) - assert tuple(result.shape) == tuple(tensor.shape) - assert not result.is_contiguous() - return result - - -def _training_weight_defect(weights, field: str, defect: str): - tensor = getattr(weights, field) - expected_shape = tensor.logical_shape - if defect == "plain_tensor": - invalid = torch.empty( - expected_shape, - dtype=torch.bfloat16, - device=tensor.device, - ) - error_type = TypeError - message = f"weights.{field} must be an MXFP8 BlockScaledTensor for " "fixed training resources" - elif defect == "logical_shape": - wrong_shape = (expected_shape[0] - 1, *expected_shape[1:]) - invalid = replace( - tensor, - data=tensor.data[: wrong_shape[0]].contiguous(), - scale=tensor.scale[: wrong_shape[0]].contiguous(), - logical_shape=wrong_shape, - ) - error_type = ValueError - message = f"weights.{field} logical shape must be {expected_shape}, " f"got {wrong_shape}" - elif defect == "axis": - invalid = _training_empty_block_scaled_like( - tensor, - axis=2, - format="mxfp8", - ) - error_type = ValueError - message = f"weights.{field} block-scaled axis must be 1, got 2" - elif defect == "format": - invalid = _training_empty_block_scaled_like( - tensor, - axis=1, - format="nvfp4", - ) - error_type = NotImplementedError - message = f"weights.{field} must use format='mxfp8', got 'nvfp4'" - elif defect == "device": - invalid = replace( - tensor, - data=torch.empty_like(tensor.data, device="meta"), - scale=torch.empty_like(tensor.scale, device="meta"), - ) - error_type = ValueError - message = f"weights.{field} must be on cpu, got meta" - else: - part = "data" if defect == "data_noncontiguous" else "scale" - invalid = replace( - tensor, - **{part: _training_same_shape_noncontiguous(getattr(tensor, part))}, - ) - error_type = ValueError - message = f"weights.{field} data and scale must be contiguous for fixed " "training weight binding" - return replace(weights, **{field: invalid}), error_type, message - - -class _TrainingResourceContractOwner: - def __init__(self, *, slot_count: int = 2, lane_count: int = 1) -> None: - self.slot_count = slot_count - self.lane_count = lane_count - self.close_calls = 0 - self.refresh_calls = 0 - self.views_calls = 0 - - def refresh_weights(self) -> None: - self.refresh_calls += 1 - - def views(self, **kwargs): - del kwargs - self.views_calls += 1 - raise AssertionError("binding rejection must happen before owner views") - - def _flat_views(self, token_count: int): - del token_count - raise AssertionError("invalid finalization must fail before workspace access") - - def finalize_overflow(self, slots, *, lane): - from cudnn.moe_ep._megamoe_backend.mxfp8._training_resources import ( - Mxfp8TrainingResourceOwner, - ) - - return Mxfp8TrainingResourceOwner.finalize_overflow( - self, - slots, - lane=lane, - ) - - def close(self) -> None: - self.close_calls += 1 - - -def _training_contract_resources( - *, - owner=None, - slot_count: int = 2, - lane_count: int = 1, -): - from cudnn.moe_ep import MoeEpTrainingResources - - if owner is None: - owner = _TrainingResourceContractOwner( - slot_count=slot_count, - lane_count=lane_count, - ) - resources = MoeEpTrainingResources( - owner=owner, - operator_token=object(), - weights=SimpleNamespace(mock_training_weights=True), - slot_count=slot_count, - lane_count=lane_count, - device=torch.device("cpu"), - ) - return resources, owner - - -def _training_staging_tensors(*, capacity: int | None = None): - activation, _, _, topk_idx, topk_weights = make_forward_inputs(torch.device("cpu")) - source = activation.dequantize(torch.bfloat16).contiguous() - token_count, hidden = source.shape - top_k = topk_idx.shape[1] - if capacity is None: - capacity = token_count - return { - "source": source, - "topk_idx": topk_idx, - "topk_weights": topk_weights.float().contiguous(), - "output": torch.empty( - (capacity, hidden), - dtype=torch.float8_e4m3fn, - ), - "output_sf": torch.empty( - (capacity, hidden // 32), - dtype=torch.float8_e8m0fnu, - ), - "output_topk_idx": torch.empty( - (capacity, top_k), - dtype=torch.int32, - ), - "output_topk_weights": torch.empty( - (capacity, top_k), - dtype=torch.float32, - ), - } - - # Forward @@ -949,7 +692,7 @@ def _dequantize_wgrad_operand( def _dense_wgrads_from_operands(operands): - """Reference grouped matmuls over the exported operand ABI.""" + """Reference grouped matmuls over the producer-native operand ABI.""" fc1_a = _dequantize_wgrad_operand( operands.fc1_a, @@ -1017,7 +760,7 @@ def _run_grouped_wgrad_kernel( if prefix not in ("fc1", "fc2"): raise ValueError(f"prefix must be 'fc1' or 'fc2', got {prefix!r}") - # Graph callers provide one persistent output per training slot. This is + # Graph callers provide one persistent output per training lane. This is # currently also the isolation key for a temporary production-WGrad # workaround: an EP2 graph with two same-signature calls produced correct # operands but corrupted the second WGrad when both calls shared the @@ -1090,16 +833,12 @@ def _assert_grouped_wgrads_match_reference( expected_fp32 = expected_dw.float() absolute_error = (actual_fp32 - expected_fp32).abs() max_absolute_error = absolute_error.max().item() - max_relative_error = ( - (absolute_error / expected_fp32.abs().clamp_min(1.0e-6)).max().item() - ) + max_relative_error = (absolute_error / expected_fp32.abs().clamp_min(1.0e-6)).max().item() torch.testing.assert_close( actual_fp32, expected_fp32, msg=lambda default, name=name: ( - f"{name} does not match {reference_name}; " - f"max_abs_error={max_absolute_error:.6g}, " - f"max_rel_error={max_relative_error:.6g}\n{default}" + f"{name} does not match {reference_name}; " f"max_abs_error={max_absolute_error:.6g}, " f"max_rel_error={max_relative_error:.6g}\n{default}" ), **close_kwargs, ) @@ -1124,9 +863,9 @@ def _reference_backward(config) -> MoeEpReference: def _fixed_training_weights(args): - """Build the four stable MXFP8 source packs required by training.""" + """Build independent source packs for allocation-free native packing.""" - from cudnn.moe_ep import MoeEpTrainingWeights + from cudnn.moe_ep import MoeEpBackwardWeights, MoeEpForwardWeights from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( _quantize_plain_mxfp8, ) @@ -1135,18 +874,110 @@ def _fixed_training_weights(args): fc2_weight = args[2] dense_fc1 = fc1_weight if isinstance(fc1_weight, torch.Tensor) else fc1_weight.dequantize() dense_fc2 = fc2_weight if isinstance(fc2_weight, torch.Tensor) else fc2_weight.dequantize() - return MoeEpTrainingWeights( - forward_fc1=(_quantize_plain_mxfp8(dense_fc1, axis=1) if isinstance(fc1_weight, torch.Tensor) else fc1_weight), - forward_fc2=(_quantize_plain_mxfp8(dense_fc2, axis=1) if isinstance(fc2_weight, torch.Tensor) else fc2_weight), - backward_w2_transpose=_quantize_plain_mxfp8( + forward = MoeEpForwardWeights( + fc1=(_quantize_plain_mxfp8(dense_fc1, axis=1) if isinstance(fc1_weight, torch.Tensor) else fc1_weight), + fc2=(_quantize_plain_mxfp8(dense_fc2, axis=1) if isinstance(fc2_weight, torch.Tensor) else fc2_weight), + ) + backward = MoeEpBackwardWeights( + w2_transpose=_quantize_plain_mxfp8( dense_fc2.transpose(1, 2).contiguous(), axis=1, ), - backward_w1_transpose=_quantize_plain_mxfp8( + w1_transpose=_quantize_plain_mxfp8( dense_fc1.transpose(1, 2).contiguous(), axis=1, ), ) + return forward, backward + + +def _allocate_training_weight_staging(weights): + """Allocate caller-owned native pack destinations for one source pair.""" + + from cudnn.moe_ep import ( + MoeEpBackwardWeightStaging, + MoeEpForwardWeightStaging, + ) + + forward, backward = weights + fc1 = forward.fc1 + fc2 = forward.fc2 + experts, hidden, gate_up = fc1.data.shape + intermediate = fc2.data.shape[1] + + def scale(elements): + return torch.empty( + (experts, elements), + dtype=torch.float8_e8m0fnu, + device=fc1.device, + ) + + def blocked_elements(rows, columns): + return ((rows + 127) // 128 * 128) * ((columns + 3) // 4 * 4) + + forward_out = MoeEpForwardWeightStaging( + fc1_payload=torch.empty_strided( + fc1.data.shape, + (hidden * gate_up, 1, hidden), + dtype=fc1.data.dtype, + device=fc1.device, + ), + fc1_scale=scale(blocked_elements(gate_up, hidden // 32)), + fc2_payload=torch.empty_strided( + fc2.data.shape, + (intermediate * hidden, 1, intermediate), + dtype=fc2.data.dtype, + device=fc2.device, + ), + fc2_scale=scale(blocked_elements(hidden, intermediate // 32)), + ) + w2t = backward.w2_transpose + w1t = backward.w1_transpose + backward_out = MoeEpBackwardWeightStaging( + w2_transpose_payload=torch.empty_like(w2t.data), + w2_transpose_scale=scale(blocked_elements(intermediate, hidden // 32)), + w1_transpose_payload=torch.empty_like(w1t.data), + w1_transpose_scale=scale(blocked_elements(hidden, gate_up // 32)), + ) + return forward_out, backward_out + + +def _allocate_stateless_training_outputs(requirements, device): + """Allocate every advertised caller-owned output contract.""" + + from cudnn.moe_ep import ( + MoeEpTrainingBackwardOutputs, + MoeEpTrainingForwardOutputs, + ) + + def allocate(name): + shape, stride, dtype, _alignment = requirements[name] + return torch.empty_strided( + shape, + stride, + dtype=dtype, + device=device, + ) + + forward = MoeEpTrainingForwardOutputs( + output=allocate("output"), + fc1_preact=allocate("fc1_preact"), + fc1_a=allocate("fc1_a"), + fc1_sfa=allocate("fc1_sfa"), + valid_route_counts=allocate("valid_route_counts"), + expert_offsets=allocate("expert_offsets"), + ) + backward = MoeEpTrainingBackwardOutputs( + grad_activation=allocate("grad_activation"), + dprob=allocate("dprob"), + fc1_b=allocate("fc1_b"), + fc1_sfb=allocate("fc1_sfb"), + fc2_a=allocate("fc2_a"), + fc2_sfa=allocate("fc2_sfa"), + fc2_b=allocate("fc2_b"), + fc2_sfb=allocate("fc2_sfb"), + ) + return forward, backward def _fixed_training_reference( @@ -1181,7 +1012,7 @@ def _fixed_training_reference( generate_c=True, backward_wgrad_mode="operands", # The standalone operand oracle's legacy ABI uses 256-row - # segments. Production fixed resources use 128-row segments; + # segments. The stateless producer ABI uses 128-row segments; # their represented dense gradients are compared below. token_padding_size=256, ) @@ -1222,14 +1053,6 @@ def _grad_output( ) -def _expected_backward(reference, grad_output, args, stash): - return reference.backward( - grad_output, - *_reference_args(args)[1:], - *stash, - ) - - def _assert_backward_matches(actual, expected, topk_idx) -> None: assert len(actual) == len(expected) == 2 for name, gradient, reference, close_kwargs in zip( @@ -1252,6 +1075,25 @@ def _assert_backward_matches(actual, expected, topk_idx) -> None: assert actual[1][dropped].eq(0).all() +def _interleave_fc1_wgrad( + tensor: torch.Tensor, + interleave_size: int = 32, +) -> torch.Tensor: + """Convert logical gate-then-up columns to producer-native strip order.""" + + out_features = tensor.shape[-1] + return ( + tensor.view( + *tensor.shape[:-1], + 2, + out_features // (2 * interleave_size), + interleave_size, + ) + .transpose(-3, -2) + .reshape(tensor.shape) + ) + + def _assert_wgrads_match_reference( actual, expected, @@ -1273,18 +1115,10 @@ def _assert_wgrads_match_reference( expected_dense = expected.dense_wgrads() if weight_interleave_size is not None: expected_fc1, expected_fc2 = expected_dense - fc1_out_features = expected_fc1.shape[-1] - expected_fc1 = ( - expected_fc1.view( - *expected_fc1.shape[:-1], - 2, - fc1_out_features // (2 * weight_interleave_size), - weight_interleave_size, - ) - .transpose(-3, -2) - .reshape(expected_fc1.shape) + expected_dense = ( + _interleave_fc1_wgrad(expected_fc1, weight_interleave_size), + expected_fc2, ) - expected_dense = (expected_fc1, expected_fc2) for name, actual_dw, expected_dw in zip( ("grad_fc1_weight", "grad_fc2_weight"), actual_dense, @@ -1296,278 +1130,3 @@ def _assert_wgrads_match_reference( msg=lambda default, name=name: (f"{name} does not match the independent reference\n{default}"), **_WGRAD_CLOSE_KWARGS, ) - - -_TRAINING_WGRAD_DATA_FIELDS = ("fc1_a", "fc1_b", "fc2_a", "fc2_b") -_TRAINING_WGRAD_SF_FIELDS = ("fc1_sfa", "fc1_sfb", "fc2_sfa", "fc2_sfb") -_TRAINING_WEIGHT_FIELDS = ( - "forward_fc1", - "forward_fc2", - "backward_w2_transpose", - "backward_w1_transpose", -) - - -def _fixed_training_case(device): - args = list(make_forward_inputs(device)) - args[0] = args[0].dequantize(torch.bfloat16) - args[4] = args[4].float() - args[3].fill_(-1) - args[4].zero_() - args[3][0, 0] = 0 - args[4][0, 0] = 1 - grad_output = _grad_output( - device, - args[0].shape[0], - seed=20260828, - ) - return args, grad_output - - -def _assert_fixed_training_matches_reference( - actual, - expected, - topk_idx, -) -> None: - actual_y, actual_dx, actual_dprob, actual_wgrads = actual - expected_y, expected_dx, expected_dprob, expected_wgrads = expected - _assert_matches_reference(actual_y, expected_y) - _assert_backward_matches( - (actual_dx, actual_dprob), - (expected_dx, expected_dprob), - topk_idx, - ) - _assert_wgrads_match_reference(actual_wgrads, expected_wgrads) - - -def _run_fixed_training_batch(resources, lane, cases): - """Run refresh, ordered forwards/backwards, and one overflow finalization.""" - - resources.refresh_weights() - outputs = [resources.forward(slot, lane, args[0], args[3], args[4]) for slot, args, _ in cases] - backwards = [resources.backward(slot, lane, grad_output) for slot, _, grad_output in cases] - overflow = resources.finalize_overflow( - tuple(slot for slot, _, _ in cases), - lane, - ) - return tuple( - SimpleNamespace( - y=output, - dx=backward[0], - dprob=backward[1], - wgrads=backward[2], - overflow=overflow, - ) - for output, backward in zip(outputs, backwards) - ) - - -def _capture_fixed_training_batch( - resources, - lane, - cases, - capture_stream, - *, - grouped_wgrad_outputs=None, -): - """Capture the shared fixed-training sequence for one or more slots.""" - - if grouped_wgrad_outputs is not None and len(grouped_wgrad_outputs) != len(cases): - raise ValueError("grouped_wgrad_outputs must match the captured case count") - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=capture_stream): - actuals = _run_fixed_training_batch(resources, lane, cases) - grouped_wgrads = ( - None - if grouped_wgrad_outputs is None - else tuple( - _dense_wgrads_from_grouped_kernel( - actual.wgrads, - wgrad_tensors=outputs, - ) - for actual, outputs in zip(actuals, grouped_wgrad_outputs) - ) - ) - capture_stream.synchronize() - return SimpleNamespace( - graph=graph, - actuals=actuals, - grouped_wgrads=grouped_wgrads, - public_pointers=tuple(_training_public_pointers(actual) for actual in actuals), - ) - - -def _training_source_pointers(case) -> dict[str, int]: - return { - name: getattr(case, name).data_ptr() - for name in ( - "activation", - "topk_idx", - "topk_weights", - "grad_output", - ) - } - - -def _training_weight_source_pointers(weights) -> dict[str, int]: - return {f"{name}.{part}": getattr(getattr(weights, name), part).data_ptr() for name in _TRAINING_WEIGHT_FIELDS for part in ("data", "scale")} - - -def _training_weight_source_values(weights) -> dict[str, torch.Tensor]: - return {f"{name}.{part}": getattr(getattr(weights, name), part).clone() for name in _TRAINING_WEIGHT_FIELDS for part in ("data", "scale")} - - -def _assert_training_weight_sources_changed(weights, previous) -> None: - for name in _TRAINING_WEIGHT_FIELDS: - for part in ("data", "scale"): - assert not torch.equal( - getattr(getattr(weights, name), part), - previous[f"{name}.{part}"], - ) - - -def _training_public_pointers(actual) -> dict[str, int]: - pointers = { - "y": actual.y.data_ptr(), - "dx": actual.dx.data_ptr(), - "dprob": actual.dprob.data_ptr(), - "overflow": actual.overflow.data_ptr(), - } - pointers.update( - { - f"wgrads.{name}": getattr(actual.wgrads, name).data_ptr() - for name in ( - *_TRAINING_WGRAD_DATA_FIELDS, - *_TRAINING_WGRAD_SF_FIELDS, - "expert_offsets", - "valid_route_counts", - ) - } - ) - return pointers - - -def _prefill_training_graph_sentinels(slot_views, actual) -> None: - """Poison every history-sensitive full-capacity destination.""" - - slot_views.routing_topk_idx.fill_(0x1A2B3C) - slot_views.routing_topk_weights.fill_(31.25) - slot_views.forward_output.fill_(29.0) - slot_views.backward_output.fill_(-27.0) - slot_views.grad_activation.fill_(23.0) - slot_views.dprob.fill_(-19.0) - slot_views.expert_offsets.fill_(-17) - slot_views.valid_route_counts.fill_(-13) - for name in _TRAINING_WGRAD_DATA_FIELDS: - getattr(actual.wgrads, name).fill_(1.0) - for name in _TRAINING_WGRAD_SF_FIELDS: - getattr(actual.wgrads, name).view(torch.uint8).fill_(0) - - -def _assert_training_graph_tails_are_reset( - slot_views, - actual, - *, - token_count: int, - capacity: int, -) -> None: - if token_count < capacity: - assert slot_views.routing_topk_idx[token_count:].eq(-1).all() - assert slot_views.routing_topk_weights[token_count:].eq(0).all() - assert slot_views.forward_output[token_count:].eq(0).all() - assert slot_views.backward_output[token_count:].eq(0).all() - assert slot_views.grad_activation[token_count:].eq(0).all() - assert slot_views.dprob[token_count:].eq(0).all() - - counts = actual.wgrads.valid_route_counts.detach().cpu().tolist() - expected_offsets = [] - offset = 0 - for count in counts: - offset += (int(count) + 127) // 128 * 128 - expected_offsets.append(offset) - assert actual.wgrads.expert_offsets.detach().cpu().tolist() == expected_offsets - - -def _copy_training_weight_sources_(destination, source) -> None: - for name in _TRAINING_WEIGHT_FIELDS: - destination_pack = getattr(destination, name) - source_pack = getattr(source, name) - destination_pack.data.copy_(source_pack.data) - destination_pack.scale.copy_(source_pack.scale) - - -def _fixed_training_drop_overflow_case(device): - base_args, base_grad_output = _fixed_training_case(device) - topk_idx = torch.tensor( - [[0, 1]], - dtype=torch.int32, - device=device, - ) - topk_weights = torch.tensor( - [[0.75, 0.25]], - dtype=torch.float32, - device=device, - ) - args = ( - base_args[0][:1].clone(), - base_args[1], - base_args[2], - topk_idx, - topk_weights, - ) - return args, base_grad_output[:1].clone() - - -def _fixed_training_drop_overflow_reference( - args, - grad_output, - *, - drop_expert1, -): - reference_topk_idx = args[3].clone() - if drop_expert1: - assert reference_topk_idx.shape == (1, 2) - assert reference_topk_idx.detach().cpu().tolist() == [[0, 1]] - reference_topk_idx[0, 1] = -1 - reference_args = ( - args[0], - args[1], - args[2], - reference_topk_idx, - args[4].clone(), - ) - return ( - _fixed_training_reference( - reference_args, - grad_output, - combine_format="bf16", - gate_up_clamp=None, - ), - reference_topk_idx, - ) - - -def _assert_fixed_training_drop_overflow_result( - actual, - expected, - reference_topk_idx, - *, - expected_overflow, -): - assert actual.overflow.eq(expected_overflow).all() - _assert_fixed_training_matches_reference( - (actual.y, actual.dx, actual.dprob, actual.wgrads), - expected, - reference_topk_idx, - ) - - if expected_overflow: - assert reference_topk_idx[0, 1].eq(-1) - assert actual.dprob[0, 1].eq(0) - assert actual.wgrads.valid_route_counts.detach().cpu().tolist() == [1, 0] - # Expert 0 owns the first 128-row padded segment. Expert 1 starts at - # pool capacity and therefore has no retained segment or dense dW. - assert actual.wgrads.expert_offsets.detach().cpu().tolist() == [128, 128] - dense_dw1, dense_dw2 = _dense_wgrads_from_operands(actual.wgrads) - assert dense_dw1[1].eq(0).all() - assert dense_dw2[1].eq(0).all() diff --git a/test/python/moe_ep/probe_moe_ep_training_graph.py b/test/python/moe_ep/probe_moe_ep_training_graph.py index 537c25998..67cdc3840 100644 --- a/test/python/moe_ep/probe_moe_ep_training_graph.py +++ b/test/python/moe_ep/probe_moe_ep_training_graph.py @@ -2,1489 +2,244 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Fixed-resource SM107 multi-rank CUDA Graph communication probe. - -Run from the project root on one node with two Rubin GPUs. The container -launcher selects an architecture-native Python/PyTorch environment:: - - data/script/run_moe_ep_bf16_combine_container.sh \ - --my-version-root "$PWD/my-version" \ - training-graph-probe - -Set ``MOE_EP_GRAPH_PROBE_NPROC=4`` (or another local EP size) on the host to -reuse the same probe beyond EP2. - -The probe exercises only the public fixed-resource ordinary/capture path, -including fixed-address staging/reset operations, forward/backward CuTeDSL -callables, FC1/FC2 production grouped WGrad, and a one-scalar NCCL overflow OR. -""" +"""Stateless SM107 multi-rank CUDA Graph training probe.""" from __future__ import annotations import argparse -import gc import os -import socket -import time -from contextlib import contextmanager from datetime import timedelta import torch import torch.distributed as dist -from cudnn import MoeEp, MoeEpTrainingWeights -from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( - _quantize_plain_mxfp8, -) -from cudnn.moe_ep._megamoe_backend._runtime import ( - _RuntimeWatchdog, - get_runtime_manager, -) +from cudnn import MoeEp from moe_ep.moe_ep_test_support import ( - _allocate_dense_grouped_wgrad_outputs, - _dense_wgrads_from_grouped_kernel, - _dense_wgrads_from_operands, + _allocate_stateless_training_outputs, + _allocate_training_weight_staging, + _fixed_training_weights, + _grad_output, + make_distributed_forward_inputs, ) -def _debug_phase(rank: int, phase: str) -> None: - if os.environ.get("MOE_EP_DEBUG_RUNTIME", "0") != "1": - return - print( - "[moe-ep-probe] " f"time={time.monotonic():.6f} host={socket.gethostname()} " f"pid={os.getpid()} rank={rank} phase={phase}", - flush=True, - ) - - -@contextmanager -def _debug_phase_scope(rank: int, phase: str): - _debug_phase(rank, f"{phase}.begin") - try: - yield - finally: - _debug_phase(rank, f"{phase}.end") - - -def _synchronize_with_watchdog( - rank: int, - device: torch.device, - phase: str, -) -> None: - watchdog = _RuntimeWatchdog(phase) - watchdog.start() - _debug_phase(rank, f"{phase}.begin") - try: - torch.cuda.synchronize(device) - finally: - watchdog.close() - _debug_phase(rank, f"{phase}.end") - - def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--diagnostic-replays", type=int, default=2) parser.add_argument("--burst-replays", type=int, default=100) - parser.add_argument( - "--multistream-replays", - type=int, - default=10, - help=("two-lane cross-stream graph replays; use a larger value such as " "100 for dedicated stress runs"), - ) - parser.add_argument( - "--max-recv-size-per-rank", - type=int, - default=1, - help=("bounded receive capacity; must remain below the forced-overflow " "route count so the probe retains overflow coverage"), - ) - parser.add_argument( - "--cycles", - type=int, - default=2, - help=("create/capture/destroy cycles; the first is exhaustive and later " "cycles use a minimal replay to verify teardown/re-init"), - ) + parser.add_argument("--multistream-replays", type=int, default=10) + parser.add_argument("--max-recv-size-per-rank", type=int, default=1) + parser.add_argument("--cycles", type=int, default=2) parser.add_argument("--timeout-seconds", type=int, default=600) - parser.add_argument( - "--skip-multistream", - action="store_true", - help="skip the two-lane ordered cross-stream resource probe", - ) - parser.add_argument( - "--wgrad-capture-mode", - choices=("both", "slot0", "slot1"), - default="both", - help=( - "capture both grouped-WGrad calls, or isolate exactly one slot to " - "diagnose same-signature graph reuse" - ), - ) - parser.add_argument( - "--expect-overflow-assert", - action="store_true", - help=("run only the fatal drop_on_overflow=False graph assertion probe; " "success requires every rank to observe the expected CUDA error"), - ) + parser.add_argument("--skip-multistream", action="store_true") + parser.add_argument("--expect-overflow-assert", action="store_true") return parser.parse_args() -def _require_positive(name: str, value: int) -> None: +def _positive(name: str, value: int) -> None: if value <= 0: raise ValueError(f"{name} must be positive, got {value}") -def _assert_replay_tensor( - name: str, - actual: torch.Tensor, - expected: torch.Tensor, -) -> None: - """Compare graph replay outputs with dtype-appropriate semantics.""" - - low_precision = { - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float8_e8m0fnu, - torch.uint8, - torch.int32, - torch.int64, - } - if actual.dtype in low_precision: - if not torch.equal(actual, expected): - _report_tensor_difference(name, actual, expected) - raise AssertionError(f"{name} is not bitwise equal after graph replay") - return - try: - torch.testing.assert_close( - actual, - expected, - rtol=1e-5, - atol=1e-6, - msg=f"{name} differs after graph replay", - ) - except AssertionError: - _report_tensor_difference(name, actual, expected) - raise - - -def _report_tensor_difference( - name: str, - actual: torch.Tensor, - expected: torch.Tensor, -) -> None: - """Print actionable mismatch statistics without changing pass criteria.""" - - rank = dist.get_rank() if dist.is_initialized() else 0 - if actual.shape != expected.shape or actual.dtype != expected.dtype: - print( - "MOE_EP_GRAPH_TENSOR_DIAGNOSTIC " - f"rank={rank} name={name} " - f"actual_shape={tuple(actual.shape)} expected_shape={tuple(expected.shape)} " - f"actual_dtype={actual.dtype} expected_dtype={expected.dtype}", - flush=True, - ) - return - - actual_fp32 = actual.float() - expected_fp32 = expected.float() - finite = torch.isfinite(actual_fp32) & torch.isfinite(expected_fp32) - absolute_error = (actual_fp32 - expected_fp32).abs() - relative_error = absolute_error / expected_fp32.abs().clamp_min(1.0e-6) - finite_absolute = absolute_error.masked_select(finite) - finite_relative = relative_error.masked_select(finite) - max_absolute = ( - float(finite_absolute.max().item()) if finite_absolute.numel() else float("nan") - ) - max_relative = ( - float(finite_relative.max().item()) if finite_relative.numel() else float("nan") - ) - exact_mismatch = actual.view(torch.uint8).ne(expected.view(torch.uint8)) - close_mismatch = ~torch.isclose( - actual_fp32, - expected_fp32, - rtol=1.0e-5, - atol=1.0e-6, - equal_nan=True, - ) - logical_mismatch = actual.ne(expected) - first_indices = logical_mismatch.nonzero() - first_description = "none" - if first_indices.numel(): - first_index = tuple(int(value) for value in first_indices[0].tolist()) - first_description = ( - f"index={first_index},actual={float(actual[first_index].float().item()):.9g}," - f"expected={float(expected[first_index].float().item()):.9g}" - ) - - print( - "MOE_EP_GRAPH_TENSOR_DIAGNOSTIC " - f"rank={rank} name={name} dtype={actual.dtype} shape={tuple(actual.shape)} " - f"byte_mismatches={int(exact_mismatch.sum().item())} " - f"logical_mismatches={int(logical_mismatch.sum().item())} " - f"close_mismatches={int(close_mismatch.sum().item())} " - f"max_abs={max_absolute:.9g} max_rel={max_relative:.9g} " - f"actual_nonfinite={int((~torch.isfinite(actual_fp32)).sum().item())} " - f"expected_nonfinite={int((~torch.isfinite(expected_fp32)).sum().item())} " - f"first_mismatch={first_description}", - flush=True, - ) - if actual.ndim == 3: - expert_dims = (1, 2) - expert_max_absolute = absolute_error.amax(dim=expert_dims) - expert_max_relative = relative_error.amax(dim=expert_dims) - expert_close_mismatches = close_mismatch.sum(dim=expert_dims) - print( - "MOE_EP_GRAPH_EXPERT_DIAGNOSTIC " - f"rank={rank} name={name} " - f"max_abs={expert_max_absolute.detach().cpu().tolist()} " - f"max_rel={expert_max_relative.detach().cpu().tolist()} " - f"close_mismatches={expert_close_mismatches.detach().cpu().tolist()}", - flush=True, - ) - - -def _report_grouped_wgrad_operand_consistency( - slot_name: str, - operands, - grouped_wgrads, -) -> None: - """Report whether replayed WGrad agrees with its replayed operand bundle.""" - - rank = dist.get_rank() if dist.is_initialized() else 0 - try: - decoded = _dense_wgrads_from_operands(operands) - except BaseException as error: - print( - "MOE_EP_GRAPH_OPERAND_DECODE_ERROR " - f"rank={rank} slot={slot_name} " - f"error={type(error).__name__}:{error}", - flush=True, - ) - return - for prefix, actual, expected in zip( - ("fc1", "fc2"), - grouped_wgrads, - decoded, - ): - _report_tensor_difference( - f"{slot_name}.{prefix}_wgrad_vs_decoded_operands", - actual, - expected.to(actual.dtype), - ) - - -def _make_inputs( - rank: int, - device: torch.device, -) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: - token_count = 8 - hidden = 128 - intermediate = 256 - experts_per_rank = 2 - top_k = 2 - generator = torch.Generator(device=device).manual_seed(20260828 + rank) - - activation = ( - torch.randn( - token_count, - hidden, - dtype=torch.bfloat16, - device=device, - generator=generator, - ) - / 8 - ) - fc1_weight = ( - torch.randn( - experts_per_rank, - hidden, - 2 * intermediate, - dtype=torch.bfloat16, - device=device, - generator=generator, - ) - / 16 - ) - fc2_weight = ( - torch.randn( - experts_per_rank, - intermediate, - hidden, - dtype=torch.bfloat16, - device=device, - generator=generator, +def _capture_training_graph( + op: MoeEp, + lane, + args, + grad_output, + native_forward, + native_backward, + forward_out, + backward_out, +) -> torch.cuda.CUDAGraph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op.training_forward( + lane, + args[0], + args[3], + args[4], + weights=native_forward, + out=forward_out, ) - / 16 - ) - topk_idx = torch.full( - (token_count, top_k), - -1, - dtype=torch.int32, - device=device, - ) - topk_weights = torch.zeros( - (token_count, top_k), - dtype=torch.float32, - device=device, - ) - # Exactly one route is received by each rank during eager warmup, so every - # positive max_recv_size_per_rank remains within capacity. - topk_idx[0, 0] = rank * experts_per_rank - topk_weights[0, 0] = 1.0 - grad_output = ( - torch.randn( - token_count, - hidden, - dtype=torch.float32, - device=device, - generator=generator, + op.training_backward( + lane, + grad_output, + args[3], + args[4], + weights=native_backward, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, ) - / 8 - ) - return ( - activation, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights, - ), grad_output + return graph -def _route_pattern( - kind: str, - rank: int, - world_size: int, +def _prepare_case( + *, device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor]: - token_count = 8 - top_k = 2 - experts_per_rank = 2 - indices = torch.full( - (token_count, top_k), - -1, - dtype=torch.int32, - device=device, - ) - weights = torch.zeros( - (token_count, top_k), - dtype=torch.float32, - device=device, - ) - if kind == "local": - indices[0, 0] = rank * experts_per_rank - weights[0, 0] = 1.0 - elif kind == "remote": - peer = (rank + 1) % world_size - indices[0, 0] = peer * experts_per_rank - weights[0, 0] = 1.0 - elif kind == "overflow": - # Every source sends all routes to rank 0. Its raw receive count is - # therefore much larger than max_recv_size_per_rank=1. - indices.fill_(0) - weights.fill_(0.5) - else: - raise ValueError(f"unknown route pattern {kind!r}") - return indices, weights - - -def _make_two_slot_inputs( rank: int, world_size: int, - device: torch.device, -) -> tuple[ - tuple[torch.Tensor, ...], - torch.Tensor, - tuple[torch.Tensor, ...], - torch.Tensor, - tuple[torch.Tensor, torch.Tensor], - tuple[torch.Tensor, torch.Tensor], -]: - args0, grad0 = _make_inputs(rank, device) - local = _route_pattern("local", rank, world_size, device) - remote = _route_pattern("remote", rank, world_size, device) - args0 = (*args0[:3], local[0].clone(), local[1].clone()) - args1 = ( - args0[0].clone(), - args0[1], - args0[2], - remote[0].clone(), - remote[1].clone(), - ) - return args0, grad0, args1, grad0.clone(), local, remote - - -def _make_training_weights( - args: tuple[torch.Tensor, ...], -) -> MoeEpTrainingWeights: - return MoeEpTrainingWeights( - forward_fc1=_quantize_plain_mxfp8(args[1], axis=1), - forward_fc2=_quantize_plain_mxfp8(args[2], axis=1), - backward_w2_transpose=_quantize_plain_mxfp8( - args[2].transpose(1, 2).contiguous(), - axis=1, - ), - backward_w1_transpose=_quantize_plain_mxfp8( - args[1].transpose(1, 2).contiguous(), - axis=1, - ), - ) - - -def _make_operator( - *, - world_size: int, - group, + lane_count: int, max_recv_size_per_rank: int, drop_on_overflow: bool, -) -> MoeEp: - return MoeEp( +): + args = make_distributed_forward_inputs(rank, world_size, device) + args = (*args[:4], args[4].float().contiguous()) + grad_output = _grad_output(device, args[0].shape[0], seed=7000 + rank) + source_weights = _fixed_training_weights(args) + op = MoeEp( num_experts=2 * world_size, hidden_size=128, intermediate_size=256, top_k=2, - ep_group=group, + ep_group=dist.group.WORLD, + # This is a collective ABI capacity, not the rank-local token count. + # make_distributed_forward_inputs intentionally varies local shapes. max_tokens_per_rank=8, max_recv_size_per_rank=max_recv_size_per_rank, drop_on_overflow=drop_on_overflow, combine_format="bf16", + weight_interleave_size=32, ) - - -def _close_probe_operator( - *, - device: torch.device, - group, - op: MoeEp, -) -> None: - torch.cuda.synchronize(device) - dist.barrier(group=group) - op.close() - gc.collect() - torch.cuda.synchronize(device) - dist.barrier(group=group) - - -def _run_single_wgrad_slot_capture_probe( - *, - rank: int, - world_size: int, - device: torch.device, - group, - max_recv_size_per_rank: int, - slot_index: int, -) -> None: - """Capture the full training chain with one grouped-WGrad invocation.""" - - args0, grad0, args1, grad1, _, _ = _make_two_slot_inputs( - rank, - world_size, - device, + requirements = op.prepare_training(lane_count=lane_count, device=device) + forward_staging, backward_staging = _allocate_training_weight_staging(source_weights) + native_forward = op.pack_forward_weights( + source_weights[0], + out=forward_staging, ) - op = _make_operator( - world_size=world_size, - group=group, - max_recv_size_per_rank=max_recv_size_per_rank, - drop_on_overflow=True, - ) - graph = None - try: - resources = op.prepare_training_resources( - _make_training_weights(args0), - slot_count=2, - lane_count=1, - ) - slot0, slot1 = resources.slots - lane = resources.lanes[0] - - resources.refresh_weights() - eager_y0 = resources.forward(slot0, lane, args0[0], args0[3], args0[4]) - eager_y1 = resources.forward(slot1, lane, args1[0], args1[3], args1[4]) - eager_dx0, eager_dp0, eager_operands0 = resources.backward( - slot0, - lane, - grad0, - ) - eager_dx1, eager_dp1, eager_operands1 = resources.backward( - slot1, - lane, - grad1, - ) - eager_operands = (eager_operands0, eager_operands1) - grouped_outputs = tuple( - _allocate_dense_grouped_wgrad_outputs(operands) - for operands in eager_operands - ) - eager_grouped_wgrads = tuple( - _dense_wgrads_from_grouped_kernel( - operands, - wgrad_tensors=outputs, - ) - for operands, outputs in zip(eager_operands, grouped_outputs) - ) - eager_overflow = resources.finalize_overflow((slot0, slot1), lane) - torch.cuda.synchronize(device) - dist.barrier(group=group) - if int(eager_overflow.item()) != 0: - raise AssertionError("single-slot WGrad eager warmup overflowed") - - operand_fields = ( - "expert_offsets", - "valid_route_counts", - "fc1_a", - "fc1_sfa", - "fc1_b", - "fc1_sfb", - "fc2_a", - "fc2_sfa", - "fc2_b", - "fc2_sfb", - ) - eager_common = tuple( - tensor.clone() - for tensor in ( - eager_y0, - eager_y1, - eager_dx0, - eager_dx1, - eager_dp0, - eager_dp1, - ) - ) - eager_operand_snapshot = { - field: getattr(eager_operands[slot_index], field).clone() - for field in operand_fields - } - eager_wgrad_snapshot = tuple( - tensor.clone() for tensor in eager_grouped_wgrads[slot_index] - ) - selected_outputs = grouped_outputs[slot_index] - selected_output_pointers = tuple( - output.data_ptr() for output in selected_outputs - ) - - stream = torch.cuda.Stream(device=device) - stream.wait_stream(torch.cuda.current_stream(device)) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - resources.refresh_weights() - graph_y0 = resources.forward( - slot0, - lane, - args0[0], - args0[3], - args0[4], - ) - graph_y1 = resources.forward( - slot1, - lane, - args1[0], - args1[3], - args1[4], - ) - graph_dx0, graph_dp0, graph_operands0 = resources.backward( - slot0, - lane, - grad0, - ) - graph_dx1, graph_dp1, graph_operands1 = resources.backward( - slot1, - lane, - grad1, - ) - graph_operands = (graph_operands0, graph_operands1)[slot_index] - graph_grouped_wgrads = _dense_wgrads_from_grouped_kernel( - graph_operands, - wgrad_tensors=selected_outputs, - ) - graph_overflow = resources.finalize_overflow((slot0, slot1), lane) - dist.barrier(group=group) - - with torch.cuda.stream(stream): - graph.replay() - stream.synchronize() - dist.barrier(group=group) - if int(graph_overflow.item()) != 0: - raise AssertionError("single-slot WGrad graph replay overflowed") - if ( - tuple(output.data_ptr() for output in graph_grouped_wgrads) - != selected_output_pointers - ): - raise AssertionError("single-slot grouped WGrad output addresses changed") - - graph_common = ( - graph_y0, - graph_y1, - graph_dx0, - graph_dx1, - graph_dp0, - graph_dp1, - ) - for name, actual, expected in zip( - ("y0", "y1", "dx0", "dx1", "dprob0", "dprob1"), - graph_common, - eager_common, - ): - _assert_replay_tensor(name, actual, expected) - for field in operand_fields: - _assert_replay_tensor( - f"slot{slot_index}.{field}", - getattr(graph_operands, field), - eager_operand_snapshot[field], - ) - try: - for prefix, actual, expected in zip( - ("fc1", "fc2"), - graph_grouped_wgrads, - eager_wgrad_snapshot, - ): - _assert_replay_tensor( - f"slot{slot_index}.{prefix}_wgrad", - actual, - expected, - ) - except BaseException: - _report_grouped_wgrad_operand_consistency( - f"slot{slot_index}", - graph_operands, - graph_grouped_wgrads, - ) - raise - - if rank == 0: - print( - f"MOE_EP_EP{world_size}_SINGLE_WGRAD_SLOT_GRAPH_PASS " - f"slot=slot{slot_index}", - flush=True, - ) - finally: - if graph is not None: - del graph - _close_probe_operator(device=device, group=group, op=op) - - -def _run_training_resource_probe( - *, - rank: int, - world_size: int, - device: torch.device, - group, - diagnostic_replays: int, - burst_replays: int, - max_recv_size_per_rank: int, - full_probe: bool, -) -> None: - """Exercise full graph behavior or a minimal teardown/re-init replay.""" - - args0, grad0, args1, grad1, local, remote = _make_two_slot_inputs( - rank, - world_size, - device, + native_backward = op.pack_backward_weights( + source_weights[1], + out=backward_staging, ) - # Keep immutable baseline patterns separate from the graph-bound input - # tensors. Overflow injection mutates the latter in place. - weights = _make_training_weights(args0) - op = _make_operator( - world_size=world_size, - group=group, - max_recv_size_per_rank=max_recv_size_per_rank, - drop_on_overflow=True, + output_pairs = tuple(_allocate_stateless_training_outputs(requirements, device) for _ in range(lane_count)) + return ( + op, + args, + grad_output, + native_forward, + native_backward, + output_pairs, ) - graph = None - try: - resources = op.prepare_training_resources( - weights, - slot_count=2, - lane_count=1, - ) - slot0, slot1 = resources.slots - lane0 = resources.lanes[0] - - # Ordinary execution is the collective warmup for all fused staging, - # MegaMoE, and fixed-capacity WGrad export compile caches. - resources.refresh_weights() - y0 = resources.forward( - slot0, - lane0, - args0[0], - args0[3], - args0[4], - ) - y1 = resources.forward( - slot1, - lane0, - args1[0], - args1[3], - args1[4], - ) - dx0, dp0, operands0 = resources.backward(slot0, lane0, grad0) - dx1, dp1, operands1 = resources.backward(slot1, lane0, grad1) - grouped_outputs0 = _allocate_dense_grouped_wgrad_outputs(operands0) - grouped_outputs1 = _allocate_dense_grouped_wgrad_outputs(operands1) - grouped_output_pointers = tuple( - output.data_ptr() for output in (*grouped_outputs0, *grouped_outputs1) - ) - grouped_wgrads0 = _dense_wgrads_from_grouped_kernel( - operands0, - wgrad_tensors=grouped_outputs0, - ) - grouped_wgrads1 = _dense_wgrads_from_grouped_kernel( - operands1, - wgrad_tensors=grouped_outputs1, - ) - overflow_status = resources.finalize_overflow((slot0, slot1)) - torch.cuda.synchronize(device) - dist.barrier(group=group) - if int(overflow_status.item()) != 0: - raise AssertionError("ordinary fixed-resource warmup overflowed") - - comparison_names = ( - "y0", - "y1", - "dx0", - "dx1", - "dprob0", - "dprob1", - "slot0.expert_offsets", - "slot0.valid_route_counts", - "slot0.fc1_a", - "slot0.fc1_sfa", - "slot0.fc1_b", - "slot0.fc1_sfb", - "slot0.fc2_a", - "slot0.fc2_sfa", - "slot0.fc2_b", - "slot0.fc2_sfb", - "slot1.expert_offsets", - "slot1.valid_route_counts", - "slot1.fc1_a", - "slot1.fc1_sfa", - "slot1.fc1_b", - "slot1.fc1_sfb", - "slot1.fc2_a", - "slot1.fc2_sfa", - "slot1.fc2_b", - "slot1.fc2_sfb", - "slot0.fc1_wgrad", - "slot0.fc2_wgrad", - "slot1.fc1_wgrad", - "slot1.fc2_wgrad", - ) - ordinary = { - name: tensor.clone() - for name, tensor in zip( - comparison_names, - ( - y0, - y1, - dx0, - dx1, - dp0, - dp1, - operands0.expert_offsets, - operands0.valid_route_counts, - operands0.fc1_a, - operands0.fc1_sfa, - operands0.fc1_b, - operands0.fc1_sfb, - operands0.fc2_a, - operands0.fc2_sfa, - operands0.fc2_b, - operands0.fc2_sfb, - operands1.expert_offsets, - operands1.valid_route_counts, - operands1.fc1_a, - operands1.fc1_sfa, - operands1.fc1_b, - operands1.fc1_sfb, - operands1.fc2_a, - operands1.fc2_sfa, - operands1.fc2_b, - operands1.fc2_sfb, - grouped_wgrads0[0], - grouped_wgrads0[1], - grouped_wgrads1[0], - grouped_wgrads1[1], - ), - ) - } - ordinary_offsets = ( - operands0.expert_offsets.clone(), - operands1.expert_offsets.clone(), - ) - ordinary_route_counts = ( - operands0.valid_route_counts.clone(), - operands1.valid_route_counts.clone(), - ) - - stream = torch.cuda.Stream(device=device) - stream.wait_stream(torch.cuda.current_stream(device)) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - resources.refresh_weights() - graph_y0 = resources.forward( - slot0, - lane0, - args0[0], - args0[3], - args0[4], - ) - graph_y1 = resources.forward( - slot1, - lane0, - args1[0], - args1[3], - args1[4], - ) - graph_dx0, graph_dp0, graph_operands0 = resources.backward( - slot0, - lane0, - grad0, - ) - graph_dx1, graph_dp1, graph_operands1 = resources.backward( - slot1, - lane0, - grad1, - ) - graph_grouped_wgrads0 = _dense_wgrads_from_grouped_kernel( - graph_operands0, - wgrad_tensors=grouped_outputs0, - ) - graph_grouped_wgrads1 = _dense_wgrads_from_grouped_kernel( - graph_operands1, - wgrad_tensors=grouped_outputs1, - ) - graph_overflow = resources.finalize_overflow((slot0, slot1)) - dist.barrier(group=group) - - with torch.cuda.stream(stream): - graph.replay() - stream.synchronize() - dist.barrier(group=group) - if int(graph_overflow.item()) != 0: - raise AssertionError("captured fixed-resource graph overflowed") - - captured = { - name: tensor - for name, tensor in zip( - comparison_names, - ( - graph_y0, - graph_y1, - graph_dx0, - graph_dx1, - graph_dp0, - graph_dp1, - graph_operands0.expert_offsets, - graph_operands0.valid_route_counts, - graph_operands0.fc1_a, - graph_operands0.fc1_sfa, - graph_operands0.fc1_b, - graph_operands0.fc1_sfb, - graph_operands0.fc2_a, - graph_operands0.fc2_sfa, - graph_operands0.fc2_b, - graph_operands0.fc2_sfb, - graph_operands1.expert_offsets, - graph_operands1.valid_route_counts, - graph_operands1.fc1_a, - graph_operands1.fc1_sfa, - graph_operands1.fc1_b, - graph_operands1.fc1_sfb, - graph_operands1.fc2_a, - graph_operands1.fc2_sfa, - graph_operands1.fc2_b, - graph_operands1.fc2_sfb, - graph_grouped_wgrads0[0], - graph_grouped_wgrads0[1], - graph_grouped_wgrads1[0], - graph_grouped_wgrads1[1], - ), - ) - } - if ( - tuple( - output.data_ptr() - for output in ( - *graph_grouped_wgrads0, - *graph_grouped_wgrads1, - ) - ) - != grouped_output_pointers - ): - raise AssertionError("captured grouped WGrad output addresses changed") - try: - for name in comparison_names: - _assert_replay_tensor(name, captured[name], ordinary[name]) - except BaseException: - _report_grouped_wgrad_operand_consistency( - "slot0", - graph_operands0, - graph_grouped_wgrads0, - ) - _report_grouped_wgrad_operand_consistency( - "slot1", - graph_operands1, - graph_grouped_wgrads1, - ) - raise - torch.testing.assert_close( - graph_operands0.expert_offsets, - ordinary_offsets[0], - rtol=0, - atol=0, - ) - torch.testing.assert_close( - graph_operands1.expert_offsets, - ordinary_offsets[1], - rtol=0, - atol=0, - ) - torch.testing.assert_close( - graph_operands0.valid_route_counts, - ordinary_route_counts[0], - rtol=0, - atol=0, - ) - torch.testing.assert_close( - graph_operands1.valid_route_counts, - ordinary_route_counts[1], - rtol=0, - atol=0, - ) - - if full_probe: - # Diagnostic mode aligns ranks after every replay and verifies that - # fixed-slot dprob reset prevents history accumulation. - dprob_reference = graph_dp0.clone() - for _ in range(diagnostic_replays): - with torch.cuda.stream(stream): - graph.replay() - stream.synchronize() - dist.barrier(group=group) - if int(graph_overflow.item()) != 0: - raise AssertionError("fixed-resource diagnostic replay overflowed") - torch.testing.assert_close( - graph_dp0, - dprob_reference, - rtol=1e-5, - atol=1e-6, - ) - for name in ( - "slot0.fc1_wgrad", - "slot0.fc2_wgrad", - "slot1.fc1_wgrad", - "slot1.fc2_wgrad", - ): - _assert_replay_tensor(name, captured[name], ordinary[name]) - for index, graph_operands in enumerate( - (graph_operands0, graph_operands1) - ): - torch.testing.assert_close( - graph_operands.expert_offsets, - ordinary_offsets[index], - rtol=0, - atol=0, - ) - torch.testing.assert_close( - graph_operands.valid_route_counts, - ordinary_route_counts[index], - rtol=0, - atol=0, - ) - - # Production-like burst: no synchronization or host collective in - # the loop. The graph contains the captured scalar overflow OR. - with torch.cuda.stream(stream): - for _ in range(burst_replays): - graph.replay() - stream.synchronize() - dist.barrier(group=group) - if int(graph_overflow.item()) != 0: - raise AssertionError("fixed-resource replay burst overflowed") - torch.testing.assert_close( - graph_dp0, - ordinary["dprob0"], - rtol=1e-5, - atol=1e-6, - ) - for name in ( - "slot0.fc1_wgrad", - "slot0.fc2_wgrad", - "slot1.fc1_wgrad", - "slot1.fc2_wgrad", - ): - _assert_replay_tensor(name, captured[name], ordinary[name]) - - # Overflow both slots, then restore their distinct valid patterns. - overflow = _route_pattern("overflow", rank, world_size, device) - with torch.cuda.stream(stream): - args0[3].copy_(overflow[0]) - args0[4].copy_(overflow[1]) - args1[3].copy_(overflow[0]) - args1[4].copy_(overflow[1]) - graph.replay() - stream.synchronize() - dist.barrier(group=group) - if int(graph_overflow.item()) != 1: - raise AssertionError("fixed-resource overflow was not global") - - with torch.cuda.stream(stream): - args0[3].copy_(local[0]) - args0[4].copy_(local[1]) - args1[3].copy_(remote[0]) - args1[4].copy_(remote[1]) - graph.replay() - stream.synchronize() - dist.barrier(group=group) - recovered_overflow = int(graph_overflow.item()) - if recovered_overflow != 0: - raise AssertionError( - "fixed-resource graph did not recover: " - f"rank={rank}, global_overflow={recovered_overflow}, " - f"slot0_routing_restored=" - f"{torch.equal(args0[3], local[0])}, " - f"slot1_routing_restored=" - f"{torch.equal(args1[3], remote[0])}" - ) - for name in ( - "slot0.fc1_wgrad", - "slot0.fc2_wgrad", - "slot1.fc1_wgrad", - "slot1.fc2_wgrad", - ): - _assert_replay_tensor(name, captured[name], ordinary[name]) - - if rank == 0: - mode = "full" if full_probe else "reinit" - effective_burst = burst_replays if full_probe else 0 - print( - f"MOE_EP_EP{world_size}_TRAINING_RESOURCES_GRAPH_PASS " f"mode={mode} burst={effective_burst}", - flush=True, - ) - print( - f"MOE_EP_EP{world_size}_GROUPED_WGRAD_GRAPH_PASS " - f"mode={mode} burst={effective_burst}", - flush=True, - ) - finally: - if graph is not None: - del graph - _close_probe_operator(device=device, group=group, op=op) -def _run_multistream_resource_probe( - *, - rank: int, - world_size: int, - device: torch.device, - group, - replays: int, - max_recv_size_per_rank: int, -) -> None: - """Capture two independent lanes with deterministic cross-rank ordering.""" - - args0, grad0, args1, grad1, _, _ = _make_two_slot_inputs( - rank, - world_size, - device, - ) - op = _make_operator( +def _run_cycle(args: argparse.Namespace, *, device: torch.device, rank: int, world_size: int) -> None: + lane_count = 1 if args.skip_multistream else 2 + case = _prepare_case( + device=device, + rank=rank, world_size=world_size, - group=group, - max_recv_size_per_rank=max_recv_size_per_rank, - drop_on_overflow=True, + lane_count=lane_count, + max_recv_size_per_rank=args.max_recv_size_per_rank, + drop_on_overflow=not args.expect_overflow_assert, ) - graph = None + op, inputs, grad_output, native_forward, native_backward, output_pairs = case try: - with _debug_phase_scope(rank, "multistream.prepare"): - resources = op.prepare_training_resources( - _make_training_weights(args0), - slot_count=2, - lane_count=2, - ) - slot0, slot1 = resources.slots - lane0, lane1 = resources.lanes - with _debug_phase_scope(rank, "multistream.refresh-weights"): - resources.refresh_weights() - - with _debug_phase_scope(rank, "multistream.lane0-forward"): - eager_y0 = resources.forward(slot0, lane0, args0[0], args0[3], args0[4]) - with _debug_phase_scope(rank, "multistream.lane0-backward"): - eager_dx0, eager_dp0, eager_operands0 = resources.backward( - slot0, - lane0, - grad0, + # Warm each lane and every kernel specialization before capture. + for lane, (forward_out, backward_out) in zip( + op.training_lanes, + output_pairs, + ): + op.training_forward( + lane, + inputs[0], + inputs[3], + inputs[4], + weights=native_forward, + out=forward_out, ) - grouped_outputs0 = _allocate_dense_grouped_wgrad_outputs(eager_operands0) - eager_grouped_wgrads0 = _dense_wgrads_from_grouped_kernel( - eager_operands0, - wgrad_tensors=grouped_outputs0, + op.training_backward( + lane, + grad_output, + inputs[3], + inputs[4], + weights=native_backward, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, ) - with _debug_phase_scope(rank, "multistream.lane0-finalize"): - resources.finalize_overflow((slot0,), lane0) - _synchronize_with_watchdog( - rank, - device, - "multistream.lane0-synchronize", - ) - with _debug_phase_scope(rank, "multistream.lane0-barrier"): - dist.barrier(group=group) + torch.cuda.synchronize(device) + dist.barrier(group=dist.group.WORLD, device_ids=[device.index]) - with _debug_phase_scope(rank, "multistream.lane1-forward"): - eager_y1 = resources.forward(slot1, lane1, args1[0], args1[3], args1[4]) - with _debug_phase_scope(rank, "multistream.lane1-backward"): - eager_dx1, eager_dp1, eager_operands1 = resources.backward( - slot1, - lane1, - grad1, - ) - grouped_outputs1 = _allocate_dense_grouped_wgrad_outputs(eager_operands1) - eager_grouped_wgrads1 = _dense_wgrads_from_grouped_kernel( - eager_operands1, - wgrad_tensors=grouped_outputs1, + graphs = tuple( + _capture_training_graph( + op, + lane, + inputs, + grad_output, + native_forward, + native_backward, + forward_out, + backward_out, ) - grouped_output_pointers = tuple( - output.data_ptr() for output in (*grouped_outputs0, *grouped_outputs1) + for lane, (forward_out, backward_out) in zip( + op.training_lanes, + output_pairs, ) - with _debug_phase_scope(rank, "multistream.lane1-finalize"): - resources.finalize_overflow((slot1,), lane1) - _synchronize_with_watchdog( - rank, - device, - "multistream.lane1-synchronize", ) - with _debug_phase_scope(rank, "multistream.lane1-barrier"): - dist.barrier(group=group) - expected = tuple( - tensor.clone() - for tensor in ( - eager_y0, - eager_dx0, - eager_dp0, - eager_y1, - eager_dx1, - eager_dp1, - eager_grouped_wgrads0[0], - eager_grouped_wgrads0[1], - eager_grouped_wgrads1[0], - eager_grouped_wgrads1[1], - ) - ) - - capture_stream = torch.cuda.Stream(device=device) - lane_stream0 = torch.cuda.Stream(device=device) - lane_stream1 = torch.cuda.Stream(device=device) - fork_event = torch.cuda.Event() - done_event0 = torch.cuda.Event() - done_event1 = torch.cuda.Event() - capture_stream.wait_stream(torch.cuda.current_stream(device)) + dist.barrier(group=dist.group.WORLD, device_ids=[device.index]) - # One outer graph visits two lane-bound streams, rejoins them, then - # emits exactly one NCCL overflow finalizer. The MegaMoE kernels use - # device-side cross-rank software synchronization and consume one CTA - # slot per SM. Launching both lanes concurrently can let different - # ranks schedule different lanes first, leaving each lane waiting for - # peers whose matching kernel cannot be scheduled. Chain lane 1 after - # lane 0 so every rank observes the same collective order while still - # validating independent per-stream lane storage and graph edges. - graph = torch.cuda.CUDAGraph() - capture_watchdog = _RuntimeWatchdog("multistream.capture") - capture_watchdog.start() - with _debug_phase_scope(rank, "multistream.capture"): - try: - with torch.cuda.graph(graph, stream=capture_stream): - fork_event.record(capture_stream) - lane_stream0.wait_event(fork_event) - with torch.cuda.stream(lane_stream0): - graph_y0 = resources.forward( - slot0, - lane0, - args0[0], - args0[3], - args0[4], - ) - graph_dx0, graph_dp0, graph_operands0 = resources.backward( - slot0, - lane0, - grad0, - ) - graph_grouped_wgrads0 = _dense_wgrads_from_grouped_kernel( - graph_operands0, - wgrad_tensors=grouped_outputs0, - ) - done_event0.record(lane_stream0) - lane_stream1.wait_event(done_event0) - with torch.cuda.stream(lane_stream1): - graph_y1 = resources.forward( - slot1, - lane1, - args1[0], - args1[3], - args1[4], - ) - graph_dx1, graph_dp1, graph_operands1 = resources.backward( - slot1, - lane1, - grad1, - ) - graph_grouped_wgrads1 = _dense_wgrads_from_grouped_kernel( - graph_operands1, - wgrad_tensors=grouped_outputs1, - ) - done_event1.record(lane_stream1) - capture_stream.wait_event(done_event1) - graph_overflow = resources.finalize_overflow( - (slot0, slot1), - lane0, - ) - finally: - capture_watchdog.close() - with _debug_phase_scope(rank, "multistream.capture-barrier"): - dist.barrier(group=group) - - with _debug_phase_scope(rank, "multistream.replay"): - with torch.cuda.stream(capture_stream): - for _ in range(replays): - graph.replay() - replay_watchdog = _RuntimeWatchdog("multistream.replay-synchronize") - replay_watchdog.start() - with _debug_phase_scope( - rank, - "multistream.replay-synchronize", - ): - try: - capture_stream.synchronize() - finally: - replay_watchdog.close() - with _debug_phase_scope(rank, "multistream.replay-barrier"): - dist.barrier(group=group) - if int(graph_overflow.item()) != 0: - raise AssertionError("multi-stream fixed-resource graph overflowed") + replay_count = args.diagnostic_replays + args.burst_replays + if lane_count > 1: + replay_count += args.multistream_replays + caught = None + try: + for replay in range(replay_count): + graphs[replay % len(graphs)].replay() + torch.cuda.synchronize(device) + except Exception as error: + caught = error - actual = ( - graph_y0, - graph_dx0, - graph_dp0, - graph_y1, - graph_dx1, - graph_dp1, - graph_grouped_wgrads0[0], - graph_grouped_wgrads0[1], - graph_grouped_wgrads1[0], - graph_grouped_wgrads1[1], - ) - if ( - tuple( - output.data_ptr() - for output in ( - *graph_grouped_wgrads0, - *graph_grouped_wgrads1, - ) - ) - != grouped_output_pointers - ): - raise AssertionError("multistream grouped WGrad output addresses changed") - for index, (value, reference) in enumerate(zip(actual, expected)): - _assert_replay_tensor( - f"multistream[{index}]", - value, - reference, - ) - if rank == 0: - print( - f"MOE_EP_EP{world_size}_MULTISTREAM_GRAPH_PASS " f"replays={replays}", - flush=True, - ) + if args.expect_overflow_assert: + if caught is None: + raise AssertionError("expected the captured overflow assertion") + elif caught is not None: + raise caught + dist.barrier(group=dist.group.WORLD, device_ids=[device.index]) finally: - if graph is not None: - del graph - _close_probe_operator(device=device, group=group, op=op) - - -def _run_error_mode_assert_probe( - *, - rank: int, - world_size: int, - device: torch.device, - group, - max_recv_size_per_rank: int, -) -> None: - """Require a captured global overflow to assert on every rank.""" - - args, grad_output = _make_inputs(rank, device) - local = _route_pattern("local", rank, world_size, device) - overflow = _route_pattern("overflow", rank, world_size, device) - route_indices = local[0].clone() - route_weights = local[1].clone() - op = _make_operator( - world_size=world_size, - group=group, - max_recv_size_per_rank=max_recv_size_per_rank, - drop_on_overflow=False, - ) - resources = op.prepare_training_resources( - _make_training_weights(args), - slot_count=1, - lane_count=1, - ) - slot = resources.slots[0] - lane = resources.lanes[0] - - # Warm every kernel and prove the assertion accepts a valid execution. - resources.refresh_weights() - resources.forward( - slot, - lane, - args[0], - route_indices, - route_weights, - ) - resources.backward(slot, lane, grad_output) - resources.finalize_overflow((slot,), lane) - torch.cuda.synchronize(device) - dist.barrier(group=group) - - stream = torch.cuda.Stream(device=device) - stream.wait_stream(torch.cuda.current_stream(device)) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - resources.refresh_weights() - resources.forward( - slot, - lane, - args[0], - route_indices, - route_weights, - ) - resources.backward(slot, lane, grad_output) - resources.finalize_overflow((slot,), lane) - dist.barrier(group=group) - - # A valid replay confirms capture before intentionally poisoning the - # context with the fatal error-mode assertion. - with torch.cuda.stream(stream): - graph.replay() - stream.synchronize() - dist.barrier(group=group) - - try: - with torch.cuda.stream(stream): - route_indices.copy_(overflow[0]) - route_weights.copy_(overflow[1]) - graph.replay() - stream.synchronize() - except BaseException as exc: - print( - f"MOE_EP_EP{world_size}_ERROR_MODE_ASSERT_PASS " f"rank={rank} error={type(exc).__name__}", - flush=True, - ) - # CUDA device assertions poison the process context. Do not run Python - # destructors, NCCL collectives, or NVSHMEM finalization afterward. - os._exit(0) - - print( - f"MOE_EP_EP{world_size}_ERROR_MODE_ASSERT_MISSING rank={rank}", - flush=True, - ) - os._exit(1) + op.close() def main() -> None: args = _parse_args() - _require_positive("diagnostic_replays", args.diagnostic_replays) - _require_positive("burst_replays", args.burst_replays) - _require_positive("multistream_replays", args.multistream_replays) - _require_positive("cycles", args.cycles) - _require_positive( + for name in ( + "diagnostic_replays", + "burst_replays", + "multistream_replays", "max_recv_size_per_rank", - args.max_recv_size_per_rank, - ) - - world_size = int(os.environ.get("WORLD_SIZE", "1")) - rank = int(os.environ.get("RANK", "0")) - local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) - if world_size < 2: - raise RuntimeError(f"this probe requires WORLD_SIZE >= 2, got {world_size}") - forced_overflow_routes = world_size * 8 * 2 - if args.max_recv_size_per_rank >= forced_overflow_routes: - raise ValueError( - "max_recv_size_per_rank must remain below the probe's forced " - f"overflow route count {forced_overflow_routes}, got " - f"{args.max_recv_size_per_rank}" - ) + "cycles", + "timeout_seconds", + ): + _positive(name, getattr(args, name)) + local_rank = int(os.environ["LOCAL_RANK"]) + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) device = torch.device("cuda", local_rank) torch.cuda.set_device(device) - capability = torch.cuda.get_device_capability(device) - if capability != (10, 7): - raise RuntimeError("this probe requires Rubin SM107; " f"rank {rank} found compute capability {capability}") - os.environ.setdefault("CUTE_DSL_ARCH", "sm_107a") - dist.init_process_group( backend="nccl", - init_method="env://", - device_id=device, timeout=timedelta(seconds=args.timeout_seconds), + device_id=device, ) try: - if args.expect_overflow_assert: - _run_error_mode_assert_probe( - rank=rank, - world_size=world_size, - device=device, - group=dist.group.WORLD, - max_recv_size_per_rank=args.max_recv_size_per_rank, - ) - raise AssertionError("fatal overflow assertion probe returned") - if args.wgrad_capture_mode == "both": - for cycle in range(args.cycles): - with _debug_phase_scope( - rank, - f"training-resources-cycle-{cycle}", - ): - _run_training_resource_probe( - rank=rank, - world_size=world_size, - device=device, - group=dist.group.WORLD, - diagnostic_replays=args.diagnostic_replays, - burst_replays=args.burst_replays, - max_recv_size_per_rank=args.max_recv_size_per_rank, - full_probe=cycle == 0, - ) - if not args.skip_multistream: - with _debug_phase_scope(rank, "multistream"): - _run_multistream_resource_probe( - rank=rank, - world_size=world_size, - device=device, - group=dist.group.WORLD, - replays=args.multistream_replays, - max_recv_size_per_rank=args.max_recv_size_per_rank, - ) - else: - slot_index = int(args.wgrad_capture_mode[-1]) - with _debug_phase_scope( - rank, - f"single-wgrad-slot{slot_index}", - ): - _run_single_wgrad_slot_capture_probe( - rank=rank, - world_size=world_size, - device=device, - group=dist.group.WORLD, - max_recv_size_per_rank=args.max_recv_size_per_rank, - slot_index=slot_index, - ) + if torch.cuda.get_device_capability(device) != (10, 7): + raise RuntimeError("stateless training graph probe requires SM107") + for _ in range(args.cycles): + _run_cycle(args, device=device, rank=rank, world_size=world_size) if rank == 0: print( - f"MOE_EP_EP{world_size}_CUDA_GRAPH_PROBE_PASS " - f"wgrad_capture_mode={args.wgrad_capture_mode}", + "stateless MoeEP training graph probe passed: " f"world_size={world_size}, cycles={args.cycles}", flush=True, ) finally: if dist.is_initialized(): - try: - with _debug_phase_scope(rank, "runtime-shutdown"): - get_runtime_manager().shutdown() - finally: - dist.destroy_process_group() + dist.destroy_process_group() if __name__ == "__main__": diff --git a/test/python/moe_ep/test_moe_ep_backward.py b/test/python/moe_ep/test_moe_ep_backward.py index eaf30b891..9b98fdd7c 100644 --- a/test/python/moe_ep/test_moe_ep_backward.py +++ b/test/python/moe_ep/test_moe_ep_backward.py @@ -1,1208 +1,786 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Fixed-resource MoE EP backward and training-graph contracts.""" +"""Stateless MoE EP training contracts.""" from __future__ import annotations -from contextlib import nullcontext -import inspect -import os -import threading from dataclasses import fields -from pathlib import Path -from types import SimpleNamespace -from unittest.mock import Mock import cudnn import pytest import torch -import torch.multiprocessing as mp from cudnn.moe_ep import ( + BlockScaledTensor, MoeEp, + MoeEpBackwardWeightStaging, + MoeEpBackwardWeights, MoeEpExecutionLane, - MoeEpTrainingResources, - MoeEpTrainingSlot, + MoeEpForwardWeights, + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, + MoeEpTrainingBackwardOutputs, + MoeEpTrainingForwardOutputs, MoeEpTrainingWgradOperands, -) -from cudnn.moe_ep._contracts import Fc1WeightLayout -from cudnn.moe_ep._validation import validate_training_weights -from cudnn.moe_ep._megamoe_backend.mxfp8._adapter import ( - _typed_k_major_view, -) -from cudnn.moe_ep._megamoe_backend._workspace import ( - BufferRegion, - WorkspaceRequirements, + pack_backward_weights, + pack_forward_weights, ) from cudnn.moe_ep._megamoe_backend.mxfp8._training_resources import ( - Mxfp8TrainingResourceOwner, + Mxfp8TrainingState, _build_training_abi_facts, _harmonize_symmetric_regions, - _verify_training_abi_across_ranks, - build_training_workspace_requirements, ) -from cudnn.moe_ep._megamoe_backend.mxfp8._training_stage import ( - Mxfp8TrainingStager, +from cudnn.moe_ep._megamoe_backend.mxfp8._training_execute import _stage_input +from cudnn.moe_ep._megamoe_backend.mxfp8._training_weights import ( + backward_native_to_kernel, + forward_native_to_kernel, ) from cudnn.moe_ep._megamoe_backend.mxfp8._training_wgrad import ( - Mxfp8TrainingWgradExporter, -) -from cudnn.moe_ep._megamoe_backend.mxfp8._fingerprint import ( - canonical_json_sha256, -) -from cudnn.moe_ep._megamoe_backend.mxfp8._training_weights import ( - Mxfp8TrainingWeightBindings, + assemble_training_wgrad_operands, ) -from moe_ep.moe_ep_reference import ( - MoeEpReference, +from cudnn.moe_ep._megamoe_backend._workspace import ( + BufferRegion, + WorkspaceRequirements, + WorkspaceViews, ) -from moe_ep.moe_ep_distributed_workers import ( - _distributed_backward_reference_worker, - _distributed_subgroup_backward_reference_worker, +from cudnn.moe_ep._megamoe_backend.mxfp8._fingerprint import canonical_json_sha256 +from cudnn.moe_ep._validation import ( + validate_native_backward_weights, + validate_native_forward_weights, + validate_training_backward_outputs, + validate_training_forward_outputs, + validate_training_forward_state, + validate_training_input, + validate_training_non_aliasing, ) +from cudnn.moe_ep.api import _resolve_training_device from moe_ep.moe_ep_test_support import ( - _allocate_dense_grouped_wgrad_outputs, - _assert_fixed_training_drop_overflow_result, - _assert_fixed_training_matches_reference, + _allocate_stateless_training_outputs, + _allocate_training_weight_staging, + _assert_backward_matches, _assert_grouped_wgrads_match_reference, - _assert_training_graph_tails_are_reset, - _assert_training_weight_sources_changed, - _capture_fixed_training_batch, - _copy_training_weight_sources_, + _assert_matches_reference, + _assert_wgrads_match_reference, _dense_wgrads_from_grouped_kernel, - _dense_wgrads_from_operands, - _fixed_training_case, - _fixed_training_drop_overflow_case, - _fixed_training_drop_overflow_reference, _fixed_training_reference, _fixed_training_weights, _grad_output, - _prefill_training_graph_sentinels, - _require_distributed_sm107, - _run_fixed_training_batch, + _interleave_fc1_wgrad, _sm107_device, - _training_public_pointers, - _training_source_pointers, - _training_weight_source_pointers, - _training_weight_source_values, - _TrainingResourceContractOwner, _training_abi_prepared, _training_config, - _training_contract_resources, - _training_inputs, _training_prepared_pair, - _training_staging_tensors, - _training_weight_defect, - _training_weights, make_forward_inputs, + quantize_mxfp8, ) -# L0 contracts +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple -@pytest.mark.L0 -@pytest.mark.parametrize( - "case", - [ - pytest.param("backward-regions", id="backward-regions"), - pytest.param("slot-lane-layout", id="slot-lane-layout"), - ], -) -def test_training_workspace_layout_contract(case): - if case == "backward-regions": - requirements = WorkspaceRequirements.for_mxfp8( - _training_config(), - kernel_local_workspace_bytes=64, - kernel_shared_workspace_bytes=128, - backward_fc1_preact_bytes=1024, - backward_dprob_bytes=32, - backward_aux_data_bytes=512, - backward_aux_scale_bytes=256, - ) - expected = ( - ("symmetric", "backward_dprob", 32, None), - ("local", "backward_fc1_preact", 1024, 128), - ("local", "backward_aux_data", 512, None), - ("local", "backward_aux_scale", 256, None), - ) - else: - config = _training_config() - forward, backward = _training_prepared_pair(config) - requirements = build_training_workspace_requirements( - config, - forward, - backward, - slot_count=2, - lane_count=1, - ) - expected = ( - ("symmetric", "lane.0.forward.symmetric.kernel_shared_workspace", None, None), - ("symmetric", "lane.0.backward.symmetric.kernel_shared_workspace", None, None), - ("symmetric", "slot.0.backward.symmetric.backward_dprob", None, None), - ("symmetric", "slot.1.backward.symmetric.backward_dprob", None, None), - ("local", "slot.0.persistent.local.fc1_preact", None, None), - ("local", "slot.1.persistent.local.fc1_preact", None, None), - ) - regions = { - "symmetric": {region.name: region for region in requirements.symmetric_regions}, - "local": {region.name: region for region in requirements.local_regions}, - } - for storage, name, nbytes, alignment in expected: - region = regions[storage][name] - if nbytes is not None: - assert region.nbytes == nbytes - if alignment is not None: - assert region.alignment == alignment - - if case == "slot-lane-layout": - assert tuple(region.name for region in requirements.symmetric_regions if region.name.startswith("slot.0.")) == ( - "slot.0.forward.symmetric.output_data", - "slot.0.backward.symmetric.backward_dprob", - "slot.0.backward.symmetric.output_data", - "slot.0.persistent.symmetric.routing_topk_weights", - ) +def _blocked_scale_elements(rows: int, columns: int) -> int: + return _round_up(rows, 128) * _round_up(columns, 4) -@pytest.mark.L0 -def test_training_workspace_harmonizes_each_symmetric_region(monkeypatch): - requirements = WorkspaceRequirements( - max_tokens_per_rank=1, - symmetric_regions=( - BufferRegion("first", 1), - BufferRegion("second", 257), +def _native_forward( + config, + *, + device: torch.device = torch.device("cpu"), +) -> MoeEpNativeForwardWeights: + e = config.experts_per_rank + h = config.hidden_size + i = config.intermediate_size + fc1 = torch.empty_strided( + (e, h, 2 * i), + (h * 2 * i, 1, h), + dtype=torch.float8_e4m3fn, + device=device, + ) + fc2 = torch.empty_strided( + (e, i, h), + (i * h, 1, i), + dtype=torch.float8_e4m3fn, + device=device, + ) + return MoeEpNativeForwardWeights( + fc1=MoeEpNativeWeight( + fc1, + torch.empty( + (e, _blocked_scale_elements(2 * i, h // 32)), + dtype=torch.float8_e8m0fnu, + device=device, + ), + MoeEpNativeWeightLayout.FORWARD_FC1_GATE_UP_INTERLEAVED_32_V1, + ), + fc2=MoeEpNativeWeight( + fc2, + torch.empty( + (e, _blocked_scale_elements(h, i // 32)), + dtype=torch.float8_e8m0fnu, + device=device, + ), + MoeEpNativeWeightLayout.FORWARD_FC2_K_MAJOR_V1, ), - local_regions=(BufferRegion("local", 1),), ) - runtime = SimpleNamespace(world_size=2, group=object()) - def all_reduce(tensor, *, op, group): - assert group is runtime.group - if tensor.numel() == 2 and op == torch.distributed.ReduceOp.MAX: - tensor.copy_(torch.tensor([257, 257], dtype=torch.int64)) - monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) - harmonized = _harmonize_symmetric_regions( - requirements, - runtime, - torch.device("cpu"), +def _native_backward( + config, + *, + device: torch.device = torch.device("cpu"), +) -> MoeEpNativeBackwardWeights: + e = config.experts_per_rank + h = config.hidden_size + i = config.intermediate_size + return MoeEpNativeBackwardWeights( + w2_transpose=MoeEpNativeWeight( + torch.empty( + (e, h, i), + dtype=torch.float8_e4m3fn, + device=device, + ), + torch.empty( + (e, _blocked_scale_elements(i, h // 32)), + dtype=torch.float8_e8m0fnu, + device=device, + ), + MoeEpNativeWeightLayout.BACKWARD_W2_TRANSPOSE_V1, + ), + w1_transpose=MoeEpNativeWeight( + torch.empty( + (e, 2 * i, h), + dtype=torch.float8_e4m3fn, + device=device, + ), + torch.empty( + (e, _blocked_scale_elements(h, 2 * i // 32)), + dtype=torch.float8_e8m0fnu, + device=device, + ), + MoeEpNativeWeightLayout.BACKWARD_W1_TRANSPOSE_GATE_UP_INTERLEAVED_32_V1, + ), ) - assert tuple(region.nbytes for region in harmonized.symmetric_regions) == (257, 257) - assert harmonized.local_regions == requirements.local_regions +def _source_weights(config): + e = config.experts_per_rank + h = config.hidden_size + i = config.intermediate_size -@pytest.mark.L0 -def test_training_abi_fingerprint_is_stable_and_structural(): - config = _training_config(ep_size=2, ep_global_ranks=(0, 1)) - forward = _training_abi_prepared("forward") - backward = _training_abi_prepared("backward") - weights = _training_weights() - requirements = WorkspaceRequirements( - max_tokens_per_rank=4, - symmetric_regions=(BufferRegion("symmetric", 256),), - local_regions=(BufferRegion("local", 128),), - ) - first = _build_training_abi_facts( - config, - forward, - backward, - weights, - requirements, - slot_count=2, - lane_count=1, - source_tree_digest="source", - ) - second = _build_training_abi_facts( - config, - forward, - backward, - weights, - requirements, - slot_count=2, - lane_count=1, - source_tree_digest="source", - ) - changed = _build_training_abi_facts( - config, - forward, - backward, - weights, - requirements, - slot_count=2, - lane_count=2, - source_tree_digest="source", - ) - changed_layout = _build_training_abi_facts( - _training_config( - ep_size=2, - ep_global_ranks=(0, 1), - weight_interleave_size=32, + def block_scaled(shape): + scale_shape = list(shape) + scale_shape[1] //= 32 + return BlockScaledTensor( + data=torch.empty(shape, dtype=torch.float8_e4m3fn), + scale=torch.empty(scale_shape, dtype=torch.float8_e8m0fnu), + format="mxfp8", + logical_shape=shape, + axis=1, + ) + + return ( + MoeEpForwardWeights( + block_scaled((e, h, 2 * i)), + block_scaled((e, i, h)), + ), + MoeEpBackwardWeights( + block_scaled((e, h, i)), + block_scaled((e, 2 * i, h)), ), - forward, - backward, - weights, - requirements, - slot_count=2, - lane_count=1, - source_tree_digest="source", ) - assert first["policy"]["fc1_weight_layout"] == "gate_then_up" - assert changed_layout["policy"]["fc1_weight_layout"] == "gate_up_interleaved_32" - assert canonical_json_sha256(first) == canonical_json_sha256(second) - assert canonical_json_sha256(first) != canonical_json_sha256(changed) - assert canonical_json_sha256(first) != canonical_json_sha256(changed_layout) - @pytest.mark.L0 -def test_training_abi_handshake_rejects_rank_mismatch(monkeypatch): - runtime = SimpleNamespace(world_size=2, group=object()) - - def all_reduce(tensor, *, op, group): - assert group is runtime.group - if op == torch.distributed.ReduceOp.MAX: - tensor.add_(1) - - def all_gather_object(output, value, *, group): - assert group is runtime.group - output[:] = [value, "different"] - - monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) - monkeypatch.setattr( - torch.distributed, - "all_gather_object", - all_gather_object, +def test_only_stateless_training_types_are_public(): + removed = ( + "MoeEpTrainingResources", + "MoeEpTrainingSlot", + "MoeEpTrainingWeights", ) - with pytest.raises(RuntimeError, match="ABI differs"): - _verify_training_abi_across_ranks( - {"schema_version": 1}, - runtime, - torch.device("cpu"), - ) + for name in removed: + assert not hasattr(cudnn, name) + for name in ( + "MoeEpForwardWeights", + "MoeEpBackwardWeights", + "MoeEpNativeWeight", + "MoeEpTrainingForwardOutputs", + "MoeEpTrainingBackwardOutputs", + "pack_forward_weights", + "pack_backward_weights", + ): + assert hasattr(cudnn, name) @pytest.mark.L0 -def test_training_resource_views_share_lane_scratch_but_not_slot_state(): - config = _training_config() - forward, backward = _training_prepared_pair(config) +def test_training_device_prefers_explicit_then_current(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 2) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 4) - class Runtime: - device = torch.device("cpu") - rank = 0 - world_size = 1 - nvshmem_enabled = False - closed = False + assert _resolve_training_device(None) == torch.device("cuda:2") + assert _resolve_training_device("cuda") == torch.device("cuda:2") + assert _resolve_training_device(1) == torch.device("cuda:1") + with pytest.raises(ValueError, match="must be CUDA"): + _resolve_training_device("cpu") - def ensure_open(self): - assert not self.closed - def close(self): - self.closed = True +@pytest.mark.L0 +def test_training_input_rejects_noncontiguous_plain_tensor(): + config = _training_config(weight_interleave_size=32) + activation = torch.empty((config.hidden_size, 2), dtype=torch.bfloat16).t() + topk_idx = torch.tensor([[0, 1], [1, 0]], dtype=torch.int32) + topk_weights = torch.ones((2, 2), dtype=torch.float32) + + with pytest.raises(ValueError, match="activation must be contiguous"): + validate_training_input( + config, + "activation", + activation, + topk_idx, + topk_weights, + device=torch.device("cpu"), + ) - runtime = Runtime() - runtime_manager = SimpleNamespace(acquire=lambda actual_config, actual_device: runtime) - weights = _training_weights() - owner = Mxfp8TrainingResourceOwner( - config, - torch.device("cpu"), - forward, - backward, - weights, - slot_count=2, - lane_count=1, - runtime_manager=runtime_manager, - ) - try: - first = owner.views(slot=0, lane=0, token_count=4) - second = owner.views(slot=1, lane=0, token_count=4) - assert first.forward.workspace.local["kernel_local_workspace"].data_ptr() == second.forward.workspace.local["kernel_local_workspace"].data_ptr() - assert first.slot.fc1_preact.data_ptr() != second.slot.fc1_preact.data_ptr() - assert first.slot.dprob.data_ptr() != second.slot.dprob.data_ptr() - assert first.forward_expert_size_snapshot is not None - assert first.forward_expert_size_snapshot.data_ptr() == (second.forward_expert_size_snapshot.data_ptr()) - finally: - owner.close() - assert runtime.closed +@pytest.mark.L0 +def test_training_bundle_fields_match_public_contracts(): + dummy = object() + assert [field.name for field in fields(MoeEpForwardWeights)] == ["fc1", "fc2"] + assert [field.name for field in fields(MoeEpBackwardWeights)] == [ + "w2_transpose", + "w1_transpose", + ] + assert MoeEpForwardWeights(dummy, dummy).fc1 is dummy + assert MoeEpBackwardWeights(dummy, dummy).w2_transpose is dummy + assert [field.name for field in fields(MoeEpTrainingForwardOutputs)] == [ + "fc1_preact", + "output", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + ] + assert [field.name for field in fields(MoeEpTrainingBackwardOutputs)] == [ + "grad_activation", + "dprob", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + ] @pytest.mark.L0 -def test_training_sources_track_adapter_grad_y2_and_dfc2_contracts(): - from cudnn.moe_ep._megamoe_backend.mxfp8 import ( - _backward_compile, - _compile, - ) +def test_native_weight_validation_and_kernel_views_are_zero_copy(): + config = _training_config(weight_interleave_size=32) + forward = _native_forward(config) + backward = _native_backward(config) - forward_source = inspect.getsource(_compile.prepare_kernel) - backward_source = inspect.getsource(_backward_compile.prepare_backward_kernel) - runtime_source = inspect.getsource(_backward_compile.build_backward_runtime_kwargs) - dglu_source = _DGLU.read_text(encoding="utf-8") - dfc2_source = _DGLU_EPILOGUE.read_text(encoding="utf-8") - - assert "gate_up_clamp=config.gate_up_clamp" in backward_source - assert "dfc2_recompute=dfc2_recompute" in backward_source - assert "enable_grad_y2_col_quant=enable_grad_y2_col_quant" in backward_source - assert '"fc1_preact":' in runtime_source - assert '"dprob":' in runtime_source - assert "generate_c=config.generate_c" in forward_source - for contract in ( - "enable_grad_y2_col_quant", - "num_ctas_grad_y2_col_quant", - "grad_y2_sizes_region", - "_snapshot_grad_y2_expert_sizes", - "grad_y2_col_quant", - "grad_y2: cute.Tensor", - "grad_y2_sf: cute.Tensor", - ): - assert contract in dglu_source - assert dglu_source.index("self._snapshot_grad_y2_expert_sizes(tidx)") < dglu_source.index("self.token_comm.reset_tail()") - assert dglu_source.index("self._topk_reduce(") < dglu_source.index("self.grad_y2_col_quant(") - assert "def _stg_col_sf_atom_value(" in dfc2_source - assert "feature_atom = feature // cutlass.Int32(128)" in dfc2_source - assert "feature_lane * cutlass.Int32(16)" in dfc2_source - assert "feature_bank * cutlass.Int32(4)" in dfc2_source - assert "real_sf[feature_atom, token_atom, atom_byte]" in dfc2_source - assert "def tma_store_dfc2_outputs(" in dfc2_source - assert dfc2_source.count("self._stg_col_sf_atom_value(") >= 2 + assert validate_native_forward_weights(config, forward) == torch.device("cpu") + assert validate_native_backward_weights(config, backward) == torch.device("cpu") + forward_kernel = forward_native_to_kernel(forward) + backward_kernel = backward_native_to_kernel(backward) + assert forward_kernel.fc1_weight.data_ptr() == forward.fc1.payload.data_ptr() + assert forward_kernel.fc1_weight_sf.data_ptr() == forward.fc1.scale.data_ptr() + assert forward_kernel.fc2_weight.data_ptr() == forward.fc2.payload.data_ptr() + assert backward_kernel.fc1_weight.data_ptr() == backward.w2_transpose.payload.data_ptr() + assert backward_kernel.fc2_weight_sf.data_ptr() == backward.w1_transpose.scale.data_ptr() -# WGrad operand contracts +@pytest.mark.L0 +def test_standalone_weight_packers_write_only_caller_staging(): + config = _training_config(weight_interleave_size=32) + source = _source_weights(config) + forward_out, backward_out = _allocate_training_weight_staging(source) -@pytest.mark.L1 -@pytest.mark.parametrize( - ("field", "defect"), - [ - pytest.param(field, "logical_shape", id=f"{field}-logical-shape") - for field in ( - "forward_fc1", - "forward_fc2", - "backward_w2_transpose", - "backward_w1_transpose", - ) - ] - + [ - pytest.param("forward_fc1", defect, id=f"forward_fc1-{defect}") - for defect in ( - "plain_tensor", - "axis", - "format", - ) - ], -) -def test_validate_training_weights_rejects_targeted_defects(field, defect): - invalid, error_type, message = _training_weight_defect( - _training_weights(), - field, - defect, - ) - with pytest.raises(error_type) as exc_info: - validate_training_weights(_training_config(), invalid) - assert str(exc_info.value) == message + native_forward = pack_forward_weights(source[0], out=forward_out) + native_backward = pack_backward_weights(source[1], out=backward_out) + assert native_forward.fc1.payload is forward_out.fc1_payload + assert native_forward.fc2.scale is forward_out.fc2_scale + assert native_backward.w2_transpose.payload is backward_out.w2_transpose_payload + assert native_backward.w1_transpose.scale is backward_out.w1_transpose_scale + validate_native_forward_weights(config, native_forward) + validate_native_backward_weights(config, native_backward) -@pytest.mark.L1 -def test_validate_training_weights_rejects_cross_field_device_mismatch(): - invalid, error_type, message = _training_weight_defect( - _training_weights(), - "backward_w1_transpose", - "device", + +@pytest.mark.L0 +def test_weight_packing_rejects_source_staging_alias(): + config = _training_config(weight_interleave_size=32) + source = _source_weights(config) + _, backward_out = _allocate_training_weight_staging(source) + aliased_out = MoeEpBackwardWeightStaging( + w2_transpose_payload=source[1].w2_transpose.data, + w2_transpose_scale=backward_out.w2_transpose_scale, + w1_transpose_payload=backward_out.w1_transpose_payload, + w1_transpose_scale=backward_out.w1_transpose_scale, ) - with pytest.raises(error_type) as exc_info: - validate_training_weights(_training_config(), invalid) - assert str(exc_info.value) == message + with pytest.raises(ValueError, match="must not alias"): + pack_backward_weights(source[1], out=aliased_out) -@pytest.mark.L1 -def test_validate_training_weights_accepts_complete_fixed_weight_set(): - assert validate_training_weights( - _training_config(), - _training_weights(), - ) == torch.device("cpu") +@pytest.mark.L0 +def test_mxfp8_training_input_bypasses_quantization_stager(): + class RejectingStager: + def stage(self, *args, **kwargs): + raise AssertionError("MXFP8 input must bypass the quantization stager") + + token_count = 2 + hidden = 32 + top_k = 2 + value = BlockScaledTensor( + data=torch.ones((token_count, hidden), dtype=torch.float8_e4m3fn), + scale=torch.ones((token_count, hidden // 32), dtype=torch.float8_e8m0fnu), + format="mxfp8", + logical_shape=(token_count, hidden), + axis=1, + ) + topk_idx = torch.tensor([[0, 1], [1, -1]], dtype=torch.int32) + topk_weights = torch.tensor([[0.75, 0.25], [1.0, 0.0]], dtype=torch.float32) + activation_data = torch.empty((4, hidden), dtype=torch.float8_e4m3fn) + activation_sf = torch.empty((4, hidden // 32), dtype=torch.float8_e8m0fnu) + routing_idx = torch.empty((4, top_k), dtype=torch.int32) + routing_weights = torch.empty((4, top_k), dtype=torch.float32) + + _stage_input( + type("Owner", (), {"stager": RejectingStager()})(), + value, + topk_idx, + topk_weights, + activation_data, + activation_sf, + routing_idx, + routing_weights, + ) -@pytest.mark.L1 -def test_interleaved_training_weights_require_direct_layouts(): - with pytest.raises(ValueError, match="requires compact K-major forward weights"): - validate_training_weights( - _training_config(weight_interleave_size=32), - _training_weights(), - ) + torch.testing.assert_close(activation_data[:token_count], value.data) + torch.testing.assert_close(activation_sf[:token_count], value.scale) + torch.testing.assert_close(routing_idx[:token_count], topk_idx) + torch.testing.assert_close(routing_weights[:token_count], topk_weights) + assert activation_data[token_count:].eq(0).all() + assert activation_sf[token_count:].view(torch.uint8).eq(0).all() + assert routing_idx[token_count:].eq(-1).all() + assert routing_weights[token_count:].eq(0).all() -@pytest.mark.L1 -@pytest.mark.parametrize("part", ["data_noncontiguous", "scale_noncontiguous"]) -def test_validate_training_weights_accepts_compact_k_major_views(part): - weights, _, _ = _training_weight_defect( - _training_weights(), - "forward_fc1", - part, +@pytest.mark.L0 +def test_native_execution_rejects_compact_or_wrong_layout_scales(): + config = _training_config(weight_interleave_size=32) + native = _native_forward(config) + bad = MoeEpNativeForwardWeights( + fc1=MoeEpNativeWeight( + native.fc1.payload, + torch.empty( + (config.experts_per_rank, config.hidden_size // 32, 2 * config.intermediate_size), + dtype=torch.float8_e8m0fnu, + ), + native.fc1.layout_id, + ), + fc2=native.fc2, ) - assert validate_training_weights(_training_config(), weights) == torch.device("cpu") - bindings = Mxfp8TrainingWeightBindings(weights) - bindings.refresh() + with pytest.raises(ValueError, match=r"weights\.fc1\.scale shape"): + validate_native_forward_weights(config, bad) -@pytest.mark.L1 +@pytest.mark.L0 @pytest.mark.parametrize( - "field", - ["backward_w2_transpose", "backward_w1_transpose"], + ("phase", "missing"), + ( + ("forward", "fc1_preact"), + ("forward", "output"), + ("forward", "fc1_a"), + ("forward", "fc1_sfa"), + ("forward", "valid_route_counts"), + ("forward", "expert_offsets"), + ("backward", "grad_activation"), + ("backward", "dprob"), + ("backward", "fc1_b"), + ("backward", "fc1_sfb"), + ("backward", "fc2_a"), + ("backward", "fc2_sfa"), + ("backward", "fc2_b"), + ("backward", "fc2_sfb"), + ), ) -@pytest.mark.parametrize("weight_interleave_size", [None, 32]) -def test_validate_training_weights_rejects_k_major_backward_data( - field, - weight_interleave_size, -): - weights, _, _ = _training_weight_defect( - _training_weights(), - field, - "data_noncontiguous", - ) - with pytest.raises( - ValueError, - match=rf"weights\.{field} data must be contiguous", - ): - validate_training_weights( - _training_config(weight_interleave_size=weight_interleave_size), - weights, - ) - +def test_training_output_requirements_reject_missing_fields(phase, missing): + requirement = ((1,), (1,), torch.float32, 1) + if phase == "forward": + names = ("output", "fc1_preact", "fc1_a", "fc1_sfa", "valid_route_counts", "expert_offsets") + values = {name: torch.empty(1) for name in names} + values[missing] = None + output = MoeEpTrainingForwardOutputs(**values) + validate = validate_training_forward_outputs + else: + names = ("grad_activation", "dprob", "fc1_b", "fc1_sfb", "fc2_a", "fc2_sfa", "fc2_b", "fc2_sfb") + values = {name: torch.empty(1) for name in names} + values[missing] = None + output = MoeEpTrainingBackwardOutputs(**values) + validate = validate_training_backward_outputs + requirements = {name: requirement for name in names} + with pytest.raises(TypeError, match=rf"out\.{missing} must be a torch.Tensor"): + validate(output, requirements, device=torch.device("cpu")) -def _operator(**overrides) -> MoeEp: - values = { - "num_experts": 2, - "hidden_size": 128, - "intermediate_size": 256, - "top_k": 2, - "max_tokens_per_rank": 4, - } - values.update(overrides) - return MoeEp( - **values, - ) +@pytest.mark.L0 +def test_training_output_types_remain_optional_before_validation(): + with pytest.raises(TypeError, match="fc1_preact"): + MoeEpTrainingForwardOutputs() + forward = MoeEpTrainingForwardOutputs(fc1_preact=torch.empty(1)) + backward = MoeEpTrainingBackwardOutputs() + assert forward.output is None + assert backward.grad_activation is None -def _install_contract_backend( - monkeypatch, - *, - weights=None, - slot_count=1, - lane_count=1, -): - import cudnn.moe_ep._backend as backend_seam - import cudnn.moe_ep.api as api_module - - weights = weights or SimpleNamespace(mock_training_weights=True) - state = SimpleNamespace( - backends=[], - validate=Mock(return_value=torch.device("cpu")), - ) - def create_backend(config, device): - del config, device - owner = _TrainingResourceContractOwner( - slot_count=slot_count, - lane_count=lane_count, +@pytest.mark.L0 +def test_training_forward_state_validation_uses_output_contract_names(): + requirement = ((1,), (1,), torch.float32, 1) + requirements = { + name: requirement + for name in ( + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", ) - backend = SimpleNamespace( - owner=owner, - prepare_training_resources=Mock(return_value=owner), - close=Mock(), + } + with pytest.raises( + TypeError, + match=r"out\.fc1_a must be a torch.Tensor", + ): + validate_training_forward_state( + fc1_preact=torch.empty(1), + fc1_a=None, + fc1_sfa=torch.empty(1), + valid_route_counts=torch.empty(1), + expert_offsets=torch.empty(1), + requirements=requirements, + device=torch.device("cpu"), ) - state.backends.append(backend) - return backend - - monkeypatch.setattr(api_module, "validate_training_weights", state.validate) - monkeypatch.setattr(backend_seam, "validate_config", lambda config: None) - monkeypatch.setattr(backend_seam, "create_backend", create_backend) - return weights, state @pytest.mark.L0 -def test_k_major_workspace_view_matches_upstream_token_major_abi(): - storage = torch.arange(12, dtype=torch.uint8) - view = _typed_k_major_view(storage, torch.uint8, (3, 4)) - - assert view.shape == (3, 4) - assert view.stride() == (1, 3) - assert torch.equal(view, storage.reshape(4, 3).transpose(0, 1)) +def test_training_backward_rejects_missing_output_bundle_after_prepare(): + op = MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=4, + max_recv_size_per_rank=4, + weight_interleave_size=32, + ) + lane = MoeEpExecutionLane(0, op._operator_token) + op._training_state = object() + op._training_requirements = {} + op._training_lanes = (lane,) + op._forward_backend_device = torch.device("cpu") + with pytest.raises(TypeError, match="out must be a MoeEpTrainingBackwardOutputs"): + op.training_backward( + lane, + torch.empty((0, 128), dtype=torch.bfloat16), + torch.empty((0, 2), dtype=torch.int32), + torch.empty((0, 2), dtype=torch.float32), + weights=None, + fc1_preact=torch.empty((0, 512), dtype=torch.bfloat16), + out=None, + ) @pytest.mark.L0 -def test_training_wgrad_data_operands_alias_backward_outputs(): - pool_rows, hidden, intermediate = 8, 4, 6 - slot = SimpleNamespace( - col_quant_data=torch.empty((pool_rows, hidden), dtype=torch.uint8), - col_quant_sf=torch.empty(1, dtype=torch.uint8), - valid_route_counts=torch.zeros(1, dtype=torch.int32), - expert_offsets=torch.zeros(1, dtype=torch.int32), - fc1_recompute=torch.empty((pool_rows, intermediate), dtype=torch.uint8), - fc1_recompute_sf=torch.empty(1, dtype=torch.uint8), - fc1_col_output=torch.empty((pool_rows, 2 * intermediate), dtype=torch.uint8), - fc1_col_output_sf=torch.empty(1, dtype=torch.uint8), - grad_y2=torch.empty((pool_rows, hidden), dtype=torch.uint8), - grad_y2_sf=torch.empty(1, dtype=torch.uint8), - wgrad_fc1_sfa=torch.empty(1, dtype=torch.uint8), - wgrad_fc1_sfb=torch.empty(1, dtype=torch.uint8), - wgrad_fc2_sfa=torch.empty(1, dtype=torch.uint8), - wgrad_fc2_sfb=torch.empty(1, dtype=torch.uint8), +def test_wgrad_assembly_returns_only_caller_owned_views(): + buffers = { + name: torch.empty(1) + for name in ( + "fc1_a", + "fc1_sfa", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + "valid_route_counts", + "expert_offsets", + ) + } + backward = MoeEpTrainingBackwardOutputs( + fc1_b=buffers["fc1_b"], + fc1_sfb=buffers["fc1_sfb"], + fc2_a=buffers["fc2_a"], + fc2_sfa=buffers["fc2_sfa"], + fc2_b=buffers["fc2_b"], + fc2_sfb=buffers["fc2_sfb"], ) - exporter = Mxfp8TrainingWgradExporter( - experts=1, - hidden=hidden, - intermediate=intermediate, - fc1_weight_layout=Fc1WeightLayout.GATE_UP_INTERLEAVED_32, + operands = assemble_training_wgrad_operands( + fc1_a=buffers["fc1_a"], + fc1_sfa=buffers["fc1_sfa"], + valid_route_counts=buffers["valid_route_counts"], + expert_offsets=buffers["expert_offsets"], + backward=backward, ) - exporter._expand_scales = Mock() - - operands = exporter.export(slot) - - assert operands.fc1_b is slot.fc1_col_output - assert operands.fc1_b.data_ptr() == slot.fc1_col_output.data_ptr() - assert operands.fc1_b.stride() == slot.fc1_col_output.stride() - assert operands.fc2_a.data_ptr() == slot.fc1_recompute.data_ptr() - assert operands.fc2_a.shape == (intermediate, pool_rows) - assert operands.fc2_a.stride() == slot.fc1_recompute.transpose(0, 1).stride() + assert isinstance(operands, MoeEpTrainingWgradOperands) + for name in buffers: + assert getattr(operands, name).data_ptr() == buffers[name].data_ptr() @pytest.mark.L0 -@pytest.mark.parametrize( - "layout", - [ - Fc1WeightLayout.GATE_THEN_UP, - Fc1WeightLayout.GATE_UP_INTERLEAVED_32, - ], -) -def test_training_wgrad_fc1_layout_matches_reference(layout): - pool_rows, hidden, intermediate = 3, 4, 64 - semantic_dc = ( - torch.arange(pool_rows * 2 * intermediate, dtype=torch.int64) - .remainder(251) - .to(torch.uint8) - .reshape(pool_rows, 2 * intermediate) - ) - interleaved_dc = ( - semantic_dc.view(pool_rows, 2, intermediate // 32, 32) - .transpose(1, 2) - .reshape_as(semantic_dc) - ) - x = torch.arange(pool_rows * hidden, dtype=torch.uint8).reshape(pool_rows, hidden) - slot = SimpleNamespace( - col_quant_data=x, - col_quant_sf=torch.empty(1, dtype=torch.uint8), - valid_route_counts=torch.zeros(1, dtype=torch.int32), - expert_offsets=torch.zeros(1, dtype=torch.int32), - fc1_recompute=torch.empty((pool_rows, intermediate), dtype=torch.uint8), - fc1_recompute_sf=torch.empty(1, dtype=torch.uint8), - fc1_col_output=interleaved_dc, - fc1_col_output_sf=torch.empty(1, dtype=torch.uint8), - grad_y2=torch.empty((pool_rows, hidden), dtype=torch.uint8), - grad_y2_sf=torch.empty(1, dtype=torch.uint8), - wgrad_fc1_b=torch.empty_like(interleaved_dc), - wgrad_fc1_sfa=torch.empty(1, dtype=torch.uint8), - wgrad_fc1_sfb=torch.empty(1, dtype=torch.uint8), - wgrad_fc2_sfa=torch.empty(1, dtype=torch.uint8), - wgrad_fc2_sfb=torch.empty(1, dtype=torch.uint8), +def test_private_training_state_has_no_bound_weights_or_wgrad_exporter(): + config = _training_config(weight_interleave_size=32) + forward, backward = _training_prepared_pair(config) + state = Mxfp8TrainingState( + config, + torch.device("cpu"), + forward, + backward, + lane_count=2, ) - exporter = Mxfp8TrainingWgradExporter( - experts=1, - hidden=hidden, - intermediate=intermediate, - fc1_weight_layout=layout, + assert not hasattr(state, "weight_bindings") + assert not hasattr(state, "wgrad_exporter") + assert not hasattr(state, "slot_count") + assert state.lane_count == 2 + assert all( + "fc1_preact" not in region.name + for region in ( + *state.requirements.symmetric_regions, + *state.requirements.local_regions, + ) ) - exporter._expand_scales = Mock() - operands = exporter.export(slot) - expected_dc = ( - semantic_dc - if layout is Fc1WeightLayout.GATE_THEN_UP - else interleaved_dc + requirements = state.public_requirements() + assert requirements["fc1_a"] == ( + (config.hidden_size, forward.pool_token_capacity), + (forward.pool_token_capacity, 1), + torch.float8_e4m3fn, + 128, ) - expected_wgrad = x.transpose(0, 1).float() @ expected_dc.float() - actual_wgrad = operands.fc1_a.float() @ operands.fc1_b.float() - - torch.testing.assert_close(actual_wgrad, expected_wgrad, atol=0, rtol=0) - if layout is Fc1WeightLayout.GATE_THEN_UP: - assert operands.fc1_b is slot.wgrad_fc1_b - assert operands.fc1_b.data_ptr() != slot.fc1_col_output.data_ptr() - assert exporter._expand_scales.call_args_list[1].kwargs[ - "deinterleave_gate_up" - ] == intermediate - else: - assert operands.fc1_b is slot.fc1_col_output - assert exporter._expand_scales.call_args_list[1].kwargs[ - "deinterleave_gate_up" - ] is None - - -@pytest.mark.L0 -def test_only_fixed_training_wgrad_types_are_public(): - expected = [f"fc{layer}_{part}" for layer in (1, 2) for part in ("a", "sfa", "b", "sfb")] - expected += ["expert_offsets", "valid_route_counts"] - assert [field.name for field in fields(MoeEpTrainingWgradOperands)] == expected - assert not hasattr(cudnn, "MoeEpWgradForwardStash") - assert not hasattr(cudnn, "MoeEpWgradOperands") + assert requirements["fc1_b"][1] == (2 * config.intermediate_size, 1) + assert requirements["fc2_a"][1] == (1, config.intermediate_size) + assert requirements["fc2_b"][1] == (1, forward.pool_token_capacity) @pytest.mark.L0 -def test_prepare_training_resources_binds_weights_and_slot_lanes(monkeypatch): - weights = _training_weights() - _, state = _install_contract_backend( - monkeypatch, - weights=weights, - slot_count=2, +def test_private_training_workspace_keeps_only_live_lane_scratch(): + config = _training_config(weight_interleave_size=32) + forward, backward = _training_prepared_pair(config) + state = Mxfp8TrainingState( + config, + torch.device("cpu"), + forward, + backward, + lane_count=2, ) - operator = _operator() - resources = operator.prepare_training_resources( - weights, - slot_count=2, - lane_count=1, + flat = WorkspaceViews( + token_count=0, + symmetric={region.name: torch.empty(region.nbytes, dtype=torch.uint8) for region in state.requirements.symmetric_regions}, + local={region.name: torch.empty(region.nbytes, dtype=torch.uint8) for region in state.requirements.local_regions}, + peer_mapping=object(), ) - assert isinstance(resources, MoeEpTrainingResources) - assert all(isinstance(slot, MoeEpTrainingSlot) for slot in resources.slots) - assert isinstance(resources.lanes[0], MoeEpExecutionLane) - resources.refresh_weights() - owner = state.backends[0].owner - assert owner.refresh_calls == 1 - operator.close() - assert resources.closed - assert owner.close_calls == 1 - - -@pytest.mark.L0 -def test_prepare_training_resources_rejects_plain_weights(): - with _operator() as operator: - with pytest.raises(TypeError, match="MoeEpTrainingWeights"): - operator.prepare_training_resources(_training_inputs()[1]) - - -@pytest.mark.L0 -def test_error_mode_requires_async_assert_before_prepare(monkeypatch): - from cudnn.moe_ep.api import _validate_training_assert_capability - - monkeypatch.setattr(torch, "_assert_async", None) - config = SimpleNamespace(drop_on_overflow=False, ep_size=1) - with pytest.raises(RuntimeError, match="callable torch._assert_async"): - _validate_training_assert_capability(config) - - _validate_training_assert_capability(SimpleNamespace(drop_on_overflow=True, ep_size=1)) + first = state._lane_scratch_views(flat, 0) + second = state._lane_scratch_views(flat, 1) + forward_workspace = state._phase_workspace( + flat, + forward.workspace_requirements, + lane=0, + phase="forward", + ) + backward_workspace = state._phase_workspace( + flat, + backward.workspace_requirements, + lane=0, + phase="backward", + ) + assert "col_quant_data" not in forward_workspace.local + assert "col_quant_sf" not in forward_workspace.local + assert "kernel_local_workspace" in forward_workspace.local + assert "backward_aux_data" in backward_workspace.local + assert "backward_aux_scale" in backward_workspace.local + for field in fields(first): + first_value = getattr(first, field.name) + second_value = getattr(second, field.name) + if isinstance(first_value, torch.Tensor): + assert first_value.data_ptr() != second_value.data_ptr() + names = {region.name for region in (*state.requirements.symmetric_regions, *state.requirements.local_regions)} + removed = ( + "valid_route_counts", + "expert_offsets", + "fc1_recompute", + "fc1_recompute_sf", + "fc1_col_output", + "fc1_col_output_sf", + "grad_y2", + "grad_y2_sf", + "col_quant_data", + "col_quant_sf", + ) + assert not any(any(name.endswith(removed_name) for removed_name in removed) for name in names) + for lane in range(2): + assert f"lane.{lane}.fallback.local.routing_topk_idx" in names + assert f"lane.{lane}.fallback.symmetric.routing_topk_weights" in names + assert f"lane.{lane}.backward.local.backward_aux_data" in names + assert f"lane.{lane}.backward.local.backward_aux_scale" in names + assert f"lane.{lane}.forward.symmetric.output_data" in names + assert f"lane.{lane}.backward.symmetric.output_data" in names + assert f"lane.{lane}.backward.symmetric.backward_dprob" in names @pytest.mark.L0 -def test_distributed_error_mode_requires_nccl(monkeypatch): - from cudnn.moe_ep.api import _validate_training_assert_capability - - monkeypatch.setattr(torch, "_assert_async", lambda *args, **kwargs: None) - monkeypatch.setattr(torch.distributed, "get_backend", lambda group: "gloo") - config = SimpleNamespace( - drop_on_overflow=False, - ep_size=2, - ep_group=object(), +def test_training_views_require_col_quant_snapshot(): + config = _training_config(weight_interleave_size=32) + forward, backward = _training_prepared_pair(config) + forward.col_quant_sizes_offset = None + state = Mxfp8TrainingState( + config, + torch.device("cpu"), + forward, + backward, + lane_count=1, ) - with pytest.raises(NotImplementedError, match="NCCL"): - _validate_training_assert_capability(config) + with pytest.raises(RuntimeError, match="persistent col-quant expert-size snapshot"): + state.views(lane=0, token_count=0) @pytest.mark.L0 -def test_training_weight_bindings_alias_data_and_stage_only_scales(): - weights = _training_weights() - from cudnn.moe_ep import BlockScaledTensor, MoeEpTrainingWeights - - def compact_k_major(tensor): - return BlockScaledTensor( - data=tensor.data.transpose(1, 2).contiguous().transpose(1, 2), - scale=tensor.scale.transpose(1, 2).contiguous().transpose(1, 2), - format=tensor.format, - logical_shape=tensor.logical_shape, - axis=tensor.axis, - ) - - weights = MoeEpTrainingWeights( - forward_fc1=compact_k_major(weights.forward_fc1), - forward_fc2=compact_k_major(weights.forward_fc2), - backward_w2_transpose=weights.backward_w2_transpose, - backward_w1_transpose=weights.backward_w1_transpose, +def test_training_abi_fingerprint_covers_lanes_and_native_layouts(): + config = _training_config( + ep_size=2, + ep_global_ranks=(0, 1), + weight_interleave_size=32, ) - assert validate_training_weights( - _training_config(weight_interleave_size=32), - weights, - ) == torch.device("cpu") - compatibility_bindings = Mxfp8TrainingWeightBindings(weights) - assert not compatibility_bindings._uses_direct_weight_bindings - assert ( - compatibility_bindings.forward.fc1_weight.data_ptr() - != weights.forward_fc1.data.data_ptr() + forward = _training_abi_prepared("forward") + backward = _training_abi_prepared("backward") + requirements = WorkspaceRequirements( + max_tokens_per_rank=4, + symmetric_regions=(BufferRegion("symmetric", 256),), + local_regions=(BufferRegion("local", 128),), ) - assert compatibility_bindings.backward.fc1_weight.is_contiguous() - assert compatibility_bindings.backward.fc2_weight.is_contiguous() - - bindings = Mxfp8TrainingWeightBindings( - weights, - fc1_weight_layout=Fc1WeightLayout.GATE_UP_INTERLEAVED_32, + first = _build_training_abi_facts( + config, + forward, + backward, + requirements, + lane_count=1, + source_tree_digest="source", ) - bindings.refresh() - data_pairs = ( - (bindings.forward.fc1_weight, weights.forward_fc1.data), - (bindings.forward.fc2_weight, weights.forward_fc2.data), - (bindings.backward.fc1_weight, weights.backward_w2_transpose.data), - (bindings.backward.fc2_weight, weights.backward_w1_transpose.data), + repeated = _build_training_abi_facts( + config, + forward, + backward, + requirements, + lane_count=1, + source_tree_digest="source", ) - scales = ( - bindings.forward.fc1_weight_sf, - bindings.forward.fc2_weight_sf, - bindings.backward.fc1_weight_sf, - bindings.backward.fc2_weight_sf, + changed_lanes = _build_training_abi_facts( + config, + forward, + backward, + requirements, + lane_count=2, + source_tree_digest="source", ) - scale_pointers = tuple(tensor.data_ptr() for tensor in scales) - scale_snapshots = tuple(tensor.clone() for tensor in scales) - weights.forward_fc1.scale.view(torch.uint8).bitwise_xor_(1) - bindings.refresh() - - assert all(bound.data_ptr() == source.data_ptr() for bound, source in data_pairs) - assert tuple(tensor.data_ptr() for tensor in scales) == scale_pointers - assert not torch.equal(bindings.forward.fc1_weight_sf, scale_snapshots[0]) + assert first["schema_version"] == 2 + assert first["native_weight_layouts"] == [layout.value for layout in MoeEpNativeWeightLayout] + assert canonical_json_sha256(first) == canonical_json_sha256(repeated) + assert canonical_json_sha256(first) != canonical_json_sha256(changed_lanes) @pytest.mark.L0 -def test_reference_wgrad_math_remains_in_test_tree(): - torch.manual_seed(20260821) - tokens, hidden, intermediate = 3, 32, 32 - topk_idx = torch.tensor([[0, 2], [2, 0], [0, 2]], dtype=torch.int32) - topk_weights = torch.tensor([[0.5, 0.25], [0.0, 0.75], [1.0, 0.125]]) - activation, fc1_weight, fc2_weight, grad_output = ( - torch.randn(shape) / 8 - for shape in ( - (tokens, hidden), - (3, hidden, 2 * intermediate), - (3, intermediate, hidden), - (tokens, hidden), - ) - ) - reference = MoeEpReference( - num_experts=3, - hidden_size=hidden, - intermediate_size=intermediate, - top_k=2, - max_tokens_per_rank=tokens, - generate_c=True, - backward_wgrad_mode="operands", - token_padding_size=256, - ) - - _, fc1_c, metadata, stash = reference(activation, fc1_weight, fc2_weight, topk_idx, topk_weights) - _, _, operands = reference.backward( - grad_output, - fc1_weight, - fc2_weight, - topk_idx, - topk_weights, - fc1_c, - metadata, - wgrad_forward_stash=stash, - ) - dw1, dw2 = operands.dense_wgrads() - - assert operands.valid_route_counts.tolist() == [3, 0, 3] - assert dw1.shape == (3, hidden, 2 * intermediate) - assert dw2.shape == (3, intermediate, hidden) - assert dw1[1].eq(0).all() - assert dw2[1].eq(0).all() - - -# Source contracts - - -_ROOT = Path(__file__).resolve().parents[3] -_CUTEDSL = _ROOT / "python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin" / "training/mega" -_DGLU = _CUTEDSL / "bwd_dglu/dglu_mxfp8_mega_moe_kernel.py" -_DGLU_EPILOGUE = _CUTEDSL / "bwd_dglu/dglu_mxfp8_fc12_epilogue.py" - - -# L1 fail-fast training-resource contracts - - -@pytest.mark.L1 -@pytest.mark.parametrize( - ("field", "value"), - [(field, value) for field in ("slot_count", "lane_count") for value in (0, True, 1.5)], -) -def test_prepare_training_resources_rejects_invalid_counts_before_backend( - monkeypatch, - field, - value, -): - import cudnn.moe_ep._backend as backend_seam - import cudnn.moe_ep.api as api_module - - def unexpected_call(*args, **kwargs): - del args, kwargs - raise AssertionError("invalid counts must fail before weight/backend work") - - monkeypatch.setattr( - api_module, - "validate_training_weights", - unexpected_call, - ) - monkeypatch.setattr(backend_seam, "create_backend", unexpected_call) - counts = {"slot_count": 1, "lane_count": 1, field: value} - - with ( - _operator() as operator, - pytest.raises( - ValueError, - match=rf"{field} must be a positive integer", - ), - ): - operator.prepare_training_resources( - SimpleNamespace(mock_training_weights=True), - **counts, - ) - - -@pytest.mark.L1 -def test_prepare_training_resources_rejects_duplicate_open_resources(monkeypatch): - weights, state = _install_contract_backend(monkeypatch) - - with _operator() as operator: - resources = operator.prepare_training_resources( - weights, - slot_count=1, - lane_count=1, - ) - with pytest.raises(RuntimeError, match="already exist"): - operator.prepare_training_resources( - weights, - slot_count=1, - lane_count=1, - ) - - assert resources.closed - backend = state.backends[0] - state.validate.assert_called_once() - backend.prepare_training_resources.assert_called_once_with( - weights, - slot_count=1, - lane_count=1, - ) - assert (backend.close.call_count, backend.owner.close_calls) == (1, 1) - - -@pytest.mark.L1 -def test_closed_training_resources_require_a_new_operator(monkeypatch): - weights, state = _install_contract_backend(monkeypatch) - - old_operator = _operator() - old_resources = old_operator.prepare_training_resources( - weights, - slot_count=1, - lane_count=1, - ) - old_resources.close() - with pytest.raises( - RuntimeError, - match="create a new MoeEp instance", - ): - old_operator.prepare_training_resources( - weights, - slot_count=1, - lane_count=1, - ) - old_operator.close() - - with _operator() as new_operator: - new_resources = new_operator.prepare_training_resources( - weights, - slot_count=1, - lane_count=1, - ) - assert not new_resources.closed - - assert len(state.backends) == 2 - assert state.validate.call_count == 2 - assert all(backend.prepare_training_resources.call_count == 1 for backend in state.backends) - - -@pytest.mark.L1 -def test_training_prepare_and_backend_close_reject_capture(monkeypatch): - from cudnn.moe_ep._megamoe_backend.mxfp8._backend import Mxfp8Backend - - monkeypatch.setattr( - torch.cuda, - "is_current_stream_capturing", - lambda: True, - ) - - owner = object.__new__(Mxfp8TrainingResourceOwner) - owner._lock = threading.RLock() - owner._closed = False - owner._runtime = None - owner._workspace = None - with pytest.raises( - RuntimeError, - match="must be prepared before CUDA graph capture", - ): - owner.prepare() - - backend = object.__new__(Mxfp8Backend) - backend._lock = threading.RLock() - backend._closed = False - backend.device = torch.device("cuda") - monkeypatch.setattr(torch.cuda, "device", lambda device: nullcontext()) - with pytest.raises(RuntimeError, match="cannot be closed during"): - backend.close() - - -@pytest.mark.L1 -def test_training_resources_reject_foreign_and_forged_slot_lane_bindings(): - resources, owner = _training_contract_resources() - foreign, _ = _training_contract_resources() - slot = resources.slots[0] - lane = resources.lanes[0] - activation = torch.empty((0, 128), dtype=torch.bfloat16) - routing = ( - torch.empty((0, 2), dtype=torch.int32), - torch.empty((0, 2), dtype=torch.float32), - ) - checks = ( - ("training slot does not belong", resources.forward, (foreign.slots[0], lane, activation, *routing)), - ("training slot does not belong", resources.backward, (MoeEpTrainingSlot(99, slot._resource_token), lane, activation.float())), - ("execution lane does not belong", resources.forward, (slot, foreign.lanes[0], activation, *routing)), - ("execution lane does not belong", resources.backward, (slot, MoeEpExecutionLane(99, lane._resource_token), activation.float())), - ) - for message, call, args in checks: - with pytest.raises(ValueError, match=message): - call(*args) - - assert owner.views_calls == 0 - - -@pytest.mark.L1 -def test_training_resources_reject_invalid_overflow_finalization(): - resources, _ = _training_contract_resources() - foreign, _ = _training_contract_resources() - slot = resources.slots[0] - lane = resources.lanes[0] - - with pytest.raises(ValueError, match="at least one slot"): - resources.finalize_overflow((), lane) - with pytest.raises(ValueError, match="slots must be unique"): - resources.finalize_overflow((slot, slot), lane) - with pytest.raises(ValueError, match="overflow slot does not belong"): - resources.finalize_overflow((foreign.slots[0],), lane) - with pytest.raises(ValueError, match="overflow execution lane does not belong"): - resources.finalize_overflow((slot,), foreign.lanes[0]) - - -@pytest.mark.L1 -def test_training_resources_reject_calls_after_close_and_close_is_idempotent(): - resources, owner = _training_contract_resources() - slot = resources.slots[0] - lane = resources.lanes[0] - activation = torch.empty((0, 128), dtype=torch.bfloat16) - routing = ( - torch.empty((0, 2), dtype=torch.int32), - torch.empty((0, 2), dtype=torch.float32), - ) - - resources.close() - resources.close() - - assert resources.closed - assert owner.close_calls == 1 - calls = ( - resources.refresh_weights, - lambda: resources.forward(slot, lane, activation, *routing), - lambda: resources.backward(slot, lane, activation.float()), - lambda: resources.finalize_overflow((slot,), lane), - ) - for call in calls: - with pytest.raises(RuntimeError, match="resources are closed"): - call() - assert owner.refresh_calls == 0 - assert owner.views_calls == 0 - - -@pytest.mark.L1 -@pytest.mark.parametrize( - ("mismatch_reduce", "message"), - [ - (2, "region counts differ"), - (4, "names, order, or alignments differ"), - ], -) -def test_harmonize_symmetric_regions_rejects_collective_metadata_mismatch( - monkeypatch, - mismatch_reduce, - message, -): +def test_training_workspace_harmonizes_symmetric_regions(monkeypatch): requirements = WorkspaceRequirements( max_tokens_per_rank=1, symmetric_regions=( - BufferRegion("first", 64, alignment=128), - BufferRegion("second", 128, alignment=256), + BufferRegion("first", 1), + BufferRegion("second", 257), ), - local_regions=(), + local_regions=(BufferRegion("local", 1),), ) - runtime = SimpleNamespace(world_size=2, group=object()) - reduce_calls = [] + runtime = type("Runtime", (), {"world_size": 2, "group": object()})() def all_reduce(tensor, *, op, group): assert group is runtime.group - reduce_calls.append(op) - if len(reduce_calls) == mismatch_reduce: - tensor.add_(1) + if tensor.numel() == 2 and op == torch.distributed.ReduceOp.MAX: + tensor.copy_(torch.tensor([257, 257], dtype=torch.int64)) monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + harmonized = _harmonize_symmetric_regions( + requirements, + runtime, + torch.device("cpu"), + ) - with pytest.raises(RuntimeError, match=message): - _harmonize_symmetric_regions( - requirements, - runtime, - torch.device("cpu"), - ) - - assert len(reduce_calls) == mismatch_reduce - - -_STAGER_FAILURES = { - "source-shape": (lambda t: t.update(source=t["source"][:, :-1].contiguous()), ValueError, r"source must have shape \(T, 128\)"), - "route-shape": (lambda t: t.update(topk_idx=t["topk_idx"][:, :-1].contiguous()), ValueError, "topk_idx shape mismatch"), - "weight-shape": (lambda t: t.update(topk_weights=t["topk_weights"][:, :-1].contiguous()), ValueError, "topk_weights shape mismatch"), - "route-dtype": (lambda t: t.update(topk_idx=t["topk_idx"].to(torch.int64)), TypeError, "contiguous Int32"), - "route-contiguity": (lambda t: t.update(topk_idx=t["topk_idx"].t().contiguous().t()), TypeError, "contiguous Int32"), - "weight-dtype": (lambda t: t.update(topk_weights=t["topk_weights"].to(torch.bfloat16)), TypeError, "contiguous FP32"), - "weight-contiguity": (lambda t: t.update(topk_weights=t["topk_weights"].t().contiguous().t()), TypeError, "contiguous FP32"), - "capacity": ( - lambda t: t.update(**{name: value[:4] for name, value in t.items() if name.startswith("output")}), - ValueError, - "token count 5 exceeds capacity 4", - ), - "device": (lambda t: t.update(source=torch.empty_like(t["source"], device="meta")), ValueError, "must share one device"), -} - - -@pytest.mark.L1 -@pytest.mark.parametrize( - ("mutator", "error_type", "message"), - [pytest.param(*case, id=name) for name, case in _STAGER_FAILURES.items()], -) -def test_training_stager_rejects_invalid_inputs(mutator, error_type, message): - tensors = _training_staging_tensors() - mutator(tensors) - with pytest.raises(error_type, match=message): - Mxfp8TrainingStager(hidden=128, top_k=2)._validate(**tensors) - - -# L1 training graph + assert tuple(region.nbytes for region in harmonized.symmetric_regions) == ( + 257, + 257, + ) + assert harmonized.local_regions == requirements.local_regions -@pytest.mark.L1 -@pytest.mark.gpu_exclusive -@pytest.mark.parametrize( - ( - "input_kind", - "combine_format", - "gate_up_clamp", - "top_k", - "all_dropped", - ), - [ - pytest.param("fixed", "bf16", None, 2, False, id="bf16-unclamped"), - pytest.param("fixed", "mxfp8", 0.5, 2, False, id="mxfp8-clamp-0.5"), - pytest.param("routed", "bf16", None, 1, False, id="topk1"), - pytest.param("routed", "bf16", None, 2, False, id="topk2"), - pytest.param("routed", "bf16", None, 2, True, id="topk2-all-dropped"), - ], -) -def test_fixed_training_resources_ep1_matches_independent_reference( - input_kind, - combine_format, - gate_up_clamp, - top_k, - all_dropped, -): - device = _sm107_device() - if input_kind == "fixed": - args, grad_output = _fixed_training_case(device) - max_recv_size = 1 - else: - base_args = make_forward_inputs(device) - args = ( - base_args[0].dequantize(torch.bfloat16), - base_args[1], - base_args[2], - base_args[3][:, :top_k].contiguous(), - base_args[4][:, :top_k].float().contiguous(), +@pytest.mark.L0 +def test_training_contract_rejects_cross_bundle_aliases(): + storage = torch.empty(16) + with pytest.raises(ValueError, match="out must not alias saved"): + validate_training_non_aliasing( + { + "saved": storage[:8], + "out": storage[4:12], + } ) - if all_dropped: - args[3].fill_(-1) - args[4].zero_() - grad_output = _grad_output(device, args[0].shape[0], seed=20260830) - max_recv_size = args[0].shape[0] * top_k - expected = _fixed_training_reference( - args, - grad_output, - combine_format=combine_format, - gate_up_clamp=gate_up_clamp, - ) - with MoeEp( + +@pytest.mark.L0 +def test_training_methods_require_prepare_and_do_not_expose_cleanup(): + op = MoeEp( num_experts=2, hidden_size=128, intermediate_size=256, - top_k=top_k, - max_tokens_per_rank=args[0].shape[0], - max_recv_size_per_rank=max_recv_size, - drop_on_overflow=True, - combine_format=combine_format, - gate_up_clamp=gate_up_clamp, - ) as op: - resources = op.prepare_training_resources( - _fixed_training_weights(args), - slot_count=1, - lane_count=1, - ) - slot = resources.slots[0] - lane = resources.lanes[0] - actual = _run_fixed_training_batch( - resources, - lane, - ((slot, args, grad_output),), - )[0] - torch.cuda.synchronize(device) - - assert actual.overflow.eq(0).all() - _assert_fixed_training_matches_reference( - (actual.y, actual.dx, actual.dprob, actual.wgrads), - expected, - args[3], + top_k=2, + max_tokens_per_rank=4, + max_recv_size_per_rank=4, + weight_interleave_size=32, + ) + assert hasattr(op, "prepare_training") + assert hasattr(op, "training_forward") + assert hasattr(op, "training_backward") + assert not hasattr(op, "prepare_training_resources") + assert not hasattr(op, "refresh_weights") + assert not hasattr(op, "finalize_overflow") + with pytest.raises(RuntimeError, match="prepare_training"): + op.training_forward( + object(), + torch.empty((0, 128), dtype=torch.bfloat16), + torch.empty((0, 2), dtype=torch.int32), + torch.empty((0, 2), dtype=torch.float32), + weights=_native_forward(_training_config(weight_interleave_size=32)), + out=MoeEpTrainingForwardOutputs( + fc1_preact=torch.empty((0, 512), dtype=torch.bfloat16), + ), ) - if all_dropped: - actual_dw1, actual_dw2 = _dense_wgrads_from_operands(actual.wgrads) - expected_dw1, expected_dw2 = expected[3].dense_wgrads() - zero_tensors = ( - actual.y, - expected[0], - actual.dx, - expected[1], - actual.dprob, - expected[2], - actual_dw1, - expected_dw1, - actual_dw2, - expected_dw2, - ) - assert all(tensor.eq(0).all() for tensor in zero_tensors) + conventional = MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=4, + max_recv_size_per_rank=4, + ) + with pytest.raises(ValueError, match="weight_interleave_size=32"): + conventional.prepare_training() @pytest.mark.L1 @pytest.mark.gpu_exclusive -def test_fixed_training_resources_ep1_grouped_wgrad_matches_independent_reference(): +def test_stateless_training_ep1_eager_and_cuda_graph_match_reference(): device = _sm107_device() base_args = make_forward_inputs(device) args = ( @@ -1210,15 +788,16 @@ def test_fixed_training_resources_ep1_grouped_wgrad_matches_independent_referenc base_args[1], base_args[2], base_args[3], - base_args[4].float(), + base_args[4].float().contiguous(), ) - grad_output = _grad_output(device, args[0].shape[0], seed=20260831) + grad_output = _grad_output(device, args[0].shape[0], seed=20260902) expected = _fixed_training_reference( args, grad_output, combine_format="bf16", gate_up_clamp=None, ) + source_weights = _fixed_training_weights(args) with MoeEp( num_experts=2, @@ -1229,665 +808,280 @@ def test_fixed_training_resources_ep1_grouped_wgrad_matches_independent_referenc max_recv_size_per_rank=args[0].shape[0] * args[3].shape[1], drop_on_overflow=True, combine_format="bf16", + weight_interleave_size=32, ) as op: - resources = op.prepare_training_resources( - _fixed_training_weights(args), - slot_count=1, - lane_count=1, + requirements = op.prepare_training(lane_count=1, device=device) + forward_staging, backward_staging = _allocate_training_weight_staging(source_weights) + native_forward = op.pack_forward_weights( + source_weights[0], + out=forward_staging, ) - actual = _run_fixed_training_batch( - resources, - resources.lanes[0], - ((resources.slots[0], args, grad_output),), - )[0] - grouped_wgrads = _dense_wgrads_from_grouped_kernel(actual.wgrads) - torch.cuda.synchronize(device) - - assert actual.overflow.eq(0).all() - _assert_fixed_training_matches_reference( - (actual.y, actual.dx, actual.dprob, actual.wgrads), - expected, - args[3], + native_backward = op.pack_backward_weights( + source_weights[1], + out=backward_staging, ) - torch.testing.assert_close( - actual.wgrads.valid_route_counts, - expected[3].valid_route_counts, - rtol=0, - atol=0, + forward_out, backward_out = _allocate_stateless_training_outputs( + requirements, + device, ) - assert actual.wgrads.valid_route_counts.gt(0).all() - expected_offsets = torch.cumsum( - torch.div( - actual.wgrads.valid_route_counts + 127, - 128, - rounding_mode="floor", + lane = op.training_lanes[0] + + def run(): + y = op.training_forward( + lane, + args[0], + args[3], + args[4], + weights=native_forward, + out=forward_out, + ) + dx, dprob, operands = op.training_backward( + lane, + grad_output, + args[3], + args[4], + weights=native_backward, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, + ) + assert operands is not None + return y, dx, dprob, operands + + def assert_matches(actual): + y, dx, dprob, operands = actual + _assert_matches_reference(y, expected[0]) + _assert_backward_matches( + (dx, dprob), + (expected[1], expected[2]), + args[3], + ) + _assert_wgrads_match_reference( + operands, + expected[3], + weight_interleave_size=32, ) - * 128, - dim=0, - dtype=actual.wgrads.expert_offsets.dtype, - ) - torch.testing.assert_close( - actual.wgrads.expert_offsets, - expected_offsets, - rtol=0, - atol=0, - ) - expected_wgrads = expected[3].dense_wgrads() + eager = run() + grouped_wgrads = _dense_wgrads_from_grouped_kernel(eager[3]) + torch.cuda.synchronize(device) + assert_matches(eager) + expected_fc1_wgrad, expected_fc2_wgrad = expected[3].dense_wgrads() _assert_grouped_wgrads_match_reference( grouped_wgrads, - expected_wgrads, + ( + _interleave_fc1_wgrad(expected_fc1_wgrad), + expected_fc2_wgrad, + ), reference_name="the independent PyTorch MXFP8 reference", ) - _assert_grouped_wgrads_match_reference( - grouped_wgrads, - _dense_wgrads_from_operands(actual.wgrads), - reference_name="the decoded production operand bundle", - close_kwargs={"rtol": 0.1, "atol": 0.1}, - ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run() + pointers = tuple(tensor.data_ptr() for bundle in (forward_out, backward_out) for tensor in vars(bundle).values() if tensor is not None) + for _ in range(2): + graph.replay() + torch.cuda.synchronize(device) + assert pointers == tuple(tensor.data_ptr() for bundle in (forward_out, backward_out) for tensor in vars(bundle).values() if tensor is not None) + assert_matches(captured) @pytest.mark.L1 @pytest.mark.gpu_exclusive -def test_fixed_training_resources_ep1_grouped_wgrad_accumulates_two_microbatches(): +def test_native_io_mxfp8_cuda_graph_replay(): + """Exercise native MXFP8 I/O and weight contracts with graph replay.""" + device = _sm107_device() base_args = make_forward_inputs(device) - args0 = ( - base_args[0].dequantize(torch.bfloat16), + activation = base_args[0] + topk_idx = base_args[3] + topk_weights = base_args[4].float().contiguous() + args = ( + activation, base_args[1], base_args[2], - base_args[3], - base_args[4].float(), - ) - args1 = ( - args0[0].mul(-0.5), - args0[1], - args0[2], - args0[3].roll(1, dims=0), - args0[4].roll(1, dims=0), + topk_idx, + topk_weights, ) - grad_outputs = ( - _grad_output(device, args0[0].shape[0], seed=20260902), - _grad_output(device, args1[0].shape[0], seed=20260903), + grad_output_plain = _grad_output( + device, + activation.shape[0], + seed=20260903, ) - references = tuple( - _fixed_training_reference( - args, - grad_output, - combine_format="bf16", - gate_up_clamp=None, - ) - for args, grad_output in zip((args0, args1), grad_outputs) + grad_output = quantize_mxfp8(grad_output_plain, axis=1) + expected = _fixed_training_reference( + args, + grad_output.dequantize(torch.float32), + combine_format="bf16", + gate_up_clamp=None, ) - with MoeEp( + op = MoeEp( num_experts=2, hidden_size=128, intermediate_size=256, top_k=2, - max_tokens_per_rank=args0[0].shape[0], - max_recv_size_per_rank=args0[0].shape[0] * args0[3].shape[1], + max_tokens_per_rank=activation.shape[0], + max_recv_size_per_rank=activation.shape[0] * topk_idx.shape[1], drop_on_overflow=True, + output_format="bf16", combine_format="bf16", - ) as op: - resources = op.prepare_training_resources( - _fixed_training_weights(args0), - slot_count=2, - lane_count=1, - ) - batch = tuple( - (slot, args, grad_output) - for slot, args, grad_output in zip( - resources.slots, - (args0, args1), - grad_outputs, - ) + weight_interleave_size=32, + ) + try: + requirements = op.prepare_training(lane_count=1, device=device) + forward_out, backward_out = _allocate_stateless_training_outputs( + requirements, + device, ) - actuals = _run_fixed_training_batch(resources, resources.lanes[0], batch) - accumulated = _allocate_dense_grouped_wgrad_outputs( - actuals[0].wgrads, - fill_value=0, + lane = op.training_lanes[0] + + # The production call below receives only native packs. The existing + # fallback packer is used once here as a test oracle to create known-good + # native contents, then copied into independent caller-owned tensors. + source_weights = _fixed_training_weights(args) + forward_staging, backward_staging = _allocate_training_weight_staging( + source_weights ) - output_pointers = tuple(output.data_ptr() for output in accumulated) - for actual in actuals: - returned = _dense_wgrads_from_grouped_kernel( - actual.wgrads, - wgrad_tensors=accumulated, - accumulate_on_output=True, - ) - assert tuple(output.data_ptr() for output in returned) == output_pointers - torch.cuda.synchronize(device) - - for actual, args, reference in zip(actuals, (args0, args1), references): - assert actual.overflow.eq(0).all() - _assert_fixed_training_matches_reference( - (actual.y, actual.dx, actual.dprob, actual.wgrads), - reference, - args[3], - ) - expected_accumulated = tuple( - reference0.float() + reference1.float() - for reference0, reference1 in zip( - references[0][3].dense_wgrads(), - references[1][3].dense_wgrads(), - ) + packed_forward = op.pack_forward_weights( + source_weights[0], + out=forward_staging, ) - _assert_grouped_wgrads_match_reference( - accumulated, - expected_accumulated, - reference_name="the sum of two independent PyTorch MXFP8 references", - close_kwargs={"rtol": 0.2, "atol": 0.25}, + packed_backward = op.pack_backward_weights( + source_weights[1], + out=backward_staging, ) - -@pytest.mark.L1 -@pytest.mark.gpu_exclusive -@pytest.mark.parametrize( - ("world_size", "combine_format", "gate_up_clamp"), - [ - pytest.param(2, "bf16", None, id="ep2-bf16"), - pytest.param(2, "mxfp8", None, id="ep2-mxfp8"), - pytest.param(4, "bf16", None, id="ep4-bf16"), - pytest.param(4, "mxfp8", None, id="ep4-mxfp8"), - pytest.param(2, "bf16", 0.5, id="ep2-bf16-clamp-0.5"), - ], -) -def test_fixed_training_resources_multi_gpu_matches_independent_reference( - world_size, - combine_format, - gate_up_clamp, - tmp_path, -): - _require_distributed_sm107(world_size) - os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") - clamp_id = "none" if gate_up_clamp is None else str(gate_up_clamp) - init_file = tmp_path / f"backward_ep{world_size}_{combine_format}_clamp_{clamp_id}.init" - mp.spawn( - _distributed_backward_reference_worker, - args=( - world_size, - str(init_file), - combine_format, - gate_up_clamp, - ), - nprocs=world_size, - join=True, - ) - - -@pytest.mark.L1 -@pytest.mark.gpu_exclusive -def test_noncontiguous_ep2_fixed_training_matches_independent_reference( - tmp_path, -): - global_world_size = 4 - _require_distributed_sm107(global_world_size) - os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") - init_file = tmp_path / "backward_two_noncontiguous_ep2.init" - mp.spawn( - _distributed_subgroup_backward_reference_worker, - args=(global_world_size, str(init_file)), - nprocs=global_world_size, - join=True, - ) - - -@pytest.mark.L1 -@pytest.mark.gpu_exclusive -@pytest.mark.parametrize( - "case", - [ - pytest.param( - SimpleNamespace( - combine_format="bf16", - drop_on_overflow=True, - max_recv_size=1, - replay_count=20, + def clone_native_tensor(tensor): + return torch.empty_strided( + tensor.shape, + tensor.stride(), + dtype=tensor.dtype, + device=tensor.device, + ).copy_(tensor) + + native_forward = MoeEpNativeForwardWeights( + fc1=MoeEpNativeWeight( + clone_native_tensor(packed_forward.fc1.payload), + clone_native_tensor(packed_forward.fc1.scale), + MoeEpNativeWeightLayout.FORWARD_FC1_GATE_UP_INTERLEAVED_32_V1, ), - id="bf16-drop", - ), - pytest.param( - SimpleNamespace( - combine_format="mxfp8", - drop_on_overflow=True, - max_recv_size=1, - replay_count=20, + fc2=MoeEpNativeWeight( + clone_native_tensor(packed_forward.fc2.payload), + clone_native_tensor(packed_forward.fc2.scale), + MoeEpNativeWeightLayout.FORWARD_FC2_K_MAJOR_V1, ), - id="mxfp8-drop", - ), - pytest.param( - SimpleNamespace( - combine_format="bf16", - drop_on_overflow=False, - max_recv_size=2, - replay_count=2, + ) + native_backward = MoeEpNativeBackwardWeights( + w2_transpose=MoeEpNativeWeight( + clone_native_tensor(packed_backward.w2_transpose.payload), + clone_native_tensor(packed_backward.w2_transpose.scale), + MoeEpNativeWeightLayout.BACKWARD_W2_TRANSPOSE_V1, ), - id="bf16-error-no-overflow", - ), - ], -) -def test_fixed_training_resources_ep1_cuda_graph_replay(case): - device = _sm107_device() - if case.drop_on_overflow: - args0, grad0 = _fixed_training_case(device) - topk_idx1 = args0[3].clone() - topk_idx1[0, 0] = 1 - inputs = ( - (args0, grad0), - ( - ( - args0[0].clone(), - args0[1], - args0[2], - topk_idx1, - args0[4].clone(), - ), - grad0.clone(), + w1_transpose=MoeEpNativeWeight( + clone_native_tensor(packed_backward.w1_transpose.payload), + clone_native_tensor(packed_backward.w1_transpose.scale), + MoeEpNativeWeightLayout.BACKWARD_W1_TRANSPOSE_GATE_UP_INTERLEAVED_32_V1, ), ) - else: - inputs = (_fixed_training_drop_overflow_case(device),) - references = tuple( - _fixed_training_reference( - args, - grad_output, - combine_format=case.combine_format, - gate_up_clamp=None, - ) - for args, grad_output in inputs - ) - - with MoeEp( - num_experts=2, - hidden_size=128, - intermediate_size=256, - top_k=2, - max_tokens_per_rank=inputs[0][0][0].shape[0], - max_recv_size_per_rank=case.max_recv_size, - drop_on_overflow=case.drop_on_overflow, - combine_format=case.combine_format, - token_padding_size=128, - ) as op: - resources = op.prepare_training_resources( - _fixed_training_weights(inputs[0][0]), - slot_count=len(inputs), - lane_count=1, - ) - lane = resources.lanes[0] - batch = tuple((slot, args, grad_output) for slot, (args, grad_output) in zip(resources.slots, inputs)) - - def assert_batch(actuals): - for actual, (args, _), reference in zip( - actuals, - inputs, - references, - ): - assert actual.overflow.shape == (1,) - assert actual.overflow.dtype == torch.int32 - assert actual.overflow.eq(0).all() - _assert_fixed_training_matches_reference( - (actual.y, actual.dx, actual.dprob, actual.wgrads), - reference, - args[3], - ) - - eager_actuals = _run_fixed_training_batch(resources, lane, batch) - torch.cuda.synchronize(device) - assert_batch(eager_actuals) - - stream = torch.cuda.Stream(device=device) - stream.wait_stream(torch.cuda.current_stream(device)) - captured = _capture_fixed_training_batch( - resources, - lane, - batch, - stream, - ) - # In error mode, each replay executes the captured torch._assert_async - # with a false overflow condition; stable public nodes prove reuse. - for _ in range(case.replay_count): - captured.graph.replay() - torch.cuda.synchronize(device) - assert captured.public_pointers == tuple(_training_public_pointers(actual) for actual in captured.actuals) - assert_batch(captured.actuals) - - -@pytest.mark.L1 -@pytest.mark.gpu_exclusive -def test_fixed_training_resources_ep1_two_shape_cuda_graph_contract(): - device = _sm107_device() - args, grad_large = _fixed_training_case(device) - max_tokens = int(args[0].shape[0]) - small_tokens = max_tokens - 2 - assert 0 < small_tokens < max_tokens - - large = SimpleNamespace( - name="large", - activation=args[0], - topk_idx=args[3], - topk_weights=args[4], - grad_output=grad_large, - ) - small = SimpleNamespace( - name="small", - activation=args[0][:small_tokens].clone(), - topk_idx=args[3][:small_tokens].clone(), - topk_weights=args[4][:small_tokens].clone(), - grad_output=grad_large[:small_tokens].clone(), - ) - assert all(getattr(large, name).data_ptr() != getattr(small, name).data_ptr() for name in ("activation", "topk_idx", "topk_weights", "grad_output")) - - weights = _fixed_training_weights(args) - weight_source_pointers = _training_weight_source_pointers(weights) - - with MoeEp( - num_experts=2, - hidden_size=128, - intermediate_size=256, - top_k=2, - max_tokens_per_rank=max_tokens, - max_recv_size_per_rank=1, - drop_on_overflow=True, - ) as op: - resources = op.prepare_training_resources( - weights, - slot_count=1, - lane_count=1, - ) - slot = resources.slots[0] - lane = resources.lanes[0] - - def case_args(case): - return ( - case.activation, - weights.forward_fc1, - weights.forward_fc2, - case.topk_idx, - case.topk_weights, - ) + assert native_forward.fc1.payload.data_ptr() != packed_forward.fc1.payload.data_ptr() + assert native_backward.w1_transpose.scale.data_ptr() != packed_backward.w1_transpose.scale.data_ptr() - def independent_reference(case): - return _fixed_training_reference( - case_args(case), - case.grad_output, - combine_format="bf16", - gate_up_clamp=None, - ) - - def warmup(case) -> None: - actual = _run_fixed_training_batch( - resources, + def run(): + output = op.training_forward( lane, - ((slot, case_args(case), case.grad_output),), - )[0] - torch.cuda.synchronize(device) - assert actual.overflow.eq(0).all(), f"{case.name} warmup overflowed" - _assert_fixed_training_matches_reference( - (actual.y, actual.dx, actual.dprob, actual.wgrads), - independent_reference(case), - case.topk_idx, + activation, + topk_idx, + topk_weights, + weights=native_forward, + out=forward_out, ) - - # Compile each static token-count specialization on the same resources, - # slot, and lane before either capture. - warmup(large) - warmup(small) - - capture_stream = torch.cuda.Stream(device=device) - capture_stream.wait_stream(torch.cuda.current_stream(device)) - - def capture(case): - # The shared sequence records refresh so replay observes in-place - # updates to all four bound source packs. - captured = _capture_fixed_training_batch( - resources, + grad_activation, dprob, operands = op.training_backward( lane, - ((slot, case_args(case), case.grad_output),), - capture_stream, - ) - return SimpleNamespace( - case=case, - graph=captured.graph, - actual=captured.actuals[0], - public_pointers=captured.public_pointers[0], - source_pointers=_training_source_pointers(case), - ) - - large_graph = capture(large) - small_graph = capture(small) - slot_views = resources._owner.views( - slot=slot.index, - lane=lane.index, - token_count=max_tokens, - ).slot - - def replay_and_check(captured): - _prefill_training_graph_sentinels(slot_views, captured.actual) - captured.graph.replay() - torch.cuda.synchronize(device) - - assert captured.actual.overflow.eq(0).all() - assert _training_public_pointers(captured.actual) == captured.public_pointers - assert _training_source_pointers(captured.case) == captured.source_pointers - assert _training_weight_source_pointers(weights) == weight_source_pointers - _assert_fixed_training_matches_reference( - ( - captured.actual.y, - captured.actual.dx, - captured.actual.dprob, - captured.actual.wgrads, - ), - independent_reference(captured.case), - captured.case.topk_idx, - ) - # The dense-dW check above decodes every expert segment and rejects - # nonzero expert padding, nonzero data capacity tails, or - # non-neutral scale tails left by the sentinels. - _assert_training_graph_tails_are_reset( - slot_views, - captured.actual, - token_count=int(captured.case.activation.shape[0]), - capacity=max_tokens, - ) - return captured.actual.y.clone() - - # The two graphs alias one persistent slot. Each replay must therefore - # fully replace the other shape's routing, gradients, and WGrad state. - for captured in (large_graph, small_graph, large_graph): - replay_and_check(captured) - - small_source_pointers = _training_source_pointers(small) - small.activation.mul_(-0.5) - small.topk_idx.fill_(-1) - small.topk_idx[0, 0] = 1 - small.topk_weights.zero_() - small.topk_weights[0, 0] = 0.625 - small.grad_output.mul_(-0.75) - assert _training_source_pointers(small) == small_source_pointers - - for captured in (small_graph, large_graph, small_graph): - replay_and_check(captured) - - old_large_y = replay_and_check(large_graph) - old_weight_values = _training_weight_source_values(weights) - generator = torch.Generator(device=device).manual_seed(20260829) - new_fc1 = ( - torch.randn( - weights.forward_fc1.logical_shape, - generator=generator, - device=device, + grad_output, + topk_idx, + topk_weights, + weights=native_backward, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, ) - / 16 - ) - new_fc2 = ( - torch.randn( - weights.forward_fc2.logical_shape, - generator=generator, - device=device, + return output, grad_activation, dprob, operands + + def assert_matches(result): + output, grad_activation, dprob, operands = result + assert output.data_ptr() == forward_out.output.data_ptr() + assert grad_activation.data_ptr() == backward_out.grad_activation.data_ptr() + assert dprob.data_ptr() == backward_out.dprob.data_ptr() + _assert_matches_reference(output, expected[0]) + _assert_backward_matches( + (grad_activation, dprob), + (expected[1], expected[2]), + topk_idx, ) - / 16 - ) - replacement = _fixed_training_weights( - ( - large.activation, - new_fc1, - new_fc2, - large.topk_idx, - large.topk_weights, + _assert_wgrads_match_reference( + operands, + expected[3], + weight_interleave_size=32, ) - ) - _copy_training_weight_sources_(weights, replacement) - - assert _training_weight_source_pointers(weights) == weight_source_pointers - _assert_training_weight_sources_changed(weights, old_weight_values) - - new_large_y = replay_and_check(large_graph) - assert not torch.equal(new_large_y, old_large_y) - - -@pytest.mark.L1 -@pytest.mark.gpu_exclusive -def test_fixed_training_resources_ep1_drop_overflow_boundary_and_graph_transitions(): - device = _sm107_device() - args, grad_output = _fixed_training_drop_overflow_case(device) - assert args[0].shape[0] == 1 - assert args[3].detach().cpu().tolist() == [[0, 1]] - - references = { - expected_overflow: _fixed_training_drop_overflow_reference( - args, - grad_output, - drop_expert1=bool(expected_overflow), - ) - for expected_overflow in (0, 1) - } - - def assert_result(actual, expected_overflow): - expected, reference_topk_idx = references[expected_overflow] - _assert_fixed_training_drop_overflow_result( - actual, - expected, - reference_topk_idx, - expected_overflow=expected_overflow, - ) - # The graph warmup below covers maxrecv=1 overflow; exercise the exact - # non-overflow boundary separately here. - with MoeEp( - num_experts=2, - hidden_size=128, - intermediate_size=256, - top_k=2, - max_tokens_per_rank=1, - max_recv_size_per_rank=2, - drop_on_overflow=True, - token_padding_size=128, - ) as op: - resources = op.prepare_training_resources( - _fixed_training_weights(args), - slot_count=1, - lane_count=1, - ) - slot = resources.slots[0] - lane = resources.lanes[0] - actual = _run_fixed_training_batch( - resources, - lane, - ((slot, args, grad_output),), - )[0] + eager = run() + grouped_wgrads = _dense_wgrads_from_grouped_kernel(eager[3]) torch.cuda.synchronize(device) - - assert_result(actual, 0) - assert args[3][0, 1].eq(1) - assert actual.wgrads.valid_route_counts.detach().cpu().tolist() == [1, 1] - assert actual.wgrads.expert_offsets.detach().cpu().tolist() == [128, 256] - - overflow_routing = args[3].clone() - expert0_only_routing = overflow_routing.clone() - expert0_only_routing[0, 1] = -1 - routing_pointer = args[3].data_ptr() - - with MoeEp( - num_experts=2, - hidden_size=128, - intermediate_size=256, - top_k=2, - max_tokens_per_rank=1, - max_recv_size_per_rank=1, - drop_on_overflow=True, - token_padding_size=128, - ) as op: - resources = op.prepare_training_resources( - _fixed_training_weights(args), - slot_count=1, - lane_count=1, - ) - slot = resources.slots[0] - lane = resources.lanes[0] - batch = ((slot, args, grad_output),) - - # Compile the fixed T=1 specialization and validate overflow eagerly - # before capturing the same forward/backward/finalize/WGrad sequence. - warmup = _run_fixed_training_batch(resources, lane, batch)[0] - grouped_outputs = _allocate_dense_grouped_wgrad_outputs(warmup.wgrads) - _dense_wgrads_from_grouped_kernel( - warmup.wgrads, - wgrad_tensors=grouped_outputs, + assert_matches(eager) + expected_fc1_wgrad, expected_fc2_wgrad = expected[3].dense_wgrads() + _assert_grouped_wgrads_match_reference( + grouped_wgrads, + ( + _interleave_fc1_wgrad(expected_fc1_wgrad), + expected_fc2_wgrad, + ), + reference_name="the independent PyTorch MXFP8 reference", ) - torch.cuda.synchronize(device) - assert_result(warmup, 1) - capture_stream = torch.cuda.Stream(device=device) - capture_stream.wait_stream(torch.cuda.current_stream(device)) - captured = _capture_fixed_training_batch( - resources, - lane, - batch, - capture_stream, - grouped_wgrad_outputs=(grouped_outputs,), + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run() + output_pointers = ( + forward_out.output.data_ptr(), + backward_out.grad_activation.data_ptr(), + backward_out.dprob.data_ptr(), ) - graph_actual = captured.actuals[0] - graph_grouped_wgrads = captured.grouped_wgrads[0] - grouped_output_pointers = tuple(output.data_ptr() for output in grouped_outputs) - - for routing, expected_overflow in ( - (overflow_routing, 1), - (expert0_only_routing, 0), - (overflow_routing, 1), - ): - args[3].copy_(routing) - for output in grouped_outputs: - output.fill_(float("nan")) - assert args[3].data_ptr() == routing_pointer - expected = _fixed_training_drop_overflow_reference( - args, - grad_output, - drop_expert1=bool(expected_overflow), - ) - captured.graph.replay() + native_weight_pointers = ( + native_forward.fc1.payload.data_ptr(), + native_forward.fc1.scale.data_ptr(), + native_forward.fc2.payload.data_ptr(), + native_forward.fc2.scale.data_ptr(), + native_backward.w2_transpose.payload.data_ptr(), + native_backward.w2_transpose.scale.data_ptr(), + native_backward.w1_transpose.payload.data_ptr(), + native_backward.w1_transpose.scale.data_ptr(), + ) + for _ in range(2): + graph.replay() torch.cuda.synchronize(device) - _assert_fixed_training_drop_overflow_result( - graph_actual, - *expected, - expected_overflow=expected_overflow, - ) - assert ( - tuple(output.data_ptr() for output in graph_grouped_wgrads) - == grouped_output_pointers + assert output_pointers == ( + forward_out.output.data_ptr(), + backward_out.grad_activation.data_ptr(), + backward_out.dprob.data_ptr(), ) - _assert_grouped_wgrads_match_reference( - graph_grouped_wgrads, - expected[0][3].dense_wgrads(), - reference_name="the independent PyTorch MXFP8 graph reference", + assert native_weight_pointers == ( + native_forward.fc1.payload.data_ptr(), + native_forward.fc1.scale.data_ptr(), + native_forward.fc2.payload.data_ptr(), + native_forward.fc2.scale.data_ptr(), + native_backward.w2_transpose.payload.data_ptr(), + native_backward.w2_transpose.scale.data_ptr(), + native_backward.w1_transpose.payload.data_ptr(), + native_backward.w1_transpose.scale.data_ptr(), ) - _assert_grouped_wgrads_match_reference( - graph_grouped_wgrads, - _dense_wgrads_from_operands(graph_actual.wgrads), - reference_name="the decoded captured production operand bundle", - close_kwargs={"rtol": 0.1, "atol": 0.1}, - ) - assert all(torch.isfinite(output).all() for output in graph_grouped_wgrads) - if expected_overflow: - assert graph_grouped_wgrads[0][1].eq(0).all() - assert graph_grouped_wgrads[1][1].eq(0).all() - assert captured.public_pointers[0] == _training_public_pointers(graph_actual) + assert_matches(captured) + finally: + op.close() diff --git a/test/python/moe_ep/test_moe_ep_cutedsl.py b/test/python/moe_ep/test_moe_ep_cutedsl.py index 2b9c58536..48800758a 100644 --- a/test/python/moe_ep/test_moe_ep_cutedsl.py +++ b/test/python/moe_ep/test_moe_ep_cutedsl.py @@ -4,17 +4,16 @@ """CUTLASS DSL version-gate tests for Rubin MegaMoE.""" import pytest +import torch @pytest.mark.L0 -@pytest.mark.parametrize("version", ["4.5.0", "4.6.1", "4.7.0"]) def test_rubin_cutedsl_gate_rejects_public_wheels_below_4_8( monkeypatch, - version, ): from cudnn.moe_ep._megamoe_backend.mxfp8 import _cutedsl - monkeypatch.setattr(_cutedsl, "_public_cutedsl_version", lambda: version) + monkeypatch.setattr(_cutedsl, "_public_cutedsl_version", lambda: "4.7.0") with pytest.raises( RuntimeError, @@ -53,6 +52,7 @@ def test_rubin_prepare_gates_before_cuda_initialization( function_name, ): module = __import__(module_name, fromlist=[function_name]) + from cudnn.moe_ep._megamoe_backend.mxfp8 import _compile_common class GateReached(RuntimeError): pass @@ -60,7 +60,60 @@ class GateReached(RuntimeError): def reject(): raise GateReached - monkeypatch.setattr(module, "require_rubin_cutedsl", reject) + monkeypatch.setattr(_compile_common, "require_rubin_cutedsl", reject) with pytest.raises(GateReached): getattr(module, function_name)(None, None, None) + + +@pytest.mark.L0 +@pytest.mark.parametrize("context", ["forward", "backward"]) +def test_rubin_environment_errors_include_compile_context( + monkeypatch, + context, +): + from cudnn.moe_ep._megamoe_backend.mxfp8 import _compile_common + + monkeypatch.setattr(_compile_common, "require_rubin_cutedsl", lambda: None) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + monkeypatch.setattr( + torch.cuda, + "get_device_capability", + lambda device: (10, 0), + ) + + with pytest.raises( + RuntimeError, + match=rf"Rubin MXFP8 {context} preparation", + ): + _compile_common._prepare_rubin_environment( + torch.device("cuda", 0), + 2, + context=context, + ) + + +@pytest.mark.L0 +def test_rubin_environment_rejects_incompatible_arch_override( + monkeypatch, +): + from cudnn.moe_ep._megamoe_backend.mxfp8 import _compile_common + + monkeypatch.setattr(_compile_common, "require_rubin_cutedsl", lambda: None) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + monkeypatch.setattr( + torch.cuda, + "get_device_capability", + lambda device: (10, 7), + ) + monkeypatch.setenv("CUTE_DSL_ARCH", "sm_100a") + + with pytest.raises( + RuntimeError, + match=r"CUTE_DSL_ARCH.*forward.*sm_100a", + ): + _compile_common._prepare_rubin_environment( + torch.device("cuda", 0), + 2, + context="forward", + ) diff --git a/test/python/moe_ep/test_moe_ep_forward.py b/test/python/moe_ep/test_moe_ep_forward.py index f03eacf1c..91b074002 100644 --- a/test/python/moe_ep/test_moe_ep_forward.py +++ b/test/python/moe_ep/test_moe_ep_forward.py @@ -459,12 +459,8 @@ def test_moe_ep_rejects_interleaved_plain_fc1_weight(): config = _forward_config() activation = torch.zeros((1, config["hidden_size"])) - fc1_weight = torch.zeros( - (config["num_experts"], config["hidden_size"], 2 * config["intermediate_size"]) - ) - fc2_weight = torch.zeros( - (config["num_experts"], config["intermediate_size"], config["hidden_size"]) - ) + fc1_weight = torch.zeros((config["num_experts"], config["hidden_size"], 2 * config["intermediate_size"])) + fc2_weight = torch.zeros((config["num_experts"], config["intermediate_size"], config["hidden_size"])) topk_idx = torch.zeros((1, config["top_k"]), dtype=torch.int32) topk_weights = torch.ones((1, config["top_k"])) with MoeEp(**config, weight_interleave_size=32) as op: @@ -891,32 +887,6 @@ def test_single_gpu_stress_and_cuda_graph_replay(): _replay_cuda_graph(op, args, original_topk_idx, expected, device) -@pytest.mark.L1 -@pytest.mark.gpu_exclusive -def test_nondefault_tuning_warmup_and_cuda_graph_replay(): - from cudnn import MoeEp, MoeEpTuningConfig - - device = _sm107_device() - args = make_forward_inputs(device) - original_topk_idx = args[3].clone() - expected = _reference_forward(args) - tuning = MoeEpTuningConfig( - token_back_mode="reuse_dispatch_warps", - epi_flag_batch=(2, 2), - token_in_flag_batch=2, - group_hint=128, - ) - - with MoeEp(**_forward_config(), tuning=tuning) as op: - _replay_cuda_graph( - op, - args, - original_topk_idx, - expected, - device, - ) - - @pytest.mark.L0 def test_forward_mxfp8_combine_is_direct_fp32(): generator = torch.Generator().manual_seed(20260819) @@ -1141,7 +1111,7 @@ def test_reference_mxfp8_inputs_bf16_combine_matches_naive( @pytest.mark.L0 def test_reference_interleaved_fc1_matches_logical_fc1(): - from cudnn import BlockScaledTensor + from moe_ep.moe_ep_reference import BlockScaledTensor torch.manual_seed(29) experts, tokens, hidden, intermediate = 2, 3, 128, 128 @@ -1163,11 +1133,7 @@ def test_reference_interleaved_fc1_matches_logical_fc1(): def interleave_last(tensor): shape = tensor.shape - return ( - tensor.view(*shape[:-1], 2, intermediate // 32, 32) - .transpose(-3, -2) - .reshape(shape) - ) + return tensor.view(*shape[:-1], 2, intermediate // 32, 32).transpose(-3, -2).reshape(shape) interleaved_fc1 = BlockScaledTensor( data=interleave_last(logical_fc1.data), diff --git a/test/python/moe_ep/test_moe_ep_multinode.py b/test/python/moe_ep/test_moe_ep_multinode.py index fe34e6171..568537af1 100644 --- a/test/python/moe_ep/test_moe_ep_multinode.py +++ b/test/python/moe_ep/test_moe_ep_multinode.py @@ -17,10 +17,7 @@ _run_backward_reference_case, _run_forward_output_case, ) -from moe_ep.moe_ep_test_support import ( - _fixed_training_weights, - make_distributed_forward_inputs, -) +from moe_ep.moe_ep_test_support import make_distributed_forward_inputs pytestmark = [ pytest.mark.L1, @@ -237,12 +234,6 @@ def test_mxfp8_forward_multinode_matches_reference( "bf16", id="backward-ep16-world16-bf16", ), - pytest.param( - 16, - 16, - "mxfp8", - id="backward-ep16-world16-mxfp8", - ), pytest.param( 32, 32, @@ -257,7 +248,7 @@ def test_mxfp8_forward_multinode_matches_reference( ), ], ) -def test_fixed_training_resources_multinode_match_independent_reference( +def test_stateless_training_multinode_matches_independent_reference( torchrun_world, ep_size, required_world_size, @@ -302,10 +293,6 @@ def test_training_prepare_multinode_rejects_rank_abi_mismatch( from cudnn import MoeEp - # A fixed helper rank gives every process byte-identical, locally valid - # weight packs. Only the locally valid lane count differs. - args = make_distributed_forward_inputs(0, 8, world.device) - weights = _fixed_training_weights(args) op = MoeEp( num_experts=16, hidden_size=128, @@ -316,15 +303,15 @@ def test_training_prepare_multinode_rejects_rank_abi_mismatch( max_recv_size_per_rank=3, drop_on_overflow=True, combine_format="bf16", + weight_interleave_size=32, ) caught_error = None try: lane_count = rank_zero_lane_count if world.rank == 0 else other_lane_count try: - op.prepare_training_resources( - weights, - slot_count=1, + op.prepare_training( lane_count=lane_count, + device=world.device, ) except Exception as error: caught_error = error From 6d7f723eb80ea4376e3e0cd8c57635b2818c0890 Mon Sep 17 00:00:00 2001 From: zhibinz Date: Wed, 2 Sep 2026 08:30:28 -0700 Subject: [PATCH 30/31] feature: sync Rubin MegaMoE CuTeDSL sources Align the vendored training kernels and peer mapping behavior with upstream while keeping compatibility tests valid across supported DSL versions. --- .../_megamoe_backend/cutedsl_src/VENDOR.md | 10 +- .../nvlink_domain/symmetric_buffer.py | 136 ++---- .../cutedsl_src/helpers/device_workspace.py | 12 +- .../cutedsl_src/helpers/ptx_helpers.py | 42 ++ .../cutedsl_src/helpers/utils.py | 12 + .../mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py | 63 ++- .../bwd_dglu/dglu_mxfp8_mega_moe_kernel.py | 18 +- .../mega/fwd_glu/glu_mxfp8_col_requant.py | 390 ++++++++++++++---- .../kernel_src/schedulers/fc12_mapping.py | 35 +- test/python/moe_ep/test_moe_ep_forward.py | 9 +- 10 files changed, 482 insertions(+), 245 deletions(-) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR.md b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR.md index a4e366e2e..db8fc6f97 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR.md +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/VENDOR.md @@ -9,14 +9,16 @@ documented in the parent backend `README.md`. - **Project**: `cutedsl_megamoe` (NVIDIA-internal repository; URL omitted). - **Source tree**: `cutedsl_megamoe/next/sources`. - **Current synchronized commit**: - `5b89819cb16069dfe20a1a0ba0778d35cb428352`. + `9b15c450e2d19472bdfaae37489317f029beb01c`. - **Earlier import points**: - base forward: `882c83e2ce4086c3cd4211fc5a2296143c5e2aea`; - selected forward updates and backward dGLU: - `92dd334af2eeedb36087834354b58ace08e880c6`. -- **Last synced**: 2026-08-28. Earlier imports occurred on 2026-08-11, - 2026-08-17, 2026-08-20, and 2026-08-24. + `92dd334af2eeedb36087834354b58ace08e880c6`; + - forward column-quantization output-layout updates: + `5b89819cb16069dfe20a1a0ba0778d35cb428352`. +- **Last synced**: 2026-09-02. Earlier imports occurred on 2026-08-11, + 2026-08-17, 2026-08-20, 2026-08-24, and 2026-08-28. - **Vendored subset**: the recursive Python import closure required by Rubin SM107 training MegaMoE forward GLU, optional forward MXFP8 column requantization, and backward dGLU. diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py index c166b817c..58186c7fd 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/communication/nvlink_domain/symmetric_buffer.py @@ -6,15 +6,13 @@ from dataclasses import dataclass from typing import Any, Optional +from packaging.version import Version + import cutlass import cutlass.cute as cute from cutlass._mlir import ir from cutlass._mlir.dialects import arith, llvm -from cutlass.base_dsl.dsl import ( - extract_mlir_values, - get_mlir_types, - new_from_mlir_values, -) +from cutlass.base_dsl.dsl import extract_mlir_values, get_mlir_types, new_from_mlir_values from cutlass.base_dsl.runtime.jit_arg_adapters import JitArgAdapterRegistry from cutlass.base_dsl.typing import get_c_pointers from cutlass.cute.typing import AddressSpace @@ -27,13 +25,14 @@ MLIR_DYNAMIC_INDEX = -(2**31) -_byval_rank_limit = 16 +_dsl_release = Version(Version(cutlass.__version__).base_version) +_grid_constant_width_is_free = _dsl_release < Version("4.0.0") or _dsl_release >= Version("4.7.0") +_byval_rank_limit = 2**31 if _grid_constant_width_is_free else 16 -def _byval_struct_type() -> Any: - return ir.Type.parse( - f"!llvm.struct<(array<{_byval_rank_limit} x i64>)>" - ) +def _byval_struct_type(rank_count: int) -> Any: + words = rank_count if _grid_constant_width_is_free else _byval_rank_limit + return ir.Type.parse(f"!llvm.struct<(array<{words} x i64>)>") @dataclass(frozen=True) @@ -46,10 +45,7 @@ class SymmetricBufferDevice: def __extract_mlir_values__(self) -> list: return [self.value] - def __new_from_mlir_values__( - self, - values: list, - ) -> "SymmetricBufferDevice": + def __new_from_mlir_values__(self, values: list) -> "SymmetricBufferDevice": return SymmetricBufferDevice(values[0], self.max_ranks) def __get_mlir_types__(self) -> list: @@ -63,19 +59,14 @@ def __extract_mlir_attributes__(self) -> list: ir.DictAttr.get( { "cute_nvgpu.grid_constant": ir.UnitAttr.get(), - "llvm.byval": ir.TypeAttr.get(_byval_struct_type()), + "llvm.byval": ir.TypeAttr.get(_byval_struct_type(self.max_ranks)), } ) ] return [ir.DictAttr.get({})] @cute.jit - def map( - self, - local_address: Int64, - destination_rank: Int32, - byte_offset: Int64 = Int64(0), - ) -> Int64: + def map(self, local_address: Int64, destination_rank: Int32, byte_offset: Int64 = Int64(0)) -> Int64: if cutlass.const_expr(self.max_ranks <= _byval_rank_limit): i64_type = ir.Type.parse("i64") offset_pointer = llvm.getelementptr( @@ -88,32 +79,17 @@ def map( ) peer_offset = Int64(llvm.load(i64_type, offset_pointer)) else: - peer_offset = Int64( - llvm.extractelement( - self.value, - destination_rank.ir_value(), - ) - ) + peer_offset = Int64(llvm.extractelement(self.value, destination_rank.ir_value())) return local_address + peer_offset + byte_offset @cute.jit - def map_pointer( - self, - pointer, - destination_rank: Int32, - byte_alignment: Optional[int] = None, - ): + def map_pointer(self, pointer, destination_rank: Int32, byte_alignment: Optional[int] = None): if cutlass.const_expr(pointer.memspace != AddressSpace.gmem): - raise ValueError( - "Only GMEM pointers can be mapped to a symmetric peer." - ) + raise ValueError("Only GMEM pointers can be mapped to a symmetric peer.") if cutlass.const_expr(byte_alignment is None): byte_alignment = pointer.max_alignment return cute.make_ptr( - pointer.dtype, - self.map(pointer.toint(), destination_rank), - pointer.memspace, - assumed_align=byte_alignment, + pointer.dtype, self.map(pointer.toint(), destination_rank), pointer.memspace, assumed_align=byte_alignment ) @@ -134,64 +110,27 @@ def _as_int64(value) -> Int64: def make_device_object(self, *, loc=None, ip=None) -> SymmetricBufferDevice: offsets = tuple(self.offsets) if len(offsets) != self.max_ranks: - raise ValueError( - f"Expected {self.max_ranks} peer offsets, got {len(offsets)}." - ) + raise ValueError(f"Expected {self.max_ranks} peer offsets, got {len(offsets)}.") if self.max_ranks <= _byval_rank_limit: pointer_type = ir.Type.parse("!llvm.ptr") - struct_type = _byval_struct_type() + struct_type = _byval_struct_type(self.max_ranks) i64_type = ir.Type.parse("i64") - one = arith.constant( - value=ir.IntegerAttr.get(i64_type, 1), - result=i64_type, - loc=loc, - ip=ip, - ) - buffer = llvm.alloca( - res=pointer_type, - elem_type=struct_type, - array_size=one, - alignment=64, - loc=loc, - ip=ip, - ) + one = arith.constant(value=ir.IntegerAttr.get(i64_type, 1), result=i64_type, loc=loc, ip=ip) + buffer = llvm.alloca(res=pointer_type, elem_type=struct_type, array_size=one, alignment=64, loc=loc, ip=ip) for index, offset in enumerate(offsets): slot = llvm.getelementptr( - pointer_type, - buffer, - [], - [index], - i64_type, - no_wrap_flags="None", - loc=loc, - ip=ip, - ) - llvm.store( - self._as_int64(offset).ir_value(), - slot, - loc=loc, - ip=ip, + pointer_type, buffer, [], [index], i64_type, no_wrap_flags="None", loc=loc, ip=ip ) + llvm.store(self._as_int64(offset).ir_value(), slot, loc=loc, ip=ip) return SymmetricBufferDevice(buffer, self.max_ranks) i32_type = ir.Type.parse("i32") vector_type = ir.Type.parse(f"vector<{self.max_ranks}xi64>") vector = llvm.mlir_zero(vector_type, loc=loc, ip=ip) for index, offset in enumerate(offsets): - element_index = arith.constant( - value=ir.IntegerAttr.get(i32_type, index), - result=i32_type, - loc=loc, - ip=ip, - ) - vector = llvm.insertelement( - vector, - self._as_int64(offset).ir_value(), - element_index, - loc=loc, - ip=ip, - ) + element_index = arith.constant(value=ir.IntegerAttr.get(i32_type, index), result=i32_type, loc=loc, ip=ip) + vector = llvm.insertelement(vector, self._as_int64(offset).ir_value(), element_index, loc=loc, ip=ip) return SymmetricBufferDevice(vector, self.max_ranks) @@ -201,15 +140,8 @@ def __init__(self, argument: SymmetricBufferHost) -> None: self._argument = argument offsets = tuple(argument.offsets) if len(offsets) != int(argument.max_ranks): - raise ValueError( - f"Expected {int(argument.max_ranks)} peer offsets, " - f"got {len(offsets)}." - ) - self._fields = ( - Int64(argument.base_address), - *(Int64(offset) for offset in offsets), - Int32(argument.rank), - ) + raise ValueError(f"Expected {int(argument.max_ranks)} peer offsets, got {len(offsets)}.") + self._fields = (Int64(argument.base_address), *(Int64(offset) for offset in offsets), Int32(argument.rank)) def __c_pointers__(self) -> list[Any]: pointers: list[Any] = [] @@ -229,25 +161,15 @@ def __extract_mlir_values__(self) -> list[ir.Value]: values.extend(extract_mlir_values(field)) return values - def __new_from_mlir_values__( - self, - values: list[ir.Value], - ) -> SymmetricBufferHost: + def __new_from_mlir_values__(self, values: list[ir.Value]) -> SymmetricBufferHost: value_index = 0 rebuilt = [] for field in self._fields: field_value_count = len(get_mlir_types(field)) - rebuilt.append( - new_from_mlir_values( - field, - values[value_index : value_index + field_value_count], - ) - ) + rebuilt.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) value_index += field_value_count if value_index != len(values): - raise ValueError( - f"Consumed {value_index} MLIR values, got {len(values)}." - ) + raise ValueError(f"Consumed {value_index} MLIR values, got {len(values)}.") result = object.__new__(SymmetricBufferHost) object.__setattr__(result, "base_address", rebuilt[0]) diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py index 252adfc16..0e75bc695 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/device_workspace.py @@ -14,6 +14,7 @@ from .utils import ( ceil_div, cosize_from_shape_stride_tuples, + flatten_shape_stride, is_nested_shape, is_power_of_two, ordered_stride, @@ -29,6 +30,15 @@ _reset_order = {"tail_reset": 0, "zero_on_first_allocate": 1, "data": 2} +def _footprint_from_shape_stride(shape: Tuple, stride: Tuple) -> int: + """Elements a region must own, as opposed to the ones its layout can address.""" + leaf_pairs = flatten_shape_stride(shape, stride) if shape else [] + claimed = max((size * step for size, step in leaf_pairs), default=1) + # Layouts whose leaves overlap can address past what any single leaf tiles, + # so the cosize stays a floor. + return int(max(claimed, cosize_from_shape_stride_tuples(shape, stride))) + + @dataclasses.dataclass(frozen=True) class DeviceRegion: """One typed region in a local or symmetric GMEM workspace.""" @@ -142,7 +152,7 @@ def finalize(self) -> None: if stride is None: raise RuntimeError(f"Region {region.name!r} stride was not resolved.") cosize = cosize_from_shape_stride_tuples(region.shape, stride) - nbytes = (cosize * int(region.dtype.width) + 7) // 8 + nbytes = (_footprint_from_shape_stride(region.shape, stride) * int(region.dtype.width) + 7) // 8 if region.reset == "tail_reset" and ( region_index == 0 or ordered[region_index - 1].reset != "tail_reset" ): diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py index 9f56235e9..ebf782ca0 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/ptx_helpers.py @@ -41,6 +41,21 @@ def nanosleep(sleep_cycles: int, *, loc: Optional[ir.Location] = None, ip: Optio ) +@dsl_user_op +def exit(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> None: + llvm.inline_asm( + res=None, + operands_=[], + asm_string="exit;", + constraints="", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + @dsl_user_op def read_clock64(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> Int64: """Read the per-SM 64-bit cycle counter.""" @@ -431,6 +446,31 @@ def red_async_add_release_gpu_s32( ) +@dsl_user_op +def red_async_add_release_sys_u32(address: Int64, value: Int32, *, loc=None, ip=None) -> None: + """Fire-and-forget cross-rank counter bump carrying its own release ordering. + + The issuing warp does not wait for the L2/HBM round trip, and the release is + enforced by the memory system, so this replaces an explicit ``membar.sys`` + followed by a relaxed reduction. ``u32`` rather than ``s32`` only because a + counter bump is the same two's-complement add either way and this is the + spelling already proven on this path. + + Operands go through uniform registers, so the address must be warp-uniform: + a fan-out where lanes target different peers cannot use this. + """ + llvm.inline_asm( + None, + [address.ir_value(), value.ir_value()], + "red.async.release.sys.global.add.u32 [$0], $1;", + "l,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + @dsl_user_op def red_add_relaxed_sys_v2_bf16x2( address, value0, value1, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None @@ -559,6 +599,7 @@ def stg_e8m0x8_from_f32( "cp_reduce_async_bulk_add_u32_s2g", "cvt_f32_to_fp8_to_f32", "cvt_f32x4_to_f8x4_pack_i32", + "exit", "lds128_v4_b32", "mbarrier_arrive_expect_tx_on_peer", "movmatrix_b16", @@ -568,6 +609,7 @@ def stg_e8m0x8_from_f32( "red_add_relaxed_sys_s32", "red_add_relaxed_sys_v2_bf16x2", "red_async_add_release_gpu_s32", + "red_async_add_release_sys_u32", "red_add_release_gpu_s32", "red_add_release_sys_s32", "store_i32_to_peer_cluster_smem_async", diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py index 59d78e743..d1349e99b 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/helpers/utils.py @@ -19,6 +19,17 @@ def ceil_div(value: IntegerType, divisor: IntegerType) -> IntegerType: return (value + divisor - 1) // divisor +def padded_expert_rows(token_count: IntegerType, padding_block: IntegerType) -> IntegerType: + """Row span one expert occupies in a block-padded pool. + + The single definition of a pool's per-expert stride. Communication components + write bases derived from it while the FC12 scheduler rebuilds the same bases + from ``expert_sizes`` at runtime; if the two ever disagree the metadata a + kernel reads no longer describes the rows it loads. + """ + return round_up(token_count, padding_block) + + def is_power_of_two(value: int) -> bool: return value > 0 and (value & (value - 1)) == 0 @@ -95,6 +106,7 @@ def product(values: Iterable[IntegerType]) -> IntegerType: "is_nested_shape", "is_power_of_two", "ordered_stride", + "padded_expert_rows", "product", "row_major_stride", "round_up", diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py index 031c753a5..142639b71 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_fc12_kernel.py @@ -731,32 +731,32 @@ def token_comm_hook_kernel_tail(self, token_comm_args, *, warp_idx, lane_idx, ti @cute.jit def __call__( self, - activation: cute.Tensor, # (token_sum_padded, hidden) = grad_out - fc1_weight: cute.Tensor, # (experts, hidden, inter_half) - activation_sf: cute.Tensor, # (token_sum_padded_sf, hidden / sf_vec_size) + activation: cute.Tensor, # (token_sum_padded, hidden) grad_out + fc1_weight: cute.Tensor, # (experts, hidden, intermediate_downproj) + activation_sf: cute.Tensor, # row-SF for activation fc1_weight_sf: cute.Tensor, # dfc2-weight SF - fc1_output: cute.Tensor, # (token_sum_padded, intermediate) - fc1_output_sf: cute.Tensor, # (token_sum_padded_sf, intermediate / sf_vec_size) - fc1_recompute: Optional[cute.Tensor], # (token_sum_padded, inter_half) - fc1_recompute_sf: Optional[cute.Tensor], # (token_sum_padded_sf, inter_half / sf_vec_size) - fc1_col_output: Optional[cute.Tensor], # (token_sum_padded, intermediate) - fc1_col_output_sf: Optional[cute.Tensor], # (sf_row_blocks, intermediate) col-SF - fc2_weight: cute.Tensor, # (experts, intermediate, hidden) - fc2_weight_sf: cute.Tensor, # dfc1-weight SF - fc2_output: cute.Tensor, # (token_sum_padded, hidden) BFloat16 = grad_x - fc1_preact: cute.Tensor, # (token_sum_padded, intermediate) BFloat16 - topk_scores: cute.Tensor, # (token_sum_padded,) Float32 - beta: cute.Tensor, # (experts,) Float32 - dprob: cute.Tensor, # (token_sum_padded,) Float32 - fc1_done_counter: cute.Tensor, # (max_token_block_per_rank,) Int32 - offs: Optional[cute.Tensor] = None, # (experts,) Int32 cumulative end offsets + fc1_output: cute.Tensor, # (token_sum_padded, intermediate_gateup) + fc1_output_sf: cute.Tensor, # (token_sum_padded_sf, intermediate_gateup_sf) + fc1_recompute: Optional[cute.Tensor], # (token_sum_padded, intermediate_downproj) + fc1_recompute_sf: Optional[cute.Tensor], # (intermediate_downproj_padded, col_sf_rows) + fc1_col_output: Optional[cute.Tensor], # (token_sum_padded, intermediate_gateup) + fc1_col_output_sf: Optional[cute.Tensor], # (intermediate_gateup_padded, col_sf_rows) + fc2_weight: cute.Tensor, # (experts, intermediate_gateup, hidden) + fc2_weight_sf: cute.Tensor, # (experts, hidden_padded * intermediate_gateup_sf_padded) + fc2_output: cute.Tensor, # (token_sum_padded, hidden) + fc1_preact: cute.Tensor, # (token_sum_padded, intermediate_gateup) BFloat16 + topk_scores: cute.Tensor, # (token_sum_padded,) Float32 + beta: cute.Tensor, # (experts,) Float32 + dprob: cute.Tensor, # (token_sum_padded,) Float32 + fc1_done_counter: cute.Tensor, # (fc1_ready_slot_count,) Int32 + offs: Optional[cute.Tensor] = None, # unsupported for dGLU; use expert_token_sizes max_active_clusters: cutlass.Constexpr = None, stream: cuda.CUstream = None, norm_const_tensor: Optional[cute.Tensor] = None, global_activation_sf: Optional[cute.Tensor] = None, global_fc1_weight_sf: Optional[cute.Tensor] = None, load_balance_counter: Optional[cute.Tensor] = None, - expert_token_sizes: Optional[cute.Tensor] = None, + expert_token_sizes: Optional[cute.Tensor] = None, # (experts,) valid token counts token_comm_args=None, overflow_flag: cute.Tensor = None, mega_peer_rank_ptr_mapper=None, @@ -857,9 +857,8 @@ def __call__( ), ) - # B_gemm (W2T): reinterpret public C-contiguous (experts, hidden, inter_half) - # as (N=inter_half, K=hidden, L=experts). The stride permutation makes this - # an N-major GEMM operand without staging or moving data. + # B_gemm (fc1 weights): (experts, hidden, intermediate_gateup) with hidden stride-1 (K-major) + # -> (N=intermediate_gateup, K=hidden, L=experts). experts, hidden_b, intermediate_gateup = fc1_weight.shape fc1_weight_gemm = cute.make_tensor( fc1_weight.iterator, @@ -925,9 +924,7 @@ def __call__( ), ) - # GEMM-domain transform for fc2 phase. W1T is public C-contiguous - # (experts, 2 * inter_half, hidden), with its reduction rows already in - # 32-wide gate/up order. Preserve that K ordering and expose hidden as N. + # GEMM-domain transform for fc2 phase ── experts2, intermediate_downproj_b2, hidden_b2 = fc2_weight.shape fc2_weight_gemm = cute.make_tensor( fc2_weight.iterator, @@ -1168,16 +1165,16 @@ def __call__( else: load_balance_counter_ptr = None - # On the MegaMoE path the per-expert sizes come from the Router (device-side), so the - # caller supplies neither offs nor expert_token_sizes. + # dGLU auxiliary layouts require per-expert token counts; cumulative offsets + # alone are insufficient. if cutlass.const_expr(not self.enable_token_comm): - if cutlass.const_expr((offs is None) == (expert_token_sizes is None)): + if cutlass.const_expr(offs is not None): raise ValueError( - "Exactly one of `offs` / `expert_token_sizes` must be " - "non-None. Got offs=" - f"{'set' if offs is not None else 'None'}, " - f"expert_token_sizes=" - f"{'set' if expert_token_sizes is not None else 'None'}." + "`offs` is not supported by dGLU; provide `expert_token_sizes`." + ) + if cutlass.const_expr(expert_token_sizes is None): + raise ValueError( + "`expert_token_sizes` must be provided for dGLU auxiliary layouts." ) self._build_scheduler( diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py index fb03d5ad3..14d0d9137 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/bwd_dglu/dglu_mxfp8_mega_moe_kernel.py @@ -141,11 +141,11 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): fake_arguments = dict( grad_out=fake_tensor(activation_dtype, (tokens, hidden), (1, 0), {0}, 16), grad_out_sf=fake_tensor(sf_dtype, (tokens, self.token_comm.activation_sf_hidden_padded), (1, 0), {0}, 16), - topk_idx=fake_tensor(cutlass.Int64, (tokens, self.num_topk), (1, 0), {0}, 16), + topk_idx=fake_tensor(cutlass.Int32, (tokens, self.num_topk), (1, 0), {0}, 16), topk_weights=fake_tensor(cutlass.Float32, (tokens, self.num_topk), (1, 0), {0}, 4), - fc1_weight=fake_tensor(self.ab_dtype, (experts, hidden, inter_half), (2, 1, 0), {0, 2}, 16), + fc1_weight=fake_tensor(self.ab_dtype, (experts, hidden, inter_half), (2, 0, 1), {0, 2}, 16), fc1_weight_sf=fake_tensor(sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16), - fc2_weight=fake_tensor(self.ab_dtype, (experts, gate_up, hidden), (2, 1, 0), {0, 2}, 16), + fc2_weight=fake_tensor(self.ab_dtype, (experts, gate_up, hidden), (2, 0, 1), {0, 2}, 16), fc2_weight_sf=fake_tensor(sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16), beta=fake_tensor(cutlass.Float32, (experts,), (0,), {0}, 4), fc1_preact=fake_tensor(cutlass.BFloat16, self.get_fc1_preact_shape(), (1, 0), set(), 128), @@ -754,7 +754,7 @@ def _snapshot_grad_y2_expert_sizes(self, tidx) -> None: def __call__( self, grad_out: cute.Tensor, # (max_tokens_per_rank, hidden) fp8 - grad_out_sf: cute.Tensor, # (max_tokens_per_rank, hidden // sf_vec_size) E8M0 + grad_out_sf: cute.Tensor, # (max_tokens_per_rank, activation_sf_hidden_padded) E8M0 topk_idx: cute.Tensor, # (max_tokens_per_rank, num_topk) topk_weights: cute.Tensor, # (max_tokens_per_rank, num_topk) Float32 (prob) fc1_weight: cute.Tensor, # W2^T: (experts_per_rank, hidden, inter_downproj) @@ -763,13 +763,13 @@ def __call__( fc2_weight_sf: cute.Tensor, beta: cute.Tensor, # (experts_per_rank,) Float32 fc1_preact: cute.Tensor, # (pool_token_capacity, intermediate_gateup) BFloat16 - output_activation: cute.Tensor, # (max_tokens_per_rank, topk, hidden) BF16 + output_activation: cute.Tensor, # (max_tokens_per_rank, hidden) BF16 overflow_flag: cute.Tensor, # (1,) Int32, per-rank FC12 overflow output dprob: cute.Tensor, # (max_tokens_per_rank, topk) Float32; symmetric, pre-zeroed - fc1_recompute: cute.Tensor, # (pool_token_capacity, inter_downproj), token-major - fc1_recompute_sf: cute.Tensor, # WGrad2 SFA: (inter_padded, col_sf_rows) - fc1_col_output: cute.Tensor, # (pool_token_capacity, gateup), token-major - fc1_col_output_sf: cute.Tensor, # WGrad1 SFB: (gateup_padded, col_sf_rows) + fc1_recompute: cute.Tensor, # (pool_token_capacity, inter_downproj) + fc1_recompute_sf: cute.Tensor, # (inter_padded, col_sf_rows) + fc1_col_output: cute.Tensor, # (pool_token_capacity, gateup) + fc1_col_output_sf: cute.Tensor, # (gateup_padded, col_sf_rows) grad_y2: cute.Tensor, # (pool_token_capacity, hidden) token-axis MXFP8 grad_y2_sf: cute.Tensor, # flat MN-major E8M0 bytes local_workspace: cute.Pointer, diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py index dbf7b9eb2..cf745ee35 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/rubin/training/mega/fwd_glu/glu_mxfp8_col_requant.py @@ -13,6 +13,8 @@ import cutlass import cutlass.cute as cute import cutlass.cute.nvgpu.cpasync as cpasync +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cutlass_dsl import ( Float32, Int32, @@ -329,6 +331,7 @@ class Mxfp8ColRequant: # --- this kernel's tile ------------------------------------------------ TILE_TOK: int = 128 # == SfAtomNonK, so a tile is one whole SF atom row NSTAGE: int = 2 # 2 is what leaves shared memory for 2 CTAs per SM + KMajorStages: int = 1 # Short-lived K-major CTAs do not amortize stage two. # Tile width and per-lane access width, per requant path. Both are tuned: # the two paths have different shared-memory budgets and different @@ -354,6 +357,7 @@ class Mxfp8ColRequant: # Grid depth, in resident waves. The curve is broad and flat above this. GridWaves: int = 24 + KMajorGridWaves: int = 13 ProducerWarps: int = 1 # ConsumerWarps is derived: @@ -444,6 +448,13 @@ def __init__( if self.scaled_cvt: _cols_choices = (self.ColsPerLaneScaled, self.ColsPerLanePortable) _preferred_tile_hid = self.TileHidScaled + # The extra K-major staging tile makes 256 columns faster on Rubin: + # it preserves two resident CTAs and wins despite twice as many + # hidden groups (84 us versus 101 us for the DS3 production case). + if self.dst_k_major: + _preferred_tile_hid = min( + _preferred_tile_hid, self.TileHidPortable + ) else: _cols_choices = (self.ColsPerLanePortable,) _preferred_tile_hid = self.TileHidPortable @@ -487,14 +498,14 @@ def __init__( self.sf_in_pad = self.SfInPad self.sf_out_pad = self.SfOutPad self.smem_capacity = _smem_capacity() + self.NumStages = self.KMajorStages if self.dst_k_major else self.NSTAGE while ( - self._smem_bytes_for(self.TILE_HID, self.NSTAGE) > self.smem_capacity + self._smem_bytes_for(self.TILE_HID, self.NumStages) > self.smem_capacity and len(self._tile_hid_choices) > 1 ): self._tile_hid_choices.pop() self.TILE_HID = self._tile_hid_choices[-1] - self.NumStages = self.NSTAGE self.TmaBoxHidU32 = self.TILE_HID // 4 self._hidden_atoms = self.hidden // self.SfAtomNonK @@ -519,17 +530,22 @@ def __init__( # --- SMEM ------------------------------------------------------------ self.smem_data_bytes = self.NumStages * self.TILE_TOK * self.TILE_HID + # K-major output staging is separate from the input pipeline. + self.smem_data_out_bytes = ( + self.TILE_TOK * self.TILE_HID if self.dst_k_major else 0 + ) self.smem_sf_in_bytes = self.NumStages * self.SfTileBytes self.SfOutStride = self.SfAtomBytes + self.sf_out_pad self.smem_sf_out_bytes = self.HidAtomsPerTile * self.SfOutStride self.smem_table_bytes = 3 * (self.num_experts + 1) * 4 self.smem_bytes = ( self.smem_data_bytes + + self.smem_data_out_bytes + self.smem_sf_in_bytes + self.smem_sf_out_bytes + self.smem_table_bytes + 2 * self.NumStages * 8 - + 256 + + (1024 if self.dst_k_major else 256) ) if self.smem_bytes > self.smem_capacity: raise ValueError( @@ -560,7 +576,12 @@ def __init__( # The wave count is tuned for 2 resident CTAs. A shape that gets only # one keeps a single-wave grid rather than extrapolating that tuning # point outside the regime it was taken in. - _want = self.GridWaves if self.CtasPerSm >= 2 else 1 + if self.CtasPerSm >= 2: + _want = ( + self.KMajorGridWaves if self.dst_k_major else self.GridWaves + ) + else: + _want = 1 _max_tiles = -(-self.max_total_tokens // self.TILE_TOK) * self.hidden_groups _waves = max(1, min(_want, _max_tiles // (_min_tiles * self.ResidentCtas))) if num_persistent_ctas > 0: @@ -588,6 +609,16 @@ def __init__( self._experts_per_lane = (self.num_experts + 31) // 32 # ------------------------------------------------------------------ host + def _k_major_tma_smem_layout(self): + """Canonical swizzled FP8 hidden-by-token TMA staging tile.""" + staged = sm100_utils.make_smem_layout_epi( + self.quant_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self.TILE_HID, self.TILE_TOK), + 1, + ) + return cute.select(staged, mode=[0, 1]) + @cute.jit def __call__( self, @@ -630,8 +661,20 @@ def __call__( ) if cutlass.const_expr(self.dst_k_major): - tma_atom_st = None - tma_tensor_st = None + k_major_smem_layout = self._k_major_tma_smem_layout() + dst_u8 = cute.make_tensor( + cute.recast_ptr(dst_data.iterator, dtype=cutlass.Uint8), + cute.make_layout( + (dst_data.shape[1], dst_data.shape[0]), + stride=(dst_data.shape[0], 1), + ), + ) + tma_atom_st, tma_tensor_st = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + dst_u8, + k_major_smem_layout, + (self.TILE_HID, BOX_T), + ) else: dst_u32 = cute.make_tensor( cute.recast_ptr(dst_data.iterator, dtype=cutlass.Uint32), @@ -778,6 +821,11 @@ def ws_kernel( smem_data = smem.allocate_array( self.quant_dtype, self.NumStages * TOK * W, byte_alignment=128 ) + smem_data_out = None + if cutlass.const_expr(self.dst_k_major): + smem_data_out = smem.allocate_array( + self.quant_dtype, TOK * W, byte_alignment=1024 + ) if tidx == Int32(0): for s in cutlass.range_constexpr(0, S, 1): @@ -794,6 +842,9 @@ def ws_kernel( total_tiles = Int32(tbl_data[self.num_experts]) // Int32(TOK) smem_data_base = smem_data.toint() + smem_data_out_base = None + if cutlass.const_expr(self.dst_k_major): + smem_data_out_base = smem_data_out.toint() smem_sf_in_base = smem_sf_in.toint() smem_sf_out_base = smem_sf_out.toint() src_sf_base = src_sf_u8.iterator.toint() @@ -808,8 +859,9 @@ def ws_kernel( ) else: self.consume_scaled( - smem_data_base, smem_sf_in_base, smem_sf_out_base, mbar_full, mbar_empty, - tbl_vend, tbl_data, tbl_sf, dst_data, dst_sf_base, + smem_data_base, smem_data_out_base, smem_sf_in_base, + smem_sf_out_base, mbar_full, mbar_empty, + tbl_vend, tbl_data, tbl_sf, dst_sf_base, bidx, grid_dim_x, total_tiles, warp_idx - Int32(self.ProducerWarps), lane_idx, tma_atom_st, tma_tensor_st, TOKPAD, @@ -909,8 +961,9 @@ def produce( # ------------------------------------------------- consumer (scaled cvt) @cute.jit def consume_scaled( - self, smem_data_base, smem_sf_in_base, smem_sf_out_base, mbar_full, mbar_empty, - tbl_vend, tbl_data, tbl_sf, dst_data, dst_sf_base, + self, smem_data_base, smem_data_out_base, smem_sf_in_base, + smem_sf_out_base, mbar_full, mbar_empty, + tbl_vend, tbl_data, tbl_sf, dst_sf_base, bidx, grid_dim_x, total_tiles, cw, lane_idx, tma_atom_st=None, tma_tensor_st=None, token_padding_block: cutlass.Constexpr = None, @@ -928,7 +981,6 @@ def consume_scaled( CONS_THREADS = cutlass.const_expr(self.ConsumerWarps * 32) HATOMS = cutlass.const_expr(self.HidAtomsPerTile) SF_PREFIX_DIFFERS = cutlass.const_expr(TOKPAD != self.sf_padding_block) - SFOUT_ONE = cutlass.const_expr(self.HidAtomsPerTile * self.SfOutStride) C = cutlass.const_expr(self.ColsPerLane) # hidden columns per lane SEGW = cutlass.const_expr(32 * C) # columns per consumer segment @@ -945,13 +997,27 @@ def consume_scaled( sf_lane_off = tb * Int32(4) if cutlass.const_expr(self.dst_k_major): - sDo_u8 = cute.make_tensor( + k_major_smem_layout = self._k_major_tma_smem_layout() + sDoT = cute.make_tensor( cute.make_ptr( - cutlass.Uint8, smem_data_base, AddressSpace.smem, assumed_align=128, + cutlass.Uint8, + smem_data_out_base, + AddressSpace.smem, + assumed_align=128, ), - cute.make_layout((TOK, W, S), stride=(W, 1, TOK * W)), + k_major_smem_layout, ) - dst_u8_pointer = cute.recast_ptr(dst_data.iterator, dtype=cutlass.Uint8) + gDo = cute.group_modes( + cute.local_tile(tma_tensor_st, (W, TOK), (None, None)), 0, 2 + ) + tDsDo, tDgDo = cpasync.tma_partition( + tma_atom_st, + 0, + cute.make_layout(1), + cute.group_modes(sDoT, 0, 2), + gDo, + ) + cpasync.prefetch_descriptor(tma_atom_st) else: BOX_HS = cutlass.const_expr(self.TmaBoxHidU32) sDo = cute.make_tensor( @@ -993,16 +1059,28 @@ def consume_scaled( if cutlass.const_expr(SF_PREFIX_DIFFERS): sf_live = cutlass.min(Int32(1), cutlass.max(Int32(0), valid_rows)) - cute.arch.mbarrier_wait(mbar_full + stage, (t // Int32(S)) % Int32(2)) - - stage_data = smem_data_base + stage * Int32(TOK * W) + data_lane_off - stage_sf = smem_sf_in_base + stage * Int32(SFB) + sf_lane_off - # A tile's sf_out store is drained before the stage is handed - # back, so one buffer is enough. + cute.arch.mbarrier_wait( + mbar_full + stage, (t // Int32(S)) % Int32(2) + ) + stage_data = ( + smem_data_base + stage * Int32(TOK * W) + data_lane_off + ) + stage_sf = ( + smem_sf_in_base + stage * Int32(SFB) + sf_lane_off + ) sfout = smem_sf_out_base self._sp_tile_body( - stage_data, stage_sf, sfout, ldsw, hb0, seg, tb, lane_idx, valid_rows, + stage_data, + stage_sf, + sfout, + smem_data_out_base, + ldsw, + hb0, + seg, + tb, + lane_idx, + valid_rows, ) # Cross-proxy ordering, and it is NOT optional. The consumer warps @@ -1017,23 +1095,18 @@ def consume_scaled( # neutralised to zero in shared memory, which is what the pool # expects to find there. if cutlass.const_expr(self.dst_k_major): - linear = cw * Int32(32) + lane_idx - while linear < Int32(TOK * W): - token_in_tile = linear // Int32(W) - feature_in_tile = linear % Int32(W) - token = data_row0 + token_in_tile - feature = hid_begin + feature_in_tile - if token < dst_data.shape[0] and feature < dst_data.shape[1]: - dst_offset = ( - Int64(feature) * Int64(dst_data.shape[0]) - + Int64(token) - ) - dst_slot = cute.make_tensor( - dst_u8_pointer + dst_offset, - cute.make_layout(1), - ) - dst_slot[0] = sDo_u8[token_in_tile, feature_in_tile, stage] - linear = linear + Int32(CONS_THREADS) + if cw == Int32(0): + cute.copy( + tma_atom_st, + tDsDo, + tDgDo[ + ( + None, + hid_begin // Int32(W), + token_tile, + ) + ], + ) else: if cw == Int32(0): cute.copy( @@ -1063,19 +1136,50 @@ def consume_scaled( ), Int32(self.SfAtomBytes), ) - cute.arch.cp_async_bulk_commit_group() - # Drain every store before the stage goes back to the producer. - cute.arch.cp_async_bulk_wait_group(0, read=True) - cute.arch.barrier(barrier_id=self.ConsumerBarrierId, number_of_threads=CONS_THREADS) - if lane_idx == Int32(0): - cute.arch.mbarrier_arrive(mbar_empty + stage) + # Only consumer warp 0 issues async stores. Its wait followed by + # the CTA barrier makes completion visible to every consumer before + # either the output buffer or the producer stage is reused. + if cw == Int32(0): + cute.arch.cp_async_bulk_commit_group() + if cutlass.const_expr(self.dst_k_major): + # K-major has a separate output tile: release the input stage + # while its async store drains, overlapping the next TMA load. + if lane_idx == Int32(0): + cute.arch.mbarrier_arrive(mbar_empty + stage) + if cw == Int32(0): + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.barrier( + barrier_id=self.ConsumerBarrierId, + number_of_threads=CONS_THREADS, + ) + else: + # Row-major stores directly from the input stage, so it cannot + # be released until the async store has finished reading it. + if cw == Int32(0): + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.barrier( + barrier_id=self.ConsumerBarrierId, + number_of_threads=CONS_THREADS, + ) + if lane_idx == Int32(0): + cute.arch.mbarrier_arrive(mbar_empty + stage) t = t + Int32(1) work_idx = work_idx + grid_dim_x @cute.jit def _sp_tile_body( - self, stage_data, stage_sf_base, sfout, ldsw, hb0, seg, tb, lane_idx, valid_rows, + self, + stage_data, + stage_sf_base, + sfout, + smem_data_out_base, + ldsw, + hb0, + seg, + tb, + lane_idx, + valid_rows, ): """The single-pass arithmetic for one lane's share of one tile.""" TOK = cutlass.const_expr(self.TILE_TOK) @@ -1114,20 +1218,23 @@ def _sp_tile_body( # instead of 0x00. The fc1 pool really does leave 0xFF in padding # scale bytes. Writing 127 is the same value the masked arm # substitutes, paid once per tile instead of once per read. - if valid_rows < Int32(TOK): - for tt in cutlass.range_constexpr(0, NB, 1): - if tb * Int32(NB) + Int32(tt) >= valid_rows: - self._store_words( - base + Int32(tt * W), ldsw, zeros, NWc, LW - ) - sf_t = cute.make_tensor( - cute.make_ptr( - cutlass.Uint8, sfb + Int32(tt * 16), - AddressSpace.smem, assumed_align=1, - ), - cute.make_layout((1,)), - ) - sf_t[0] = Uint8(127) + dead_tt = cutlass.min( + Int32(NB), + cutlass.max(Int32(0), valid_rows - tb * Int32(NB)), + ) + while dead_tt < Int32(NB): + self._store_words( + base + dead_tt * Int32(W), ldsw, zeros, NWc, LW + ) + sf_t = cute.make_tensor( + cute.make_ptr( + cutlass.Uint8, sfb + dead_tt * Int32(16), + AddressSpace.smem, assumed_align=1, + ), + cute.make_layout((1,)), + ) + sf_t[0] = Uint8(127) + dead_tt = dead_tt + Int32(1) # ---- the one scan: unpack-with-scale, keep BF16, accumulate amax -- d = [[None] * NPc for _ in range(NB)] @@ -1135,7 +1242,9 @@ def _sp_tile_body( for tt in cutlass.range_constexpr(0, NB, 1): # Both loads are issued before either is consumed, so the scale # load overlaps the data load. - words = self._load_words(base + Int32(tt * W), ldsw, NWc, LW) + words = self._load_words( + base + Int32(tt * W), ldsw, NWc, LW + ) raw_sf = self._src_scale_raw(sfb + Int32(tt * 16)) s16 = raw_sf | (raw_sf << Int32(8)) for w in cutlass.range_constexpr(0, NWc, 1): @@ -1187,16 +1296,158 @@ def _sp_tile_body( ) out_t[0] = Uint8(raws[j]) - for tt in cutlass.range_constexpr(0, NB, 2): - out0, out1 = self._requant_token_pair( - d[tt], d[tt + 1], scs, invs, NPc, NWc + if cutlass.const_expr(self.dst_k_major): + for tt in cutlass.range_constexpr(0, NB, 16): + # Keep the two-token result column-major and stage it + # directly. The previous implementation first stored + # out0/out1 back to row-major SMEM and then reread every + # byte in a separate transpose pass. + rot = (lane_idx >> Int32(1)) & Int32(3) + swap_adjacent = (rot & Int32(1)) != Int32(0) + swap_halves = (rot & Int32(2)) != Int32(0) + adjacent_rotated = [ + cute.make_rmem_tensor((4,), cutlass.Int32) + for _ in range(LW) + ] + + # Produce only two token pairs at a time and immediately + # combine them into one four-token word. This avoids + # keeping all eight pair tensors live until a later pass. + for q in cutlass.range_constexpr(0, 4, 1): + pair0 = self._requant_token_pair_kmajor( + d[tt + 4 * q], + d[tt + 4 * q + 1], + scs, + invs, + NPc, + NWc, + LW, + ) + pair1 = self._requant_token_pair_kmajor( + d[tt + 4 * q + 2], + d[tt + 4 * q + 3], + scs, + invs, + NPc, + NWc, + LW, + ) + packed = [None] * LW + for j in cutlass.range_constexpr(0, LW, 1): + packed[j] = Int32( + cute.arch.prmt( + pair0[j], + pair1[j], + Int32(0x5410), + ) + ) + for j in cutlass.range_constexpr(0, LW, 1): + adjacent_rotated[j][q] = Int32( + arith.select( + swap_adjacent.ir_value(), + packed[j ^ 1].ir_value(), + packed[j].ir_value(), + ) + ) + + token = tb * Int32(NB) + Int32(tt) + # The adjacent-column exchange above applies rot bit 0. + # This final exchange applies rot bit 1, preserving the + # same XOR-rotated store order with one select per word. + for phase in cutlass.range_constexpr(0, LW, 1): + j_rot = Int32(phase) ^ rot + token_words = cute.make_rmem_tensor( + (4,), cutlass.Int32 + ) + for q in cutlass.range_constexpr(0, 4, 1): + token_words[q] = Int32( + arith.select( + swap_halves.ir_value(), + Int32( + adjacent_rotated[phase ^ 2][q] + ).ir_value(), + Int32( + adjacent_rotated[phase][q] + ).ir_value(), + ) + ) + self._store_kmajor_token_16( + smem_data_out_base, + col0 + j_rot, + token, + token_words, + ) + else: + for tt in cutlass.range_constexpr(0, NB, 2): + out0, out1 = self._requant_token_pair( + d[tt], d[tt + 1], scs, invs, NPc, NWc + ) + self._store_words(base + Int32(tt * W), ldsw, out0, NWc, LW) + self._store_words( + base + Int32((tt + 1) * W), ldsw, out1, NWc, LW + ) + + def _requant_token_pair_kmajor(self, d0, d1, scs, invs, NPc, NWc, LW): + """Return one packed token pair per hidden column for K-major staging.""" + QT = self.quant_dtype + pairs = cute.make_rmem_tensor((LW,), cutlass.Int32) + if cutlass.const_expr(self.scaled_cvt): + for k in range(0, NPc, 1): + lo = Int32(cute.arch.prmt(d0[k], d1[k], Int32(0x5410))) + hi = Int32(cute.arch.prmt(d0[k], d1[k], Int32(0x7632))) + pairs[2 * k] = cvt_scaled_dn_fp8x2(lo, scs[2 * k], QT) + pairs[2 * k + 1] = cvt_scaled_dn_fp8x2( + hi, scs[2 * k + 1], QT + ) + else: + for w in range(0, NWc, 1): + k0 = 2 * w + k1 = 2 * w + 1 + p00 = cvt_dn_fp8x2_portable( + d0[k0], invs[2 * k0], invs[2 * k0 + 1], QT + ) + p01 = cvt_dn_fp8x2_portable( + d0[k1], invs[2 * k1], invs[2 * k1 + 1], QT ) - self._store_words(base + Int32(tt * W), ldsw, out0, NWc, LW) - self._store_words(base + Int32((tt + 1) * W), ldsw, out1, NWc, LW) + p10 = cvt_dn_fp8x2_portable( + d1[k0], invs[2 * k0], invs[2 * k0 + 1], QT + ) + p11 = cvt_dn_fp8x2_portable( + d1[k1], invs[2 * k1], invs[2 * k1 + 1], QT + ) + pair01 = Int32(cute.arch.prmt(p00, p10, Int32(0x5140))) + pair23 = Int32(cute.arch.prmt(p01, p11, Int32(0x5140))) + pairs[4 * w] = pair01 & Int32(0xFFFF) + pairs[4 * w + 1] = (pair01 >> Int32(16)) & Int32(0xFFFF) + pairs[4 * w + 2] = pair23 & Int32(0xFFFF) + pairs[4 * w + 3] = (pair23 >> Int32(16)) & Int32(0xFFFF) + return pairs + + @cute.jit + def _store_kmajor_token_16( + self, smem_data_out_base, col, token, token_words + ): + """Store 16 adjacent tokens with one 128-bit R2S copy.""" + linear = col * Int32(self.TILE_TOK) + token + offset = linear ^ ((linear & Int32(0x380)) >> Int32(3)) + st128 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Int32, + num_bits_per_copy=128, + ) + dst = cute.make_tensor( + cute.make_ptr( + cutlass.Int32, + smem_data_out_base + offset, + AddressSpace.smem, + assumed_align=16, + ), + cute.make_layout((4,)), + ) + cute.copy(st128, token_words, dst) def _requant_token_pair(self, d0, d1, scs, invs, NPc, NWc): - """Requantise one lane's two consecutive tokens. THE ONLY PLACE THIS - KERNEL DEPENDS ON THE TARGET ARCHITECTURE. + """Requantise one lane's two consecutive tokens for row-major output. ``d0``/``d1`` are lists of NPc Int32, each a token-major BF16x2: register ``k`` is hidden columns ``(2k, 2k+1)`` of one token. ``out0``/``out1`` @@ -1282,11 +1533,12 @@ def _smem_bytes_for(self, tile_hid: int, stages: int) -> int: hid_atoms = tile_hid // self.SfAtomNonK return ( stages * self.TILE_TOK * tile_hid + + (self.TILE_TOK * tile_hid if self.dst_k_major else 0) + stages * hid_atoms * (self.SfAtomBytes + self.sf_in_pad) + hid_atoms * (self.SfAtomBytes + self.sf_out_pad) + 3 * (self.num_experts + 1) * 4 + 2 * stages * 8 - + 256 + + (1024 if self.dst_k_major else 256) ) @cute.jit diff --git a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py index af270cf14..ab3797072 100644 --- a/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py +++ b/python/cudnn/moe_ep/_megamoe_backend/cutedsl_src/kernel_src/schedulers/fc12_mapping.py @@ -13,6 +13,7 @@ from cutlass.cutlass_dsl import Boolean, Int32, extract_mlir_values, new_from_mlir_values from ...helpers.iket_compat import iket +from ...helpers.utils import padded_expert_rows from .base import SchedulerWorkTileBase @@ -494,12 +495,8 @@ def _load_expert_batch_metrics( token_blocks = (token_count + Int32(mapping_state.mapping_cluster_tile_m - 1)) // Int32( mapping_state.mapping_cluster_tile_m ) - data_rows = ( - (token_count + Int32(mapping_state.token_padding_block - 1)) // Int32(mapping_state.token_padding_block) - ) * Int32(mapping_state.token_padding_block) - sf_rows = ( - (token_count + Int32(mapping_state.sf_padding_block - 1)) // Int32(mapping_state.sf_padding_block) - ) * Int32(mapping_state.sf_padding_block) + data_rows = padded_expert_rows(token_count, Int32(mapping_state.token_padding_block)) + sf_rows = padded_expert_rows(token_count, Int32(mapping_state.sf_padding_block)) fc1_tiles = token_blocks * mapping_state.num_fc1_intermediate_blocks fc2_tiles = token_blocks * mapping_state.num_fc2_hidden_blocks return expert_idx, token_count, token_blocks, data_rows, sf_rows, fc1_tiles, fc2_tiles @@ -660,13 +657,12 @@ def _seek_expert_for_work_id(linear_work_id: Int32, mapping_state: Fc12TaskMappi base_token_block_cumulative = cursor.current_token_block_cumulative if cursor.current_expert_idx >= cursor.current_group_first_expert: current_token_count = cursor.current_expert_token_count - base_data_cumulative = base_data_cumulative + ( - (current_token_count + Int32(mapping_state.token_padding_block - 1)) - // Int32(mapping_state.token_padding_block) - ) * Int32(mapping_state.token_padding_block) - base_sf_cumulative = base_sf_cumulative + ( - (current_token_count + Int32(mapping_state.sf_padding_block - 1)) // Int32(mapping_state.sf_padding_block) - ) * Int32(mapping_state.sf_padding_block) + base_data_cumulative = base_data_cumulative + padded_expert_rows( + current_token_count, Int32(mapping_state.token_padding_block) + ) + base_sf_cumulative = base_sf_cumulative + padded_expert_rows( + current_token_count, Int32(mapping_state.sf_padding_block) + ) base_token_block_cumulative = base_token_block_cumulative + cursor.current_token_block_count search_begin = cutlass.max(cursor.current_expert_idx + Int32(1), cursor.current_group_first_expert) @@ -1034,13 +1030,12 @@ def _advance_phase_cursor( cursor: _PhaseFc12CursorState, mapping_state: PhaseInterleavedFc12MappingState ) -> _PhaseFc12CursorState: previous_token_count = cursor.current_expert_token_count - cursor.data_cumulative = cursor.data_cumulative + ( - (previous_token_count + Int32(mapping_state.token_padding_block - 1)) - // Int32(mapping_state.token_padding_block) - ) * Int32(mapping_state.token_padding_block) - cursor.sf_cumulative = cursor.sf_cumulative + ( - (previous_token_count + Int32(mapping_state.sf_padding_block - 1)) // Int32(mapping_state.sf_padding_block) - ) * Int32(mapping_state.sf_padding_block) + cursor.data_cumulative = cursor.data_cumulative + padded_expert_rows( + previous_token_count, Int32(mapping_state.token_padding_block) + ) + cursor.sf_cumulative = cursor.sf_cumulative + padded_expert_rows( + previous_token_count, Int32(mapping_state.sf_padding_block) + ) cursor.token_block_cumulative = cursor.token_block_cumulative + cursor.current_token_block_count cursor.expert_idx = cursor.expert_idx + Int32(1) diff --git a/test/python/moe_ep/test_moe_ep_forward.py b/test/python/moe_ep/test_moe_ep_forward.py index 91b074002..60a6dd1d6 100644 --- a/test/python/moe_ep/test_moe_ep_forward.py +++ b/test/python/moe_ep/test_moe_ep_forward.py @@ -1262,8 +1262,10 @@ def test_megamoe_capability_and_kernel_config_accept_ep_above_16(): @pytest.mark.L0 -def test_ep32_peer_mapping_selects_vector_payload(): +def test_ep32_peer_mapping_selects_version_compatible_payload(): + import cutlass from cutlass._mlir import ir + from packaging.version import Version from cudnn.moe_ep._megamoe_backend._comm import PeerMapping from cudnn.moe_ep._megamoe_backend.cutedsl_src.communication.nvlink_domain.symmetric_buffer import ( @@ -1286,7 +1288,10 @@ def test_ep32_peer_mapping_selects_vector_payload(): assert host.offsets == offsets assert int(host.max_ranks) == 32 - assert device_type_text == "vector<32xi64>" + dsl_release = Version(Version(cutlass.__version__).base_version) + grid_constant_width_is_free = dsl_release < Version("4.0.0") or dsl_release >= Version("4.7.0") + expected_type = "!llvm.ptr" if grid_constant_width_is_free else "vector<32xi64>" + assert device_type_text == expected_type @pytest.fixture From e3c3a22c2e75d095201a9d56443d096259e24a0f Mon Sep 17 00:00:00 2001 From: zhibinz Date: Wed, 2 Sep 2026 10:57:59 -0700 Subject: [PATCH 31/31] feature: add explicit MoeEP sweep autotuning Add rank-consistent inference and training sweeps so callers can select and apply a measured performance configuration before graph capture. --- docs/fe-oss-apis/moe_ep.md | 51 ++ docs/operations/moe_ep.md | 18 + python/cudnn/__init__.py | 228 ++++-- python/cudnn/moe_ep/__init__.py | 8 +- python/cudnn/moe_ep/_autotune.py | 243 +++++++ python/cudnn/moe_ep/_tuning.py | 81 ++- python/cudnn/moe_ep/api.py | 661 +++++++++++++++++- .../moe_ep/moe_ep_distributed_workers.py | 65 +- test/python/moe_ep/test_moe_ep_autotune.py | 577 +++++++++++++++ test/python/moe_ep/test_moe_ep_forward.py | 58 +- 10 files changed, 1892 insertions(+), 98 deletions(-) create mode 100644 python/cudnn/moe_ep/_autotune.py create mode 100644 test/python/moe_ep/test_moe_ep_autotune.py diff --git a/docs/fe-oss-apis/moe_ep.md b/docs/fe-oss-apis/moe_ep.md index 1e40d278e..5519bbb8a 100644 --- a/docs/fe-oss-apis/moe_ep.md +++ b/docs/fe-oss-apis/moe_ep.md @@ -46,6 +46,57 @@ op = MoeEp( Native training requires `weight_interleave_size=32`. FC1 payloads then use alternating 32-element gate/up strips. +## Explicit sweep autotuning + +`MoeEp.autotune` measures inference forward. `MoeEp.autotune_training` +measures one training forward immediately followed by its matching backward: + +```python +result = op.autotune( + activation, fc1_weight, fc2_weight, topk_idx, topk_weights, + candidates=candidates, + warmup_iters=3, + timed_iters=10, +) + +training_result = op.autotune_training( + activation, grad_output, topk_idx, topk_weights, + forward_weights=native_fw, + backward_weights=native_bw, + candidates=candidates, +) +``` + +Both calls are collective over `ep_group`, must use the same ordered candidate +list on every rank, and must run outside CUDA Graph capture. +`autotune_training` must run before `prepare_training`. It accepts only native +weights; source packing and allocation are intentionally outside its measured +region. + +The current `MoeEpTuningConfig` is prepended as a baseline, duplicate values are +removed, and the normalized list is limited to 32 candidates. Autotuning keeps +`reduce_topk_in_kernel` fixed because that flag changes where top-k reduction +is performed. Each timed iteration is reduced with rank MAX and the candidate +score is the median of those slow-rank samples. Equal scores select the earlier +candidate. `MoeEpAutotuneResult` reports `winner`, per-candidate `latency_ms` +and `samples_ms`, and `evaluated_candidates`. + +The sweep is fail-fast. Any validation, allocation, compile, launch, timing, +synchronization, or teardown error ends the whole sweep. An error after +runtime/collective entry poisons the operator, and later execution is rejected; +close it and create a new instance. Compiled candidate kernels remain in the +process JIT cache. The production sweep does not compare candidate outputs at +runtime; supported candidates are covered by the separate correctness suite. + +Autotuning commits one active winner per instance. A later inference or +training sweep replaces it. Existing CUDA Graph executables are invalid after +the winner changes. Use these sequences: + +- inference: `autotune` → eager winner launch (performed by `autotune`) → + capture; +- training: `autotune_training` → `prepare_training` → allocate outputs → + eager forward/backward → rank synchronization → capture. + ## Stateless training preparation Preparation is collective over `ep_group` and must run outside CUDA Graph diff --git a/docs/operations/moe_ep.md b/docs/operations/moe_ep.md index d2ae0ba87..b4c2f48ea 100644 --- a/docs/operations/moe_ep.md +++ b/docs/operations/moe_ep.md @@ -91,6 +91,24 @@ output = op( For inference CUDA Graph capture, call `op.warmup(...)` with the exact bindings before capture. `MoeEp` supports `close()` and context-manager use. +Explicit sweep autotuning is available before capture: + +```python +from cudnn import MoeEpTuningConfig + +result = op.autotune( + activation, fc1_weight, fc2_weight, topk_idx, topk_weights, + candidates=[ + MoeEpTuningConfig(token_in_flag_batch=2), + MoeEpTuningConfig(group_hint=256), + ], +) +``` + +The current tuning is always included as the baseline. Candidates are +de-duplicated and limited to 32 including that baseline. The winner is applied +only to this operator instance. + Stateless training prepares only private execution lanes. Every invocation receives independent native weights and caller-owned outputs: diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index cddc99535..863ca5879 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -108,7 +108,9 @@ def set_stream(handle, stream): per stream regardless.) """ if not isinstance(handle, Handle): - raise TypeError(f"cudnn.set_stream expects a cudnn.Handle (from cudnn.create_handle()), got {type(handle).__name__}") + raise TypeError( + f"cudnn.set_stream expects a cudnn.Handle (from cudnn.create_handle()), got {type(handle).__name__}" + ) if handle.stream == stream: return if handle.backend_handle is not None: @@ -120,7 +122,9 @@ def get_stream(handle): """The CUDA stream a :class:`cudnn.Handle` runs on -- the cached ``Handle.stream``, no backend round-trip.""" if not isinstance(handle, Handle): - raise TypeError(f"cudnn.get_stream expects a cudnn.Handle (from cudnn.create_handle()), got {type(handle).__name__}") + raise TypeError( + f"cudnn.get_stream expects a cudnn.Handle (from cudnn.create_handle()), got {type(handle).__name__}" + ) return handle.stream @@ -129,7 +133,9 @@ def destroy_handle(handle): after destruction so a reused Handle object cannot pass a released ``cudnnHandle_t`` back to C++ (a double-destroy or a later set_stream).""" if not isinstance(handle, Handle): - raise TypeError(f"cudnn.destroy_handle expects a cudnn.Handle (from cudnn.create_handle()), got {type(handle).__name__}") + raise TypeError( + f"cudnn.destroy_handle expects a cudnn.Handle (from cudnn.create_handle()), got {type(handle).__name__}" + ) backend = handle.backend_handle if backend is None: handle.stream = None @@ -202,10 +208,14 @@ def _set_data_type( def load_cudnn(): # First look at python site packages - lib_path = glob.glob(os.path.join(sysconfig.get_path("purelib"), "nvidia/cudnn/bin/cudnn64_9.dll")) + lib_path = glob.glob( + os.path.join(sysconfig.get_path("purelib"), "nvidia/cudnn/bin/cudnn64_9.dll") + ) if lib_path: - assert len(lib_path) == 1, f"Found {len(lib_path)} libcudnn.dll.x in nvidia-cudnn-cuXX." + assert ( + len(lib_path) == 1 + ), f"Found {len(lib_path)} libcudnn.dll.x in nvidia-cudnn-cuXX." lib = ctypes.windll.LoadLibrary(lib_path[0]) else: # Fallback lib = ctypes.windll.LoadLibrary("cudnn64_9.dll") @@ -230,13 +240,23 @@ def _dlopen_cudnn(): return # Then look at python site packages - lib_path = glob.glob(os.path.join(sysconfig.get_path("purelib"), "nvidia/cudnn/lib/libcudnn.so.*[0-9]")) + lib_path = glob.glob( + os.path.join( + sysconfig.get_path("purelib"), "nvidia/cudnn/lib/libcudnn.so.*[0-9]" + ) + ) if not lib_path: - lib_path = glob.glob(os.path.join(sysconfig.get_path("purelib"), "nvidia/cudnn_jit/lib/libcudnn.so.*[0-9]")) + lib_path = glob.glob( + os.path.join( + sysconfig.get_path("purelib"), "nvidia/cudnn_jit/lib/libcudnn.so.*[0-9]" + ) + ) if lib_path: - assert len(lib_path) == 1, f"Found {len(lib_path)} libcudnn.so.x in nvidia-cudnn-cuXX." + assert ( + len(lib_path) == 1 + ), f"Found {len(lib_path)} libcudnn.so.x in nvidia-cudnn-cuXX." lib = ctypes.CDLL(lib_path[0]) else: # Fallback try: @@ -306,11 +326,16 @@ def _dlopen_cudnn(): __all__ = [*_EAGER_PUBLIC_NAMES, "Graph", "wrapper"] _CUTEDSL_INSTALL_HINT = "Install with 'pip install nvidia-cudnn-frontend[cutedsl]'" -_MOE_EP_INSTALL_HINT = "Install with 'pip install " '"nvidia-cudnn-frontend[cutedsl,comm]" torch torch-c-dlpack-ext\'' +_MOE_EP_INSTALL_HINT = ( + "Install with 'pip install " + '"nvidia-cudnn-frontend[cutedsl,comm]" torch torch-c-dlpack-ext\'' +) _MOE_EP_OPTIONAL_IMPORTS = { "moe_ep", "BlockScaledTensor", "MoeEp", + "MoeEpAutotuneCandidateResult", + "MoeEpAutotuneResult", "MoeEpBackwardWeightStaging", "MoeEpBackwardWeights", "MoeEpExecutionLane", @@ -335,6 +360,11 @@ def _dlopen_cudnn(): "moe_ep": (".moe_ep", None), "BlockScaledTensor": (".moe_ep", "BlockScaledTensor"), "MoeEp": (".moe_ep", "MoeEp"), + "MoeEpAutotuneCandidateResult": ( + ".moe_ep", + "MoeEpAutotuneCandidateResult", + ), + "MoeEpAutotuneResult": (".moe_ep", "MoeEpAutotuneResult"), "MoeEpBackwardWeightStaging": (".moe_ep", "MoeEpBackwardWeightStaging"), "MoeEpBackwardWeights": (".moe_ep", "MoeEpBackwardWeights"), "MoeEpExecutionLane": (".moe_ep", "MoeEpExecutionLane"), @@ -356,9 +386,18 @@ def _dlopen_cudnn(): "pack_backward_weights": (".moe_ep", "pack_backward_weights"), "pack_forward_weights": (".moe_ep", "pack_forward_weights"), "BSA": (".block_sparse_attention", "BSA"), - "block_sparse_attention_forward": (".block_sparse_attention", "block_sparse_attention_forward"), - "block_sparse_attention_fp8_forward": (".block_sparse_attention", "block_sparse_attention_fp8_forward"), - "block_sparse_attention_backward": (".block_sparse_attention", "block_sparse_attention_backward"), + "block_sparse_attention_forward": ( + ".block_sparse_attention", + "block_sparse_attention_forward", + ), + "block_sparse_attention_fp8_forward": ( + ".block_sparse_attention", + "block_sparse_attention_fp8_forward", + ), + "block_sparse_attention_backward": ( + ".block_sparse_attention", + "block_sparse_attention_backward", + ), "DSA": (".deepseek_sparse_attention", "DSA"), "CSA": (".csa", "CSA"), "CSACompressorForward": (".csa", "CSACompressorForward"), @@ -367,66 +406,159 @@ def _dlopen_cudnn(): "csa_compressor_backward_wrapper": (".csa", "csa_compressor_backward_wrapper"), "NSA": (".native_sparse_attention", "NSA"), "GemmSwigluSm100": (".gemm.cutedsl.dense.swiglu", "GemmSwigluSm100"), - "gemm_swiglu_wrapper_sm100": (".gemm.cutedsl.dense.swiglu", "gemm_swiglu_wrapper_sm100"), + "gemm_swiglu_wrapper_sm100": ( + ".gemm.cutedsl.dense.swiglu", + "gemm_swiglu_wrapper_sm100", + ), "gemm_swiglu_jax_sm100": (".gemm.cutedsl.dense.swiglu", "gemm_swiglu_jax_sm100"), "gemm_srelu_jax_sm100": (".gemm.cutedsl.dense.srelu", "gemm_srelu_jax_sm100"), "gemm_dsrelu_jax_sm100": (".gemm.cutedsl.dense.dsrelu", "gemm_dsrelu_jax_sm100"), "GemmSreluSm100": (".gemm.cutedsl.dense.srelu", "GemmSreluSm100"), - "gemm_srelu_wrapper_sm100": (".gemm.cutedsl.dense.srelu", "gemm_srelu_wrapper_sm100"), + "gemm_srelu_wrapper_sm100": ( + ".gemm.cutedsl.dense.srelu", + "gemm_srelu_wrapper_sm100", + ), "GemmDsreluSm100": (".gemm.cutedsl.dense.dsrelu", "GemmDsreluSm100"), - "gemm_dsrelu_wrapper_sm100": (".gemm.cutedsl.dense.dsrelu", "gemm_dsrelu_wrapper_sm100"), + "gemm_dsrelu_wrapper_sm100": ( + ".gemm.cutedsl.dense.dsrelu", + "gemm_dsrelu_wrapper_sm100", + ), "GemmAmaxSm100": (".gemm.cutedsl.dense.amax", "GemmAmaxSm100"), "gemm_amax_wrapper_sm100": (".gemm.cutedsl.dense.amax", "gemm_amax_wrapper_sm100"), "gemm_amax_jax_sm100": (".gemm.cutedsl.dense.amax", "gemm_amax_jax_sm100"), - "GemmProjRopeMxfp8Bf16InSm100": (".gemm.cutedsl.dense.proj_rope_mxfp8", "GemmProjRopeMxfp8Bf16InSm100"), - "GemmProjRopeMxfp8Mxfp8InSm100": (".gemm.cutedsl.dense.proj_rope_mxfp8", "GemmProjRopeMxfp8Mxfp8InSm100"), - "gemm_proj_rope_mxfp8_wrapper_sm100": (".gemm.cutedsl.dense.proj_rope_mxfp8", "gemm_proj_rope_mxfp8_wrapper_sm100"), - "gemm_proj_rope_mxfp8_jax_sm100": (".gemm.cutedsl.dense.proj_rope_mxfp8", "gemm_proj_rope_mxfp8_jax_sm100"), + "GemmProjRopeMxfp8Bf16InSm100": ( + ".gemm.cutedsl.dense.proj_rope_mxfp8", + "GemmProjRopeMxfp8Bf16InSm100", + ), + "GemmProjRopeMxfp8Mxfp8InSm100": ( + ".gemm.cutedsl.dense.proj_rope_mxfp8", + "GemmProjRopeMxfp8Mxfp8InSm100", + ), + "gemm_proj_rope_mxfp8_wrapper_sm100": ( + ".gemm.cutedsl.dense.proj_rope_mxfp8", + "gemm_proj_rope_mxfp8_wrapper_sm100", + ), + "gemm_proj_rope_mxfp8_jax_sm100": ( + ".gemm.cutedsl.dense.proj_rope_mxfp8", + "gemm_proj_rope_mxfp8_jax_sm100", + ), "RmsNormRhtAmaxSm100": (".rmsnorm_rht_amax", "RmsNormRhtAmaxSm100"), - "rmsnorm_rht_amax_wrapper_sm100": (".rmsnorm_rht_amax", "rmsnorm_rht_amax_wrapper_sm100"), + "rmsnorm_rht_amax_wrapper_sm100": ( + ".rmsnorm_rht_amax", + "rmsnorm_rht_amax_wrapper_sm100", + ), "grouped_gemm": (".gemm.cutedsl.grouped", None), "GroupedGemmSm100": (".gemm.cutedsl.grouped", "GroupedGemmSm100"), - "grouped_gemm_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_wrapper_sm100"), + "grouped_gemm_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_wrapper_sm100", + ), "grouped_gemm_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_jax_sm100"), - "grouped_gemm_glu_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_glu_jax_sm100"), - "grouped_gemm_dglu_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dglu_jax_sm100"), - "grouped_gemm_dsrelu_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dsrelu_jax_sm100"), - "grouped_gemm_wgrad_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_wgrad_jax_sm100"), - "discrete_grouped_gemm_swiglu_jax_sm100": (".gemm.cutedsl.discrete_grouped", "discrete_grouped_gemm_swiglu_jax_sm100"), - "discrete_grouped_gemm_dswiglu_jax_sm100": (".gemm.cutedsl.discrete_grouped", "discrete_grouped_gemm_dswiglu_jax_sm100"), + "grouped_gemm_glu_jax_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_glu_jax_sm100", + ), + "grouped_gemm_dglu_jax_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_dglu_jax_sm100", + ), + "grouped_gemm_dsrelu_jax_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_dsrelu_jax_sm100", + ), + "grouped_gemm_wgrad_jax_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_wgrad_jax_sm100", + ), + "discrete_grouped_gemm_swiglu_jax_sm100": ( + ".gemm.cutedsl.discrete_grouped", + "discrete_grouped_gemm_swiglu_jax_sm100", + ), + "discrete_grouped_gemm_dswiglu_jax_sm100": ( + ".gemm.cutedsl.discrete_grouped", + "discrete_grouped_gemm_dswiglu_jax_sm100", + ), "GroupedGemmSwigluSm100": (".gemm.cutedsl.grouped", "GroupedGemmSwigluSm100"), - "grouped_gemm_swiglu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_swiglu_wrapper_sm100"), + "grouped_gemm_swiglu_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_swiglu_wrapper_sm100", + ), "GroupedGemmDswigluSm100": (".gemm.cutedsl.grouped", "GroupedGemmDswigluSm100"), - "grouped_gemm_dswiglu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dswiglu_wrapper_sm100"), + "grouped_gemm_dswiglu_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_dswiglu_wrapper_sm100", + ), "GroupedGemmSreluSm100": (".gemm.cutedsl.grouped", "GroupedGemmSreluSm100"), - "grouped_gemm_srelu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_srelu_wrapper_sm100"), + "grouped_gemm_srelu_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_srelu_wrapper_sm100", + ), "GroupedGemmDsreluSm100": (".gemm.cutedsl.grouped", "GroupedGemmDsreluSm100"), - "grouped_gemm_dsrelu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dsrelu_wrapper_sm100"), + "grouped_gemm_dsrelu_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_dsrelu_wrapper_sm100", + ), "HSTUFwdSm100": (".hstu_attention", "HSTUFwdSm100"), "HSTUBwdSm100": (".hstu_attention", "HSTUBwdSm100"), "hstu_attention_forward": (".hstu_attention", "hstu_attention_forward"), "hstu_attention_backward": (".hstu_attention", "hstu_attention_backward"), "GroupedGemmQuantSm100": (".gemm.cutedsl.grouped", "GroupedGemmQuantSm100"), - "grouped_gemm_quant_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_quant_wrapper_sm100"), + "grouped_gemm_quant_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_quant_wrapper_sm100", + ), "GroupedGemmGluSm100": (".gemm.cutedsl.grouped", "GroupedGemmGluSm100"), - "grouped_gemm_glu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_glu_wrapper_sm100"), - "GroupedGemmGluHadamardSm100": (".gemm.cutedsl.grouped", "GroupedGemmGluHadamardSm100"), - "grouped_gemm_glu_hadamard_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_glu_hadamard_wrapper_sm100"), - "GroupedGemmGluHadamardQuantSm100": (".gemm.cutedsl.grouped", "GroupedGemmGluHadamardQuantSm100"), - "grouped_gemm_glu_hadamard_quant_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_glu_hadamard_quant_wrapper_sm100"), + "grouped_gemm_glu_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_glu_wrapper_sm100", + ), + "GroupedGemmGluHadamardSm100": ( + ".gemm.cutedsl.grouped", + "GroupedGemmGluHadamardSm100", + ), + "grouped_gemm_glu_hadamard_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_glu_hadamard_wrapper_sm100", + ), + "GroupedGemmGluHadamardQuantSm100": ( + ".gemm.cutedsl.grouped", + "GroupedGemmGluHadamardQuantSm100", + ), + "grouped_gemm_glu_hadamard_quant_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_glu_hadamard_quant_wrapper_sm100", + ), "GroupedGemmDgluSm100": (".gemm.cutedsl.grouped", "GroupedGemmDgluSm100"), - "grouped_gemm_dglu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dglu_wrapper_sm100"), + "grouped_gemm_dglu_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_dglu_wrapper_sm100", + ), "GroupedGemmWgradSm100": (".gemm.cutedsl.grouped", "GroupedGemmWgradSm100"), "get_grouped_gemm_wgrad_workspace_size_sm100": ( ".gemm.cutedsl.grouped", "get_grouped_gemm_wgrad_workspace_size_sm100", ), - "grouped_gemm_wgrad_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_wgrad_wrapper_sm100"), + "grouped_gemm_wgrad_wrapper_sm100": ( + ".gemm.cutedsl.grouped", + "grouped_gemm_wgrad_wrapper_sm100", + ), "discrete_grouped_gemm": (".gemm.cutedsl.discrete_grouped", None), - "DiscreteGroupedGemmSwigluSm100": (".gemm.cutedsl.discrete_grouped", "DiscreteGroupedGemmSwigluSm100"), - "discrete_grouped_gemm_swiglu_wrapper_sm100": (".gemm.cutedsl.discrete_grouped", "discrete_grouped_gemm_swiglu_wrapper_sm100"), - "DiscreteGroupedGemmDswigluSm100": (".gemm.cutedsl.discrete_grouped", "DiscreteGroupedGemmDswigluSm100"), - "discrete_grouped_gemm_dswiglu_wrapper_sm100": (".gemm.cutedsl.discrete_grouped", "discrete_grouped_gemm_dswiglu_wrapper_sm100"), + "DiscreteGroupedGemmSwigluSm100": ( + ".gemm.cutedsl.discrete_grouped", + "DiscreteGroupedGemmSwigluSm100", + ), + "discrete_grouped_gemm_swiglu_wrapper_sm100": ( + ".gemm.cutedsl.discrete_grouped", + "discrete_grouped_gemm_swiglu_wrapper_sm100", + ), + "DiscreteGroupedGemmDswigluSm100": ( + ".gemm.cutedsl.discrete_grouped", + "DiscreteGroupedGemmDswigluSm100", + ), + "discrete_grouped_gemm_dswiglu_wrapper_sm100": ( + ".gemm.cutedsl.discrete_grouped", + "discrete_grouped_gemm_dswiglu_wrapper_sm100", + ), } @@ -436,8 +568,14 @@ def _load_optional_symbol(name: str) -> Any: module = importlib.import_module(module_name, package=__name__) value = module if attr_name is None else getattr(module, attr_name) except Exception as e: - install_hint = _MOE_EP_INSTALL_HINT if name in _MOE_EP_OPTIONAL_IMPORTS else _CUTEDSL_INSTALL_HINT - raise ImportError(f"{name} requires optional dependencies. {install_hint}: {e}") from e + install_hint = ( + _MOE_EP_INSTALL_HINT + if name in _MOE_EP_OPTIONAL_IMPORTS + else _CUTEDSL_INSTALL_HINT + ) + raise ImportError( + f"{name} requires optional dependencies. {install_hint}: {e}" + ) from e globals()[name] = value return value diff --git a/python/cudnn/moe_ep/__init__.py b/python/cudnn/moe_ep/__init__.py index e918a2b37..c267cba51 100644 --- a/python/cudnn/moe_ep/__init__.py +++ b/python/cudnn/moe_ep/__init__.py @@ -1,7 +1,11 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from ._tuning import MoeEpTuningConfig +from ._tuning import ( + MoeEpAutotuneCandidateResult, + MoeEpAutotuneResult, + MoeEpTuningConfig, +) from ._types import ( BlockScaledTensor, MoeEpBackwardWeightStaging, @@ -24,6 +28,8 @@ __all__ = [ "BlockScaledTensor", "MoeEp", + "MoeEpAutotuneCandidateResult", + "MoeEpAutotuneResult", "MoeEpBackwardWeightStaging", "MoeEpBackwardWeights", "MoeEpExecutionLane", diff --git a/python/cudnn/moe_ep/_autotune.py b/python/cudnn/moe_ep/_autotune.py new file mode 100644 index 000000000..6830e9422 --- /dev/null +++ b/python/cudnn/moe_ep/_autotune.py @@ -0,0 +1,243 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Collective coordination and timing helpers for explicit MoeEP sweeps.""" + +from __future__ import annotations + +import math +import statistics +from collections.abc import Callable, Sequence +from typing import TypeVar + +import torch +import torch.distributed as dist + +from ._tuning import MoeEpAutotuneCandidateResult, MoeEpTuningConfig +from ._types import MoeEpTrainingBackwardOutputs, MoeEpTrainingForwardOutputs + +_T = TypeVar("_T") +_MAX_AUTOTUNE_CANDIDATES = 32 + + +def normalize_candidates( + baseline: MoeEpTuningConfig, + candidates: Sequence[MoeEpTuningConfig], + *, + warmup_iters: int, + timed_iters: int, + max_candidates: int, +) -> tuple[MoeEpTuningConfig, ...]: + """Validate, de-duplicate, and prepend the current configuration.""" + + if ( + isinstance(warmup_iters, bool) + or not isinstance(warmup_iters, int) + or warmup_iters < 0 + ): + raise ValueError( + f"warmup_iters must be a non-negative integer, got {warmup_iters!r}" + ) + if ( + isinstance(timed_iters, bool) + or not isinstance(timed_iters, int) + or timed_iters <= 0 + ): + raise ValueError(f"timed_iters must be a positive integer, got {timed_iters!r}") + if ( + isinstance(max_candidates, bool) + or not isinstance(max_candidates, int) + or not 1 <= max_candidates <= _MAX_AUTOTUNE_CANDIDATES + ): + raise ValueError( + f"max_candidates must be an integer in [1, {_MAX_AUTOTUNE_CANDIDATES}], " + f"got {max_candidates!r}" + ) + if not isinstance(candidates, Sequence) or isinstance(candidates, (str, bytes)): + raise TypeError("candidates must be a sequence of MoeEpTuningConfig values") + if not candidates: + raise ValueError("candidates must not be empty") + + ordered: list[MoeEpTuningConfig] = [baseline] + seen = {baseline} + for index, candidate in enumerate(candidates): + if not isinstance(candidate, MoeEpTuningConfig): + raise TypeError( + "candidates must contain only MoeEpTuningConfig values; " + f"candidates[{index}] is {type(candidate).__name__}" + ) + if candidate.reduce_topk_in_kernel != baseline.reduce_topk_in_kernel: + raise ValueError( + "autotune does not sweep reduce_topk_in_kernel; " + f"candidate {index} has {candidate.reduce_topk_in_kernel}, " + f"baseline has {baseline.reduce_topk_in_kernel}" + ) + if candidate not in seen: + ordered.append(candidate) + seen.add(candidate) + + if len(ordered) > max_candidates: + raise ValueError( + f"autotune has {len(ordered)} unique candidates including the baseline, " + f"exceeding max_candidates={max_candidates}" + ) + return tuple(ordered) + + +def verify_candidates_across_ranks( + candidates: tuple[MoeEpTuningConfig, ...], + group: dist.ProcessGroup | None, +) -> None: + """Fail before runtime allocation when EP ranks supplied different lists.""" + + if group is None: + return + gathered: list[object] = [None] * dist.get_world_size(group) + dist.all_gather_object(gathered, candidates, group=group) + if any(value != candidates for value in gathered): + raise RuntimeError( + f"MoeEp autotune candidates must match on every EP rank; " + f"rank candidate lists: {gathered}" + ) + + +def verify_state_across_ranks( + state: tuple[object, ...], + group: dist.ProcessGroup | None, +) -> None: + """Require matching operator lifecycle state before collective teardown.""" + + if group is None: + return + gathered: list[object] = [None] * dist.get_world_size(group) + dist.all_gather_object(gathered, state, group=group) + if any(value != state for value in gathered): + raise RuntimeError( + f"MoeEp autotune requires matching lifecycle state on every EP rank; " + f"rank states: {gathered}" + ) + + +def raise_preflight_errors( + error: BaseException | None, + *, + phase: str, + group: dist.ProcessGroup | None, +) -> None: + """Turn rank-local preflight failures into one collective failure.""" + + if group is None: + if error is not None: + raise error + return + local = None if error is None else (type(error).__name__, str(error)) + gathered: list[object] = [None] * dist.get_world_size(group) + dist.all_gather_object(gathered, local, group=group) + failures = [ + (rank, value) for rank, value in enumerate(gathered) if value is not None + ] + if failures: + raise RuntimeError( + f"MoeEp autotune {phase} failed before runtime entry; rank errors: {failures}" + ) from error + + +def synchronize_candidate( + device: torch.device, + group: dist.ProcessGroup | None, +) -> None: + """Drain device work and align ranks at a healthy candidate boundary.""" + + torch.cuda.synchronize(device) + if group is not None: + dist.barrier(group=group) + + +def benchmark_candidate( + run: Callable[[], _T], + *, + device: torch.device, + group: dist.ProcessGroup | None, + timed_iters: int, +) -> tuple[float, tuple[float, ...]]: + """Return median(per-iteration rank-MAX) in milliseconds.""" + + stream = torch.cuda.current_stream(device) + local_samples: list[float] = [] + for _ in range(timed_iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record(stream) + run() + end.record(stream) + end.synchronize() + local_samples.append(float(start.elapsed_time(end))) + + slow_rank_samples = torch.tensor(local_samples, dtype=torch.float64, device=device) + if group is not None: + dist.all_reduce(slow_rank_samples, op=dist.ReduceOp.MAX, group=group) + samples = tuple(float(value) for value in slow_rank_samples.cpu().tolist()) + latency_ms = float(statistics.median(samples)) + if not math.isfinite(latency_ms): + raise RuntimeError( + f"MoeEp autotune produced a non-finite latency: {latency_ms}" + ) + return latency_ms, samples + + +def select_winner( + results: Sequence[MoeEpAutotuneCandidateResult], +) -> MoeEpAutotuneCandidateResult: + """Choose the first minimum-latency candidate for stable tie-breaking.""" + + if not results: + raise ValueError("cannot select an autotune winner without results") + return min(results, key=lambda result: result.latency_ms) + + +def allocate_training_outputs( + requirements, + device: torch.device, +) -> tuple[MoeEpTrainingForwardOutputs, MoeEpTrainingBackwardOutputs]: + """Allocate private one-lane outputs from the production ABI contract.""" + + def allocate(name: str) -> torch.Tensor: + shape, stride, dtype, alignment = requirements[name] + tensor = torch.empty_strided(shape, stride, dtype=dtype, device=device) + if tensor.data_ptr() % alignment: + raise RuntimeError( + f"autotune output {name} is not {alignment}-byte aligned" + ) + return tensor + + forward = MoeEpTrainingForwardOutputs( + fc1_preact=allocate("fc1_preact"), + output=allocate("output"), + fc1_a=allocate("fc1_a"), + fc1_sfa=allocate("fc1_sfa"), + valid_route_counts=allocate("valid_route_counts"), + expert_offsets=allocate("expert_offsets"), + ) + backward = MoeEpTrainingBackwardOutputs( + grad_activation=allocate("grad_activation"), + dprob=allocate("dprob"), + fc1_b=allocate("fc1_b"), + fc1_sfb=allocate("fc1_sfb"), + fc2_a=allocate("fc2_a"), + fc2_sfa=allocate("fc2_sfa"), + fc2_b=allocate("fc2_b"), + fc2_sfb=allocate("fc2_sfb"), + ) + return forward, backward + + +__all__ = [ + "allocate_training_outputs", + "benchmark_candidate", + "normalize_candidates", + "raise_preflight_errors", + "select_winner", + "synchronize_candidate", + "verify_candidates_across_ranks", + "verify_state_across_ranks", +] diff --git a/python/cudnn/moe_ep/_tuning.py b/python/cudnn/moe_ep/_tuning.py index 1d7bfa779..a52897f7f 100644 --- a/python/cudnn/moe_ep/_tuning.py +++ b/python/cudnn/moe_ep/_tuning.py @@ -8,6 +8,8 @@ from dataclasses import dataclass from typing import Literal +AutotuneMode = Literal["inference", "training"] + TokenBackMode = Literal[ "epi_warps", "standalone_warps", @@ -56,18 +58,75 @@ class MoeEpTuningConfig: reduce_topk_in_kernel: bool = False def __post_init__(self) -> None: - if not isinstance(self.token_back_mode, str) or self.token_back_mode not in _TOKEN_BACK_MODES: - raise ValueError("token_back_mode must be one of " f"{tuple(sorted(_TOKEN_BACK_MODES))}, got " f"{self.token_back_mode!r}") - if not isinstance(self.epi_flag_batch, tuple) or self.epi_flag_batch not in _EPI_FLAG_BATCHES: - raise ValueError("epi_flag_batch must be one of " f"{tuple(sorted(_EPI_FLAG_BATCHES))}, got " f"{self.epi_flag_batch!r}") - if isinstance(self.token_in_flag_batch, bool) or self.token_in_flag_batch not in _TOKEN_IN_FLAG_BATCHES: - raise ValueError("token_in_flag_batch must be one of " f"{tuple(sorted(_TOKEN_IN_FLAG_BATCHES))}, got " f"{self.token_in_flag_batch!r}") - if self.group_hint is not None and (isinstance(self.group_hint, bool) or self.group_hint not in _GROUP_HINTS): - raise ValueError("group_hint must be None or one of " f"{tuple(sorted(_GROUP_HINTS))}, got {self.group_hint!r}") + if ( + not isinstance(self.token_back_mode, str) + or self.token_back_mode not in _TOKEN_BACK_MODES + ): + raise ValueError( + "token_back_mode must be one of " + f"{tuple(sorted(_TOKEN_BACK_MODES))}, got " + f"{self.token_back_mode!r}" + ) + if ( + not isinstance(self.epi_flag_batch, tuple) + or self.epi_flag_batch not in _EPI_FLAG_BATCHES + ): + raise ValueError( + "epi_flag_batch must be one of " + f"{tuple(sorted(_EPI_FLAG_BATCHES))}, got " + f"{self.epi_flag_batch!r}" + ) + if ( + isinstance(self.token_in_flag_batch, bool) + or self.token_in_flag_batch not in _TOKEN_IN_FLAG_BATCHES + ): + raise ValueError( + "token_in_flag_batch must be one of " + f"{tuple(sorted(_TOKEN_IN_FLAG_BATCHES))}, got " + f"{self.token_in_flag_batch!r}" + ) + if self.group_hint is not None and ( + isinstance(self.group_hint, bool) or self.group_hint not in _GROUP_HINTS + ): + raise ValueError( + "group_hint must be None or one of " + f"{tuple(sorted(_GROUP_HINTS))}, got {self.group_hint!r}" + ) if not isinstance(self.reduce_topk_in_kernel, bool): - raise ValueError("reduce_topk_in_kernel must be a bool, got " f"{self.reduce_topk_in_kernel!r}") + raise ValueError( + "reduce_topk_in_kernel must be a bool, got " + f"{self.reduce_topk_in_kernel!r}" + ) if self.reduce_topk_in_kernel and self.token_back_mode != "epi_warps": - raise ValueError("reduce_topk_in_kernel requires " "token_back_mode='epi_warps'") + raise ValueError( + "reduce_topk_in_kernel requires " "token_back_mode='epi_warps'" + ) + + +@dataclass(frozen=True) +class MoeEpAutotuneCandidateResult: + """Measured slow-rank latency for one successfully evaluated candidate.""" + + tuning: MoeEpTuningConfig + latency_ms: float + samples_ms: tuple[float, ...] + +@dataclass(frozen=True) +class MoeEpAutotuneResult: + """Winner and measurements produced by one explicit tuning sweep.""" -__all__ = ["MoeEpTuningConfig"] + mode: AutotuneMode + winner: MoeEpTuningConfig + candidates: tuple[MoeEpAutotuneCandidateResult, ...] + + @property + def evaluated_candidates(self) -> int: + return len(self.candidates) + + +__all__ = [ + "MoeEpAutotuneCandidateResult", + "MoeEpAutotuneResult", + "MoeEpTuningConfig", +] diff --git a/python/cudnn/moe_ep/api.py b/python/cudnn/moe_ep/api.py index 19cb47a91..c37c91427 100644 --- a/python/cudnn/moe_ep/api.py +++ b/python/cudnn/moe_ep/api.py @@ -14,14 +14,19 @@ import math import threading import warnings +from dataclasses import replace from numbers import Real -from typing import Mapping, Optional, Union +from typing import Mapping, Optional, Sequence, Union import torch import torch.distributed as dist from ._contracts import Fc1WeightLayout, ForwardConfig, normalize_fc1_weight_layout -from ._tuning import MoeEpTuningConfig +from ._tuning import ( + MoeEpAutotuneCandidateResult, + MoeEpAutotuneResult, + MoeEpTuningConfig, +) from ._types import ( BlockScaledTensor, MoeEpBackwardWeightStaging, @@ -60,18 +65,24 @@ def _resolve_ep_topology( if ep_group is None: return 1, 0, () if not dist.is_available() or not dist.is_initialized(): - raise RuntimeError("ep_group requires an initialized torch.distributed process group") + 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 ep_size <= 0 or ep_rank < 0 or ep_rank >= ep_size: raise ValueError("the current process must be a member of ep_group") - ep_global_ranks = tuple(dist.get_global_rank(ep_group, group_rank) for group_rank in range(ep_size)) + ep_global_ranks = tuple( + dist.get_global_rank(ep_group, group_rank) for group_rank in range(ep_size) + ) if len(set(ep_global_ranks)) != ep_size: raise RuntimeError("ep_group returned duplicate global ranks") if ep_global_ranks[ep_rank] != dist.get_rank(): - raise RuntimeError("ep_group rank mapping is inconsistent with the current global rank") + raise RuntimeError( + "ep_group rank mapping is inconsistent with the current global rank" + ) return ep_size, ep_rank, ep_global_ranks @@ -81,12 +92,18 @@ def _validate_training_assert_capability(config: ForwardConfig) -> None: if config.drop_on_overflow: return if not callable(getattr(torch, "_assert_async", None)): - raise RuntimeError("drop_on_overflow=False training requires callable " "torch._assert_async before CUDA Graph capture") + raise RuntimeError( + "drop_on_overflow=False training requires callable " + "torch._assert_async before CUDA Graph capture" + ) if config.ep_size <= 1: return backend = dist.get_backend(config.ep_group) if backend != dist.Backend.NCCL and str(backend).lower() != "nccl": - raise NotImplementedError("drop_on_overflow=False EP2+ training requires an NCCL " "process group for the captured scalar global overflow OR") + raise NotImplementedError( + "drop_on_overflow=False EP2+ training requires an NCCL " + "process group for the captured scalar global overflow OR" + ) def _resolve_training_device( @@ -106,7 +123,11 @@ def _resolve_training_device( resolved = torch.device("cuda", torch.cuda.current_device()) if resolved.type != "cuda": raise ValueError(f"training device must be CUDA, got {resolved}") - if resolved.index is None or resolved.index < 0 or resolved.index >= torch.cuda.device_count(): + if ( + resolved.index is None + or resolved.index < 0 + or resolved.index >= torch.cuda.device_count() + ): raise ValueError(f"CUDA device {resolved} is not available") return resolved @@ -213,13 +234,25 @@ def __init__( if isinstance(value, bool) or 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 (isinstance(max_tokens_per_rank, bool) or not isinstance(max_tokens_per_rank, int) or max_tokens_per_rank < 0): - raise ValueError("max_tokens_per_rank must be a non-negative integer or None") + raise ValueError( + f"top_k ({top_k}) cannot exceed num_experts ({num_experts})" + ) + if max_tokens_per_rank is not None and ( + isinstance(max_tokens_per_rank, bool) + or not isinstance(max_tokens_per_rank, int) + or max_tokens_per_rank < 0 + ): + raise ValueError( + "max_tokens_per_rank must be a non-negative integer or None" + ) if max_recv_size_per_rank is not None and ( - isinstance(max_recv_size_per_rank, bool) or not isinstance(max_recv_size_per_rank, int) or max_recv_size_per_rank <= 0 + isinstance(max_recv_size_per_rank, bool) + or not isinstance(max_recv_size_per_rank, int) + or max_recv_size_per_rank <= 0 ): - raise ValueError("max_recv_size_per_rank must be a positive integer or None") + raise ValueError( + "max_recv_size_per_rank must be a positive integer or None" + ) if not isinstance(drop_on_overflow, bool): raise ValueError("drop_on_overflow must be a bool") if not isinstance(apply_topk_in_fc1, bool): @@ -232,9 +265,15 @@ def __init__( if isinstance(value, bool) or not isinstance(value, int) or value <= 0: raise ValueError(f"{name} must be a positive integer, got {value!r}") if sf_padding_size % 128: - raise ValueError("sf_padding_size must be a positive multiple of 128, " f"got {sf_padding_size}") + raise ValueError( + "sf_padding_size must be a positive multiple of 128, " + f"got {sf_padding_size}" + ) if tuning is not None and not isinstance(tuning, MoeEpTuningConfig): - raise TypeError("tuning must be a MoeEpTuningConfig or None, " f"got {type(tuning).__name__}") + raise TypeError( + "tuning must be a MoeEpTuningConfig or None, " + f"got {type(tuning).__name__}" + ) if gate_up_clamp is not None: if isinstance(gate_up_clamp, bool) or not isinstance(gate_up_clamp, Real): raise ValueError("gate_up_clamp must be a finite real number or None") @@ -243,10 +282,15 @@ def __init__( raise ValueError("gate_up_clamp must be a finite real number or None") if ep_group is not None and not isinstance(ep_group, dist.ProcessGroup): - raise ValueError(f"ep_group must be a torch.distributed.ProcessGroup or None, " f"got {type(ep_group).__name__}") + raise ValueError( + f"ep_group must be a torch.distributed.ProcessGroup or None, " + f"got {type(ep_group).__name__}" + ) ep_size, ep_rank, ep_global_ranks = _resolve_ep_topology(ep_group) if num_experts % ep_size != 0: - raise ValueError(f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})") + 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 @@ -270,17 +314,27 @@ def __init__( self.sf_padding_size = sf_padding_size self.tuning = MoeEpTuningConfig() if tuning is None else tuning if self.tuning.reduce_topk_in_kernel and ( - self.combine_format is not MoeFormat.BF16 or self.output_format is not MoeFormat.BF16 or not self.apply_topk_in_fc1 + self.combine_format is not MoeFormat.BF16 + or self.output_format is not MoeFormat.BF16 + or not self.apply_topk_in_fc1 ): - raise ValueError("reduce_topk_in_kernel requires BF16 combine/output and " "apply_topk_in_fc1=True") + raise ValueError( + "reduce_topk_in_kernel requires BF16 combine/output and " + "apply_topk_in_fc1=True" + ) 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 + 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}") + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by " + f"{required_multiple} for {name}={fmt.value}" + ) self._forward_config = ForwardConfig( num_experts=self.num_experts, @@ -320,6 +374,7 @@ def __init__( ] | None ) = None + self._poisoned = False self._closed = False @staticmethod @@ -337,10 +392,20 @@ def _get_backend(self, request): with self._lifecycle_lock: if self._closed: raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError( + "MoeEp is unusable after an autotune runtime failure" + ) from . import _backend - if self._forward_backend is not None and request.device != self._forward_backend_device: - raise ValueError(f"MoeEp backend is bound to {self._forward_backend_device}; " f"create a separate MoeEp instance for {request.device}") + if ( + self._forward_backend is not None + and request.device != self._forward_backend_device + ): + raise ValueError( + f"MoeEp backend is bound to {self._forward_backend_device}; " + f"create a separate MoeEp instance for {request.device}" + ) _backend.validate_config(self._forward_config) _backend.validate_request(request) @@ -374,8 +439,16 @@ def __call__( with self._lifecycle_lock: if self._closed: raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError( + "MoeEp is unusable after an autotune runtime failure" + ) topk_version = self._tensor_version(topk_idx) - validate_expert_ids = not (self._validated_topk_idx is topk_idx and topk_version is not None and topk_version == self._validated_topk_version) + validate_expert_ids = not ( + self._validated_topk_idx is topk_idx + and topk_version is not None + and topk_version == self._validated_topk_version + ) request = validate_forward( self._forward_config, activation, @@ -394,6 +467,499 @@ def __call__( self._validated_topk_version = None return self._get_backend(request).forward(request) + def autotune( + self, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + candidates: Sequence[MoeEpTuningConfig], + warmup_iters: int = 3, + timed_iters: int = 10, + max_candidates: int = 32, + ) -> MoeEpAutotuneResult: + """Collectively sweep inference configurations and apply the winner. + + Candidate compilation, allocation, and warmup are excluded from CUDA + Event timing. The measured region includes input/weight staging, the + MegaMoE launch, and the output copy performed by a normal forward. + """ + + from . import _backend + from ._autotune import ( + benchmark_candidate, + normalize_candidates, + raise_preflight_errors, + select_winner, + synchronize_candidate, + verify_candidates_across_ranks, + verify_state_across_ranks, + ) + + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError( + "MoeEp is unusable after an autotune runtime failure" + ) + if self._training_state is not None: + raise RuntimeError("autotune must be called before prepare_training()") + + normalized = normalize_candidates( + self.tuning, + candidates, + warmup_iters=warmup_iters, + timed_iters=timed_iters, + max_candidates=max_candidates, + ) + verify_candidates_across_ranks(normalized, self._forward_config.ep_group) + verify_state_across_ranks( + ( + self._forward_backend is not None, + self._training_state is not None, + ( + None + if self._forward_backend_device is None + else str(self._forward_backend_device) + ), + ), + self._forward_config.ep_group, + ) + + candidate_requests = [] + preflight_error: BaseException | None = None + try: + for index, tuning in enumerate(normalized): + try: + config = replace(self._forward_config, tuning=tuning) + request = validate_forward( + config, + activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ) + if request.device.type != "cuda": + raise ValueError( + f"autotune requires CUDA inputs, got {request.device}" + ) + with torch.cuda.device(request.device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "autotune cannot run during CUDA Graph capture" + ) + _backend.validate_config(config) + _backend.validate_request(request) + candidate_requests.append(request) + except BaseException as exc: + raise RuntimeError( + f"MoeEp autotune candidate {index} {tuning!r} " + f"failed during preflight: {exc}" + ) from exc + except BaseException as exc: + preflight_error = exc + raise_preflight_errors( + preflight_error, + phase="inference preflight", + group=self._forward_config.ep_group, + ) + assert candidate_requests + device = candidate_requests[0].device + + if self._forward_backend is not None: + try: + synchronize_candidate(device, self._forward_config.ep_group) + self._forward_backend.close() + self._forward_backend = None + self._forward_backend_device = None + if self._forward_config.ep_group is not None: + dist.barrier(group=self._forward_config.ep_group) + except BaseException as exc: + self._poisoned = True + raise RuntimeError( + f"MoeEp autotune failed during active backend teardown: {exc}" + ) from exc + + results: list[MoeEpAutotuneCandidateResult] = [] + for index, (tuning, request) in enumerate( + zip(normalized, candidate_requests) + ): + backend = None + runtime_entered = False + phase = "backend creation" + try: + backend = _backend.create_backend(request.config, device) + phase = "compile/prime" + runtime_entered = True + with torch.cuda.device(device): + output = backend.forward(request) + del output + phase = "warmup" + for _ in range(warmup_iters): + output = backend.forward(request) + del output + phase = "pre-timing synchronize" + synchronize_candidate(device, self._forward_config.ep_group) + phase = "timing" + latency_ms, samples_ms = benchmark_candidate( + lambda: backend.forward(request), + device=device, + group=self._forward_config.ep_group, + timed_iters=timed_iters, + ) + phase = "post-timing synchronize" + synchronize_candidate(device, self._forward_config.ep_group) + results.append( + MoeEpAutotuneCandidateResult( + tuning=tuning, + latency_ms=latency_ms, + samples_ms=samples_ms, + ) + ) + phase = "teardown" + backend.close() + backend = None + if self._forward_config.ep_group is not None: + dist.barrier(group=self._forward_config.ep_group) + except BaseException as exc: + if backend is not None and self._forward_config.ep_size == 1: + with contextlib.suppress(Exception): + backend.close() + if runtime_entered: + self._poisoned = True + raise RuntimeError( + f"MoeEp autotune candidate {index} {tuning!r} failed during {phase}: {exc}" + ) from exc + + winner = select_winner(results) + winner_request = candidate_requests[normalized.index(winner.tuning)] + winner_backend = None + try: + winner_backend = _backend.create_backend( + winner_request.config, + device, + ) + with torch.cuda.device(device): + output = winner_backend.forward(winner_request) + del output + synchronize_candidate(device, self._forward_config.ep_group) + except BaseException as exc: + if winner_backend is not None and self._forward_config.ep_size == 1: + with contextlib.suppress(Exception): + winner_backend.close() + self._poisoned = True + raise RuntimeError( + f"MoeEp autotune winner {winner.tuning!r} failed final validation: {exc}" + ) from exc + + self.tuning = winner.tuning + self._forward_config = winner_request.config + self._forward_backend = winner_backend + self._forward_backend_device = device + self._validated_topk_idx = None + self._validated_topk_version = None + return MoeEpAutotuneResult( + mode="inference", + winner=winner.tuning, + candidates=tuple(results), + ) + + def autotune_training( + self, + activation: MoeTensor, + grad_output: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + forward_weights: MoeEpNativeForwardWeights, + backward_weights: MoeEpNativeBackwardWeights, + candidates: Sequence[MoeEpTuningConfig], + warmup_iters: int = 3, + timed_iters: int = 10, + max_candidates: int = 32, + ) -> MoeEpAutotuneResult: + """Sweep complete training forward+backward latency and apply the winner. + + This collective API uses private one-lane temporary resources. It must + run before :meth:`prepare_training` and accepts kernel-native weights + so packing allocation and source-layout conversion are not timed. + """ + + from . import _backend + from ._autotune import ( + allocate_training_outputs, + benchmark_candidate, + normalize_candidates, + raise_preflight_errors, + select_winner, + synchronize_candidate, + verify_candidates_across_ranks, + verify_state_across_ranks, + ) + from ._megamoe_backend.mxfp8._training_execute import ( + launch_training_backward, + launch_training_forward, + ) + + with self._lifecycle_lock: + if self._closed: + raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError( + "MoeEp is unusable after an autotune runtime failure" + ) + if self._training_state is not None: + raise RuntimeError( + "autotune_training must be called before prepare_training()" + ) + + normalized = normalize_candidates( + self.tuning, + candidates, + warmup_iters=warmup_iters, + timed_iters=timed_iters, + max_candidates=max_candidates, + ) + verify_candidates_across_ranks(normalized, self._forward_config.ep_group) + verify_state_across_ranks( + ( + self._forward_backend is not None, + self._training_state is not None, + ( + None + if self._forward_backend_device is None + else str(self._forward_backend_device) + ), + ), + self._forward_config.ep_group, + ) + + device: torch.device | None = None + preflight_error: BaseException | None = None + candidate_configs: list[ForwardConfig] = [] + token_count = -1 + try: + device = torch.device(activation.device) + if device.type != "cuda": + raise ValueError( + f"autotune_training requires CUDA inputs, got {device}" + ) + with torch.cuda.device(device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "autotune_training cannot run during CUDA Graph capture" + ) + for index, tuning in enumerate(normalized): + try: + config = replace(self._forward_config, tuning=tuning) + _validate_training_assert_capability(config) + _backend.validate_config(config) + activation_tokens = validate_training_input( + config, + "activation", + activation, + topk_idx, + topk_weights, + device=device, + ) + grad_tokens = validate_training_input( + config, + "grad_output", + grad_output, + topk_idx, + topk_weights, + device=device, + ) + if activation_tokens != grad_tokens: + raise ValueError( + "activation and grad_output must have the same token " + f"count, got {activation_tokens} and {grad_tokens}" + ) + validate_native_forward_weights( + config, forward_weights, device=device + ) + validate_native_backward_weights( + config, backward_weights, device=device + ) + token_count = activation_tokens + candidate_configs.append(config) + except BaseException as exc: + raise RuntimeError( + f"MoeEp autotune_training candidate {index} {tuning!r} " + f"failed during preflight: {exc}" + ) from exc + except BaseException as exc: + preflight_error = exc + raise_preflight_errors( + preflight_error, + phase="training preflight", + group=self._forward_config.ep_group, + ) + assert device is not None and candidate_configs and token_count >= 0 + + if self._forward_backend is not None: + try: + synchronize_candidate(device, self._forward_config.ep_group) + self._forward_backend.close() + self._forward_backend = None + self._forward_backend_device = None + if self._forward_config.ep_group is not None: + dist.barrier(group=self._forward_config.ep_group) + except BaseException as exc: + self._poisoned = True + raise RuntimeError( + "MoeEp autotune_training failed during active backend " + f"teardown: {exc}" + ) from exc + + results: list[MoeEpAutotuneCandidateResult] = [] + for index, (tuning, config) in enumerate( + zip(normalized, candidate_configs) + ): + backend = None + runtime_entered = False + phase = "backend creation" + try: + backend = _backend.create_backend(config, device) + runtime_entered = True + phase = "training preparation" + with torch.cuda.device(device): + state = backend.prepare_training(lane_count=1) + requirements = state.public_requirements() + forward_out, backward_out = allocate_training_outputs( + requirements, + device, + ) + forward_names = ( + "output", + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + ) + backward_names = ( + "grad_activation", + "dprob", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + ) + validate_training_forward_outputs( + forward_out, + {name: requirements[name] for name in forward_names}, + device=device, + ) + validate_training_backward_outputs( + backward_out, + {name: requirements[name] for name in backward_names}, + device=device, + ) + validate_training_forward_state( + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + requirements={ + name: requirements[name] + for name in ( + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + ) + }, + device=device, + ) + execution = state.views(lane=0, token_count=token_count) + + def run_training_pair(): + launch_training_forward( + state, + execution, + activation, + topk_idx, + topk_weights, + weights=forward_weights, + out=forward_out, + ) + return launch_training_backward( + state, + execution, + grad_output, + topk_idx, + topk_weights, + weights=backward_weights, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, + ) + + phase = "compile/prime" + run_training_pair() + phase = "warmup" + for _ in range(warmup_iters): + run_training_pair() + phase = "pre-timing synchronize" + synchronize_candidate(device, self._forward_config.ep_group) + phase = "timing" + latency_ms, samples_ms = benchmark_candidate( + run_training_pair, + device=device, + group=self._forward_config.ep_group, + timed_iters=timed_iters, + ) + phase = "post-timing synchronize" + synchronize_candidate(device, self._forward_config.ep_group) + results.append( + MoeEpAutotuneCandidateResult( + tuning=tuning, + latency_ms=latency_ms, + samples_ms=samples_ms, + ) + ) + phase = "teardown" + backend.close() + backend = None + if self._forward_config.ep_group is not None: + dist.barrier(group=self._forward_config.ep_group) + except BaseException as exc: + if backend is not None and self._forward_config.ep_size == 1: + with contextlib.suppress(Exception): + backend.close() + if runtime_entered: + self._poisoned = True + raise RuntimeError( + f"MoeEp autotune_training candidate {index} {tuning!r} " + f"failed during {phase}: {exc}" + ) from exc + + winner = select_winner(results) + winner_config = candidate_configs[normalized.index(winner.tuning)] + self.tuning = winner.tuning + self._forward_config = winner_config + self._forward_backend = None + self._forward_backend_device = None + self._validated_topk_idx = None + self._validated_topk_version = None + return MoeEpAutotuneResult( + mode="training", + winner=winner.tuning, + candidates=tuple(results), + ) + def warmup( self, activation: MoeTensor, @@ -417,6 +983,10 @@ def warmup( with self._lifecycle_lock: if self._closed: raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError( + "MoeEp is unusable after an autotune runtime failure" + ) output = self( activation, fc1_weight, @@ -453,8 +1023,18 @@ def prepare_training( with self._lifecycle_lock: if self._closed: raise RuntimeError("MoeEp is closed") - if isinstance(lane_count, bool) or not isinstance(lane_count, int) or lane_count <= 0: - raise ValueError(f"lane_count must be a positive integer, got {lane_count!r}") + if self._poisoned: + raise RuntimeError( + "MoeEp is unusable after an autotune runtime failure" + ) + if ( + isinstance(lane_count, bool) + or not isinstance(lane_count, int) + or lane_count <= 0 + ): + raise ValueError( + f"lane_count must be a positive integer, got {lane_count!r}" + ) if self._training_state is not None: raise RuntimeError("MoeEp training is already prepared") if self._fc1_weight_layout is not Fc1WeightLayout.GATE_UP_INTERLEAVED_32: @@ -464,8 +1044,14 @@ def prepare_training( from . import _backend _backend.validate_config(self._forward_config) - if self._forward_backend is not None and resolved_device != self._forward_backend_device: - raise ValueError(f"MoeEp backend is bound to {self._forward_backend_device}; " f"got {resolved_device}") + if ( + self._forward_backend is not None + and resolved_device != self._forward_backend_device + ): + raise ValueError( + f"MoeEp backend is bound to {self._forward_backend_device}; " + f"got {resolved_device}" + ) if self._forward_backend is None: self._forward_backend = _backend.create_backend( self._forward_config, @@ -477,7 +1063,10 @@ def prepare_training( lane_count=lane_count, ) self._training_state = state - self._training_lanes = tuple(MoeEpExecutionLane(index, self._operator_token) for index in range(lane_count)) + self._training_lanes = tuple( + MoeEpExecutionLane(index, self._operator_token) + for index in range(lane_count) + ) self._training_requirements = state.public_requirements() return self._training_requirements @@ -487,9 +1076,15 @@ def _require_training_lane( ) -> None: if self._closed: raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError("MoeEp is unusable after an autotune runtime failure") if self._training_state is None or self._training_requirements is None: raise RuntimeError("prepare_training() must be called first") - if not isinstance(lane, MoeEpExecutionLane) or lane._operator_token is not self._operator_token or lane not in self._training_lanes: + if ( + not isinstance(lane, MoeEpExecutionLane) + or lane._operator_token is not self._operator_token + or lane not in self._training_lanes + ): raise ValueError("execution lane does not belong to this MoeEp") def _training_requirement_subset( @@ -760,6 +1355,10 @@ def __enter__(self) -> "MoeEp": with self._lifecycle_lock: if self._closed: raise RuntimeError("MoeEp is closed") + if self._poisoned: + raise RuntimeError( + "MoeEp is unusable after an autotune runtime failure" + ) return self def __exit__(self, exc_type, exc_value, traceback) -> bool: @@ -787,6 +1386,8 @@ def __del__(self) -> None: __all__ = [ "BlockScaledTensor", "MoeEp", + "MoeEpAutotuneCandidateResult", + "MoeEpAutotuneResult", "MoeEpBackwardWeightStaging", "MoeEpBackwardWeights", "MoeEpExecutionLane", diff --git a/test/python/moe_ep/moe_ep_distributed_workers.py b/test/python/moe_ep/moe_ep_distributed_workers.py index 0acf2ef2e..72b5bb25e 100644 --- a/test/python/moe_ep/moe_ep_distributed_workers.py +++ b/test/python/moe_ep/moe_ep_distributed_workers.py @@ -31,6 +31,7 @@ ) __all__ = [ + "_distributed_autotune_worker", "_distributed_output_worker", "_distributed_subgroup_output_worker", "_run_backward_reference_case", @@ -38,6 +39,61 @@ ] +def _distributed_autotune_worker( + rank: int, + world_size: int, + init_file: str, +) -> None: + """Run an EP sweep and verify one rank-consistent applied winner.""" + + device = torch.device("cuda", rank) + torch.cuda.set_device(device) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + device_id=device, + timeout=timedelta(seconds=180), + ) + try: + from cudnn import MoeEp, MoeEpTuningConfig + + args = make_distributed_forward_inputs(rank, world_size, device) + config = _forward_config( + num_experts=2 * world_size, + ep_group=dist.group.WORLD, + max_tokens_per_rank=8, + ) + expected = _reference_forward(args, **config) + candidate = MoeEpTuningConfig(token_in_flag_batch=2) + op = MoeEp(**config) + try: + result = op.autotune( + *args, + candidates=[candidate], + warmup_iters=1, + timed_iters=2, + ) + actual = op(*args) + torch.cuda.synchronize(device) + winners = [None] * world_size + dist.all_gather_object(winners, result.winner) + assert all(winner == result.winner for winner in winners) + assert op.tuning == result.winner + _assert_matches_reference(actual, expected) + dist.barrier() + op.close() + op = None + dist.barrier() + finally: + if op is not None: + op.close() + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + def _run_forward_output_case( *, device: torch.device, @@ -141,12 +197,17 @@ def _distributed_subgroup_output_worker( ) try: subgroup_memberships = ((0, 2), (1, 3)) - subgroups = [dist.new_group(list(members), backend="nccl") for members in subgroup_memberships] + subgroups = [ + dist.new_group(list(members), backend="nccl") + for members in subgroup_memberships + ] subgroup_index = global_rank % 2 ep_group = subgroups[subgroup_index] ep_rank = dist.get_rank(ep_group) ep_size = dist.get_world_size(ep_group) - actual_global_ranks = tuple(dist.get_global_rank(ep_group, group_rank) for group_rank in range(ep_size)) + actual_global_ranks = tuple( + dist.get_global_rank(ep_group, group_rank) for group_rank in range(ep_size) + ) _run_forward_output_case( device=device, diff --git a/test/python/moe_ep/test_moe_ep_autotune.py b/test/python/moe_ep/test_moe_ep_autotune.py new file mode 100644 index 000000000..c9a801313 --- /dev/null +++ b/test/python/moe_ep/test_moe_ep_autotune.py @@ -0,0 +1,577 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts and smoke coverage for the explicit MoeEP sweep autotuner.""" + +from __future__ import annotations + +import contextlib +from types import SimpleNamespace + +import pytest +import torch + +from moe_ep.moe_ep_test_support import ( + _allocate_stateless_training_outputs, + _allocate_training_weight_staging, + _assert_backward_matches, + _assert_matches_reference, + _fixed_training_reference, + _fixed_training_weights, + _forward_config, + _grad_output, + _reference_forward, + _replay_cuda_graph, + _sm107_device, + make_forward_inputs, +) + + +def _validated_request(config, *args, **kwargs): + del args, kwargs + return SimpleNamespace(config=config, device=torch.device("cuda", 0)) + + +def _patch_common_inference_dependencies( + patch, + *, + api_module, + backend_module, + create_backend, +) -> None: + patch.setattr(api_module, "validate_forward", _validated_request) + patch.setattr(backend_module, "validate_config", lambda config: None) + patch.setattr(backend_module, "validate_request", lambda request: None) + patch.setattr(backend_module, "create_backend", create_backend) + patch.setattr(torch.cuda, "device", lambda device: contextlib.nullcontext()) + patch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + +@pytest.mark.L0 +def test_autotune_core_contracts(monkeypatch): + import cudnn.moe_ep._autotune as autotune_module + from cudnn import ( + MoeEpAutotuneCandidateResult, + MoeEpAutotuneResult, + MoeEpTuningConfig, + ) + + baseline = MoeEpTuningConfig() + candidate = MoeEpTuningConfig(token_in_flag_batch=2) + normalize = autotune_module.normalize_candidates + assert normalize( + baseline, + [candidate, baseline, candidate], + warmup_iters=0, + timed_iters=1, + max_candidates=2, + ) == (baseline, candidate) + + invalid_limits = ( + ({"warmup_iters": -1}, "warmup_iters"), + ({"timed_iters": 0}, "timed_iters"), + ({"max_candidates": 33}, "max_candidates"), + ) + for overrides, message in invalid_limits: + arguments = { + "warmup_iters": 0, + "timed_iters": 1, + "max_candidates": 32, + **overrides, + } + with pytest.raises(ValueError, match=message): + normalize(baseline, [baseline], **arguments) + + with pytest.raises(ValueError, match="does not sweep reduce_topk_in_kernel"): + normalize( + baseline, + [MoeEpTuningConfig(reduce_topk_in_kernel=True)], + warmup_iters=0, + timed_iters=1, + max_candidates=32, + ) + with pytest.raises(ValueError, match="exceeding max_candidates=1"): + normalize( + baseline, + [candidate], + warmup_iters=0, + timed_iters=1, + max_candidates=1, + ) + + first = MoeEpAutotuneCandidateResult(baseline, 1.0, (1.0,)) + second = MoeEpAutotuneCandidateResult(candidate, 1.0, (1.0,)) + assert autotune_module.select_winner((first, second)) is first + result = MoeEpAutotuneResult("inference", first.tuning, (first, second)) + assert result.evaluated_candidates == 2 + + with monkeypatch.context() as patch: + remote = (candidate,) + patch.setattr(autotune_module.dist, "get_world_size", lambda group: 2) + + def gather(output, value, *, group): + del group + output[:] = [value, remote] + + patch.setattr(autotune_module.dist, "all_gather_object", gather) + with pytest.raises(RuntimeError, match="must match on every EP rank"): + autotune_module.verify_candidates_across_ranks((baseline,), object()) + + with monkeypatch.context() as patch: + local_samples = iter((1.0, 7.0, 3.0)) + + class Event: + def record(self, stream): + del stream + + def synchronize(self): + pass + + def elapsed_time(self, end): + del end + return next(local_samples) + + patch.setattr(torch.cuda, "current_stream", lambda device: object()) + patch.setattr(torch.cuda, "Event", lambda enable_timing: Event()) + patch.setattr( + autotune_module.dist, + "all_reduce", + lambda values, **kwargs: values.copy_( + torch.tensor([5.0, 8.0, 4.0], dtype=values.dtype) + ), + ) + latency, samples = autotune_module.benchmark_candidate( + lambda: None, + device=torch.device("cpu"), + group=object(), + timed_iters=3, + ) + assert samples == (5.0, 8.0, 4.0) + assert latency == 5.0 + + +@pytest.mark.L0 +def test_autotune_api_transactions(monkeypatch): + import cudnn.moe_ep._autotune as autotune_module + import cudnn.moe_ep._backend as backend_module + import cudnn.moe_ep._megamoe_backend.mxfp8._training_execute as execute_module + import cudnn.moe_ep.api as api_module + from cudnn import MoeEp, MoeEpTuningConfig + + baseline = MoeEpTuningConfig() + candidate = MoeEpTuningConfig(token_in_flag_batch=2) + + # Validation failures happen before teardown and preserve active state. + op = MoeEp(**_forward_config()) + active_backend = object() + op._forward_backend = active_backend + with pytest.raises(ValueError, match="does not sweep reduce_topk_in_kernel"): + op.autotune( + None, + None, + None, + None, + None, + candidates=[MoeEpTuningConfig(reduce_topk_in_kernel=True)], + warmup_iters=0, + timed_iters=1, + ) + assert op.tuning == baseline + assert op._forward_backend is active_backend + op._forward_backend = None + op.close() + + # A runtime failure is fail-fast and permanently poisons the instance. + with monkeypatch.context() as patch: + calls = [] + + class FailingBackend: + def forward(self, request): + calls.append(request.config.tuning) + raise RuntimeError("launch failed") + + def close(self): + pass + + _patch_common_inference_dependencies( + patch, + api_module=api_module, + backend_module=backend_module, + create_backend=lambda config, device: FailingBackend(), + ) + patch.setattr( + autotune_module, + "verify_state_across_ranks", + lambda state, group: None, + ) + op = MoeEp(**_forward_config()) + with pytest.raises(RuntimeError, match="candidate 0.*compile/prime"): + op.autotune( + None, + None, + None, + None, + None, + candidates=[candidate], + warmup_iters=0, + timed_iters=1, + ) + assert calls == [baseline] + with pytest.raises(RuntimeError, match="unusable"): + op.autotune( + None, + None, + None, + None, + None, + candidates=[candidate], + warmup_iters=0, + timed_iters=1, + ) + op._forward_backend = None + op.close() + + # Inference commits only the measured winner and retains its rebuilt backend. + with monkeypatch.context() as patch: + active_backends = [] + + class InferenceBackend: + def __init__(self, config): + self.config = config + self.closed = False + + def forward(self, request): + assert request.config == self.config + return object() + + def close(self): + self.closed = True + + def create_inference_backend(config, device): + del device + active_backends.append(InferenceBackend(config)) + return active_backends[-1] + + def benchmark_inference(run, *, device, group, timed_iters): + del device, group, timed_iters + run() + tuning = active_backends[-1].config.tuning + latency = 1.0 if tuning == candidate else 2.0 + return latency, (latency,) + + _patch_common_inference_dependencies( + patch, + api_module=api_module, + backend_module=backend_module, + create_backend=create_inference_backend, + ) + patch.setattr( + autotune_module, + "benchmark_candidate", + benchmark_inference, + ) + patch.setattr( + autotune_module, + "synchronize_candidate", + lambda device, group: None, + ) + op = MoeEp(**_forward_config()) + result = op.autotune( + None, + None, + None, + None, + None, + candidates=[candidate], + warmup_iters=0, + timed_iters=1, + ) + assert result.winner == candidate + assert result.evaluated_candidates == 2 + assert op.tuning == op._forward_config.tuning == candidate + assert op._forward_backend is active_backends[-1] + assert not active_backends[-1].closed + op.close() + + # Training times forward/backward pairs and leaves preparation to the caller. + with monkeypatch.context() as patch: + launches = [] + active_backends = [] + requirement_names = ( + "output", + "fc1_preact", + "fc1_a", + "fc1_sfa", + "valid_route_counts", + "expert_offsets", + "grad_activation", + "dprob", + "fc1_b", + "fc1_sfb", + "fc2_a", + "fc2_sfa", + "fc2_b", + "fc2_sfb", + ) + + class TrainingState: + def public_requirements(self): + return {name: None for name in requirement_names} + + def views(self, *, lane, token_count): + return lane, token_count + + class TrainingBackend: + def __init__(self, config): + self.config = config + + def prepare_training(self, *, lane_count): + assert lane_count == 1 + return TrainingState() + + def close(self): + pass + + forward_outputs = SimpleNamespace( + fc1_preact=object(), + output=object(), + fc1_a=object(), + fc1_sfa=object(), + valid_route_counts=object(), + expert_offsets=object(), + ) + backward_outputs = SimpleNamespace() + + patch.setattr( + api_module, + "_validate_training_assert_capability", + lambda config: None, + ) + for validation_name in ( + "validate_native_forward_weights", + "validate_native_backward_weights", + "validate_training_forward_outputs", + "validate_training_backward_outputs", + "validate_training_forward_state", + ): + patch.setattr( + api_module, + validation_name, + lambda *args, **kwargs: None, + ) + patch.setattr( + api_module, + "validate_training_input", + lambda *args, **kwargs: 2, + ) + patch.setattr(backend_module, "validate_config", lambda config: None) + + def create_training_backend(config, device): + del device + active_backends.append(TrainingBackend(config)) + return active_backends[-1] + + patch.setattr(backend_module, "create_backend", create_training_backend) + patch.setattr( + autotune_module, + "allocate_training_outputs", + lambda requirements, device: (forward_outputs, backward_outputs), + ) + patch.setattr( + autotune_module, + "synchronize_candidate", + lambda device, group: None, + ) + + def benchmark_training(run, *, device, group, timed_iters): + del device, group, timed_iters + run() + tuning = active_backends[-1].config.tuning + latency = 1.0 if tuning == candidate else 2.0 + return latency, (latency,) + + patch.setattr( + autotune_module, + "benchmark_candidate", + benchmark_training, + ) + patch.setattr( + torch.cuda, + "device", + lambda device: contextlib.nullcontext(), + ) + patch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + patch.setattr( + execute_module, + "launch_training_forward", + lambda *args, **kwargs: launches.append("forward"), + ) + patch.setattr( + execute_module, + "launch_training_backward", + lambda *args, **kwargs: launches.append("backward"), + ) + + op = MoeEp(**_forward_config(), weight_interleave_size=32) + value = SimpleNamespace(device=torch.device("cuda", 0)) + result = op.autotune_training( + value, + value, + None, + None, + forward_weights=object(), + backward_weights=object(), + candidates=[candidate], + warmup_iters=0, + timed_iters=1, + ) + assert result.mode == "training" + assert result.winner == candidate + assert launches == ["forward", "backward"] * 4 + assert op._training_state is None + assert op._forward_backend is None + op.close() + + +def _print_candidate_timings(label, result) -> None: + print(f"\n{label} autotune timings:", flush=True) + for index, measurement in enumerate(result.candidates): + samples = ", ".join(f"{sample:.4f}" for sample in measurement.samples_ms) + print( + f" [{index}] median={measurement.latency_ms:.4f} ms " + f"samples=[{samples}] tuning={measurement.tuning}", + flush=True, + ) + + +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_autotune_sm107_inference_training_and_graph(): + from cudnn import MoeEp, MoeEpTuningConfig + + device = _sm107_device() + candidates = [ + MoeEpTuningConfig(), + MoeEpTuningConfig(token_in_flag_batch=2), + MoeEpTuningConfig(token_in_flag_batch=4), + MoeEpTuningConfig(epi_flag_batch=(2, 1)), + MoeEpTuningConfig(group_hint=64), + ] + + inference_args = make_forward_inputs(device) + inference_expected = _reference_forward(inference_args) + original_topk_idx = inference_args[3].clone() + with MoeEp(**_forward_config()) as op: + result = op.autotune( + *inference_args, + candidates=candidates, + warmup_iters=1, + timed_iters=2, + ) + _print_candidate_timings("inference", result) + actual = op(*inference_args) + torch.cuda.synchronize(device) + assert result.evaluated_candidates == len(candidates) == 5 + assert result.winner in candidates + assert op.tuning == result.winner + assert op._forward_backend is not None + _assert_matches_reference(actual, inference_expected) + _replay_cuda_graph( + op, + inference_args, + original_topk_idx, + inference_expected, + device, + ) + + base_args = make_forward_inputs(device) + training_args = ( + base_args[0].dequantize(torch.bfloat16), + base_args[1], + base_args[2], + base_args[3], + base_args[4].float().contiguous(), + ) + grad_output = _grad_output( + device, + training_args[0].shape[0], + seed=20260903, + ) + training_expected = _fixed_training_reference( + training_args, + grad_output, + combine_format="bf16", + gate_up_clamp=None, + ) + source_weights = _fixed_training_weights(training_args) + with MoeEp( + num_experts=2, + hidden_size=128, + intermediate_size=256, + top_k=2, + max_tokens_per_rank=training_args[0].shape[0], + max_recv_size_per_rank=(training_args[0].shape[0] * training_args[3].shape[1]), + drop_on_overflow=True, + combine_format="bf16", + weight_interleave_size=32, + ) as op: + forward_staging, backward_staging = _allocate_training_weight_staging( + source_weights + ) + native_forward = op.pack_forward_weights( + source_weights[0], + out=forward_staging, + ) + native_backward = op.pack_backward_weights( + source_weights[1], + out=backward_staging, + ) + result = op.autotune_training( + training_args[0], + grad_output, + training_args[3], + training_args[4], + forward_weights=native_forward, + backward_weights=native_backward, + candidates=candidates, + warmup_iters=1, + timed_iters=2, + ) + _print_candidate_timings("training", result) + requirements = op.prepare_training(lane_count=1, device=device) + forward_out, backward_out = _allocate_stateless_training_outputs( + requirements, + device, + ) + lane = op.training_lanes[0] + actual_y = op.training_forward( + lane, + training_args[0], + training_args[3], + training_args[4], + weights=native_forward, + out=forward_out, + ) + actual_dx, actual_dprob, _ = op.training_backward( + lane, + grad_output, + training_args[3], + training_args[4], + weights=native_backward, + fc1_preact=forward_out.fc1_preact, + fc1_a=forward_out.fc1_a, + fc1_sfa=forward_out.fc1_sfa, + valid_route_counts=forward_out.valid_route_counts, + expert_offsets=forward_out.expert_offsets, + out=backward_out, + ) + torch.cuda.synchronize(device) + assert result.evaluated_candidates == len(candidates) == 5 + assert result.winner in candidates + assert op.tuning == result.winner + _assert_matches_reference(actual_y, training_expected[0]) + _assert_backward_matches( + (actual_dx, actual_dprob), + (training_expected[1], training_expected[2]), + training_args[3], + ) diff --git a/test/python/moe_ep/test_moe_ep_forward.py b/test/python/moe_ep/test_moe_ep_forward.py index 60a6dd1d6..4e128ebb9 100644 --- a/test/python/moe_ep/test_moe_ep_forward.py +++ b/test/python/moe_ep/test_moe_ep_forward.py @@ -17,6 +17,7 @@ import torch.multiprocessing as mp from moe_ep.moe_ep_distributed_workers import ( + _distributed_autotune_worker, _distributed_output_worker, _distributed_subgroup_output_worker, ) @@ -120,7 +121,9 @@ def test_moe_ep_tuning_public_contract_mapping_and_cache_key(): ) with MoeEp(**_forward_config()) as default_op: - default_config = Mxfp8KernelConfig.from_forward_config(default_op._forward_config) + default_config = Mxfp8KernelConfig.from_forward_config( + default_op._forward_config + ) key_args = ( torch.device("cuda", 0), (10, 7), @@ -147,7 +150,11 @@ def test_internal_column_requant_config_is_disabled_by_default_and_cache_distinc ) assert default_config.enable_col_quant is False - assert default_config.max_recv_size_per_rank == (default_forward.ep_size * default_forward.max_tokens_per_rank * default_forward.top_k) + assert default_config.max_recv_size_per_rank == ( + default_forward.ep_size + * default_forward.max_tokens_per_rank + * default_forward.top_k + ) assert enabled_config.enable_col_quant is True assert enabled_config.col_quant_num_ctas == 512 with pytest.raises(ValueError, match="max_recv_size_per_rank"): @@ -155,7 +162,9 @@ def test_internal_column_requant_config_is_disabled_by_default_and_cache_distinc with pytest.raises(ValueError, match="col_quant_num_ctas"): replace(default_config, col_quant_num_ctas=0) key_args = (torch.device("cuda", 0), (10, 7), 123, ()) - assert default_config.compile_key(*key_args) != enabled_config.compile_key(*key_args) + assert default_config.compile_key(*key_args) != enabled_config.compile_key( + *key_args + ) @pytest.mark.L0 @@ -459,8 +468,12 @@ def test_moe_ep_rejects_interleaved_plain_fc1_weight(): config = _forward_config() activation = torch.zeros((1, config["hidden_size"])) - fc1_weight = torch.zeros((config["num_experts"], config["hidden_size"], 2 * config["intermediate_size"])) - fc2_weight = torch.zeros((config["num_experts"], config["intermediate_size"], config["hidden_size"])) + fc1_weight = torch.zeros( + (config["num_experts"], config["hidden_size"], 2 * config["intermediate_size"]) + ) + fc2_weight = torch.zeros( + (config["num_experts"], config["intermediate_size"], config["hidden_size"]) + ) topk_idx = torch.zeros((1, config["top_k"]), dtype=torch.int32) topk_weights = torch.ones((1, config["top_k"])) with MoeEp(**config, weight_interleave_size=32) as op: @@ -737,7 +750,9 @@ def test_nondefault_moe_ep_tuning_matches_reference_and_reuses_plan(): assert backend._compiled is compiled assert backend._plan._workspace is workspace - assert backend.kernel_config.tuning_signature(backend._prepared_kernel.launch_cluster_count) == ("standalone_warps", (4, 2), 4, 64, False) + assert backend.kernel_config.tuning_signature( + backend._prepared_kernel.launch_cluster_count + ) == ("standalone_warps", (4, 2), 4, 64, False) _assert_matches_reference(first, expected) _assert_matches_reference(second, expected) @@ -938,6 +953,21 @@ def test_mxfp8_forward_multi_gpu_matches_reference( ) +@pytest.mark.L1 +@pytest.mark.gpu_exclusive +def test_mxfp8_forward_ep2_autotune_is_rank_consistent(tmp_path): + world_size = 2 + _require_distributed_sm107(world_size) + os.environ.setdefault("NVIDIA_IMEX_CHANNELS", "0") + init_file = tmp_path / "autotune_ep2.init" + mp.spawn( + _distributed_autotune_worker, + args=(world_size, str(init_file)), + nprocs=world_size, + join=True, + ) + + @pytest.mark.L1 @pytest.mark.gpu_exclusive def test_mxfp8_forward_noncontiguous_ep_subgroups(tmp_path): @@ -1024,7 +1054,11 @@ def test_activation_scale_rows_are_padded_to_16_bytes(): kernel_local_workspace_bytes=128, kernel_shared_workspace_bytes=128, ) - activation_scale = next(region for region in requirements.symmetric_regions if region.name == "activation_scale") + activation_scale = next( + region + for region in requirements.symmetric_regions + if region.name == "activation_scale" + ) assert activation_scale.nbytes == 5 * 16 @@ -1133,7 +1167,11 @@ def test_reference_interleaved_fc1_matches_logical_fc1(): def interleave_last(tensor): shape = tensor.shape - return tensor.view(*shape[:-1], 2, intermediate // 32, 32).transpose(-3, -2).reshape(shape) + return ( + tensor.view(*shape[:-1], 2, intermediate // 32, 32) + .transpose(-3, -2) + .reshape(shape) + ) interleaved_fc1 = BlockScaledTensor( data=interleave_last(logical_fc1.data), @@ -1289,7 +1327,9 @@ def test_ep32_peer_mapping_selects_version_compatible_payload(): assert host.offsets == offsets assert int(host.max_ranks) == 32 dsl_release = Version(Version(cutlass.__version__).base_version) - grid_constant_width_is_free = dsl_release < Version("4.0.0") or dsl_release >= Version("4.7.0") + grid_constant_width_is_free = dsl_release < Version( + "4.0.0" + ) or dsl_release >= Version("4.7.0") expected_type = "!llvm.ptr" if grid_constant_width_is_free else "vector<32xi64>" assert device_type_text == expected_type