diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 631e01ecdb63..a8dcaf1aec25 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -264,13 +264,8 @@ common-files: &common_files | tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py | tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py | tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py | tensorrt_llm/_torch/modules/fused_moe/interface.py | tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py | - tensorrt_llm/_torch/modules/fused_moe/ops/__init__.py | - tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py | - tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_cutlass.py | - tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py | tensorrt_llm/_torch/modules/fused_moe/quantization.py | tensorrt_llm/_torch/modules/fused_moe/routing.py | tensorrt_llm/_torch/modules/gated_mlp.py | @@ -1038,13 +1033,8 @@ legacy-files: &legacy_files | tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py | tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py | tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py | - tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py | tensorrt_llm/_torch/modules/fused_moe/interface.py | tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py | - tensorrt_llm/_torch/modules/fused_moe/ops/__init__.py | - tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py | - tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_cutlass.py | - tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py | tensorrt_llm/_torch/modules/fused_moe/quantization.py | tensorrt_llm/_torch/modules/fused_moe/routing.py | tensorrt_llm/_torch/modules/gated_mlp.py | diff --git a/examples/layer_wise_benchmarks/run.py b/examples/layer_wise_benchmarks/run.py index c3596267f9dc..cd240fed603f 100644 --- a/examples/layer_wise_benchmarks/run.py +++ b/examples/layer_wise_benchmarks/run.py @@ -244,7 +244,7 @@ def comma_separated_floats(s): ctx_attn_workspace = torch.empty((0,), device="cuda", dtype=torch.int8) with mock.patch.dict( os.environ, - {"TRTLLM_FORCE_ALLTOALL_METHOD": "NotEnabled", "TRTLLM_FORCE_COMM_METHOD": "ALLGATHER"}, + {"TRTLLM_FORCE_COMM_METHOD": "ALLGATHER"}, clear=False, ): ctx_runner = Runner( diff --git a/legacy-files.txt b/legacy-files.txt index b01ca7fbe20c..69de2c5b22bb 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -256,13 +256,8 @@ tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py -tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py tensorrt_llm/_torch/modules/fused_moe/interface.py tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py -tensorrt_llm/_torch/modules/fused_moe/ops/__init__.py -tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py -tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_cutlass.py -tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py tensorrt_llm/_torch/modules/fused_moe/quantization.py tensorrt_llm/_torch/modules/fused_moe/routing.py tensorrt_llm/_torch/modules/gated_mlp.py diff --git a/pyproject.toml b/pyproject.toml index 725161fec289..1c26fd19c1ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -313,13 +313,8 @@ exclude = [ "tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py", "tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py", "tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py", "tensorrt_llm/_torch/modules/fused_moe/interface.py", "tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py", - "tensorrt_llm/_torch/modules/fused_moe/ops/__init__.py", - "tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py", - "tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_cutlass.py", - "tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py", "tensorrt_llm/_torch/modules/fused_moe/quantization.py", "tensorrt_llm/_torch/modules/fused_moe/routing.py", "tensorrt_llm/_torch/modules/gated_mlp.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index ff58435cf25a..bf435dce35bb 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -273,13 +273,8 @@ include = [ "tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py", "tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py", "tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py", - "tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py", "tensorrt_llm/_torch/modules/fused_moe/interface.py", "tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py", - "tensorrt_llm/_torch/modules/fused_moe/ops/__init__.py", - "tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py", - "tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_cutlass.py", - "tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py", "tensorrt_llm/_torch/modules/fused_moe/quantization.py", "tensorrt_llm/_torch/modules/fused_moe/routing.py", "tensorrt_llm/_torch/modules/gated_mlp.py", diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index f528145ff93e..b30c713d69f0 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -63,7 +63,6 @@ from ..modules.fused_moe import (DeepSeekV3MoeRoutingMethod, MoEWeightLoadingMode, create_moe, is_moe_weight_owner) -from ..modules.fused_moe.fused_moe_wide_ep import WideEPMoE from ..modules.mla import MLA # isort: off @@ -1151,9 +1150,6 @@ def compute_routed_output(self, hidden_states, hidden_states_fp4, output_dtype=hidden_states.dtype, all_rank_num_tokens=all_rank_num_tokens, use_dp_padding=use_dp_padding, - **({ - "alltoall_result_do_sum": False - } if isinstance(self.experts, WideEPMoE) else {}), ) return routed_output diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index ca53aa6f0856..fa0823b3c7f7 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -76,7 +76,6 @@ resolve_moe_cls, ) from ..modules.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE -from ..modules.fused_moe.fused_moe_wide_ep import WideEPMoE from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode, WeightsLoadingConfig from ..modules.mhc.hyper_connection import HCHead, HCState, mHC @@ -1697,7 +1696,6 @@ def compute_routed_output( output_dtype=hidden_states.dtype, all_rank_num_tokens=all_rank_num_tokens, use_dp_padding=use_dp_padding, - **({"alltoall_result_do_sum": False} if isinstance(self.experts, WideEPMoE) else {}), ) return routed_output diff --git a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md index 86b9e97adf99..b03b26a9bb49 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -112,15 +112,16 @@ The codebase is transitioning between two architectures: | Entry | `XXFusedMoE` (e.g., `CutlassFusedMoE`) | `ConfigurableMoE` + `XXBackend` + `MoEScheduler` | | Communication | Embedded inside each backend | Separated into `communication/` (or fused into kernel for `FUSED_COMM`) | | Forward execution | Inline in backend | `MoEScheduler` (`moe_scheduler.py`) | -| EPLB | Only in WideEPMoE | Available to all backends | +| EPLB | Not supported | Available on EPLB-capable backends | | Status | Being replaced | Active development | ConfigurableMoE currently supports these backends (`create_moe.py`): -- `CutlassFusedMoE`, `TRTLLMGenFusedMoE`, `DeepGemmFusedMoE`, `CuteDslFusedMoE`, `DenseGEMMFusedMoE`, `MegaMoEDeepGemm`, `MegaMoECuteDsl` +- `CutlassFusedMoE`, `TRTLLMGenFusedMoE`, `DeepGemmFusedMoE`, `CuteDslFusedMoE`, + `CuteDslB12xFusedMoE`, `DenseGEMMFusedMoE`, `MegaMoEDeepGemm`, + `MegaMoECuteDsl`, `MarlinFusedMoE` Still on old path (standalone, with embedded communication): - `TritonFusedMoE`, `VanillaMoE` -- `WideEPMoE` — deprecated, the `WIDEEP` backend can no longer be selected **Rule: All new features should target ConfigurableMoE + Backend + Scheduler architecture.** @@ -137,7 +138,7 @@ Still on old path (standalone, with embedded communication): | `impl_contract.py` | Selection vocabulary — `MoEProblem`, `MoEDeployment`, `MoEEnvironment`, `MoEEligibility`, `MoERejectReason`, `MoEResolutionReport` | | `impl_environment.py` | The only place that probes the machine (SM, optional wheels, env flags) and freezes the result | | `impl_identity.py` | `MoEImplId` / `MoEImplDescriptor` / registry — the stable one-id-per-leaf-class mechanism used after an implementation migrates | -| `interface.py` | Base class `MoE` and enums (`MoEWeightLoadingMode`, `MoESchedulerKind`, `AlltoallMethodType`) | +| `interface.py` | Base class `MoE` and enums (`MoEWeightLoadingMode`, `MoESchedulerKind`) | | `quantization.py` | Quantization method implementations (`FusedMoEMethod` subclasses: weight creation, loading, quant/dequant ops per quant mode) | | `routing.py` | Routing methods (`TopKRouting`, etc.) | | `moe_load_balancer.py` | EPLB implementation | @@ -157,7 +158,6 @@ Still on old path (standalone, with embedded communication): | `mega_moe/mega_moe_cute_dsl.py` | `MegaMoECuteDsl` | SM100/SM103 | NVFP4 via ported CuteDSL `Sm100MegaMoEKernel` fused dispatch+FC1+act+FC2+combine kernel; requires CUDA 13 Cutlass DSL runtime (PR #14354) and NVSHMEM provider (hard gate); threads per-expert `fc31_alpha`/`fc2_alpha`/`fc1_norm_const` through the kernel ABI and supports SwiGLU clamp via `swiglu_limit`; default deepgemm graph (topk score folded before fc1-out quant, host `combine_output.sum(dim=1)`) | `FUSED_COMM` | | `fused_moe_marlin.py` | `MarlinFusedMoE` | SM89-SM99 | W4A16 NVFP4 on Ada/Hopper (BF16 activations + FP4 weights, fused single-launch `marlin_nvfp4_moe_gemm` kernel); supports attention-DP + EP via external comm (scheduler precomputes routing; dispatch payload is plain BF16, no activation scales); non-NVFP4 layers (e.g. unquantized MTP draft layers) degrade to Cutlass in `resolve_moe_impl`, recorded in the layer's `MoEResolutionReport`; no dynamic EPLB | `EXTERNAL_COMM` | | `fused_moe_triton.py` | `TritonFusedMoE` | SM90 only | GPT-OSS on Hopper (requires `swiglu_gptoss_style=True`) | (legacy path) | -| `fused_moe_wide_ep.py` | `WideEPMoE` | All GPUs | Deprecated — `moe_resolution.py` raises on the `WIDEEP` backend literal. Wide EP and EPLB are available on the other backends: use `DEEPGEMM` for FP8 block-scale checkpoints, or `TRTLLM` / `CUTEDSL` / `CUTLASS` otherwise. Class kept for reference only | (legacy path) | | `fused_moe_vanilla.py` | `VanillaMoE` | All devices | Reference / debugging only | (legacy path) | ### Communication (`fused_moe/communication/`) @@ -310,26 +310,21 @@ unconditionally, not here. Each backend's `can_implement(p, d)` classmethod declares what it supports. Source of truth: the `can_implement` classmethod in each backend file. -| Quantization | Cutlass | TRTLLMGen | DeepGemm | DenseGEMM | CuteDSL | MegaMoE-DG | MegaMoE-CuteDSL | Triton | Marlin | WideEP (retired)† | Vanilla | -|---|---|---|---|---|---|---|---|---|---|---|---| -| Unquantized (BF16/FP16) | Y (SM80+) | Y (SM100/103, BF16, needs FlashInfer `trtllm_bf16_moe`)§ | N | N | N | N | N | Y (SM90, BF16) | N | Y | Y | -| FP8 QDQ | Y (SM89+) | N | N | N | N | N | N | Y (SM90) | N | Y | Y | -| FP8 Block Scales | Y (SM90, SM120) | Y (SM100/103) | Y (SM100/103) | N | N‡ | N | N | N | N | Y | Y | -| NVFP4 | Y (SM100/103/120/121) | Y (SM100/103) | N | Y (SM100/103) | Y (SM100/103/120/121) | N | Y (SM100/103, cu13 cutlass-dsl + NVSHMEM provider; per-expert alpha/norm_const + SwiGLU clamp) | N | Y (SM89-SM99) | Y | Y | -| W4A16 NVFP4 | Y (SM80+, dequant-on-the-fly) | N | N | N | Y (SM120/121 via `CuteDslB12xFusedMoE`, needs flashinfer) | N | N | N | Y (SM89-SM99, BF16) | N | Y | -| W4A8 NVFP4 FP8 | N | Y (SM100/103) | N | N | N | N | N | N | N | N | N | -| W4A16 MXFP4 | Y (SM90) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | N | -| W4A8 MXFP4 FP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | N | -| W4A8 MXFP4 MXFP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | Y (SM100/103, requires `hidden_size % 512 == 0`) | N | N | N | N | N | -| W8A8 MXFP8 MXFP8 | Y (SM100/103) | N | N | N | N | N | N | N | N | N | N | -| W4A8 AWQ | Y (SM89/90) | N | N | N | N | N | N | N | N | N | N | -| W8A16 | Y (SM80+) | N | N | N | N | N | N | N | N | N | N | -| INT4 WoQ (W4AFP8) | N | N | N | N | N | N | N | N | N | Y | N | - -† The `WideEP` column is history, not an option. `resolve_moe_impl` raises on the -`WIDEEP` literal, so nothing in that column can be selected — read a `Y` there as -"the retired class implemented this", never as "you may request this". For -`INT4 WoQ (W4AFP8)` that leaves the row without a selectable backend. +| Quantization | Cutlass | TRTLLMGen | DeepGemm | DenseGEMM | CuteDSL | MegaMoE-DG | MegaMoE-CuteDSL | Triton | Marlin | Vanilla | +|---|---|---|---|---|---|---|---|---|---|---| +| Unquantized (BF16/FP16) | Y (SM80+) | Y (SM100/103, BF16, needs FlashInfer `trtllm_bf16_moe`)§ | N | N | N | N | N | Y (SM90, BF16) | N | Y | +| FP8 QDQ | Y (SM89+) | N | N | N | N | N | N | Y (SM90) | N | Y | +| FP8 Block Scales | Y (SM90, SM120) | Y (SM100/103) | Y (SM100/103) | N | N‡ | N | N | N | N | Y | +| NVFP4 | Y (SM100/103/120/121) | Y (SM100/103) | N | Y (SM100/103) | Y (SM100/103/120/121) | N | Y (SM100/103, cu13 cutlass-dsl + NVSHMEM provider; per-expert alpha/norm_const + SwiGLU clamp) | N | Y (SM89-SM99) | Y | +| W4A16 NVFP4 | Y (SM80+, dequant-on-the-fly) | N | N | N | Y (SM120/121 via `CuteDslB12xFusedMoE`, needs flashinfer) | N | N | N | Y (SM89-SM99, BF16) | Y | +| W4A8 NVFP4 FP8 | N | Y (SM100/103) | N | N | N | N | N | N | N | N | +| W4A16 MXFP4 | Y (SM90) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | +| W4A8 MXFP4 FP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | +| W4A8 MXFP4 MXFP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | Y (SM100/103, requires `hidden_size % 512 == 0`) | N | N | N | N | +| W8A8 MXFP8 MXFP8 | Y (SM100/103) | N | N | N | N | N | N | N | N | N | +| W4A8 AWQ | Y (SM89/90) | N | N | N | N | N | N | N | N | N | +| W8A16 | Y (SM80+) | N | N | N | N | N | N | N | N | N | +| INT4 WoQ (W4AFP8) | N | N | N | N | N | N | N | N | N | N | § The unquantized `TRTLLMGenFusedMoE` path is not a TRTLLM-Gen kernel at all: it calls FlashInfer's `trtllm_bf16_moe` / `trtllm_bf16_routed_moe`, which is why it diff --git a/tensorrt_llm/_torch/modules/fused_moe/__init__.py b/tensorrt_llm/_torch/modules/fused_moe/__init__.py index 1702599e5680..c036d55f4381 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/__init__.py +++ b/tensorrt_llm/_torch/modules/fused_moe/__init__.py @@ -7,7 +7,6 @@ from .fused_moe_triton import TritonFusedMoE from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE from .fused_moe_vanilla import VanillaMoE -from .fused_moe_wide_ep import WideEPMoE from .impl_base import MoEImplBase from .interface import MoE, MoEWeightLoadingMode from .moe_load_balancer import (MoeLoadBalancer, @@ -61,5 +60,4 @@ "TritonFusedMoE", "TRTLLMGenFusedMoE", "VanillaMoE", - "WideEPMoE", ] diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 011f77500322..1986e81e9e95 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -18,7 +18,6 @@ from .fused_moe_triton import TritonFusedMoE from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE from .fused_moe_vanilla import VanillaMoE -from .fused_moe_wide_ep import WideEPMoE from .interface import MoE, MoEWeightLoadingMode from .mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm from .moe_load_balancer import get_moe_load_balancer @@ -98,9 +97,6 @@ def create_moe_backend( Returns: MoE: MoE backend instance """ - if moe_cls is WideEPMoE: - raise ValueError(WIDEEP_DEPRECATION_MESSAGE) - shapes = derive_moe_layer_shapes( model_config, num_experts=num_experts, @@ -499,7 +495,7 @@ def create_moe( communication_method=communication_method, ) - # WideEPMoE, TritonFusedMoE and VanillaMoE are not wrapped by ConfigurableMoE + # TritonFusedMoE and VanillaMoE are not wrapped by ConfigurableMoE # and own their communication and forward paths. if communication_method is not None: raise ValueError("communication_method requires ConfigurableMoE.") diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py deleted file mode 100755 index 6ef70fadd28d..000000000000 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py +++ /dev/null @@ -1,1006 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. - -import inspect -import os -from typing import Dict, List, Optional, Tuple, Union - -import torch - -from tensorrt_llm._mnnvl_utils import MnnvlMemory, MnnvlMoe, MoEAlltoallInfo -from tensorrt_llm._utils import is_sm_100f, local_mpi_size -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator - -from ...distributed import allgather, reducescatter -from ...expert_statistic import ExpertStatistic -from ...model_config import ModelConfig -from ...utils import (ActivationType, AuxStreamType, EventType, - Fp4QuantizedTensor) -from .deep_ep_utils import buffer_pool, deep_ep_installed -from .interface import AlltoallMethodType, MoE -from .ops import MoEOp, MoEOpSelector -from .quantization import (DeepSeekFP8BlockScalesFusedMoEMethod, - DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm, - FP8QDQFusedMoEMethod, FusedMoEQuantScalesW4A8, - MoEWeightLoadingMode, NVFP4CutlassFusedMoEMethod, - UnquantizedFusedMoEMethod, WInt4AFP8FusedMoEMethod) -from .routing import BaseMoeRoutingMethod - - -class WideEPMoE(MoE): - """ - Fused Mixture of Experts (MoE) Layer with for wide EP. - - Args: - num_experts (int): Number of experts in the MoE layer. - top_k (int): Number of top experts to select for each input token. - hidden_size (int): Size of the hidden state. - intermediate_size (int): Size of the intermediate state. - aux_stream_dict (Optional[Dict[AuxStreamType, torch.cuda.Stream]]): Auxiliary CUDA streams for overlapping. - dtype (Optional[torch.dtype]): Data type for the weights. - reduce_results (bool): Whether to reduce the results across devices. - model_config (ModelConfig): Configuration object for the model. - - Large-scale EP: - When we have redundant expert, we have more weight slots than `num_experts`, in that case, we separate the concepts of expert and slot. - Expert is the concept from model's perspective while slot is the concept from model engine's perspective. - There should be at least `num_experts` slots in the model engine. More than that is OK, in that case, some experts may have multiple replicas. - """ - - def __init__( - self, - *, - routing_method: BaseMoeRoutingMethod, - num_experts: int, - hidden_size: int, - intermediate_size: int, - dtype: Optional[torch.dtype] = None, - reduce_results: bool = False, - model_config: ModelConfig = ModelConfig(), - aux_stream_dict: Optional[Dict[AuxStreamType, - torch.cuda.Stream]] = None, - weight_loading_mode: MoEWeightLoadingMode = MoEWeightLoadingMode. - VANILLA, - apply_router_weight_on_input: bool = False, - layer_idx: Optional[int] = None, - swiglu_limit: Optional[torch.Tensor] = None, - swiglu_limit_scalar: Optional[float] = None, - activation_type: ActivationType = ActivationType.Swiglu, - ): - - super().__init__( - routing_method=routing_method, - num_experts=num_experts, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - dtype=dtype, - reduce_results=reduce_results, - model_config=model_config, - aux_stream_dict=aux_stream_dict, - weight_loading_mode=weight_loading_mode, - swiglu_limit=swiglu_limit, - swiglu_limit_scalar=swiglu_limit_scalar, - layer_idx=layer_idx, - activation_type=activation_type, - ) - - assert self.use_dp, "Attention DP should be used with WideEP." - assert self.parallel_size > 1, "WideEP should only be enabled with parallel_size > 1" - # If True, the router weight will be multiplied on the input rather than at the end of FC2 - self.apply_router_weight_on_input = apply_router_weight_on_input - - # Store original hidden size before any potential padding - self.unpadded_hidden_size = self.hidden_size - - self.use_cuda_graph = model_config.use_cuda_graph - - # moe_max_num_tokens is set in ModelConfig.__post_init__ if not specified - # The default value is max_num_tokens * dp_size - self.moe_max_num_tokens = model_config.moe_max_num_tokens - # The auxiliary CUDA stream and CUDA events are only used when MoE chunking is applied - default_moe_max_num_tokens = model_config.max_num_tokens * model_config.mapping.dp_size - if self.moe_max_num_tokens < default_moe_max_num_tokens: - self.aux_stream = aux_stream_dict[ - AuxStreamType. - MoeChunkingOverlap] if aux_stream_dict is not None else torch.cuda.Stream( - ) - self.event_dict = { - key: torch.cuda.Event() - for key in [EventType.Main, EventType.MoeChunkingOverlap] - } - else: - self.aux_stream = None - self.event_dict = None - - # The profiler converges on the same best tactic when the number of tokens is large enough. - # To avoid long profiling time, the max number of tokens used in the profiling is capped to - # around 16k tokens per expert, which is well into the compute bound domain. - self.tune_max_num_tokens = min( - self.moe_max_num_tokens, - 16384 * self.num_slots // routing_method.get_experts_per_token(), - ) - self.has_been_profiled = False - - self.alltoall_method_type = self.select_alltoall_method_type( - model_config.mapping, routing_method.experts_per_token, dtype, - model_config.use_cuda_graph) - logger.info_once( - f"{self.__class__.__name__} selects alltoall_method_type {self.alltoall_method_type!r}", - key="alltoall_method_type") - self.use_postquant_alltoall = False - self.use_low_precision_combine = False - if self.enable_alltoall: - self.use_postquant_alltoall = (os.environ.get( - "TRTLLM_MOE_POST_QUANT_ALLTOALLV", "1") == "1") - self.use_low_precision_combine = model_config.use_low_precision_moe_combine - - if self.alltoall_method_type == AlltoallMethodType.NVLinkTwoSided: - MnnvlMemory.initialize() - self.alltoall_workspace = MnnvlMoe.get_moe_workspaces( - model_config.mapping) - self.alltoall_prepare_workspace = MnnvlMoe.get_moe_prepare_workspace( - model_config.mapping) - elif self.alltoall_method_type == AlltoallMethodType.DeepEP: - self.deep_ep_buffer = buffer_pool.get_buffer( - model_config.mapping) - self.deep_ep_buffer.reserve(hidden_size, dtype) - elif self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: - self.deep_ep_max_num_tokens = int( - os.environ.get( - "TRTLLM_DEEP_EP_TOKEN_LIMIT", - str( - min(model_config.max_num_tokens, - self.moe_max_num_tokens)))) - # Set nvshmem queue pair depth larger than the number of on-flight WRs (ref: https://github.com/deepseek-ai/DeepEP/issues/427) - os.environ['NVSHMEM_QP_DEPTH'] = str( - 2 * (self.deep_ep_max_num_tokens + 1)) - self.deep_ep_buffer = buffer_pool.get_low_latency_buffer( - model_config.mapping) - self.deep_ep_buffer.reserve(self.deep_ep_max_num_tokens, - hidden_size, self.num_slots) - else: - raise NotImplementedError( - f"Unsupported alltoall method type: {self.alltoall_method_type!r}" - ) - - self.use_fused_finalize = not model_config.moe_disable_finalize_fusion - - self._weights_created = False - if not model_config.skip_create_weights_in_init: - self.create_weights() - - # MoE op will be lazily initialized when first accessed (see moe_op_impl property) - self._moe_op_impl = None - - def _check_configs(self): - assert self._weights_created - - if self.apply_router_weight_on_input: - assert self.routing_method.top_k == 1, "Current walkaround only supports top-1 routing" - - if self.quant_config and self.quant_config.quant_mode.has_any_quant( - exclude_kv_cache=True): - if not (self.quant_config.quant_mode.has_nvfp4() - | self.quant_config.quant_mode.has_fp8_block_scales() - | self.quant_config.quant_mode.has_fp8_qdq() - | self.quant_config.quant_mode. - is_int4_weight_only_per_group()): - raise ValueError( - f"unsupported quantization mode: {self.quant_config.quant_mode}" - ) - - @staticmethod - def select_alltoall_method_type(mapping: Mapping, top_k: int, - dtype: torch.dtype, - use_cuda_graph: bool) -> AlltoallMethodType: - - # Check if DeepEP is feasible for the given number of ranks - # DeepEP supports two modes: - # 1. Intranode: Single node with 2, 4, or 8 ranks - # 2. Internode: 2, 4, 8, or 16 nodes with 8 ranks per node - def is_deepep_feasible(num_ranks: int) -> bool: - NUM_INTRANODE_SUPPORTED_RANKS = {2, 4, 8} - REQUIRED_LOCAL_MPI_SIZE = 8 - NUM_INTERNODE_SUPPORTED_RDMA_RANKS = {2, 4, 8, 16} - mpi_size = local_mpi_size() - # Intranode cases - if num_ranks == mpi_size and num_ranks in NUM_INTRANODE_SUPPORTED_RANKS: - return True - # Internode cases - if mpi_size != REQUIRED_LOCAL_MPI_SIZE: - return False - num_rdma_nodes = num_ranks // mpi_size - return num_rdma_nodes in NUM_INTERNODE_SUPPORTED_RDMA_RANKS - - all2all_method_type_env = os.environ.get("TRTLLM_FORCE_ALLTOALL_METHOD") - if all2all_method_type_env is not None: - alltoall_method_type = AlltoallMethodType[all2all_method_type_env] - if alltoall_method_type == AlltoallMethodType.NVLinkOneSided: - raise NotImplementedError( - "NVLinkOneSided is not supported for WideEPMoE. Please use NVLinkTwoSided or switch to CutlassFusedMoE." - ) - return alltoall_method_type - - if not mapping.enable_attention_dp: - return AlltoallMethodType.NotEnabled - - if mapping.tp_size == 1: - return AlltoallMethodType.NotEnabled - - if mapping.moe_ep_size <= top_k: - return AlltoallMethodType.NotEnabled - - if MnnvlMemory.supports_mnnvl(): - return AlltoallMethodType.NVLinkTwoSided - - if os.environ.get("TRTLLM_CAN_USE_DEEP_EP", "0") == "1": - if deep_ep_installed and dtype == torch.bfloat16: - # Choose DeepEP if feasible - if is_deepep_feasible(mapping.moe_ep_size): - return AlltoallMethodType.DeepEP - return AlltoallMethodType.DeepEPLowLatency - - return AlltoallMethodType.NotEnabled - - @property - def has_w4afp8(self): - assert self._weights_created - return self.quant_config and self.quant_config.quant_mode.is_int4_weight_only_per_group( - ) - - @property - def enable_alltoall(self): - """ enable_alltoall (bool): whether to enable alltoall instead of allgather/reducescatter - """ - return self.alltoall_method_type != AlltoallMethodType.NotEnabled - - def calculate_num_chunks(self, all_rank_num_tokens: List[int]) -> int: - num_rows = sum(all_rank_num_tokens) - return (num_rows + self.moe_max_num_tokens - - 1) // self.moe_max_num_tokens - - def can_use_alltoall(self, all_rank_num_tokens, all_rank_max_num_tokens): - if self.alltoall_method_type == AlltoallMethodType.NVLinkTwoSided: - return True - - # Disable alltoall when chunking is used - if self.calculate_num_chunks(all_rank_num_tokens) > 1: - return False - - # For DeepEPLowLatency, check if tokens exceed the threshold - if (self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency - and all_rank_max_num_tokens > self.deep_ep_max_num_tokens): - return False - - return self.enable_alltoall - - def deep_ep_low_latency_dispatch_modify_output_to_adapt_fused_moe( - self, x: torch.Tensor, x_sf: Optional[torch.Tensor], - recv_expert_count: torch.Tensor, final_scales_dtype: torch.dtype - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, - torch.Tensor]: - # x shape: [#local experts, EP size * all_rank_max_num_tokens, hidden_size] - # recv_expert_count shape: [#local experts] - - # Adapter between `torch.ops.trtllm.fused_moe` and DeepEP - # TODO: remove the adapter by changing `torch.ops.trtllm.fused_moe` API - mask = torch.arange(x.shape[1], - dtype=torch.int32, device=x.device).expand( - x.shape[0], - x.shape[1]) < recv_expert_count.unsqueeze(1) - token_selected_slots = torch.where( - mask, - torch.arange(x.shape[0] * self.mapping.moe_ep_rank, - x.shape[0] * (self.mapping.moe_ep_rank + 1), - dtype=torch.int32, - device=x.device).unsqueeze(1), self.num_slots) - x = x.reshape(x.shape[0] * x.shape[1], x.shape[2]) - if x_sf is not None: - x_sf = x_sf.reshape(x_sf.shape[0] * x_sf.shape[1], x_sf.shape[2]) - # Cheat the fused_moe API with fake top_k=1 - token_selected_slots = token_selected_slots.view(x.shape[0], 1) - token_final_scales = torch.ones_like(token_selected_slots, - dtype=final_scales_dtype) - return x, x_sf, token_selected_slots, token_final_scales - - def _get_quant_method(self): - if self.quant_config is not None and self.quant_config.layer_quant_mode.has_any_quant( - exclude_kv_cache=True): - if self.quant_config.layer_quant_mode.has_fp8_qdq(): - return FP8QDQFusedMoEMethod() - elif self.quant_config.layer_quant_mode.has_fp8_block_scales(): - if is_sm_100f(): - return DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm() - else: - return DeepSeekFP8BlockScalesFusedMoEMethod() - elif self.quant_config.layer_quant_mode.has_nvfp4(): - return NVFP4CutlassFusedMoEMethod() - elif self.quant_config.layer_quant_mode.is_int4_weight_only_per_group( - ): - return WInt4AFP8FusedMoEMethod() - else: - raise ValueError( - f"Unsupported quantization mode: {self.quant_config.quant_mode}" - ) - else: - return UnquantizedFusedMoEMethod() - - def create_weights(self): - if self._weights_created: - return - - self.quant_method = self._get_quant_method() - self.quant_method.create_weights(self) - - self._weights_created = True - self._check_configs() - - @property - def moe_op_impl(self) -> MoEOp: - """ - Lazily initialize and return the MoE op. - - The op is selected based on hardware capabilities and quantization - configuration, which are only available after weights are created. - """ - if self._moe_op_impl is None: - assert self._weights_created, "Weights must be created before accessing moe_op" - self._moe_op_impl = MoEOpSelector.select_op(self) - return self._moe_op_impl - - def reducescatter_or_allreduce( - self, - inputs, - use_all_to_all: bool, - all_rank_num_tokens: Optional[List[int]] = None, - use_dp_padding: Optional[bool] = None, - ): - outputs = inputs - if not use_all_to_all: - if self.enable_dummy_allreduce: - self.dummy_allreduce() - outputs = reducescatter( - inputs, - self.mapping, - dim=0, - sizes=None if use_dp_padding else all_rank_num_tokens) - return outputs - - def is_post_quant_all2all_supported(self): - if not self.use_postquant_alltoall: - return False - if self.alltoall_method_type == AlltoallMethodType.NVLinkTwoSided: - return True - elif self.alltoall_method_type == AlltoallMethodType.DeepEP: - return self.has_nvfp4 - elif self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: - return self.has_fp8_qdq or self.has_nvfp4 or self.has_w4afp8 - else: - return False - - def is_low_precision_combine_supported(self): - if not self.use_low_precision_combine: - return False - if self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: - return self.has_fp8_qdq or self.has_nvfp4 or self.has_w4afp8 - return False - - def forward_chunk( - self, - x: Union[torch.Tensor, Fp4QuantizedTensor], - router_logits: torch.Tensor, - use_all_to_all: bool, - input_ids: Optional[torch.IntTensor] = None, - output_dtype: Optional[torch.dtype] = None, - all_rank_num_tokens: Optional[List[int]] = None, - use_dp_padding: Optional[bool] = None, - repeating_info: Tuple = (True, True), - alltoall_result_do_sum: bool = True, - ) -> torch.Tensor: - all_rank_max_num_tokens = max(all_rank_num_tokens) - if isinstance(x, Fp4QuantizedTensor): - assert output_dtype is not None - else: - output_dtype = x.dtype - - is_first_call, is_last_call = repeating_info - - self._load_balancer_start_wait_gpu_stage(is_first_call) - - if not use_all_to_all or self.alltoall_method_type != AlltoallMethodType.NVLinkTwoSided: - alltoall_result_do_sum = True - - weight_dtype = self.w3_w1_weight.dtype - - token_selected_experts, token_final_scales = self.routing_method.apply( - router_logits, input_ids) - - assert token_selected_experts.shape[ - 1] == self.routing_method.experts_per_token - assert token_selected_experts.shape == token_final_scales.shape - assert token_selected_experts.shape[0] == router_logits.shape[0] - assert token_final_scales.dtype == torch.float32 - assert token_selected_experts.dtype == torch.int32 - - if self.apply_router_weight_on_input: - assert x.dtype != torch.float8_e4m3fn, "Current workaround for apply_router_weight_on_input does not support fp8 input" - x = x * token_final_scales.to(x.dtype) - # TODO: remove this once we have correct fusedmoe kernel ready - if self.alltoall_method_type in ( - AlltoallMethodType.DeepEP, - AlltoallMethodType.DeepEPLowLatency): - # DeepEP doesn't support token_final_scales is None - token_final_scales = torch.ones_like(token_final_scales) - else: - token_final_scales = None - - if self.layer_load_balancer: - self._load_balancer_done_wait_gpu_stage(is_first_call) - ignore_allreduce = self.alltoall_method_type == AlltoallMethodType.NVLinkTwoSided - self._load_balancer_update_statistic(token_selected_experts, - is_first_call, is_last_call, - ignore_allreduce) - token_selected_slots = self._load_balancer_route( - token_selected_experts, self.use_dp) - else: - token_selected_slots = token_selected_experts - - # If load balancer is disabled, the statistics are collected from expert IDs. - # If load balancer is enabled, the statistics are collected from expert slot IDs. - ExpertStatistic.set_layer(self.layer_idx) - ExpertStatistic.maybe_add_info(self.num_slots, token_selected_slots) - token_selected_slots = get_calibrator().maybe_collect_or_replay_slots( - self.num_slots, token_selected_slots) - - use_allgather = not use_all_to_all - - # If alltoall is disabled, we need also disable use_postquant_alltoall - use_postquant_alltoall = use_all_to_all and self.is_post_quant_all2all_supported( - ) - - # Prepare additional information for profiling in case padding is applied when using alltoall. - # Only the non-alltoall case is considered for profiling in the warmup phase. - # Therefore, to get the correct tactics during the actual inference, the inputs to the tuner should be the same as when not using alltoall. - if use_all_to_all: - if all_rank_num_tokens is not None: - tuner_num_tokens = sum(all_rank_num_tokens) - else: - tuner_num_tokens = x.shape[0] * self.mapping.tp_size - tuner_top_k = token_selected_slots.shape[1] - else: - tuner_num_tokens = None - tuner_top_k = None - alltoall_info = None - if use_all_to_all: - if self.alltoall_method_type == AlltoallMethodType.NVLinkTwoSided: - if self.enable_dummy_allreduce: - self.dummy_allreduce() - token_count = x.shape[0] - if is_last_call: - loadbalancer_local_statistic_info = self._load_balancer_get_local_statistic_tensor( - ) - else: - loadbalancer_local_statistic_info = None - token_selected_slots, gathered_loadbalancer_local_statistic_info, alltoall_info = \ - self.alltoall_prepare(all_rank_max_num_tokens, - token_selected_slots, - loadbalancer_local_statistic_info) - - if gathered_loadbalancer_local_statistic_info is not None: - gathered_loadbalancer_local_statistic_info = gathered_loadbalancer_local_statistic_info.view( - (self.mapping.moe_ep_size, self.num_experts)) - self._load_balancer_update_statistic_with_gathered_statistic( - gathered_loadbalancer_local_statistic_info) - elif self.alltoall_method_type == AlltoallMethodType.DeepEP: - if not use_postquant_alltoall: - x, recv_topk_idx, token_final_scales, num_recv_tokens_per_expert_list, deep_ep_handle = \ - self.deep_ep_buffer.dispatch(x, token_selected_slots, token_final_scales, self.num_slots, - self.expert_size_per_partition * self.mapping.moe_ep_rank, all_rank_max_num_tokens, self.ep_size, self.use_cuda_graph) - padded, x, _, token_selected_slots, token_final_scales = self.pad_empty_recv_tensors( - x, None, recv_topk_idx, token_final_scales) - elif self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: - if not use_postquant_alltoall: - deep_ep_topk_idx = token_selected_slots - deep_ep_topk_weights = token_final_scales - assert all_rank_max_num_tokens <= self.deep_ep_max_num_tokens - x, recv_expert_count, deep_ep_handle = \ - self.deep_ep_buffer.low_latency_dispatch(x, deep_ep_topk_idx, all_rank_max_num_tokens, self.num_slots) - x, _, token_selected_slots, token_final_scales = self.deep_ep_low_latency_dispatch_modify_output_to_adapt_fused_moe( - x, None, recv_expert_count, token_final_scales.dtype) - - x_sf = None - x_row = x.shape[0] - x_col = x.shape[1] - if self.has_any_quant: - if self.has_fp8_qdq: - x, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( - x, self.fc31_input_dequant) - elif self.has_nvfp4: - if isinstance(x, Fp4QuantizedTensor): - assert not x.is_sf_swizzled, "Fp4QuantizedTensor should not be swizzled before allgather or postquant alltoall" - x, x_sf = x.fp4_tensor, x.scaling_factor - x_row = x.shape[0] - # note: we use uint8 to store 2 fp4 values - x_col = x.shape[1] * 2 - else: - # for both postquant alltoall and allgather, we need non swizzle layout - x_row = x.shape[0] - x_col = x.shape[1] - x, x_sf = torch.ops.trtllm.fp4_quantize( - x, - self.fc31_input_scale, - self.scaling_vector_size, - sfUseUE8M0=False, - isSfSwizzledLayout=False) - x_sf = x_sf.view((x_row, -1)) - - elif self.has_deepseek_fp8_block_scales: - pass - elif self.has_w4afp8: - weight_dtype = torch.quint4x2 - else: - raise ValueError( - f"unsupported quantization mode: {self.quant_config.quant_mode}" - ) - - if use_allgather: - # using allgather case. - if self.enable_dummy_allreduce: - self.dummy_allreduce() - x, x_sf, token_selected_slots, token_final_scales = allgather( - [ - x, - x_sf, - token_selected_slots, - token_final_scales, - ], - self.mapping, - dim=0, - sizes=None if use_dp_padding else all_rank_num_tokens) - x_row = x.shape[0] - - w3_w1_weight = self.w3_w1_weight - w2_weight = self.w2_weight - quant_scales = self.quant_scales - - if self.alltoall_method_type == AlltoallMethodType.NVLinkTwoSided: - top_k = self.routing_method.experts_per_token - x, x_sf, token_selected_slots, token_final_scales = self.alltoall_dispatch( - x, x_sf, token_selected_slots, token_final_scales, - all_rank_max_num_tokens, top_k, alltoall_info) - - if use_postquant_alltoall: - if self.alltoall_method_type == AlltoallMethodType.NVLinkTwoSided: - pass - elif self.alltoall_method_type == AlltoallMethodType.DeepEP: - assert self.has_nvfp4, "DeepEP postquant alltoall should have nvfp4" - if x_sf is not None: - # Adapter between `x_sf` and DeepEP - # TODO: remove the adapter by adding dtype support to DeepEP - x_sf_dtype = x_sf.dtype - x_sf = x_sf.view(torch.float32) - (x, x_sf), recv_topk_idx, token_final_scales, num_recv_tokens_per_expert_list, deep_ep_handle = \ - self.deep_ep_buffer.dispatch((x, x_sf), token_selected_slots, token_final_scales, self.num_slots, - self.expert_size_per_partition * self.mapping.moe_ep_rank, all_rank_max_num_tokens, self.ep_size, self.use_cuda_graph) - padded, x, x_sf, token_selected_slots, token_final_scales = self.pad_empty_recv_tensors( - x, x_sf, recv_topk_idx, token_final_scales) - if x_sf is not None: - x_sf = x_sf.view(x_sf_dtype) - elif self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: - assert self.has_any_quant, "DeepEPLowLatency postquant alltoall should have quantization" - assert all_rank_max_num_tokens <= self.deep_ep_max_num_tokens - deep_ep_topk_idx = token_selected_slots - deep_ep_topk_weights = token_final_scales - if self.has_fp8_qdq: - assert x.dtype == torch.float8_e4m3fn and x_sf is None, "x should be torch.float8_e4m3fn and x_sf should be None in fp8 postquant alltoall" - x = x.view(torch.bfloat16) - x, recv_expert_count, deep_ep_handle = \ - self.deep_ep_buffer.low_latency_dispatch(x, deep_ep_topk_idx, all_rank_max_num_tokens, self.num_slots) - x = x.view(torch.float8_e4m3fn) - elif self.has_nvfp4: - token_num = x_row - hidden_size = x_col - assert x.dtype == torch.uint8 and x_sf is not None and x_sf.dtype == torch.uint8 - assert hidden_size % 32 == 0, "HiddenSize should be divisible by 32 in nvfp4 postquant alltoall" - assert x_sf.shape[0] == token_num and x_sf.shape[ - 1] == hidden_size // 16 - assert x.shape[0] == token_num and x.shape[ - 1] == hidden_size // 2 - - x, x_sf, recv_expert_count, deep_ep_handle = \ - self.deep_ep_buffer.low_latency_dispatch_fp4(x, x_sf, deep_ep_topk_idx, all_rank_max_num_tokens, self.num_slots) - assert x.dtype == torch.uint8 and x_sf.dtype == torch.uint8 - assert x.dim() == 3 and x_sf.dim() == 3 - assert x.shape[2] == hidden_size // 2 and x_sf.shape[ - 2] == hidden_size // 16 - elif self.has_w4afp8: - assert isinstance(quant_scales, FusedMoEQuantScalesW4A8) - pre_quant_scales = quant_scales.pre_quant_scale_1 - assert pre_quant_scales.shape == ( - 1, x.shape[1]) and pre_quant_scales.dtype == x.dtype - x = (x * pre_quant_scales).to(torch.float8_e4m3fn).view( - torch.bfloat16) - x, recv_expert_count, deep_ep_handle = \ - self.deep_ep_buffer.low_latency_dispatch(x, deep_ep_topk_idx, all_rank_max_num_tokens, self.num_slots) - x = x.view(torch.float8_e4m3fn) - else: - raise ValueError( - f"unsupported quantization mode in postquant alltoall: {self.quant_config.quant_mode}" - ) - x, x_sf, token_selected_slots, token_final_scales = self.deep_ep_low_latency_dispatch_modify_output_to_adapt_fused_moe( - x, x_sf, recv_expert_count, token_final_scales.dtype) - else: - raise NotImplementedError( - f"Unsupported alltoall method type: {self.alltoall_method_type!r}" - ) - - final_hidden_states = self.moe_op_impl.run_moe( - self, - x, - token_selected_slots, - token_final_scales, - w3_w1_weight.view(weight_dtype), - None, # w3_w1_bias - w2_weight.view(weight_dtype), - None, # w2_bias - output_dtype, - quant_scales=quant_scales, - use_all_to_all=use_all_to_all, - input_sf=x_sf, - swizzled_input_sf=False, - min_latency_mode=False, - use_fused_finalize=self.use_fused_finalize, - tuner_num_tokens=tuner_num_tokens, - tuner_top_k=tuner_top_k, - ) - - self._load_balancer_start_set_cpu_stage(is_last_call) - - # Only in cutlass_min_latency_mode, the output is a list of tensors. - # Otherwise, the output should be unpacked as a single tensor. - final_hidden_states = final_hidden_states[0] - - if use_all_to_all: - if self.alltoall_method_type == AlltoallMethodType.NVLinkTwoSided: - if self.enable_dummy_allreduce: - self.dummy_allreduce() - final_hidden_states = self.alltoall_combine( - final_hidden_states, alltoall_info, token_count, - alltoall_result_do_sum) - elif self.alltoall_method_type == AlltoallMethodType.DeepEP: - final_hidden_states = self.unpad_tensors( - padded, final_hidden_states) - final_hidden_states = self.deep_ep_buffer.combine( - final_hidden_states, deep_ep_handle) - elif self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: - num_tokens_per_expert_for_fused_moe = self.mapping.moe_ep_size * all_rank_max_num_tokens - final_hidden_states = final_hidden_states.view( - self.expert_size_per_partition, - num_tokens_per_expert_for_fused_moe, self.hidden_size) - if self.is_low_precision_combine_supported(): - precision = "fp8" - global_scales = None - if self.has_nvfp4: - precision = "nvfp4" - global_scales = torch.ops.trtllm.calculate_nvfp4_global_scale( - final_hidden_states, recv_expert_count) - final_hidden_states = self.deep_ep_buffer.low_latency_combine_low_precision( - precision, final_hidden_states, global_scales, - deep_ep_topk_idx, deep_ep_topk_weights, deep_ep_handle) - else: - final_hidden_states = self.deep_ep_buffer.low_latency_combine( - final_hidden_states, deep_ep_topk_idx, - deep_ep_topk_weights, deep_ep_handle) - else: - raise NotImplementedError( - f"Unsupported alltoall method type: {self.alltoall_method_type!r}" - ) - - self._load_balancer_done_set_cpu_stage(is_last_call) - - return final_hidden_states - - def forward_impl( - self, - x: Union[torch.Tensor, Fp4QuantizedTensor], - router_logits: torch.Tensor, - *, - input_ids: Optional[torch.IntTensor] = None, - do_finalize: bool = True, - output_dtype: Optional[torch.dtype] = None, - all_rank_num_tokens: Optional[List[int]] = None, - use_dp_padding: Optional[bool] = None, - alltoall_result_do_sum: bool = True, - **kwargs, - ) -> torch.Tensor: - assert all_rank_num_tokens is not None - assert use_dp_padding is not None - - all_rank_max_num_tokens = max(all_rank_num_tokens) - - if use_dp_padding: - all_rank_num_tokens_padded = [all_rank_max_num_tokens - ] * len(all_rank_num_tokens) - else: - all_rank_num_tokens_padded = all_rank_num_tokens - - # in case of num_rows is larger than max_chunk_size, we need to split the input into multiple chunks - num_chunks = self.calculate_num_chunks(all_rank_num_tokens_padded) - use_all_to_all = self.can_use_alltoall(all_rank_num_tokens_padded, - all_rank_max_num_tokens) - if num_chunks == 1: - is_first_call = self.repeat_idx == 0 - is_last_call = self.repeat_idx == self.repeat_count - 1 - outputs = self.forward_chunk( - x, - router_logits, - use_all_to_all, - input_ids=input_ids, - output_dtype=output_dtype, - all_rank_num_tokens=all_rank_num_tokens_padded, - use_dp_padding=use_dp_padding, - repeating_info=(is_first_call, is_last_call), - alltoall_result_do_sum=alltoall_result_do_sum) - outputs = self.reducescatter_or_allreduce( - outputs, - use_all_to_all, - all_rank_num_tokens=all_rank_num_tokens_padded, - use_dp_padding=use_dp_padding) - else: - - def split_chunk(split_token_num: int, split_num_chunks: int): - val_div = split_token_num // split_num_chunks - val_mod = split_token_num % split_num_chunks - split_chunk_size_list = [val_div + 1] * val_mod + [val_div] * ( - split_num_chunks - val_mod) - return split_chunk_size_list - - all_rank_chunk_size_list = [ - split_chunk(val, num_chunks) - for val in all_rank_num_tokens_padded - ] - all_rank_num_tokens_list = [[ - val[idx_chunk] for val in all_rank_chunk_size_list - ] for idx_chunk in range(num_chunks)] - all_rank_max_num_tokens_list = split_chunk(all_rank_max_num_tokens, - num_chunks) - chunk_size_list = all_rank_chunk_size_list[self.parallel_rank] - if use_all_to_all: - all_rank_num_tokens_list = [[ - 1 if val == 0 else val for val in val_list - ] for val_list in all_rank_num_tokens_list] - all_rank_max_num_tokens_list = [ - 1 if val == 0 else val - for val in all_rank_max_num_tokens_list - ] - - x_list = x.split(chunk_size_list) - router_logits_list = router_logits.split(chunk_size_list) - input_ids_list = input_ids.split( - chunk_size_list) if input_ids is not None else [None - ] * num_chunks - - if not use_all_to_all: - self.event_dict[EventType.Main].record() - with torch.cuda.stream(self.aux_stream): - self.event_dict[EventType.Main].wait() - - outputs_list = [] - # Postpone reduce-scatter/all-reduce to the next iteration to achieve better overlap - for idx_chunk, (x, router_logits, input_ids_chunk) in enumerate( - zip(x_list, router_logits_list, input_ids_list)): - is_first_call = idx_chunk == 0 and self.repeat_idx == 0 - is_last_call = idx_chunk == num_chunks - 1 and self.repeat_idx == self.repeat_count - 1 - if not use_all_to_all: - if idx_chunk % 2 == 0: - with torch.cuda.stream(self.aux_stream): - outputs = self.forward_chunk( - x, - router_logits, - use_all_to_all, - input_ids=input_ids_chunk, - all_rank_num_tokens=all_rank_num_tokens_list[ - idx_chunk], - use_dp_padding=use_dp_padding, - repeating_info=(is_first_call, is_last_call), - alltoall_result_do_sum=alltoall_result_do_sum) - if idx_chunk > 0: - outputs_list[-1] = self.reducescatter_or_allreduce( - outputs_list[-1], - use_all_to_all, - all_rank_num_tokens=all_rank_num_tokens_list[ - idx_chunk - 1], - use_dp_padding=use_dp_padding) - else: - outputs = self.forward_chunk( - x, - router_logits, - use_all_to_all, - input_ids=input_ids_chunk, - all_rank_num_tokens=all_rank_num_tokens_list[ - idx_chunk], - use_dp_padding=use_dp_padding, - repeating_info=(is_first_call, is_last_call), - alltoall_result_do_sum=alltoall_result_do_sum) - with torch.cuda.stream(self.aux_stream): - outputs_list[-1] = self.reducescatter_or_allreduce( - outputs_list[-1], - use_all_to_all, - all_rank_num_tokens=all_rank_num_tokens_list[ - idx_chunk - 1], - use_dp_padding=use_dp_padding) - else: - outputs = self.forward_chunk( - x, - router_logits, - use_all_to_all, - input_ids=input_ids_chunk, - all_rank_num_tokens=all_rank_num_tokens_list[idx_chunk], - repeating_info=(is_first_call, is_last_call), - alltoall_result_do_sum=alltoall_result_do_sum) - - outputs_list.append(outputs) - if not use_all_to_all: - if num_chunks % 2 == 0: - outputs_list[-1] = self.reducescatter_or_allreduce( - outputs_list[-1], - use_all_to_all, - all_rank_num_tokens=all_rank_num_tokens_list[-1], - use_dp_padding=use_dp_padding) - else: - with torch.cuda.stream(self.aux_stream): - outputs_list[-1] = self.reducescatter_or_allreduce( - outputs_list[-1], - use_all_to_all, - all_rank_num_tokens=all_rank_num_tokens_list[-1], - use_dp_padding=use_dp_padding) - with torch.cuda.stream(self.aux_stream): - self.event_dict[EventType.MoeChunkingOverlap].record() - self.event_dict[EventType.MoeChunkingOverlap].wait() - outputs = torch.cat(outputs_list) - rank = self.parallel_rank - outputs = outputs[:all_rank_num_tokens[rank]] - self.repeat_idx = 0 if self.repeat_idx == self.repeat_count - 1 else self.repeat_idx + 1 - return outputs - - def alltoall_prepare(self, all_rank_max_num_tokens: int, - token_selected_slots: torch.Tensor, - local_statistic_tensor: Optional[torch.Tensor]): - top_k = self.routing_method.experts_per_token - - alltoall_info, gathered_local_statistic_tensor = MnnvlMoe.mnnvl_moe_alltoallv_prepare_without_allgather( - token_selected_slots, local_statistic_tensor, - self.alltoall_prepare_workspace, all_rank_max_num_tokens, - self.ep_rank, self.ep_size, self.num_experts, self.num_slots, top_k) - - return token_selected_slots, gathered_local_statistic_tensor, alltoall_info - - def alltoall_dispatch(self, x: torch.Tensor, x_sf: Optional[torch.Tensor], - token_selected_slots: torch.Tensor, - token_final_scales: Optional[torch.Tensor], - all_rank_max_num_tokens: int, top_k: int, - alltoall_info: MoEAlltoallInfo): - - x, x_sf, token_selected_slots, token_final_scales = MnnvlMoe.mnnvl_moe_alltoallv( - [x, x_sf, token_selected_slots, token_final_scales], alltoall_info, - self.alltoall_workspace, self.ep_rank, self.ep_size) - - torch.ops.trtllm.memset_expert_ids(token_selected_slots, - alltoall_info.recv_rank_count_cumsum, - all_rank_max_num_tokens, top_k, - self.num_slots, self.ep_size) - - return x, x_sf, token_selected_slots, token_final_scales - - def alltoall_combine(self, final_hidden_states: torch.Tensor, - alltoall_info: MoEAlltoallInfo, token_count: int, - alltoall_result_do_sum: bool): - top_k = self.routing_method.experts_per_token - if isinstance(final_hidden_states, list): - final_hidden_states = final_hidden_states[0] - final_hidden_states = MnnvlMoe.mnnvl_moe_alltoallv_combine( - final_hidden_states, - alltoall_info, - self.alltoall_workspace, - ep_rank=self.ep_rank, - ep_size=self.ep_size, - top_k=top_k, - token_count=token_count, - use_low_precision_combine=self.use_low_precision_combine, - do_reduce=alltoall_result_do_sum) - - return final_hidden_states - - def pad_empty_recv_tensors( - self, x: torch.Tensor, x_sf: Optional[torch.Tensor], - recv_topk_idx: torch.Tensor, token_final_scales: torch.Tensor - ) -> Tuple[bool, torch.Tensor, Optional[torch.Tensor], torch.Tensor, - torch.Tensor]: - """ - Pad the output of DeepEP `dispatch` if the output length is zero. - We can remove the adapter if both `fused_moe` op and `swizzle_sf` - accept zero-length inputs. - """ - if x.shape[0] == 0: - padded = True - x = torch.zeros((1, x.shape[1]), dtype=x.dtype, device=x.device) - if x_sf is not None: - x_sf = torch.zeros((1, x_sf.shape[1]), - dtype=x_sf.dtype, - device=x_sf.device) - recv_topk_idx = torch.full((1, recv_topk_idx.shape[1]), - self.num_slots, - dtype=recv_topk_idx.dtype, - device=recv_topk_idx.device) - token_final_scales = torch.ones((1, token_final_scales.shape[1]), - dtype=token_final_scales.dtype, - device=token_final_scales.device) - else: - padded = False - return padded, x, x_sf, recv_topk_idx, token_final_scales - - def unpad_tensors(self, padded: bool, - final_hidden_states: torch.Tensor) -> torch.Tensor: - if padded: - final_hidden_states = final_hidden_states[:0] - return final_hidden_states - - def _supports_load_balancer(self) -> bool: - """WideEPMoE supports load balancer.""" - return True - - def load_weights(self, - weights: List[Dict], - allow_partial_loading: bool = False): - assert self._weights_created - assert len(weights) == 1 - weights = weights[0] - - kargs = {} - if "allow_partial_loading" in inspect.getfullargspec( - self.quant_method.load_weights).args: - kargs["allow_partial_loading"] = allow_partial_loading - self.quant_method.load_weights(self, weights, self.weight_loading_mode, - **kargs) - - def forward_fake( - self, - x: Union[torch.Tensor, Fp4QuantizedTensor], - router_logits: torch.Tensor, - *, - do_finalize: bool = True, - output_dtype: Optional[torch.dtype] = None, - all_rank_num_tokens: Optional[List[int]] = None, - use_dp_padding: Optional[bool] = None, - alltoall_result_do_sum: bool = True, - **kwargs, - ) -> Union[torch.Tensor, List[torch.Tensor]]: - moe_output = super().forward_fake( - x, - router_logits, - do_finalize=do_finalize, - output_dtype=torch.bfloat16, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=use_dp_padding, - **kwargs) - if self.alltoall_method_type == AlltoallMethodType.NVLinkTwoSided and not alltoall_result_do_sum: - shape = moe_output.shape - top_k = self.routing_method.experts_per_token - new_shape = [shape[0], top_k, shape[1]] - return moe_output.new_empty(new_shape) - else: - return moe_output diff --git a/tensorrt_llm/_torch/modules/fused_moe/interface.py b/tensorrt_llm/_torch/modules/fused_moe/interface.py index 6babb96f8db9..c70d88aca832 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/interface.py +++ b/tensorrt_llm/_torch/modules/fused_moe/interface.py @@ -16,7 +16,7 @@ import os import weakref from abc import abstractmethod -from enum import Enum, IntEnum +from enum import Enum from typing import Dict, List, Optional, Tuple, Union, final import torch @@ -91,22 +91,6 @@ class MoEWeightLoadingMode(Enum): W4A8_CUSTOM = 2 -# The type of alltoall method -class AlltoallMethodType(IntEnum): - # Not available - NotEnabled = 0 - # NVLink One-Sided - NVLinkOneSided = 1 - # NVLink Two-Sided - NVLinkTwoSided = 2 - # DeepEP intranode or internode: CUDA Graphs are supported, IBGDA is required by internode - DeepEP = 3 - # DeepEP low latency: CUDA Graphs are supported, IBGDA is required - DeepEPLowLatency = 4 - # NCCL EP: Low-latency expert parallelism via NCCL EP library - NcclEP = 5 - - class MoESchedulerKind(Enum): """Selects which forward-execution scheduler ConfigurableMoE picks for a backend. diff --git a/tensorrt_llm/_torch/modules/fused_moe/ops/__init__.py b/tensorrt_llm/_torch/modules/fused_moe/ops/__init__.py deleted file mode 100644 index 3c6f3bc3a799..000000000000 --- a/tensorrt_llm/_torch/modules/fused_moe/ops/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# 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. -"""MoE ops module for different computation implementations.""" - -from .moe_op import MoEOp, MoEOpSelector -from .moe_op_cutlass import CutlassMoEOp -from .moe_op_deepgemm import DeepGemmMoEOp - -__all__ = ['MoEOp', 'MoEOpSelector', 'CutlassMoEOp', 'DeepGemmMoEOp'] diff --git a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py b/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py deleted file mode 100644 index c9771996d55b..000000000000 --- a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py +++ /dev/null @@ -1,235 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# 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. -""" -MoE Op abstraction for supporting different MoE computation implementations. -This module provides a unified interface for different MoE ops (Cutlass, DeepGemm, etc.) -""" - -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, List, Optional - -import torch - -from tensorrt_llm._utils import is_sm_100f - -if TYPE_CHECKING: - from ..interface import MoE - - -class MoEOp(ABC): - """Abstract base class for MoE computation ops. - - This class provides a strategy pattern for different MoE computation implementations. - It is used by MoE modules (like WideEPMoE) to delegate the actual computation. - - Note: MoEOp is NOT a MoE module itself, but a computation strategy. - The actual MoE module (e.g., WideEPMoE) inherits from MoE and uses MoEOp - for the computation implementation. - """ - - # Op-specific abstract methods - @abstractmethod - def finalize_tactic( - self, - module: 'MoE', - tuner_input: torch.Tensor, - output_dtype: torch.dtype, - min_latency_mode: bool = False, - use_fused_finalize: bool = True, - tuner_top_k: Optional[int] = None, - ) -> None: - """ - Finalize tactics for the MoE computation. - For Cutlass op, this includes profiling and tactic selection. - For DeepGemm op, this can be a no-op. - - Args: - module: The MoE module containing MoE configurations - tuner_input: Real input used for tuning (same shape/layout as non-alltoall) - output_dtype: Output dtype for tuner run - min_latency_mode: Whether to profile for min-latency path - use_fused_finalize: Whether to use fused finalize - tuner_top_k: Top-k value for tuning (Cutlass specific) - """ - - @abstractmethod - def compute_moe( - self, - module: 'MoE', - # Input tensors - x: torch.Tensor, - token_selected_slots: torch.Tensor, - token_final_scales: Optional[torch.Tensor], - # Weight tensors - w3_w1_weight: torch.Tensor, - w3_w1_bias: Optional[torch.Tensor], - w2_weight: torch.Tensor, - w2_bias: Optional[torch.Tensor], - # Output configuration - output_dtype: torch.dtype, - # Quantization parameters - quant_scales: List[torch.Tensor], - use_all_to_all: bool, - input_sf: Optional[torch.Tensor] = None, - swizzled_input_sf: bool = True, - # Performance tuning (only runtime-variable parameters) - min_latency_mode: bool = False, - use_fused_finalize: bool = True, - tuner_num_tokens: Optional[int] = None, - tuner_top_k: Optional[int] = None, - **kwargs) -> torch.Tensor: - """ - Perform the actual MoE computation. - - Configuration parameters (tp_size, ep_size, swiglu params, etc.) are - automatically extracted from the module parameter. - - Args: - module: MoE module containing configuration and parameters. - x: Input tensor - token_selected_slots: Selected expert slots - token_final_scales: Scaling factors - w3_w1_weight: Fused gate and up projection weights - w3_w1_bias: Optional bias - w2_weight: Down projection weights - w2_bias: Optional bias - output_dtype: Output data type - quant_scales: Quantization scales - use_all_to_all: Whether to use all-to-all communication - input_sf: Input scaling factor - swizzled_input_sf: Whether input_sf is swizzled - min_latency_mode: Use minimum latency optimizations - use_fused_finalize: Use fused finalization - tuner_num_tokens: Number of tokens for tuning - tuner_top_k: Top-k value for tuning - - Returns: - Computed MoE output tensor - """ - - def run_moe( - self, - module: 'MoE', - # Input tensors - input: torch.Tensor, - token_selected_slots: torch.Tensor, - token_final_scales: torch.Tensor, - w3_w1_weight: torch.Tensor, - w3_w1_bias: Optional[torch.Tensor], - w2_weight: torch.Tensor, - w2_bias: Optional[torch.Tensor], - output_dtype: torch.dtype, - # Quantization parameters - quant_scales: List[torch.Tensor], - use_all_to_all: bool, - input_sf: Optional[torch.Tensor] = None, - swizzled_input_sf: bool = True, - # Performance tuning (only runtime-variable parameters) - min_latency_mode: bool = False, - use_fused_finalize: bool = True, - tuner_num_tokens: Optional[int] = None, - tuner_top_k: Optional[int] = None, - **kwargs) -> torch.Tensor: - """ - Run the complete MoE computation pipeline. - - Configuration parameters are automatically extracted from the module. - - Args: - module: MoE module containing configuration - input: Input tensor to the MoE layer - token_selected_slots: Selected expert slots for each token - token_final_scales: Final scaling factors for each token - w3_w1_weight: Concatenated weights for w3 and w1 projections - w3_w1_bias: Optional bias for w3/w1 projections - w2_weight: Weight for w2 projection - w2_bias: Optional bias for w2 projection - output_dtype: Desired output data type - quant_scales: Quantization scales for weights - use_all_to_all: Whether to use all-to-all communication - input_sf: Optional input scale factors for quantization - swizzled_input_sf: Whether input scale factors are swizzled - min_latency_mode: Use minimum latency optimizations - use_fused_finalize: Use fused finalization - tuner_num_tokens: Number of tokens for tuner input - tuner_top_k: Top-k value for tuning - - Returns: - Computed MoE output tensor - """ - self.finalize_tactic(module, input, output_dtype, min_latency_mode, - use_fused_finalize, tuner_top_k) - - # Call compute_moe with module - return self.compute_moe(module=module, - x=input, - token_selected_slots=token_selected_slots, - token_final_scales=token_final_scales, - w3_w1_weight=w3_w1_weight, - w3_w1_bias=w3_w1_bias, - w2_weight=w2_weight, - w2_bias=w2_bias, - output_dtype=output_dtype, - quant_scales=quant_scales, - use_all_to_all=use_all_to_all, - input_sf=input_sf, - swizzled_input_sf=swizzled_input_sf, - min_latency_mode=min_latency_mode, - use_fused_finalize=use_fused_finalize, - tuner_num_tokens=tuner_num_tokens, - tuner_top_k=tuner_top_k, - **kwargs) - - -class MoEOpSelector: - """ - Utility class for selecting the appropriate MoE op based on - hardware capabilities and quantization configuration. - - This class implements the strategy pattern for op selection, - choosing between Cutlass and DeepGemm implementations based on: - - Hardware capabilities (SM version) - - Quantization configuration (block FP8 support) - """ - - @staticmethod - def select_op(module: 'MoE') -> MoEOp: - """ - Select the appropriate MoE op based on module configuration. - - Selection criteria: - - Blackwell (SM100) with block FP8 quantization -> DeepGemm op - - All other configurations -> Cutlass op - - Args: - module: The MoE module containing configuration information - - Returns: - MoEOp: Selected op instance (CutlassMoEOp or DeepGemmMoEOp) - - Example: - >>> op = MoEOpSelector.select_op(moe_module) - >>> output = op.run_moe(input, ...) - """ - from .moe_op_cutlass import CutlassMoEOp - from .moe_op_deepgemm import DeepGemmMoEOp - - # Check if we should use DeepGemm op - # Blackwell has SM version 100 - is_blackwell = is_sm_100f() - has_block_fp8 = module.has_deepseek_fp8_block_scales - - if is_blackwell and has_block_fp8: - # Use DeepGemm op for Blackwell with block FP8 - return DeepGemmMoEOp() - else: - # Use Cutlass op for all other cases - return CutlassMoEOp() diff --git a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_cutlass.py deleted file mode 100644 index 4a166a6905fe..000000000000 --- a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_cutlass.py +++ /dev/null @@ -1,333 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# 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. -""" -Cutlass-based MoE op implementation. -""" - -from typing import TYPE_CHECKING, List, Optional - -import torch - -from .moe_op import MoEOp - -if TYPE_CHECKING: - from ..interface import MoE - - -class CutlassMoEOp(MoEOp): - """Cutlass-based MoE op using torch.ops.trtllm.fused_moe.""" - - def __init__(self): - """Initialize the Cutlass op.""" - super().__init__() - self.moe_runner = None - self.gemm_tactics = None - - def finalize_tactic( - self, - module: 'MoE', - tuner_input: torch.Tensor, - output_dtype: torch.dtype, - min_latency_mode: bool = False, - use_fused_finalize: bool = True, - tuner_top_k: Optional[int] = None, - ) -> None: - """ - Finalize tactics for Cutlass MoE by profiling and selecting optimal GEMM tactics. - """ - - # Import necessary modules for profiling - from ....custom_ops.torch_custom_ops import AutoTuner, MoERunner - - # Use real tuner_input rather than dummy input - assert tuner_input is not None, "tuner_input must be provided to finalize_tactic" - if tuner_top_k is None: - tuner_top_k = getattr(module.routing_method, 'experts_per_token', 1) - - # Determine view dtype for weights to match runtime quantization layout - weight_view_dtype = module.w3_w1_weight.dtype - if getattr(module, 'has_w4afp8', False): - weight_view_dtype = torch.quint4x2 - elif module.has_w4a16_mxfp4: - weight_view_dtype = torch.uint8 - - # Create MoERunner for profiling - if self.moe_runner is None: - self.moe_runner = MoERunner( - x_dtype=tuner_input.dtype, - weight_dtype=weight_view_dtype, - output_dtype=output_dtype, - top_k=tuner_top_k, - tp_size=module.tp_size, - tp_rank=module.tp_rank, - ep_size=module.ep_size, - ep_rank=module.ep_rank, - cluster_size=module.cluster_size, - cluster_rank=module.cluster_rank, - use_deepseek_fp8_block_scale=module. - has_deepseek_fp8_block_scales, - use_w4_group_scaling=getattr(module, 'has_w4afp8', False), - use_int8_woq_per_channel=getattr(module, - 'has_int8_woq_per_channel', - False), - use_mxfp8_act_scaling=getattr(module, 'has_mxfp8_act_scaling', - False), - min_latency_mode=min_latency_mode, - use_fused_finalize=use_fused_finalize, - activation_type=module.activation_type, - ) - - # Set tuning configuration - MoERunner.tuning_config.tune_max_num_tokens = getattr( - module, 'tune_max_num_tokens', 8192) - - # Get AutoTuner for tactic selection - tuner = AutoTuner.get() - - # Profile and select tactics (GEMM1) - _, gemm_tactic_1 = tuner.choose_one( - "trtllm::fused_moe::gemm1", - [self.moe_runner], - MoERunner.tuning_config, - [ - tuner_input, - module.w3_w1_weight.view(weight_view_dtype), - getattr(module, 'w3_w1_bias', None), - module.w2_weight.view(weight_view_dtype), - getattr(module, 'w2_bias', None), - ], - gemm_idx=1, - ) - - # Profile and select tactics (GEMM2) - _, gemm_tactic_2 = tuner.choose_one( - "trtllm::fused_moe::gemm2", - [self.moe_runner], - MoERunner.tuning_config, - [ - tuner_input, - module.w3_w1_weight.view(weight_view_dtype), - getattr(module, 'w3_w1_bias', None), - module.w2_weight.view(weight_view_dtype), - getattr(module, 'w2_bias', None), - ], - gemm_idx=2, - ) - - # Store selected tactics - self.gemm_tactics = [gemm_tactic_1, gemm_tactic_2] - - def compute_moe( - self, - module: 'MoE', # Now required as first parameter - # Input tensors - x: torch.Tensor, - token_selected_slots: torch.Tensor, - token_final_scales: Optional[torch.Tensor], - # Weight tensors - w3_w1_weight: torch.Tensor, - w3_w1_bias: Optional[torch.Tensor], - w2_weight: torch.Tensor, - w2_bias: Optional[torch.Tensor], - # Output configuration - output_dtype: torch.dtype, - # Quantization parameters - quant_scales: List[torch.Tensor], - use_all_to_all: bool, - input_sf: Optional[torch.Tensor] = None, - swizzled_input_sf: bool = True, - # Performance tuning (only runtime-variable parameters) - min_latency_mode: bool = False, - use_fused_finalize: bool = True, - tuner_num_tokens: Optional[int] = None, - tuner_top_k: Optional[int] = None, - **kwargs) -> torch.Tensor: - """ - Compute MoE using Cutlass op with MoERunner. - """ - # Extract parameters from module - tp_size = module.tp_size - tp_rank = module.tp_rank - ep_size = module.ep_size - ep_rank = module.ep_rank - cluster_size = module.cluster_size - cluster_rank = module.cluster_rank - use_all_to_all = use_all_to_all - swiglu_alpha = module.swiglu_alpha - swiglu_beta = module.swiglu_beta - swiglu_limit = module.swiglu_limit - use_w4_group_scaling = getattr(module, 'has_w4afp8', False) - activation_type = module.activation_type - - # Determine weight dtype for view operation if needed - weight_dtype = w3_w1_weight.dtype - if use_w4_group_scaling and weight_dtype != torch.quint4x2: - weight_dtype = torch.quint4x2 - - # Validate that tactics have been finalized - if self.gemm_tactics is None or len(self.gemm_tactics) == 0: - raise RuntimeError( - "GEMM tactics have not been finalized. " - "Call finalize_tactic() before compute_moe() or use run_moe() instead." - ) - - if self.moe_runner is None: - raise RuntimeError( - "MoERunner has not been initialized. " - "Call finalize_tactic() before compute_moe() or use run_moe() instead." - ) - - # Select the appropriate run method based on latency mode - run_moe = self.moe_runner.fused_moe_runner.run_moe_min_latency if min_latency_mode else self.moe_runner.fused_moe_runner.run_moe - - # Get unpadded_hidden_size from module if available, otherwise use hidden_size - # For now it is the user's responsibility to set unpadded_hidden_size. - # DeepGemmFusedMoE and WideEPMoE both have unpadded_hidden_size. - unpadded_hidden_size = getattr(module, 'unpadded_hidden_size', - x.shape[1]) - - use_dynamic_fc2_scale = ( - getattr(module, 'has_nvfp4', False) - and getattr(module, 'force_dynamic_quantization', False) - and hasattr(module, 'fc2_weight_scale_2')) - - # Ensure quant_scales is a plain list for C++ ArrayRef binding. - # NamedTuple (e.g. FusedMoEQuantScalesW4A8) may not convert correctly. - if not isinstance(quant_scales, list): - quant_scales = list(quant_scales) - - # The C++ run_moe / run_moe_min_latency TorchBind methods share this - # positional prefix. The TorchBind schemas do not honor C++ default - # arguments, so every positional argument must be supplied. This - # profiling/tuning path never applies MoE LoRA (LoRA is fused into - # torch.ops.trtllm.fused_moe), so all eager (per-request) and - # slot-indexed LoRA args are passed as None/0. - run_moe_args = [ - x, token_selected_slots, token_final_scales, - w3_w1_weight.view(weight_dtype), w3_w1_bias, - w2_weight.view(weight_dtype), w2_bias, quant_scales, input_sf, - swizzled_input_sf, swiglu_alpha, swiglu_beta, swiglu_limit, tp_size, - tp_rank, ep_size, ep_rank, cluster_size, cluster_rank, - use_all_to_all, min_latency_mode, self.gemm_tactics, - activation_type, unpadded_hidden_size, tuner_num_tokens, None - ] - if not min_latency_mode: - run_moe_args += [ - use_dynamic_fc2_scale, - # Eager (per-request) LoRA args. - None, - None, - None, - None, - None, - None, - None, - None, - 0, - # Slot-indexed (CUDA-graph) LoRA args. - None, - None, - None, - None, - None, - None, - None - ] - output = run_moe(*run_moe_args) - - # Return output based on latency mode - return output if min_latency_mode else [output] - - def run_moe( - self, - module: 'MoE', - # Input tensors - input: torch.Tensor, - token_selected_slots: torch.Tensor, - token_final_scales: torch.Tensor, - w3_w1_weight: torch.Tensor, - w3_w1_bias: Optional[torch.Tensor], - w2_weight: torch.Tensor, - w2_bias: Optional[torch.Tensor], - output_dtype: torch.dtype, - # Quantization parameters - quant_scales: List[torch.Tensor], - use_all_to_all: bool, - input_sf: Optional[torch.Tensor] = None, - swizzled_input_sf: bool = True, - # Performance tuning (only runtime-variable parameters) - min_latency_mode: bool = False, - use_fused_finalize: bool = True, - tuner_num_tokens: Optional[int] = None, - tuner_top_k: Optional[int] = None, - **kwargs) -> torch.Tensor: - """ - Run the complete MoE computation pipeline for Cutlass op. - - This override handles the specific tuner_input logic needed for Cutlass. - - Args: - module: MoE module containing configuration - input: Input tensor to the MoE layer - token_selected_slots: Selected expert slots for each token - token_final_scales: Final scaling factors for each token - w3_w1_weight: Concatenated weights for w3 and w1 projections - w3_w1_bias: Optional bias for w3/w1 projections - w2_weight: Weight for w2 projection - w2_bias: Optional bias for w2 projection - output_dtype: Desired output data type - quant_scales: Quantization scales for weights - input_sf: Optional input scale factors for quantization - swizzled_input_sf: Whether input scale factors are swizzled - min_latency_mode: Use minimum latency optimizations - use_fused_finalize: Use fused finalization - tuner_num_tokens: Number of tokens for tuner input - tuner_top_k: Top-k value for tuning - use_all_to_all: Whether to use all-to-all communication - - Returns: - Computed MoE output tensor - """ - use_all_to_all = use_all_to_all - - # Compute tuner_input per fused_moe logic - if use_all_to_all: - assert tuner_num_tokens is not None - assert tuner_top_k is not None - tuner_input = input[:tuner_num_tokens] - else: - assert tuner_num_tokens is None - assert tuner_top_k is None - tuner_input = input - tuner_top_k = token_selected_slots.size(1) - - self.finalize_tactic(module, tuner_input, output_dtype, - min_latency_mode, use_fused_finalize, tuner_top_k) - - # Call compute_moe with module - return self.compute_moe(module=module, - x=input, - token_selected_slots=token_selected_slots, - token_final_scales=token_final_scales, - w3_w1_weight=w3_w1_weight, - w3_w1_bias=w3_w1_bias, - w2_weight=w2_weight, - w2_bias=w2_bias, - output_dtype=output_dtype, - quant_scales=quant_scales, - use_all_to_all=use_all_to_all, - input_sf=input_sf, - swizzled_input_sf=swizzled_input_sf, - min_latency_mode=min_latency_mode, - use_fused_finalize=use_fused_finalize, - tuner_num_tokens=tuner_num_tokens, - tuner_top_k=tuner_top_k, - **kwargs) diff --git a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py deleted file mode 100644 index 9b86c02673f2..000000000000 --- a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py +++ /dev/null @@ -1,326 +0,0 @@ -# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. -# 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. -""" -DeepGemm-based MoE op implementation for GB200 block FP8. -""" - -from typing import TYPE_CHECKING, Dict, List, Optional - -import torch - -from ....memory_buffer_utils import get_memory_buffers -from .moe_op import MoEOp - -if TYPE_CHECKING: - from ..interface import MoE - - -class DeepGemmMoEOp(MoEOp): - """DeepGemm-based MoE op for GB200 block FP8.""" - buffers = get_memory_buffers() - - def __init__(self): - """Initialize DeepGemm op.""" - super().__init__() - import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils - self.fp8_utils = fp8_utils - - from ..fused_moe_deepgemm import deepgemm_fp8_group_blockwise_gemm - self.deepgemm_fp8_group_blockwise_gemm = deepgemm_fp8_group_blockwise_gemm - - def finalize_tactic( - self, - module: 'MoE', - tuner_input: torch.Tensor, - output_dtype: torch.dtype, - min_latency_mode: bool = False, - use_fused_finalize: bool = True, - tuner_top_k: Optional[int] = None, - ) -> None: - """ - No-op for DeepGemm op as it doesn't require tactic profiling. - - Args: - module: The MoE module - tuner_input: Input tensor for tuning - output_dtype: Output dtype - min_latency_mode: Whether to use min-latency mode - use_fused_finalize: Whether to use fused finalize - tuner_top_k: Top-k value for tuning - """ - - def _get_deepgemm_workspace(self, module: 'MoE', m_max: int, - group_size: int) -> Dict[str, torch.Tensor]: - """ - Get workspace for DeepGemm op operations. - - Args: - module: The MoE module containing configuration - m_max: Maximum number of tokens (aligned) - group_size: Group size for quantization - - Returns: - Dictionary containing workspace tensors - """ - import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils - - # Get dimensions from module - hidden_size = module.hidden_size - intermediate_size = module.intermediate_size - expert_size_per_partition = module.expert_size_per_partition - - # Calculate aligned dimensions - m_padded = fp8_utils.align(m_max, 4) - fp8_dim = max(hidden_size, intermediate_size) - scale_k = fp8_utils.ceil_div(fp8_dim, group_size) - scale_k_padded = fp8_utils.align(scale_k, 4) - - # Allocate workspace tensors - workspace = {} - - # Workspace for FP8 activations - capture_graph = torch.cuda.is_current_stream_capturing() - workspace["workspace_0"] = DeepGemmMoEOp.buffers.get_buffer( - [expert_size_per_partition, m_max, fp8_dim], - dtype=torch.float8_e4m3fn, - buffer_name='workspace_0', - reserve_buffer=capture_graph) - - # Workspace for intermediate results - workspace["workspace_1"] = DeepGemmMoEOp.buffers.get_buffer( - [ - expert_size_per_partition, m_max, - max(intermediate_size * 2, hidden_size) - ], - dtype=torch.bfloat16, - buffer_name='workspace_1', - reserve_buffer=capture_graph) - - # Workspace for scaling factors - workspace["workspace_sf"] = DeepGemmMoEOp.buffers.get_buffer( - [expert_size_per_partition, (scale_k_padded // 4), m_padded], - dtype=torch.int32, - buffer_name='workspace_sf', - reserve_buffer=capture_graph) - - return workspace - - def compute_moe( - self, - module: 'MoE', - # Input tensors - x: torch.Tensor, - token_selected_slots: torch.Tensor, - token_final_scales: Optional[torch.Tensor], - # Weight tensors - w3_w1_weight: torch.Tensor, - w3_w1_bias: Optional[torch.Tensor], - w2_weight: torch.Tensor, - w2_bias: Optional[torch.Tensor], - # Output configuration - output_dtype: torch.dtype, - # Quantization parameters - quant_scales: List[torch.Tensor], - use_all_to_all: bool, - input_sf: Optional[torch.Tensor] = None, - swizzled_input_sf: bool = True, - # Performance tuning (only runtime-variable parameters) - min_latency_mode: bool = False, - use_fused_finalize: bool = True, - tuner_num_tokens: Optional[int] = None, - tuner_top_k: Optional[int] = None, - **kwargs) -> torch.Tensor: - """ - Compute MoE using DeepGemm op with block FP8 quantization. - - Note: This assumes the data has already been gathered/alltoall'd - by the WideEP forward_chunk method. - """ - - # Import necessary functions for DeepGemm - from ..fused_moe_deepgemm import (fused_expand_group_quant_fp8, - preprocess_after_permute, set_strides, - triton_masked_index_gather) - - # Extract parameters from module - tp_size = module.tp_size - tp_rank = module.tp_rank - ep_size = module.ep_size - ep_rank = module.ep_rank - cluster_size = module.cluster_size - cluster_rank = module.cluster_rank - use_all_to_all = use_all_to_all - - # Not supported: min_latency_mode. Raise error if enabled. - if min_latency_mode: - raise NotImplementedError( - "DeepGemm op does not support min_latency_mode=True") - - # Get expert configuration from module - expert_size_per_partition = module.expert_size_per_partition - intermediate_size = module.intermediate_size - hidden_size = x.shape[1] - - # Permute the data for expert-parallel processing. - # Unlike DeepGemmFusedMoE (which fuses gather+finalize and never touches - # permuted_data_tensor), this op reuses permuted_data_tensor as a - # write-before-read scratch buffer in the gather+finalize tail below, so - # it is kept; only the genuinely unused outputs are discarded with `_`. - ( - permuted_row_to_unpermuted_row_tensor, - _, # permuted_token_selected_experts_tensor (unused) - permuted_data_tensor, - expert_first_token_offset_tensor, - _, # permuted_token_final_scales_tensor (uninitialized under skip_data_expand) - unpermuted_row_to_permuted_row_tensor, - ) = torch.ops.trtllm.moe_permute_op( - x, - token_selected_slots, - token_final_scales, - None, # w3_w1_weight - None, # w2_weight - None, # quant_scales - input_sf=input_sf, - num_experts_on_rank=expert_size_per_partition, - tp_size=tp_size, - tp_rank=tp_rank, - ep_size=ep_size, - ep_rank=ep_rank, - cluster_size=cluster_size, - cluster_rank=cluster_rank, - min_latency_mode=min_latency_mode, - use_fp8_block_scaling=True, # Always use block scaling for DeepGemm - skip_data_expand=True, - ) - - # Take the expanded-token count from the populated permutation map (one - # entry per permuted token) rather than the uninitialized data tensor. - num_permuted_tokens = permuted_row_to_unpermuted_row_tensor.shape[0] - if num_permuted_tokens == 0: - return torch.zeros_like(x) - - # Preprocess for masked operations - masked_m, token_to_expert_map = preprocess_after_permute( - expert_first_token_offset_tensor, num_permuted_tokens) - - expected_m = (token_selected_slots.numel() + expert_size_per_partition - - 1) // expert_size_per_partition - - # Get workspace for DeepGemm operations - m_max = self.fp8_utils.align(x.shape[0], 128) - workspace = self._get_deepgemm_workspace(module, m_max, 128) - - # Padding and quantization for first GEMM input - m_padded = self.fp8_utils.align(m_max, 4) - scale_k = self.fp8_utils.ceil_div(hidden_size, 128) - scale_k_padded = self.fp8_utils.align(scale_k, 4) - - act_input_fp8 = set_strides(workspace["workspace_0"], - expert_size_per_partition, m_max, - hidden_size) - act_input_sf = set_strides(workspace["workspace_sf"], - expert_size_per_partition, - scale_k_padded // 4, m_padded) - - # Fused expand + quantize (reads from original input via perm map) - act_input_sf = fused_expand_group_quant_fp8( - act_input_fp8, - act_input_sf, - x, - permuted_row_to_unpermuted_row_tensor, - expert_first_token_offset_tensor, - token_to_expert_map, - experts_per_token=token_selected_slots.shape[1], - group_size=128) - - # First grouped GEMM (w3 and w1) - h1 = set_strides(workspace["workspace_1"], expert_size_per_partition, - m_max, intermediate_size * 2) - - self.deepgemm_fp8_group_blockwise_gemm( - d=h1, - a=act_input_fp8, - b=w3_w1_weight, - sfa=act_input_sf, - sfb=quant_scales[0] if quant_scales else None, - masked_m=masked_m, - expected_m=expected_m, - ) - - # SiLU activation and quantization for second GEMM - act_input_fp8 = set_strides(workspace["workspace_0"], - expert_size_per_partition, m_max, - intermediate_size) - - scale_k = self.fp8_utils.ceil_div(intermediate_size, 128) - scale_k_padded = self.fp8_utils.align(scale_k, 4) - act_input_sf = set_strides(workspace["workspace_sf"], - expert_size_per_partition, - scale_k_padded // 4, m_padded) - - act_input_sf = self.fp8_utils.silu_and_mul_masked_post_quant_fwd( - output=act_input_fp8, - output_scale=act_input_sf, - input=h1, - quant_group_size=128, - masked_m=masked_m, - scale_ue8m0=True, - swiglu_limit=getattr(module, "swiglu_limit_scalar", None)) - - # Second grouped GEMM (w2) - h3 = set_strides(workspace["workspace_1"], expert_size_per_partition, - m_max, hidden_size) - - self.deepgemm_fp8_group_blockwise_gemm( - d=h3, - a=act_input_fp8, - b=w2_weight, - sfa=act_input_sf, - sfb=quant_scales[1] if quant_scales else None, - masked_m=masked_m, - expected_m=expected_m, - ) - - # Gather results back to original token order - triton_masked_index_gather(permuted_data_tensor, h3, - expert_first_token_offset_tensor, - token_to_expert_map) - - # Finalize and scale the output - # Get unpadded_hidden_size from module if available, otherwise use hidden_size - # For now it is the user's responsibility to set unpadded_hidden_size. - # DeepGemmFusedMoE and WideEPMoE both have unpadded_hidden_size. - unpadded_hidden_size = getattr(module, 'unpadded_hidden_size', - x.shape[1]) - - final_hidden_states = torch.ops.trtllm.moe_finalize_scale_op( - permuted_data_tensor, - None, # biases (w2_bias could be added here if needed) - token_final_scales, - unpermuted_row_to_permuted_row_tensor, - permuted_row_to_unpermuted_row_tensor, - token_selected_slots, - expert_first_token_offset_tensor, - use_all_to_all, - x.shape[0], # num_rows - x.shape[1], # hidden_size - unpadded_hidden_size, # unpadded_hidden_size (may be different from hidden_size if padding was applied) - module.routing_method.top_k if module else 1, # experts_per_token - expert_size_per_partition, # num_experts_per_node - tp_size, - tp_rank, - ep_size, - ep_rank, - ) - - return final_hidden_states if min_latency_mode else [ - final_hidden_states - ] diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index f1f88a59e953..a48110e5fecf 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1418,8 +1418,8 @@ def get_layer_initial_global_assignments( class MoeConfig(StrictBaseModel): """Configuration for MoE.""" backend: Literal[ - "AUTO", "CUTLASS", "CUTEDSL", "WIDEEP", "TRTLLM", "DEEPGEMM", - "DENSEGEMM", "VANILLA", "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM", + "AUTO", "CUTLASS", "CUTEDSL", "TRTLLM", "DEEPGEMM", "DENSEGEMM", + "VANILLA", "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"] = Field( default='AUTO', description="MoE backend to use. " diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index 2b3d61e4a802..2d263b171afe 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -18,7 +18,6 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_utils import PostInitCaller, skip_forward from tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE -from tensorrt_llm._torch.modules.fused_moe.fused_moe_wide_ep import WideEPMoE from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from tensorrt_llm._torch.pyexecutor._util import get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.config_utils import ( @@ -463,7 +462,6 @@ def make_load_pretrained_config(mapping, load_pretrained_config_orig): # To run the problem size of $B$ GPUs on $A$ GPUs, we need: # (1) Attention: If TP, reduce the number of attention heads; If DP, nothing to change. # (2) MoE: If EP, reduce the number of experts; If TP, reduce head size. - # Maintain the result of AllToAll method selection because it is affected by EP size. def load_pretrained_config(*args, **kwargs): pretrained_config = load_pretrained_config_orig(*args, **kwargs) if not mapping.enable_attention_dp: @@ -484,33 +482,13 @@ def load_pretrained_config(*args, **kwargs): return load_pretrained_config - def make_select_alltoall_method_type(select_alltoall_method_type_orig): - def select_alltoall_method_type( - cls: type, mapping: Mapping, top_k: int, *args, **kwargs - ): - # Replace the condition `mapping.moe_ep_size <= top_k` with `scaled_from <= top_k` - # by replacing `top_k` with `fake_top_k` - if scaled_from <= top_k: - fake_top_k = mapping.moe_ep_size + 1 - else: - fake_top_k = mapping.moe_ep_size - 1 - assert (mapping.moe_ep_size <= fake_top_k) == (scaled_from <= top_k) - return select_alltoall_method_type_orig(mapping, fake_top_k, *args, **kwargs) - - return select_alltoall_method_type - - select_alltoall_method_type_wide_ep = WideEPMoE.select_alltoall_method_type tensorrt_llm._torch.model_config.load_pretrained_config = make_load_pretrained_config( mapping, load_pretrained_config ) - WideEPMoE.select_alltoall_method_type = make_select_alltoall_method_type( - select_alltoall_method_type_wide_ep - ) try: yield finally: tensorrt_llm._torch.model_config.load_pretrained_config = load_pretrained_config - WideEPMoE.select_alltoall_method_type = select_alltoall_method_type_wide_ep @staticmethod @contextlib.contextmanager diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index fb92bfba7ad2..670d741a96b1 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1064,7 +1064,6 @@ "AUTO", "CUTLASS", "CUTEDSL", - "WIDEEP", "TRTLLM", "DEEPGEMM", "DENSEGEMM", @@ -1074,7 +1073,7 @@ "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL" ], - "annotation": "Literal['AUTO', 'CUTLASS', 'CUTEDSL', 'WIDEEP', 'TRTLLM', 'DEEPGEMM', 'DENSEGEMM', 'VANILLA', 'TRITON', 'MARLIN', 'MEGAMOE_DEEPGEMM', 'MEGAMOE_CUTEDSL']", + "annotation": "Literal['AUTO', 'CUTLASS', 'CUTEDSL', 'TRTLLM', 'DEEPGEMM', 'DENSEGEMM', 'VANILLA', 'TRITON', 'MARLIN', 'MEGAMOE_DEEPGEMM', 'MEGAMOE_CUTEDSL']", "converter": "", "kind": "categorical", "path": "moe_config.backend" diff --git a/tests/microbenchmarks/bench_moe/search.py b/tests/microbenchmarks/bench_moe/search.py index 491ecc2fb256..6aec8ec4fed8 100644 --- a/tests/microbenchmarks/bench_moe/search.py +++ b/tests/microbenchmarks/bench_moe/search.py @@ -45,7 +45,7 @@ def _is_deepep_feasible(num_ranks: int) -> bool: Intranode: num_ranks in {2, 4, 8} and num_ranks == local_mpi_size(). Internode: exactly 8 ranks per node, with 2/4/8/16 RDMA nodes. - Mirrors the feasibility check in fused_moe_wide_ep.py::select_alltoall_method_type. + Mirrors the feasibility check in communication/deep_ep.py::DeepEP._is_deepep_feasible. """ _INTRANODE_RANKS = {2, 4, 8} _REQUIRED_LOCAL_SIZE = 8 diff --git a/tests/unittest/_torch/modules/moe/test_cutlass_moe_op_smoke.py b/tests/unittest/_torch/modules/moe/test_cutlass_moe_op_smoke.py deleted file mode 100644 index d592f44b6d3d..000000000000 --- a/tests/unittest/_torch/modules/moe/test_cutlass_moe_op_smoke.py +++ /dev/null @@ -1,231 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Non-LoRA smoke test for the CutlassMoEOp op path. - -CutlassMoEOp.run_moe calls the TorchBind C++ method fused_moe_runner.run_moe. -That method takes trailing routed-expert LoRA arguments, and TorchBind schemas -do not honor C++ default-argument values, so every positional argument must be -supplied by the caller. This op is what the WideEP and ConfigurableMoE paths -select via MoEOpSelector, so a missing or extra argument here is a silent -regression for a feature unrelated to LoRA. - -This test exercises that callsite on a single GPU (tp/ep/cluster = 1, the -degenerate WideEP configuration) with unquantized bf16 weights and no LoRA, -confirming the call succeeds and produces finite output. -""" - -import pytest -import torch - -_TRTLLM_AVAILABLE = hasattr(torch.ops, "trtllm") and hasattr(torch.ops.trtllm, "fused_moe") - -requires_cuda_and_op = pytest.mark.skipif( - not torch.cuda.is_available() or not _TRTLLM_AVAILABLE, - reason="Requires CUDA and built TensorRT-LLM C++ extension (torch.ops.trtllm.fused_moe).", -) - - -class _RoutingStub: - def __init__(self, top_k): - self.experts_per_token = top_k - - -class _ModuleStub: - """Minimal stand-in for an MoE module exposing only the attributes that - CutlassMoEOp.finalize_tactic and compute_moe read.""" - - def __init__(self, w3_w1_weight, w2_weight, top_k, hidden_size, activation_type): - self.w3_w1_weight = w3_w1_weight - self.w2_weight = w2_weight - self.routing_method = _RoutingStub(top_k) - - # Parallelism: single rank (degenerate WideEP). - self.tp_size = 1 - self.tp_rank = 0 - self.ep_size = 1 - self.ep_rank = 0 - self.cluster_size = 1 - self.cluster_rank = 0 - - # SwiGLU params, unused here. - self.swiglu_alpha = None - self.swiglu_beta = None - self.swiglu_limit = None - - # Quantization and layout flags, all off for unquantized bf16. - self.has_w4afp8 = False - self.has_w4a16_mxfp4 = False - self.has_deepseek_fp8_block_scales = False - self.has_int8_woq_per_channel = False - self.has_mxfp8_act_scaling = False - self.has_nvfp4 = False - self.force_dynamic_quantization = False - - self.activation_type = activation_type - self.tune_max_num_tokens = 8192 - self.unpadded_hidden_size = hidden_size - - -@requires_cuda_and_op -def test_cutlass_moe_op_run_moe_no_lora_smoke(): - from tensorrt_llm._torch.modules.fused_moe.ops.moe_op_cutlass import CutlassMoEOp - from tensorrt_llm._torch.utils import ActivationType - - device = torch.device("cuda") - dtype = torch.bfloat16 - num_tokens, hidden_size, inter_size = 8, 128, 256 - num_experts, top_k = 4, 2 - - torch.manual_seed(0) - x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) - # FC1 packs [up(w3); gate(w1)] -> [E, 2 * inter, hidden]. - w3_w1_weight = ( - torch.randn(num_experts, 2 * inter_size, hidden_size, dtype=dtype, device=device) * 0.02 - ) - w2_weight = torch.randn(num_experts, hidden_size, inter_size, dtype=dtype, device=device) * 0.02 - - logits = torch.randn(num_tokens, num_experts, dtype=torch.float32, device=device) - topk_scores, topk_ids = torch.topk(logits, k=top_k, dim=-1) - topk_scores = torch.softmax(topk_scores, dim=-1) - - module = _ModuleStub( - w3_w1_weight=w3_w1_weight, - w2_weight=w2_weight, - top_k=top_k, - hidden_size=hidden_size, - activation_type=int(ActivationType.Swiglu), - ) - - op = CutlassMoEOp() - # No LoRA arguments: guards the run_moe callsite in moe_op_cutlass.py. - out = op.run_moe( - module=module, - input=x, - token_selected_slots=topk_ids.to(torch.int32), - token_final_scales=topk_scores.to(torch.float32), - w3_w1_weight=w3_w1_weight, - w3_w1_bias=None, - w2_weight=w2_weight, - w2_bias=None, - output_dtype=dtype, - quant_scales=[], - use_all_to_all=False, - ) - - # Non-min-latency run_moe returns a single-element list. - assert isinstance(out, list) and len(out) == 1 - result = out[0] - assert result.shape == (num_tokens, hidden_size) - assert torch.isfinite(result).all() - - -@requires_cuda_and_op -def test_cutlass_moe_op_run_moe_no_lora_matches_fused_moe_op(): - """The CutlassMoEOp path and the direct torch.ops.trtllm.fused_moe path - should produce the same result for an unquantized bf16 MoE with no LoRA, - since both ultimately call the same C++ run_moe.""" - from tensorrt_llm._torch.modules.fused_moe.ops.moe_op_cutlass import CutlassMoEOp - from tensorrt_llm._torch.utils import ActivationType - - device = torch.device("cuda") - dtype = torch.bfloat16 - num_tokens, hidden_size, inter_size = 8, 128, 256 - num_experts, top_k = 4, 2 - - torch.manual_seed(1) - x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) - w3_w1_weight = ( - torch.randn(num_experts, 2 * inter_size, hidden_size, dtype=dtype, device=device) * 0.02 - ) - w2_weight = torch.randn(num_experts, hidden_size, inter_size, dtype=dtype, device=device) * 0.02 - logits = torch.randn(num_tokens, num_experts, dtype=torch.float32, device=device) - topk_scores, topk_ids = torch.topk(logits, k=top_k, dim=-1) - topk_scores = torch.softmax(topk_scores, dim=-1).to(torch.float32) - topk_ids = topk_ids.to(torch.int32) - - # Direct op reference. - ref = torch.ops.trtllm.fused_moe( - input=x, - token_selected_experts=topk_ids, - token_final_scales=topk_scores, - fc1_expert_weights=w3_w1_weight, - fc1_expert_biases=None, - fc2_expert_weights=w2_weight, - fc2_expert_biases=None, - output_dtype=dtype, - quant_scales=[], - )[0] - - module = _ModuleStub( - w3_w1_weight=w3_w1_weight, - w2_weight=w2_weight, - top_k=top_k, - hidden_size=hidden_size, - activation_type=int(ActivationType.Swiglu), - ) - out = CutlassMoEOp().run_moe( - module=module, - input=x, - token_selected_slots=topk_ids, - token_final_scales=topk_scores, - w3_w1_weight=w3_w1_weight, - w3_w1_bias=None, - w2_weight=w2_weight, - w2_bias=None, - output_dtype=dtype, - quant_scales=[], - use_all_to_all=False, - )[0] - - torch.testing.assert_close(out, ref, rtol=5e-2, atol=1e-2) - - -@requires_cuda_and_op -def test_cutlass_fp8_block_scale_ep_rank_with_no_tokens_returns_zeros(): - """An EP rank with no selected local experts must not enter block-scale GEMM empty.""" - device = torch.device("cuda") - dtype = torch.bfloat16 - num_tokens, hidden_size, inter_size = 8, 128, 128 - local_experts, top_k = 32, 8 - ep_size, ep_rank = 8, 7 - - torch.manual_seed(2) - x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) - w3_w1_weight = torch.randn( - local_experts, 2 * inter_size, hidden_size, dtype=dtype, device=device - ).to(torch.float8_e4m3fn) - w2_weight = torch.randn(local_experts, hidden_size, inter_size, dtype=dtype, device=device).to( - torch.float8_e4m3fn - ) - fc1_scales = torch.ones(local_experts, 2 * inter_size // 128, hidden_size // 128, device=device) - fc2_scales = torch.ones(local_experts, hidden_size // 128, inter_size // 128, device=device) - - # Rank 7 owns global experts [224, 256); route every token to [0, 8). - topk_ids = ( - torch.arange(top_k, dtype=torch.int32, device=device).expand(num_tokens, -1).contiguous() - ) - topk_scores = torch.full((num_tokens, top_k), 1.0 / top_k, dtype=torch.float32, device=device) - - def run_moe(ep_size: int, ep_rank: int): - return torch.ops.trtllm.fused_moe( - input=x, - token_selected_experts=topk_ids, - token_final_scales=topk_scores, - fc1_expert_weights=w3_w1_weight, - fc1_expert_biases=None, - fc2_expert_weights=w2_weight, - fc2_expert_biases=None, - output_dtype=dtype, - quant_scales=[fc1_scales, fc2_scales], - ep_size=ep_size, - ep_rank=ep_rank, - use_deepseek_fp8_block_scale=True, - tune_max_num_tokens=num_tokens, - )[0] - - # Copying the offsets for GEMM must preserve a rank with local work. - torch.testing.assert_close(run_moe(ep_size=8, ep_rank=0), run_moe(ep_size=1, ep_rank=0)) - - for _ in range(3): - out = run_moe(ep_size=ep_size, ep_rank=ep_rank) - torch.testing.assert_close(out, torch.zeros_like(out)) diff --git a/tests/unittest/_torch/modules/test_fused_moe.py b/tests/unittest/_torch/modules/test_fused_moe.py index d33c0c5cf083..5f9d23016243 100644 --- a/tests/unittest/_torch/modules/test_fused_moe.py +++ b/tests/unittest/_torch/modules/test_fused_moe.py @@ -4,7 +4,6 @@ from contextlib import contextmanager from itertools import product from typing import Dict, List, Optional -from unittest import mock import _torch.helpers import cloudpickle @@ -26,17 +25,16 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl import \ CuteDslFusedMoE -from tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm import \ - DeepGemmFusedMoE -from tensorrt_llm._torch.modules.fused_moe.interface import ( - AlltoallMethodType, MoEWeightLoadingMode) +from tensorrt_llm._torch.modules.fused_moe.interface import MoEWeightLoadingMode # isort and yapf will fight against each other here, so we disable isort # isort: off -from tensorrt_llm._torch.modules.fused_moe import ( - BaseMoeRoutingMethod, CutlassFusedMoE, TRTLLMGenFusedMoE, - DefaultMoeRoutingMethod, RenormalizeMoeRoutingMethod, TritonFusedMoE, - create_moe, WideEPMoE) +from tensorrt_llm._torch.modules.fused_moe import (BaseMoeRoutingMethod, + CutlassFusedMoE, + TRTLLMGenFusedMoE, + DefaultMoeRoutingMethod, + RenormalizeMoeRoutingMethod, + TritonFusedMoE, create_moe) from tensorrt_llm._torch.modules.fused_moe.quantization import \ NVFP4CutlassFusedMoEMethod # isort: on @@ -227,309 +225,6 @@ def test_fused_moe_multi_gpu(moe_cls, ep_size): assert r is None -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@pytest.mark.skipif(torch.cuda.device_count() < 4, - reason="needs 4 GPUs to run this test") -@pytest.mark.parametrize("alltoall_method_type", [ - AlltoallMethodType.NVLinkTwoSided, AlltoallMethodType.DeepEP, - AlltoallMethodType.DeepEPLowLatency -], - ids=lambda s: s.name) -def test_fused_moe_alltoall(alltoall_method_type): - world_size = 4 - dtype = torch.bfloat16 - HIDDEN_SIZE = 2560 - INTERMEDIATE_SIZE = 1536 - NUM_EXPERTS = 72 - TOP_K = 6 - MAX_NUM_TOKENS = 2048 - - def per_rank_test_fused_moe_alltoall(job_id): - routing_method = DefaultMoeRoutingMethod(top_k=TOP_K) - mapping = Mapping(world_size=world_size, - rank=mpi_rank(), - tp_size=world_size, - moe_ep_size=world_size, - moe_tp_size=1, - enable_attention_dp=True) - torch.cuda.set_device(mapping.rank) - torch.manual_seed(mapping.rank) - - weights = {} - for expert_id in range(NUM_EXPERTS): - w1_weight = torch.empty((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=dtype, - device="cuda") - w2_weight = torch.empty((HIDDEN_SIZE, INTERMEDIATE_SIZE), - dtype=dtype, - device="cuda") - w3_weight = torch.empty((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=dtype, - device="cuda") - torch.nn.init.xavier_uniform_(w1_weight) - torch.nn.init.xavier_uniform_(w2_weight) - torch.nn.init.xavier_uniform_(w3_weight) - weights[f"{expert_id}.w1.weight"] = w1_weight - weights[f"{expert_id}.w2.weight"] = w2_weight - weights[f"{expert_id}.w3.weight"] = w3_weight - with mock.patch.object(WideEPMoE, - "select_alltoall_method_type", - return_value=alltoall_method_type): - alltoall_model = WideEPMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - reduce_results=True, - model_config=ModelConfig(mapping=mapping, - max_num_tokens=MAX_NUM_TOKENS, - moe_max_num_tokens=MAX_NUM_TOKENS), - ) - alltoall_model.to("cuda") - alltoall_model.load_weights([weights]) - - ref_model = CutlassFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - reduce_results=True, - model_config=ModelConfig(mapping=mapping, - max_num_tokens=MAX_NUM_TOKENS), - ) - ref_model.to("cuda") - ref_model.load_weights([weights]) - - # Evaluate the outputs on a variant sequence length to verify the robustness of alltoall methods - m = MAX_NUM_TOKENS - while m >= 1: - x = torch.randn((m, HIDDEN_SIZE), dtype=dtype, device="cuda") - router_logits = torch.randn((m, NUM_EXPERTS), - dtype=dtype, - device="cuda") - all_rank_num_tokens = [m] * mapping.world_size - - with torch.inference_mode(): - output = alltoall_model.forward( - x, - router_logits, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False) - ref_output = ref_model.forward( - x, - router_logits, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False) - - if output.ndim == 3: - output = output.sum(dim=1) - print(f"output: {output.shape}") - print(f"ref_output: {ref_output.shape}") - # Evaluate outputs - torch.testing.assert_close(output, - ref_output, - rtol=0.05, - atol=0.003) - m //= 2 - - with MPIPoolExecutor(max_workers=world_size) as executor: - results = executor.map(per_rank_test_fused_moe_alltoall, - range(world_size)) - for r in results: - assert r is None - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@pytest.mark.skipif(torch.cuda.device_count() < 4, - reason="needs 4 GPUs to run this test") -@pytest.mark.parametrize("alltoall_method_type", [ - AlltoallMethodType.NVLinkTwoSided, AlltoallMethodType.DeepEP, - AlltoallMethodType.DeepEPLowLatency -], - ids=lambda s: s.name) -def test_fused_moe_alltoall_fp4(alltoall_method_type): - world_size = 4 - dtype = torch.bfloat16 - HIDDEN_SIZE = 4096 - INTERMEDIATE_SIZE = 1536 - NUM_EXPERTS = 72 - TOP_K = 6 - MAX_NUM_TOKENS = 2048 - - torch.manual_seed(0) - torch.cuda.manual_seed(0) - - x_list_world = [] - weights_world = [] - - for i in range(world_size): - x_list = [] - m = MAX_NUM_TOKENS - while m >= 1: - x = torch.randn((m, HIDDEN_SIZE), dtype=dtype) - x_list.append(x) - m //= 2 - - x_abs_max = torch.cat([x.flatten() for x in x_list]).abs().max().float() - x_sf_global = (448 * 6) / x_abs_max - - weights = {} - for expert_id in range(NUM_EXPERTS): - - w1_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=dtype, - device="cuda") - w1_sf_global = (448 * 6) / w1_weight.abs().max().float() - - w2_weight = torch.randn((HIDDEN_SIZE, INTERMEDIATE_SIZE), - dtype=dtype, - device="cuda") - w2_sf_global = (448 * 6) / w2_weight.abs().max().float() - - w3_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=dtype, - device="cuda") - w3_sf_global = (448 * 6) / w3_weight.abs().max().float() - - w3_w1_global = min( - w1_sf_global, - w3_sf_global) # w3 global and w1 global must be the same - - SCALING_VECTOR_SIZE = 16 - - w1_weight_nvfp4, w1_sf_block = torch.ops.trtllm.fp4_quantize( - w1_weight, w3_w1_global, SCALING_VECTOR_SIZE, False) - w1_sf_block_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - w1_sf_block.cpu().view(INTERMEDIATE_SIZE, -1)) - - w2_weight_nvfp4, w2_sf_block = torch.ops.trtllm.fp4_quantize( - w2_weight, w2_sf_global, SCALING_VECTOR_SIZE, False) - w2_sf_block_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - w2_sf_block.cpu().view(HIDDEN_SIZE, -1)) - - w3_weight_nvfp4, w3_sf_block = torch.ops.trtllm.fp4_quantize( - w3_weight, w3_w1_global, SCALING_VECTOR_SIZE, False) - w3_sf_block_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - w3_sf_block.cpu().view(INTERMEDIATE_SIZE, -1)) - - weights[f"{expert_id}.w1.weight"] = w1_weight_nvfp4.cpu() - weights[f"{expert_id}.w2.weight"] = w2_weight_nvfp4.cpu() - weights[f"{expert_id}.w3.weight"] = w3_weight_nvfp4.cpu() - weights[f"{expert_id}.w1.weight_scale"] = w1_sf_block_unswizzled - weights[f"{expert_id}.w2.weight_scale"] = w2_sf_block_unswizzled - weights[f"{expert_id}.w3.weight_scale"] = w3_sf_block_unswizzled - - weights[f"{expert_id}.w1.input_scale"] = 1.0 / x_sf_global - weights[f"{expert_id}.w2.input_scale"] = 1.0 / x_sf_global - weights[f"{expert_id}.w3.input_scale"] = 1.0 / x_sf_global - weights[f"{expert_id}.w1.weight_scale_2"] = 1.0 / w3_w1_global.cpu() - weights[f"{expert_id}.w2.weight_scale_2"] = 1.0 / w2_sf_global.cpu() - weights[f"{expert_id}.w3.weight_scale_2"] = 1.0 / w3_w1_global.cpu() - - x_list_world.append(x_list) - weights_world.append(weights) - torch.cuda.synchronize() - - def per_rank_test_fused_moe_alltoall(job_id, weights, x_list): - routing_method = DefaultMoeRoutingMethod(top_k=TOP_K) - mapping = Mapping(world_size=world_size, - rank=job_id, - tp_size=world_size, - moe_ep_size=world_size, - moe_tp_size=1, - enable_attention_dp=True) - torch.cuda.set_device(mapping.rank) - torch.manual_seed(mapping.rank) - - weights = {k: v.cuda() for k, v in weights.items()} - x_list = [x.cuda() for x in x_list] - - quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) - with mock.patch.object(WideEPMoE, - "select_alltoall_method_type", - return_value=alltoall_method_type): - alltoall_model = WideEPMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - reduce_results=True, - model_config=ModelConfig(mapping=mapping, - max_num_tokens=MAX_NUM_TOKENS, - quant_config=quant_config), - ) - alltoall_model.to("cuda") - alltoall_model.load_weights([weights]) - - ref_model = CutlassFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - reduce_results=True, - model_config=ModelConfig(mapping=mapping, - max_num_tokens=MAX_NUM_TOKENS, - quant_config=quant_config), - ) - ref_model.to("cuda") - ref_model.load_weights([weights]) - - # Evaluate the outputs on a variant sequence length to verify the robustness of alltoall methods - m = MAX_NUM_TOKENS - i = 0 - while m >= 1: - x = x_list[i] - i += 1 - router_logits = torch.randn((m, NUM_EXPERTS), - dtype=dtype, - device="cuda") - all_rank_num_tokens = [m] * mapping.world_size - - with torch.inference_mode(): - output = alltoall_model.forward( - x, - router_logits, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False) - # Verify the fake impl is correct. - output_fake = alltoall_model.forward_fake( - x, - router_logits, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False) - assert output_fake.shape == output.shape - assert output_fake.dtype == output.dtype - if len(output.shape) == 3: - output = torch.sum(output, dim=1, keepdim=False) - ref_output = ref_model.forward( - x, - router_logits, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False) - - # Evaluate outputs - torch.testing.assert_close(output, ref_output, rtol=0.05, atol=0.5) - m //= 2 - - with MPIPoolExecutor(max_workers=world_size) as executor: - results = executor.map( - per_rank_test_fused_moe_alltoall, - *zip(*[(i, weights_world[i], x_list_world[i]) - for i in range(world_size)])) - for r in results: - assert r is None - - @pytest.mark.skip( reason= "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." @@ -722,165 +417,6 @@ def set_tensor_value_4(x, num_row, num_cols): x.copy_(repeated) -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@skip_pre_blackwell -@pytest.mark.skipif(torch.cuda.device_count() < 4, - reason="needs 4 GPUs to run this test") -@pytest.mark.parametrize( - "alltoall_method_type", - [AlltoallMethodType.NVLinkTwoSided, AlltoallMethodType.NotEnabled], - ids=lambda s: s.name) -def test_fused_moe_fp8_blockwise_wide_ep(alltoall_method_type): - """Test WideEPMoE with FP8 block-wise quantization using DeepGemmFusedMoE as reference.""" - - world_size = 4 - dtype = torch.bfloat16 - # Reduce model size to avoid MPI int32 overflow - HIDDEN_SIZE = 768 - INTERMEDIATE_SIZE = 512 - NUM_EXPERTS = 16 - TOP_K = 2 - MAX_NUM_TOKENS = 256 - - # The MPI can not support FP8, so create weights on each rank - def per_rank_test_fused_moe_alltoall_fp8_blockwise(job_id): - routing_method = DefaultMoeRoutingMethod(top_k=TOP_K) - mapping = Mapping(world_size=world_size, - rank=mpi_rank(), - tp_size=world_size, - moe_ep_size=world_size, - moe_tp_size=1, - enable_attention_dp=True) - torch.cuda.set_device(mapping.rank) - # Use same seed for all ranks to ensure consistency - torch.manual_seed(0) - torch.cuda.manual_seed(0) - - # Generate test data locally on each rank - x_list = [] - m = MAX_NUM_TOKENS - while m >= 1: - x = torch.randn((m, HIDDEN_SIZE), dtype=dtype, device="cuda") - set_tensor_value_2(x, m, HIDDEN_SIZE) - x_list.append(x) - m //= 2 - - # Generate weights locally on each rank (same weights due to same seed) - weights = {} - for expert_id in range(NUM_EXPERTS): - w1_weight = torch.randn( - (INTERMEDIATE_SIZE, HIDDEN_SIZE), dtype=dtype, - device="cuda") / HIDDEN_SIZE - w2_weight = torch.randn((HIDDEN_SIZE, INTERMEDIATE_SIZE), - dtype=dtype, - device="cuda") - w3_weight = torch.randn( - (INTERMEDIATE_SIZE, HIDDEN_SIZE), dtype=dtype, - device="cuda") / HIDDEN_SIZE - - set_tensor_value_3(w1_weight, INTERMEDIATE_SIZE, HIDDEN_SIZE) - set_tensor_value_4(w2_weight, HIDDEN_SIZE, INTERMEDIATE_SIZE) - set_tensor_value_3(w3_weight, INTERMEDIATE_SIZE, HIDDEN_SIZE) - - # FP8 block-wise quantization - w1_weight_fp8, w1_weight_scale = per_block_cast_to_fp8_e8m0( - w1_weight) - w1_weight_fp8 = w1_weight_fp8.view(torch.float8_e4m3fn).cuda() - - w2_weight_fp8, w2_weight_scale = per_block_cast_to_fp8_e8m0( - w2_weight) - w2_weight_fp8 = w2_weight_fp8.view(torch.float8_e4m3fn).cuda() - - w3_weight_fp8, w3_weight_scale = per_block_cast_to_fp8_e8m0( - w3_weight) - w3_weight_fp8 = w3_weight_fp8.view(torch.float8_e4m3fn).cuda() - - weights[f"{expert_id}.w1.weight"] = w1_weight_fp8 - weights[f"{expert_id}.w2.weight"] = w2_weight_fp8 - weights[f"{expert_id}.w3.weight"] = w3_weight_fp8 - weights[f"{expert_id}.w1.weight_scale_inv"] = w1_weight_scale - weights[f"{expert_id}.w2.weight_scale_inv"] = w2_weight_scale - weights[f"{expert_id}.w3.weight_scale_inv"] = w3_weight_scale - weights[f"{expert_id}.w1.weight_scale"] = w1_weight_scale - weights[f"{expert_id}.w2.weight_scale"] = w2_weight_scale - weights[f"{expert_id}.w3.weight_scale"] = w3_weight_scale - - quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) - - # Test WideEPMoE with alltoall method - with mock.patch.object(WideEPMoE, - "select_alltoall_method_type", - return_value=alltoall_method_type): - alltoall_model = WideEPMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - reduce_results=True, - model_config=ModelConfig(mapping=mapping, - max_num_tokens=MAX_NUM_TOKENS, - quant_config=quant_config), - ) - alltoall_model.to("cuda") - alltoall_model.load_weights([weights]) - alltoall_model.post_load_weights() - - # Use DeepGemmFusedMoE as reference - ref_model = DeepGemmFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - reduce_results=True, - model_config=ModelConfig(mapping=mapping, - max_num_tokens=MAX_NUM_TOKENS, - quant_config=quant_config), - ) - ref_model.to("cuda") - ref_model.load_weights([weights]) - ref_model.post_load_weights() - - # Evaluate the outputs on variant sequence lengths - m = MAX_NUM_TOKENS - i = 0 - while m >= 1: - x = x_list[i] - i += 1 - router_logits = torch.randn((m, NUM_EXPERTS), - dtype=dtype, - device="cuda") - all_rank_num_tokens = [m] * mapping.world_size - with torch.inference_mode(): - output = alltoall_model.forward( - x, - router_logits, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False) - ref_output = ref_model.forward( - x, - router_logits, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=False) - - # Evaluate outputs with relaxed tolerance for FP8 - # If WideEPMoE output has TOP_K dimension, reduce it to match DeepGemmFusedMoE - if output.dim() == 3 and output.shape[1] == TOP_K: - output = output.sum(dim=1) - torch.testing.assert_close(output, ref_output, rtol=1e-2, atol=0.1) - m //= 2 - - with MPIPoolExecutor(max_workers=world_size) as executor: - results = executor.map(per_rank_test_fused_moe_alltoall_fp8_blockwise, - range(world_size)) - for r in results: - assert r is None - - @pytest.mark.skip( reason= "Deprecated: covered by tests/unittest/_torch/modules/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." diff --git a/tests/unittest/tools/test_layer_wise_benchmarks.py b/tests/unittest/tools/test_layer_wise_benchmarks.py index c3e4eee83030..edb68422956c 100644 --- a/tests/unittest/tools/test_layer_wise_benchmarks.py +++ b/tests/unittest/tools/test_layer_wise_benchmarks.py @@ -244,9 +244,10 @@ def test_deepseek_v32_ctx_dep(llm_root, world_size): # The pinned DeepSeek FP4 checkpoint requires SM100+. @pytest.mark.skip( - reason="--scaled-from rewrites WideEPMoE.select_alltoall_method_type, which is " - "the only alltoall-selection hook it patches. The WIDEEP backend is deprecated, " - "so weak scaling has no equivalent backend until the hook is generalized." + reason="--scaled-from makes the CTX prefill pack come out all-NaN, independently of " + "the MoE backend: on 4x B200 every combination of gen backend (CUTEDSL, CUTLASS) and " + "prefill backend (CUTLASS, DEEPGEMM) fails the NaN check, while the same command " + "without --scaled-from passes. Re-enable once weak scaling yields finite activations." ) @skip_pre_blackwell @pytest.mark.parametrize("world_size", [4]) @@ -264,7 +265,7 @@ def test_deepseek_r1_gen_scaled_from_16_dep(llm_root, world_size): model_root / "DeepSeek-R1" / "DeepSeek-R1-0528-FP4-v2", "--layer-indices=5,6", "--scaled-from=16", - "--moe-backend=WIDEEP", + "--moe-backend=CUTEDSL", ], cwd=llm_root / "examples" / "layer_wise_benchmarks", env={