diff --git a/cpp/tensorrt_llm/deep_gemm/CMakeLists.txt b/cpp/tensorrt_llm/deep_gemm/CMakeLists.txt index 88f8b6ac143c..14f35e4c6f2c 100644 --- a/cpp/tensorrt_llm/deep_gemm/CMakeLists.txt +++ b/cpp/tensorrt_llm/deep_gemm/CMakeLists.txt @@ -49,6 +49,17 @@ foreach(SOURCE_FILE ${DEEP_GEMM_ALL_FILES}) string(REPLACE "_C." "tensorrt_llm.deep_gemm_cpp_tllm." _content "${_content}") + if(REL_PATH STREQUAL "__init__.py") + string( + PREPEND + _content + "# Use the PyTorch-owned cuBLASLt handle by default. DeepGEMM's\n" + "# standalone handle destructor can run after PyTorch/cuBLASLt teardown\n" + "# in MPI worker shutdown and trigger a double free.\n" + "import os\n" + "os.environ.setdefault(\"DG_USE_PYTORCH_CUBLASLT_HANDLE\", \"1\")\n") + endif() + # Add adaptation header string( PREPEND 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 71c8703f2e4f..5b7bbe49e4c1 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -33,35 +33,75 @@ Input Hidden States ### ConfigurableMoE: The Orchestrator -ConfigurableMoE composes independent components via composition (not inheritance): +`ConfigurableMoE` composes independent components via composition (not inheritance) and **owns module lifecycle** (backend construction, weight loading, comm-strategy creation, `repeat_idx` advancement, DWDP record). Forward-time execution is delegated to a **scheduler**: ```text ConfigurableMoE -├── Backend (pure computation): routing → quantize → FC1 → activation → FC2 -├── Communication (distributed): dispatch tokens → compute → combine results -├── EPLB (optional): dynamic expert migration across GPUs -└── Multi-chunk: splits tokens into chunks to reduce peak memory usage +├── Backend (pure computation: routing → quantize → FC1 → act → FC2) +├── Communication (distributed, optional: dispatch tokens → compute → combine) +├── EPLB (optional: dynamic expert migration across GPUs) +└── MoEScheduler (forward-execution strategy: chunking, EPLB hook ordering, + comm orchestration; selected by backend.scheduler_kind) ``` -Execution flow within ConfigurableMoE (`_forward_chunk_impl`): +`forward_impl` is thin — it resolves `output_dtype`, delegates to `self.scheduler.forward(...)`, then runs wrapper-level bookkeeping that both schedulers share: + +```python +def forward_impl(self, x, router_logits, ...): + outputs = self.scheduler.forward(x, router_logits, ...) + if self.enable_dwdp: + self.dwdp_manager.record_compute_and_prefetch_next(self.layer_idx) + self.repeat_idx = (self.repeat_idx + 1) % self.repeat_count + return outputs +``` + +### Scheduler Selection (`MoESchedulerKind`) + +Each backend declares one of two scheduler kinds via the `scheduler_kind` class attribute (defined on `MoE` base, default `EXTERNAL_COMM`): + +| Kind | Scheduler class | Used by | Cross-rank EP exchange | +|------|-----------------|---------|------------------------| +| `EXTERNAL_COMM` | `ExternalCommMoEScheduler` | Cutlass, DeepGemm, CuteDSL, DenseGEMM, TRTLLMGen | Host issues `Communication.dispatch` / `.combine` outside the MoE kernel; supports per-chunk EPLB hooks and multi-stream chunk overlap | +| `FUSED_COMM` | `FusedCommMoEScheduler` | MegaMoEDeepGemm | Comm is fused into the backend kernel via NVLink SymmBuffer; no host comm; lockstep chunk launches; EPLB stats AllReduced internally | + +The two paths have *deliberately opposite* invariants (`use_dp_padding` honored vs ignored, ADP padding kept vs stripped, empty-chunk substituted vs zero-token kernel launch, multi-stream overlap allowed vs forbidden). See `moe_scheduler.py` class docstrings and `MOE_SCHEDULER_DESIGN.md` for the full contract. + +### External-comm execution flow (most backends) + +`ExternalCommMoEScheduler._forward_chunk_impl` runs per chunk: ```text -routing() → [EPLB] → quantize/dispatch (adaptive order) → backend.run_moe() → combine() - │ │ - Communication Communication +[EPLB start_wait_gpu] → routing → [EPLB done_wait_gpu + update_statistic + route] + → [comm.prepare_dispatch (NVLink2-sided)] → quantize/dispatch (adaptive order) + → backend.run_moe → [EPLB start_set_cpu] → comm.combine → [EPLB done_set_cpu] -Adaptive order (based on comm.supports_post_quant_dispatch()): +Adaptive quantize/dispatch order (gated by comm.supports_post_quant_dispatch()): Post-quant flow: quantize_input() → comm.dispatch() (send quantized data) Pre-quant flow: comm.dispatch() → quantize_input() (send raw, quantize locally) ``` +EPLB hooks fire only at the first/last chunk of the first/last `repeat_idx`. Multi-stream chunk overlap is enabled when `not enable_alltoall and aux_stream is not None`. + +### Fused-comm execution flow (MegaMoE-style) + +`FusedCommMoEScheduler._forward_chunk` runs per chunk: + +```text +[EPLB start_wait_gpu] → routing → [EPLB done_wait_gpu + update_statistic + route] + → backend.quantize_input → backend.run_moe (fused dispatch+GEMM+act+GEMM+combine) + → [EPLB start_set_cpu + done_set_cpu] +``` + +No external `Communication.dispatch` / `.combine`. Zero-token chunks still launch the kernel so peer EP ranks can cross the in-kernel NVLink barrier. + ### Core Design Principles -1. **Composition over inheritance** — Backend, Communication, and EPLB are independent, composable components -2. **Any Backend × Any Communication × EPLB On/Off** — All valid combinations should work +1. **Composition over inheritance** — Backend, Communication, EPLB, and Scheduler are independent, composable components +2. **Any Backend × Any Communication × EPLB On/Off** — All valid combinations should work (subject to `can_implement` and `scheduler_kind`) 3. **Backend = pure computation** — No communication logic, no EPLB logic inside backends -4. **Communication is pluggable** — Auto-selected at runtime by `CommunicationFactory` based on hardware and workload -5. **Backend declares capabilities** — `can_implement()` declares what it supports; ConfigurableMoE adapts flow accordingly +4. **Communication is pluggable** — `EXTERNAL_COMM` backends pick a strategy via `CommunicationFactory` based on hardware/workload; `FUSED_COMM` backends bypass external comm entirely +5. **Backend declares capabilities** — `can_implement()` declares supported quant/dtype; ConfigurableMoE adapts flow accordingly +6. **Backend declares scheduler** — `scheduler_kind` class attribute selects the forward path; lifecycle code stays generic, forward path stays specialized ## Architecture Transition (IMPORTANT) @@ -69,18 +109,19 @@ The codebase is transitioning between two architectures: | | Old Path | New Path | |---|---|---| -| Entry | `XXFusedMoE` (e.g., `CutlassFusedMoE`) | `ConfigurableMoE` + `XXBackend` | -| Communication | Embedded inside each backend | Separated into `communication/` | +| 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 | | Status | Being replaced | Active development | ConfigurableMoE currently supports these backends (`create_moe.py`): -- CutlassFusedMoE, TRTLLMGenFusedMoE, DeepGemmFusedMoE, CuteDslFusedMoE +- `CutlassFusedMoE`, `TRTLLMGenFusedMoE`, `DeepGemmFusedMoE`, `CuteDslFusedMoE`, `DenseGEMMFusedMoE`, `MegaMoEDeepGemm` Still on old path (standalone, with embedded communication): -- TritonFusedMoE, WideEPMoE, VanillaMoE +- `TritonFusedMoE`, `WideEPMoE`, `VanillaMoE` -**Rule: All new features should target ConfigurableMoE + Backend architecture.** +**Rule: All new features should target ConfigurableMoE + Backend + Scheduler architecture.** ## File Map @@ -88,9 +129,10 @@ Still on old path (standalone, with embedded communication): | File | Role | |------|------| -| `configurable_moe.py` | Orchestrator — wires Backend + Communication + EPLB + multi-chunk | +| `configurable_moe.py` | Orchestrator — wires Backend + Communication + EPLB + Scheduler; owns lifecycle and `forward_impl` | +| `moe_scheduler.py` | Forward-execution strategies (`MoEScheduler` ABC, `ExternalCommMoEScheduler`, `FusedCommMoEScheduler`, `create_moe_scheduler` factory) | | `create_moe.py` | Factory — selects MoE class based on `model_config.moe_backend` | -| `interface.py` | Base class `MoE` and enums (`MoEWeightLoadingMode`, `AlltoallMethodType`) | +| `interface.py` | Base class `MoE` and enums (`MoEWeightLoadingMode`, `MoESchedulerKind`, `AlltoallMethodType`) | | `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 | @@ -98,68 +140,99 @@ Still on old path (standalone, with embedded communication): ### Backends (`fused_moe/`) -| File | Backend | Hardware | Scenario | -|------|---------|----------|----------| -| `fused_moe_cutlass.py` | CutlassFusedMoE | SM80+ | High throughput, most comprehensive quant support | -| `fused_moe_trtllm_gen.py` | TRTLLMGenFusedMoE | SM100/SM103 | Min-latency and high-throughput on Blackwell | -| `fused_moe_deepgemm.py` | DeepGemmFusedMoE | SM100/SM103 | FP8 Block Scales on Blackwell | -| `fused_moe_triton.py` | TritonFusedMoE | SM90 only | GPT-OSS on Hopper (requires `swiglu_gptoss_style=True`) | -| `fused_moe_cute_dsl.py` | CuteDslFusedMoE | SM100/SM103 | High throughput NVFP4, generally faster than Cutlass | -| `fused_moe_wide_ep.py` | WideEPMoE | All GPUs | Deprecating — use ConfigurableMoE instead | -| `fused_moe_vanilla.py` | VanillaMoE | All devices | Reference / debugging only | +| File | Backend | Hardware | Scenario | Scheduler | +|------|---------|----------|----------|-----------| +| `fused_moe_cutlass.py` | `CutlassFusedMoE` | SM80+ | High throughput, most comprehensive quant support | `EXTERNAL_COMM` | +| `fused_moe_trtllm_gen.py` | `TRTLLMGenFusedMoE` | SM100/SM103 | Min-latency and high-throughput on Blackwell | `EXTERNAL_COMM` | +| `fused_moe_deepgemm.py` | `DeepGemmFusedMoE` | SM100/SM103 | FP8 Block Scales on Blackwell | `EXTERNAL_COMM` | +| `fused_moe_densegemm.py` | `DenseGEMMFusedMoE` | SM100/SM103 | NVFP4 min-latency; CuTe DSL dense GEMM packs all experts into one matrix (vs Cutlass per-expert scatter), efficient for small token counts | `EXTERNAL_COMM` | +| `fused_moe_cute_dsl.py` | `CuteDslFusedMoE` | SM100/SM103 | High throughput NVFP4, generally faster than Cutlass | `EXTERNAL_COMM` | +| `mega_moe/mega_moe_deepgemm.py` | `MegaMoEDeepGemm` | SM100 only | W4A8_MXFP4_MXFP8 via DeepGEMM `fp8_fp4_mega_moe` fused dispatch+GEMM+act+GEMM+combine kernel; requires `hidden_size % 512 == 0` | `FUSED_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 | Deprecating — use ConfigurableMoE instead | (legacy path) | +| `fused_moe_vanilla.py` | `VanillaMoE` | All devices | Reference / debugging only | (legacy path) | ### Communication (`fused_moe/communication/`) -Communication strategies are auto-selected at runtime by `CommunicationFactory` based on hardware and configuration. See `communication_factory.py` for selection logic and `base.py` for the `Communication` ABC. +Communication strategies are auto-selected at runtime by `CommunicationFactory` based on hardware and configuration. Skipped for `FUSED_COMM` backends. See `communication_factory.py` for selection logic and `base.py` for the `Communication` ABC. + +### MegaMoE (`fused_moe/mega_moe/`) + +| File | Role | +|------|------| +| `mega_moe_deepgemm.py` | `MegaMoEDeepGemm` backend (DeepGEMM `fp8_fp4_mega_moe` wrapper) | +| `CHUNKING_DESIGN.md` | Chunking design for MegaMoE (sequential multi-chunk, in-kernel barrier semantics) | +| `COMMUNICATION_COMPARISON.md` | Comparison of fused-comm SymmBuffer vs external comm strategies | +| `KERNEL_INTERNALS.html` | Reference for the underlying DeepGEMM kernel layout | + +### Design Documents + +| File | Topic | +|------|-------| +| `MOE_SCHEDULER_DESIGN.md` | Scheduler refactor design + `MoEScheduler` contract | +| `mega_moe/CHUNKING_DESIGN.md` | MegaMoE chunking invariants | ### Tests | File | Tests | Status | |------|-------|--------| -| `test_moe_backend.py` | Backend unit tests (run_moe, can_implement) | Active | +| `test_moe_backend.py` | Backend unit tests (`run_moe`, `can_implement`) | Active | | `test_moe_module.py` | ConfigurableMoE integration tests (Backend × Comm × EPLB) | Active | -| `test_fused_moe.py` | Legacy moe tests | Being replaced, do NOT add new tests here | +| `test_fused_moe.py` | Legacy MoE tests | Being replaced, do NOT add new tests here | | `test_moe.py` | Legacy TRTLLM backend tests | Being replaced, do NOT add new tests here | ## Backend Capability Matrix ### Quantization Support -Each backend's `can_implement(quant_algo, dtype_activation, swiglu_gptoss_style)` method declares supported quantizations. Source of truth: the `can_implement` classmethod in each backend file. - -| Quantization | Cutlass | TRTLLMGen | DeepGemm | Triton | CuteDSL | WideEP | Vanilla | -|---|---|---|---|---|---|---|---| -| Unquantized (BF16/FP16) | Y (SM80+) | N | N | Y (SM90, BF16) | N | Y | Y | -| FP8 QDQ | Y (SM89+) | N | N | Y (SM90) | N | Y | Y | -| FP8 Block Scales | Y (SM90, SM120) | Y (SM100/103) | Y (SM100/103) | N | N | Y | Y | -| NVFP4 | Y (SM100/103/120/121) | Y (SM100/103) | N | N | Y (SM100/103) | Y | Y | -| W4A8 NVFP4 FP8 | N | Y (SM100/103) | N | N | N | N | N | -| W4A16 MXFP4 | Y (SM90) | Y (SM100/103) | N | Y (SM90) | N | N | N | -| W4A8 MXFP4 FP8 | Y (SM100/103) | Y (SM100/103) | N | Y (SM90) | N | N | N | -| W4A8 MXFP4 MXFP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | N | N | -| W4A8 AWQ | Y (SM89/90) | N | N | N | N | N | N | -| W8A16 | Y (SM80+) | N | N | N | N | N | N | -| INT4 WoQ (W4AFP8) | N | N | N | N | N | Y | N | - - +Each backend's `can_implement(quant_algo, dtype_activation, swiglu_gptoss_style, ...)` method declares supported quantizations. Source of truth: the `can_implement` classmethod in each backend file. + +| Quantization | Cutlass | TRTLLMGen | DeepGemm | DenseGEMM | CuteDSL | MegaMoE-DG | Triton | WideEP | Vanilla | +|---|---|---|---|---|---|---|---|---|---| +| Unquantized (BF16/FP16) | Y (SM80+) | N | N | N | N | N | Y (SM90, BF16) | Y | Y | +| FP8 QDQ | Y (SM89+) | N | N | N | N | N | Y (SM90) | Y | Y | +| FP8 Block Scales | Y (SM90, SM120) | Y (SM100/103) | Y (SM100/103) | N | Y (SM100/103) | N | N | Y | Y | +| NVFP4 | Y (SM100/103/120/121) | Y (SM100/103) | N | Y (SM100/103) | Y (SM100/103) | N | N | Y | Y | +| W4A8 NVFP4 FP8 | N | Y (SM100/103) | N | N | N | N | N | N | N | +| W4A16 MXFP4 | Y (SM90) | Y (SM100/103) | N | N | N | N | Y (SM90) | N | N | +| W4A8 MXFP4 FP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | N | Y (SM90) | N | N | +| W4A8 MXFP4 MXFP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | Y (SM100, requires `hidden_size % 512 == 0`) | N | N | N | +| W4A8 AWQ | Y (SM89/90) | N | N | N | N | N | N | N | N | +| W8A16 | Y (SM80+) | N | N | N | N | N | N | N | N | +| INT4 WoQ (W4AFP8) | N | N | N | N | N | N | N | Y | N | + +### Scheduler / EPLB Constraints + +- `FUSED_COMM` backends (`MegaMoEDeepGemm`) **must not** layer host-side `Communication.dispatch` / `.combine` on top of the fused kernel — `ConfigurableMoE._create_comm_strategy_auto` returns `None` for them. +- Dynamic EPLB requires backend and quantization-method support. Backends gate + wrapper-level constraints via `validate_configurable_moe`; `MegaMoEDeepGemm` + supports dynamic EPLB by routing to slot IDs and migrating transformed DG + weight tensors registered by its quantization method, with the constraint + `num_slots % ep_size == 0`. +- `FUSED_COMM` backends use `ignore_allreduce=False` for EPLB statistic update because the fused kernel AllReduces routing stats internally. ## Canonical Examples When adding new components, use these reference implementations: | Task | Reference | Key methods to implement | -|------|-----------|------------------------| -| New Backend | `fused_moe_cutlass.py` (CutlassFusedMoE) | `can_implement`, `run_moe`, `create_weights`, `load_weights` | +|------|-----------|--------------------------| +| New `EXTERNAL_COMM` Backend | `fused_moe_cutlass.py` (`CutlassFusedMoE`) | `can_implement`, `run_moe`, `create_weights`, `load_weights` | +| New `FUSED_COMM` Backend | `mega_moe/mega_moe_deepgemm.py` (`MegaMoEDeepGemm`) | Same as above + override `scheduler_kind = MoESchedulerKind.FUSED_COMM` and `validate_configurable_moe` for backend-specific constraints | | New Quantization Method | `quantization.py` → `FP8QDQFusedMoEMethod` | Subclass `FusedMoEMethod`, implement quant/dequant ops | -| New Communication Strategy | `communication/nvlink_one_sided.py` (NVLinkOneSided) | Subclass `Communication`, implement `prepare_dispatch`, `dispatch`, `combine` | +| New Communication Strategy | `communication/nvlink_one_sided.py` (`NVLinkOneSided`) | Subclass `Communication`, implement `prepare_dispatch`, `dispatch`, `combine` | +| New Scheduler | `moe_scheduler.py` (`ExternalCommMoEScheduler` / `FusedCommMoEScheduler`) | Subclass `MoEScheduler`, implement `forward`; add new `MoESchedulerKind` value and wire into `create_moe_scheduler` factory | | Backend Tests | `test_moe_backend.py` | Follow existing parametrize patterns | | Integration Tests | `test_moe_module.py` | Test Backend × Communication × EPLB combinations | -**Note on backend inheritance:** New backends should inherit from `MoE` (in `interface.py`), NOT from `CutlassFusedMoE`. Current backends inherit from `CutlassFusedMoE` as a historical shortcut to reuse infrastructure (load balancer, weight management, TP/EP). This will be refactored — a dedicated `MoEBackend` interface will be extracted. +**Note on backend inheritance:** New backends should inherit from `MoE` (in `interface.py`), NOT from `CutlassFusedMoE`. Current backends inherit from `CutlassFusedMoE` as a historical shortcut to reuse infrastructure (load balancer, weight management, TP/EP). This will be refactored — a dedicated `MoEBackend` interface will be extracted. `MegaMoEDeepGemm` and `DenseGEMMFusedMoE` already inherit directly from `MoE`. ## Anti-Patterns -- **Do NOT add communication logic inside backends** — Communication belongs in `communication/`, backends do pure computation -- **Do NOT modify old `XXFusedMoE` files for new features** — Use ConfigurableMoE + Backend architecture +- **Do NOT add communication logic inside backends** — Communication belongs in `communication/`, backends do pure computation (exception: `FUSED_COMM` backends own the SymmBuffer collective inside their fused kernel) +- **Do NOT add forward-execution policy inside backends** — chunking, EPLB hook ordering, dispatch/combine sequencing belong in `MoEScheduler` +- **Do NOT modify old `XXFusedMoE` files for new features** — Use ConfigurableMoE + Backend + Scheduler architecture - **Do NOT add new tests to `test_fused_moe.py` or `test_moe.py`** — Use `test_moe_backend.py` and `test_moe_module.py` - **Do NOT skip `can_implement()` checks** — Every backend must declare what it supports; unsupported combos must return `(False, reason)` +- **Do NOT pick `scheduler_kind` opportunistically** — Use `EXTERNAL_COMM` (default) unless your backend's fused kernel genuinely owns cross-rank exchange via SymmBuffer / equivalent in-kernel collective; `FUSED_COMM` brings hard invariants (no host comm, lockstep launches, no multi-stream overlap) +- **Schedulers MUST NOT write `moe.repeat_idx`** — `repeat_idx` is wrapper state advanced once per `forward_impl` regardless of chunk count diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py index e37d5db10819..95870b69072a 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -25,7 +25,7 @@ """ import os -from typing import List, Optional, Tuple +from typing import Dict, List, Optional, Tuple import torch @@ -55,7 +55,9 @@ class NVLinkOneSided(Communication): MAX_TOP_K = 8 MAX_PAYLOADS = 8 - # Single shared workspace/memory across the process + # Shared workspaces/memory across the process, keyed by payload layout. + _WORKSPACES: Dict[Tuple[object, ...], dict] = {} + _WORKSPACE_REFCOUNTS: Dict[Tuple[object, ...], int] = {} _WORKSPACE: dict | None = None # MetaInfo indices - initialized from C++ constants @@ -221,10 +223,27 @@ def __init__( ) self.workspace_size_per_rank = 2048 * 1024 * 1024 - # Initialize or reuse workspace + # Initialize or reuse workspace. The C++ op computes payload offsets + # from the current tensors at dispatch time, while the Python singleton + # owns the symmetric memory backing those offsets. Keep separate + # workspaces for different payload layouts so one test/layer cannot + # reuse stale one-sided state from another shape. MnnvlMemory.initialize() + self._workspace_key = ( + self.workspace_size_per_rank, + self.max_num_tokens_per_rank, + self.ep_rank, + self.ep_size, + self.eplb_stats_num_experts, + self.num_experts, + self.top_k, + hidden_size, + dtype, + self.use_low_precision_combine, + ) - if self._WORKSPACE is None: + workspace_state = NVLinkOneSided._WORKSPACES.get(self._workspace_key) + if workspace_state is None: tllm_logger.info( f"NVLinkOneSided: Allocating workspace with size {self.workspace_size_per_rank} bytes." f"ep_rank: {self.ep_rank}, ep_size: {self.ep_size}, top_k: {self.top_k}, max_num_tokens_per_rank: {self.max_num_tokens_per_rank}" @@ -238,37 +257,49 @@ def __init__( self.max_num_tokens_per_rank, self.eplb_stats_num_experts, ) - NVLinkOneSided._WORKSPACE = { + workspace_state = { "workspace_size_per_rank": self.workspace_size_per_rank, "max_num_tokens_per_rank": self.max_num_tokens_per_rank, "ep_rank": self.ep_rank, "ep_size": self.ep_size, "eplb_stats_num_experts": self.eplb_stats_num_experts, + "num_experts": self.num_experts, + "top_k": self.top_k, + "hidden_size": hidden_size, + "dtype": dtype, + "use_low_precision_combine": self.use_low_precision_combine, "mnnvl_mem": mnnvl_mem, "workspace": workspace, "metainfo": metainfo, } + NVLinkOneSided._WORKSPACES[self._workspace_key] = workspace_state else: - assert self._WORKSPACE["workspace_size_per_rank"] == self.workspace_size_per_rank, ( - "reuse workspace with different workspace_size_per_rank" - ) - assert self._WORKSPACE["max_num_tokens_per_rank"] == self.max_num_tokens_per_rank, ( - "reuse workspace with different max_num_tokens_per_rank" - ) - assert self._WORKSPACE["ep_rank"] == self.ep_rank, ( - "reuse workspace with different ep_rank" - ) - assert self._WORKSPACE["ep_size"] == self.ep_size, ( - "reuse workspace with different ep_size" - ) - assert self._WORKSPACE["eplb_stats_num_experts"] == self.eplb_stats_num_experts, ( - "reuse workspace with different eplb_stats_num_experts" - ) - - self.mnnvl_mem = self._WORKSPACE["mnnvl_mem"] - self.workspace = self._WORKSPACE["workspace"] - self.moe_a2a_metainfo = self._WORKSPACE["metainfo"] - self.max_num_tokens_per_rank = self._WORKSPACE["max_num_tokens_per_rank"] + expected_workspace_state = { + "workspace_size_per_rank": self.workspace_size_per_rank, + "max_num_tokens_per_rank": self.max_num_tokens_per_rank, + "ep_rank": self.ep_rank, + "ep_size": self.ep_size, + "eplb_stats_num_experts": self.eplb_stats_num_experts, + "num_experts": self.num_experts, + "top_k": self.top_k, + "hidden_size": hidden_size, + "dtype": dtype, + "use_low_precision_combine": self.use_low_precision_combine, + } + for key, expected_value in expected_workspace_state.items(): + assert workspace_state[key] == expected_value, ( + f"reuse workspace with different {key}" + ) + + NVLinkOneSided._WORKSPACE = workspace_state + NVLinkOneSided._WORKSPACE_REFCOUNTS[self._workspace_key] = ( + NVLinkOneSided._WORKSPACE_REFCOUNTS.get(self._workspace_key, 0) + 1 + ) + self._destroyed = False + self.mnnvl_mem = workspace_state["mnnvl_mem"] + self.workspace = workspace_state["workspace"] + self.moe_a2a_metainfo = workspace_state["metainfo"] + self.max_num_tokens_per_rank = workspace_state["max_num_tokens_per_rank"] # Initialize dispatch state self._dispatch_state = {"phase": "idle"} @@ -289,6 +320,35 @@ def supports_post_quant_dispatch(self) -> bool: """ return True + def destroy(self): + """Release this instance's reference to the shared symmetric workspace.""" + if getattr(self, "_destroyed", False): + return + + self._destroyed = True + workspace_key = getattr(self, "_workspace_key", None) + if workspace_key is None: + return + + if torch.cuda.is_available(): + torch.cuda.synchronize() + + refcount = NVLinkOneSided._WORKSPACE_REFCOUNTS.get(workspace_key, 0) - 1 + if refcount > 0: + NVLinkOneSided._WORKSPACE_REFCOUNTS[workspace_key] = refcount + else: + NVLinkOneSided._WORKSPACE_REFCOUNTS.pop(workspace_key, None) + workspace_state = NVLinkOneSided._WORKSPACES.pop(workspace_key, None) + if NVLinkOneSided._WORKSPACE is workspace_state: + NVLinkOneSided._WORKSPACE = None + if workspace_state is not None: + workspace_state.clear() + + self.mnnvl_mem = None + self.workspace = None + self.moe_a2a_metainfo = None + self._dispatch_state = {"phase": "destroyed"} + def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool: """ Check if NVLINK one-sided comm is feasible for the given workload at runtime. diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index d45283c4e01c..bbc454426c63 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -28,34 +28,41 @@ 4. Unified EPLB integration for backends that support it """ -from typing import Dict, List, Optional, Tuple, Union +from contextlib import contextmanager +from typing import Dict, List, Optional, Union import torch -from tensorrt_llm._torch.expert_statistic import ExpertStatistic from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe.interface import MoE +from tensorrt_llm._torch.modules.fused_moe.interface import MoE, MoESchedulerKind from tensorrt_llm._torch.modules.fused_moe.routing import BaseMoeRoutingMethod from tensorrt_llm._torch.pyexecutor.dwdp import get_global_dwdp_manager from tensorrt_llm._torch.utils import AuxStreamType, EventType, Fp4QuantizedTensor from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator - -from .communication import ( - AllGatherReduceScatter, - Communication, - CommunicationFactory, - DeepEP, - DeepEPLowLatency, - NVLinkOneSided, - NVLinkTwoSided, -) + +from .communication import AllGatherReduceScatter, Communication, CommunicationFactory from .fused_moe_cute_dsl import CuteDslFusedMoE -from .fused_moe_cutlass import CutlassFusedMoE -from .fused_moe_deepgemm import DeepGemmFusedMoE -from .fused_moe_densegemm import DenseGEMMFusedMoE -from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE +from .moe_scheduler import MoEScheduler, create_moe_scheduler + +# Attributes that ConfigurableMoE owns (computed in MoE.__init__ from real +# layer_idx + load balancer) and must be mirrored onto the backend after +# the backend was constructed with layer_idx=None / init_load_balancer=False. +# Adding a new EPLB-derived attribute? Append it here so the sync stays +# in one place and __init__ does not silently drift. +_BACKEND_SYNC_ATTRS = ( + "layer_idx", + "layer_idx_str", + "num_slots", + "layer_load_balancer", + "repeat_count", + "repeat_idx", + "initial_local_expert_ids", + "initial_global_assignments", + "slot_start", + "slot_end", + "expert_size_per_partition", +) class ConfigurableMoE(MoE): @@ -63,7 +70,10 @@ class ConfigurableMoE(MoE): Configurable MoE layer using composition pattern with automatic configuration This class orchestrates the MoE execution flow by composing: - - moe_backend: Existing FusedMoE implementation (CutlassFusedMoE, CuteDslFusedMoE, etc.) + - moe_backend: Existing FusedMoE implementation used as a pluggable backend. + Currently supported backends (see ``create_moe.get_moe_cls``): + CutlassFusedMoE, TRTLLMGenFusedMoE, DeepGemmFusedMoE, + CuteDslFusedMoE, DenseGEMMFusedMoE, MegaMoEDeepGemm. Note: Current FusedMoE implementations are used as backends (transitional). Future will have dedicated MoEBackend interface. - Communication: Handles distributed communication (auto-selected) @@ -81,8 +91,6 @@ class ConfigurableMoE(MoE): weight_loading_mode: Weight loading mode layer_idx: Layer index **kwargs: Additional arguments - - backend_type: Backend type ('cutlass', 'trtllm_gen_min_latency', etc.) - Default: 'cutlass' - tune_max_num_tokens: Max tokens for profiling (passed to backend) - Other backend-specific arguments @@ -93,8 +101,11 @@ class ConfigurableMoE(MoE): Auto-Detection: - EPLB: Enabled if get_moe_load_balancer() is not None - - Backend: Defaults to CutlassMoEBackend, override via backend_type - - Communication: Auto-selected based on hardware (NVLINK > DeepEP > AllGather) + - Backend: Selected by ``model_config.moe_backend`` via ``create_moe.get_moe_cls``; + defaults to CutlassFusedMoE when the requested backend is unsupported + for the active quant/SM config. + - Communication: Auto-selected based on hardware (NVLINK > DeepEP > AllGather); + skipped entirely for FUSED_COMM backends (e.g. MegaMoEDeepGemm). """ @classmethod @@ -163,72 +174,14 @@ def __init__( # 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 - # ========== Create MoE Backend (Default: Cutlass) ========== - from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend, get_moe_cls - - # Get MoE backend class based on override_quant_config or model_config - moe_cls = get_moe_cls(model_config, override_quant_config=override_quant_config) - - # Call create_moe_backend with all necessary parameters - # init_load_balancer=False: Prevents backend from registering itself with load balancer - # without_comm=True: Prevents backend from initializing communication (ConfigurableMoE handles it) - # skip_create_weights_in_init=True: Prevents backend from creating weights in __init__ - # because backend uses layer_idx=None and may have different expert assignments - # We will create weights after syncing attributes from ConfigurableMoE - tmp_skip_create_weights_in_init = model_config.skip_create_weights_in_init - model_config._frozen = False - model_config.skip_create_weights_in_init = True - model_config._frozen = True - - backend = create_moe_backend( - moe_cls=moe_cls, - routing_method=routing_method, - num_experts=self.num_experts, - hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size, - dtype=self.dtype, - reduce_results=self.reduce_results, + # ========== Create MoE Backend (selected by model_config.moe_backend) ========== + self._create_and_sync_backend( model_config=model_config, - aux_stream_dict=self.aux_stream_dict, - weight_loading_mode=self.weight_loading_mode, - bias=kwargs.get("bias", False), - apply_router_weight_on_input=self.apply_router_weight_on_input, - layer_idx=None, - swiglu_alpha=kwargs.get("swiglu_alpha"), - swiglu_beta=kwargs.get("swiglu_beta"), - swiglu_limit=kwargs.get("swiglu_limit"), - init_load_balancer=False, - without_comm=True, - activation_type=self.activation_type, + routing_method=routing_method, + override_quant_config=override_quant_config, + **kwargs, ) - self.validate_backend(backend) - self.backend = backend - self.use_flashinfer = getattr(self.backend, "use_flashinfer", False) - # Sync critical attributes from ConfigurableMoE to backend - # ConfigurableMoE's super().__init__() was called with real layer_idx and initialized load balancer. - # Backend was created with init_load_balancer=False and without_comm=True to avoid - # duplicate initialization. Now sync all attributes from ConfigurableMoE to backend. - if self.backend is not None: - self.backend.layer_idx = self.layer_idx - self.backend.layer_idx_str = self.layer_idx_str - self.backend.num_slots = self.num_slots - self.backend.layer_load_balancer = self.layer_load_balancer - self.backend.repeat_count = self.repeat_count - self.backend.repeat_idx = self.repeat_idx - self.backend.initial_local_expert_ids = self.initial_local_expert_ids - self.backend.initial_global_assignments = self.initial_global_assignments - self.backend.slot_start = self.slot_start - self.backend.slot_end = self.slot_end - self.backend.expert_size_per_partition = self.expert_size_per_partition - - # Create weights here, because the backend needs the layer_load_balancer info to create weights - model_config._frozen = False - model_config.skip_create_weights_in_init = tmp_skip_create_weights_in_init - model_config._frozen = True - if not model_config.skip_create_weights_in_init: - self.backend.create_weights() - # ========== Create Communication Strategy ========== self.comm = self._create_comm_strategy_auto() @@ -274,10 +227,113 @@ def __init__( # TODO: in the future, all the weights related work should be done only in backend. self._weights_removed = True + # ========== Create forward scheduler (ExternalComm / FusedComm) ========== + # Constructed last so the scheduler may safely read any wrapper state + # (comm, aux_stream, event_dict, moe_max_num_tokens, dwdp_*) at init + # time without ordering surprises. Selection is based on + # ``backend.scheduler_kind`` set on the backend class. + self.scheduler: MoEScheduler = create_moe_scheduler(self) + + @staticmethod + @contextmanager + def _temporarily_skip_weight_creation(model_config: ModelConfig): + """Force ``model_config.skip_create_weights_in_init = True`` for the duration. + + The backend is constructed with ``layer_idx=None`` and an unset load + balancer, so weight allocation must be deferred until ConfigurableMoE + has synced the real EPLB-derived attributes onto the backend (see + ``_BACKEND_SYNC_ATTRS``). The flag is also flipped through the + ``_frozen`` Pydantic guard, hence the bracketing dance. Using a + contextmanager guarantees the original state is restored even if + backend construction raises. + """ + previous = model_config.skip_create_weights_in_init + model_config._frozen = False + model_config.skip_create_weights_in_init = True + model_config._frozen = True + try: + yield + finally: + model_config._frozen = False + model_config.skip_create_weights_in_init = previous + model_config._frozen = True + + def _create_and_sync_backend( + self, + *, + model_config: ModelConfig, + routing_method: BaseMoeRoutingMethod, + override_quant_config: Optional["QuantConfig"], + **kwargs, + ) -> None: + """Build the MoE backend, mirror EPLB attrs, then create weights. + + Why this dance: + - ``init_load_balancer=False`` / ``without_comm=True``: the backend + would otherwise re-register itself with the load balancer and + initialize its own communication; ConfigurableMoE owns both. + - ``layer_idx=None``: the wrapper passes the real ``layer_idx`` to + ``MoE.__init__`` to drive load-balancer setup. The backend + receives ``None`` so its own EPLB hooks no-op until we sync the + real values via ``_BACKEND_SYNC_ATTRS`` below. + - ``skip_create_weights_in_init=True`` (via contextmanager): weights + depend on ``layer_load_balancer`` / ``initial_local_expert_ids`` + / etc., which only become known after the sync. Defer weight + creation to the explicit ``backend.create_weights()`` call below. + """ + from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend, get_moe_cls + + moe_cls = get_moe_cls(model_config, override_quant_config=override_quant_config) + + with self._temporarily_skip_weight_creation(model_config): + backend = create_moe_backend( + moe_cls=moe_cls, + routing_method=routing_method, + num_experts=self.num_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + dtype=self.dtype, + reduce_results=self.reduce_results, + model_config=model_config, + aux_stream_dict=self.aux_stream_dict, + weight_loading_mode=self.weight_loading_mode, + bias=kwargs.get("bias", False), + apply_router_weight_on_input=self.apply_router_weight_on_input, + layer_idx=None, + swiglu_alpha=kwargs.get("swiglu_alpha"), + swiglu_beta=kwargs.get("swiglu_beta"), + swiglu_limit=kwargs.get("swiglu_limit"), + init_load_balancer=False, + without_comm=True, + activation_type=self.activation_type, + ) + + self.validate_backend(backend) + self.backend = backend + self.use_flashinfer = getattr(self.backend, "use_flashinfer", False) + + # Mirror wrapper-owned EPLB / layer-id state onto the backend so any + # backend code path that reads e.g. ``self.layer_load_balancer`` or + # ``self.num_slots`` sees the real values resolved by MoE.__init__. + if self.backend is not None: + for attr in _BACKEND_SYNC_ATTRS: + setattr(self.backend, attr, getattr(self, attr)) + + # Sync done -- now the backend has enough info to allocate weight + # tensors with the right shard / slot count. + if not model_config.skip_create_weights_in_init: + self.backend.create_weights() + def _supports_load_balancer(self) -> bool: - """Check if this MoE implementation supports load balancer.""" - # During initialization, backend might not be created yet - # Return True by default (most backends support it), backend will validate later + """Check if this MoE implementation supports load balancer. + + ``MoE.__init__`` can query this before ``ConfigurableMoE`` has + created ``self.backend``. In that initialization window, fall back to + the wrapper-level DP/parallelism condition; ``validate_backend`` runs + after backend construction and enforces the backend-specific answer. + """ + # During initialization, backend might not be created yet. + # Backend-specific support is checked later by validate_backend. if not hasattr(self, "backend") or self.backend is None: return self.use_dp and self.parallel_size > 1 return self.backend._supports_load_balancer() @@ -310,19 +366,6 @@ def _should_enable_dwdp(self) -> bool: quant_mode is not None and hasattr(quant_mode, "has_nvfp4") and quant_mode.has_nvfp4() ) - def _create_comm_strategy(self, model_config: ModelConfig) -> Optional[Communication]: - """ - Create communication strategy based on configuration - - Default: None (will use factory to auto-select when needed) - Auto-selects best strategy based on hardware and configuration - - """ - # Communication strategy is None by default - # Will be created lazily in determine_communication_method() when first needed - # For now, return None and create on-demand - return None - def _get_quant_config_dict(self, model_config: ModelConfig) -> Optional[Dict]: """ Extract quantization configuration from model_config @@ -440,9 +483,14 @@ def _create_comm_strategy_auto(self) -> Communication: """ Auto-create the best communication strategy based on hardware and configuration - Uses factory to select optimal strategy. - + Uses factory to select optimal strategy. Backends whose fused kernel + owns cross-rank exchange (``scheduler_kind=FUSED_COMM``) skip + host-side comm entirely; layering Communication.dispatch / combine + on top of the fused exchange would double-count traffic and break + the in-kernel NVLink barrier semantics. """ + if self.backend.scheduler_kind == MoESchedulerKind.FUSED_COMM: + return None return CommunicationFactory.create_strategy( model_config=self.model_config, num_experts=self.num_experts, @@ -468,714 +516,66 @@ def forward_impl( use_dp_padding: Optional[bool] = None, **kwargs, ) -> torch.Tensor: - """ - Universal forward implementation framework + """Forward entry point. - Flow: - 1. Handle padding - 2. Calculate chunk count and determine communication method - 3. Execute MoE computation (single or multiple chunks) - 4. Handle output truncation and EPLB repeat + Acts as a thin wrapper that: + + 1. Validates / fills ``output_dtype``. + 2. Delegates the per-pass execution to ``self.scheduler`` (chosen + once at init time from ``backend.scheduler_kind``). + 3. Records DWDP compute/prefetch (per layer, not per chunk). + 4. Advances the EPLB ``repeat_idx``. + + DP-padding handling and chunking live in the scheduler. """ - # TODO: to clarify whether the output_dtype is needed. + del kwargs + if isinstance(x, Fp4QuantizedTensor): assert output_dtype is not None else: output_dtype = x.dtype - # ========== Mega MoE fast-path ========== - # DeepGEMM's ``fp8_fp4_mega_moe`` subsumes dispatch + GEMM1 + SwiGLU - # + GEMM2 + combine into one collective launch and owns routing - # weight application + EP exchange. It does not fit the pipeline - # below (EPLB, chunk-quant, padded-broadcast); short-circuit - # BEFORE the ``all_rank_num_tokens_padded`` computation so the - # backend receives the raw per-rank unpadded token counts it - # needs to slice off ADP padding. EPLB is also rejected here as - # a hard check (``validate_backend`` enforces this at __init__, - # but keep the assert defensively in case LB state changes). - from .mega_moe import MegaMoEDeepGemmFusedMoE - - if isinstance(self.backend, MegaMoEDeepGemmFusedMoE): - assert not self._using_load_balancer(), ( - "MegaMoEDeepGemmFusedMoE does not support EPLB; disable the load " - "balancer or pick a different backend." - ) - return self._forward_chunk_mega_impl( - x, - router_logits, - output_dtype=output_dtype, - all_rank_num_tokens=all_rank_num_tokens, - do_finalize=do_finalize, - ) - - # ========== Step 1: Handle padding ========== - if all_rank_num_tokens is None: - all_rank_num_tokens = [x.shape[0]] - - 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 - - # ========== Step 2: Determine communication method ========== - num_chunks = self.calculate_num_chunks(all_rank_num_tokens_padded) - - # Determine and setup communication strategy (may fallback to AllGather) - self.determine_communication_method(all_rank_num_tokens_padded, num_chunks) - - # ========== Step 3: Execute MoE computation ========== - if num_chunks == 1: - # Single chunk case - outputs = self._forward_single_chunk( - x, - router_logits, - output_dtype, - all_rank_num_tokens_padded, - use_dp_padding, - do_finalize, - ) - else: - # Multiple chunks case - outputs = self._forward_multiple_chunks( - x, - router_logits, - num_chunks, - output_dtype, - all_rank_num_tokens_padded, - use_dp_padding, - do_finalize, - ) - - # DWDP: record compute and trigger next prefetch (per-layer, not per-chunk) - if self.enable_dwdp: - self.dwdp_manager.record_compute_and_prefetch_next(self.layer_idx) - - # ========== Step 4: Handle output truncation and EPLB repeat ========== - if self.use_dp and self.parallel_size > 1: - outputs = outputs[: all_rank_num_tokens[self.mapping.tp_rank]] - - # EPLB repeat logic - self.repeat_idx = (self.repeat_idx + 1) % self.repeat_count - - return outputs - - def _prepare_workspace_deepgemm( - self, - x: Union[torch.Tensor, Fp4QuantizedTensor], - all_rank_num_tokens: List[int], - ) -> Optional[torch.Tensor]: - """ - Prepare workspace for DeepGemmFusedMoE backend. - - Args: - x: Input tensor - all_rank_num_tokens: List of token counts for all ranks (used when use_dp is True) - - Returns: - Workspace tensor or None if not using DeepGemmFusedMoE - """ - if not isinstance(self.backend, DeepGemmFusedMoE): - return None - - # Calculate the number of rows - num_rows = x.shape[0] - if self.use_dp and self.comm is not None: - # When using communication, dispatch will create tensors with shape: - # [ep_size * max_tokens_per_rank, ...] due to padding for balanced distribution - # So we need to allocate workspace based on this size - if isinstance(self.comm, DeepEPLowLatency): - # deeptplowlatency dispatch outputs shape is - # [#local_experts * moe_ep_size * max_tokens_per_rank, hidden size] - # local_experts = self.num_slots / moe_ep_size - num_rows = self.num_slots * max(all_rank_num_tokens) - else: - num_rows = self.mapping.moe_ep_size * max(all_rank_num_tokens) - - workspaces = self.backend.get_workspaces([num_rows]) - return workspaces[0] - - def _forward_single_chunk( - self, - x: Union[torch.Tensor, Fp4QuantizedTensor], - router_logits: torch.Tensor, - output_dtype: Optional[torch.dtype], - all_rank_num_tokens: List[int], - use_dp_padding: Optional[bool], - do_finalize: bool = True, - ) -> torch.Tensor: - """ - Single chunk execution path - - """ - # Calculate EPLB flags (first call or last call) - is_first_call = self.repeat_idx == 0 - is_last_call = self.repeat_idx == self.repeat_count - 1 - - # ========== Create workspace for DeepGemmFusedMoE ========== - workspace = self._prepare_workspace_deepgemm(x, all_rank_num_tokens) - - # Execute unified flow (handles both separated and fused routing) - outputs = self._forward_chunk_impl( + outputs = self.scheduler.forward( x, router_logits, - output_dtype, - all_rank_num_tokens, - use_dp_padding, - is_first_call, - is_last_call, - do_finalize, - workspace=workspace, - ) - - return outputs - - def _forward_chunk_mega_impl( - self, - x: Union[torch.Tensor, Fp4QuantizedTensor], - router_logits: torch.Tensor, - *, - output_dtype: Optional[torch.dtype], - all_rank_num_tokens: Optional[List[int]], - do_finalize: bool, - ) -> torch.Tensor: - """Run ``MegaMoEDeepGemmFusedMoE`` by the separated-routing-and-quant path. - - Mirrors the CUTLASS / CUTEDSL pipeline: compute routing + BF16→FP8 - pre-quant once in ConfigurableMoE, then hand the pre-quantized - tensors to ``backend.run_with_prequant`` which only does buffer - copies + the fused kernel launch. This keeps pre-processing out - of the inner-loop backend call so the GPU work visible to DG's - kernel matches what DG's own benchmarks measure (~330 us at - DSV3 seq=32 ep=4) instead of the 890 us we saw when routing + - Python ``per_token_cast_to_fp8`` ran inside the backend. - - Quant uses a ``torch.compile``-fused variant (see - ``mega_moe.backend._get_fused_per_token_cast_to_fp8``) so the ~8 - Python-launched elementwise / reduction kernels DG's helper - decomposes into collapse to a single Inductor-generated Triton - kernel. - - Phase 1: EPLB disabled (see ``MegaMoEDeepGemmFusedMoE._supports_load_balancer``); - ``apply_router_weight_on_input`` also rejected at backend init. - """ - assert not self.apply_router_weight_on_input, ( - "ConfigurableMoE with MegaMoEDeepGemmFusedMoE does not support apply_router_weight_on_input" - ) - assert do_finalize, ( - "MegaMoE's fused kernel always finalizes — do_finalize=False is not supported" - ) - - # Resolve the raw (unpadded) token count under ADP. MegaMoE's - # SymmBuffer collective requires every rank to enter the kernel, - # even on zero-token ranks, so we *don't* short-circuit when - # num_tokens == 0 — we still run quant/route on the empty slice - # (cheap) and let the kernel launch proceed. - # ``all_rank_num_tokens`` is one entry per EP rank (Phase 1 - # asserts ``ep_size == parallel_size``), so index by ``moe_ep_rank``. - if all_rank_num_tokens is not None: - num_tokens = int(all_rank_num_tokens[self.mapping.moe_ep_rank]) - else: - num_tokens = x.shape[0] - assert num_tokens <= x.shape[0] - - if output_dtype is None: - output_dtype = x.dtype if not isinstance(x, Fp4QuantizedTensor) else torch.bfloat16 - - # Slice to real tokens (skip DP padding rows if any) and compute - # routing + pre-quant. These used to live inside - # ``backend.forward_impl``; hoisting them up collapses the Python - # call stack by two frames and — more importantly — lets the - # backend's ``run_with_prequant`` match DG's ``run_fused`` shape - # contract exactly (4 x buf.copy_ + kernel). - x_real = x[:num_tokens] - router_logits_real = router_logits[:num_tokens] - - if num_tokens > 0: - topk_idx, topk_weights = self.routing_method.apply(router_logits_real) - topk_idx = topk_idx.to(torch.int64) - topk_weights = topk_weights.to(torch.float32) - x_fp8, x_sf = self.backend.quantize_input(x_real) - else: - # Zero-token rank: fabricate empty tensors so - # ``run_with_prequant`` still takes the same shape contract. - # x_sf is (m, hidden_size // 128) int32 — packed UE8M0 stores - # 4 u8 scales per int32 over a 32-element block, i.e. 128 - # input elements per int32 stride. - device = x.device - x_fp8 = torch.empty((0, self.hidden_size), dtype=torch.float8_e4m3fn, device=device) - x_sf = torch.empty((0, self.hidden_size // 128), dtype=torch.int32, device=device) - topk_idx = torch.empty( - (0, self.routing_method.experts_per_token), dtype=torch.int64, device=device - ) - topk_weights = torch.empty( - (0, self.routing_method.experts_per_token), dtype=torch.float32, device=device - ) - - return self.backend.run_with_prequant( - x_fp8=x_fp8, - x_sf=x_sf, - topk_idx=topk_idx, - topk_weights=topk_weights, - num_tokens=num_tokens, + do_finalize=do_finalize, output_dtype=output_dtype, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding, ) - def _forward_chunk_impl( - self, - x: Union[torch.Tensor, Fp4QuantizedTensor], - router_logits: torch.Tensor, - output_dtype: Optional[torch.dtype], - all_rank_num_tokens: List[int], - use_dp_padding: bool, - is_first_call: bool, - is_last_call: bool, - do_finalize: bool = True, - workspace: Optional[dict] = None, - ) -> torch.Tensor: - """ - Unified execution flow for all backends - - Flow (based on EPLB_in_MOE[1].html): - 1. [EPLB] Start wait GPU stage (first call only, if enabled) - 2. Apply routing (only if backend supports routing separation) - 3. [EPLB] Update statistics and route (only if EPLB enabled) - 4. Quantization and Communication (adaptive ordering) - 5. MoE computation (backend) - 6. [EPLB] Start CPU stage (last call only, if enabled) - 7. Communication combine - 8. [EPLB] Done CPU stage (last call only, if enabled) - - - Separated routing: fused_moe_wide_ep.py:456-780, fused_moe_cutlass.py:236-443 - - Fused routing: fused_moe_trtllm_gen.py - """ - - # Mega MoE never reaches here: ``forward_impl`` short-circuits to - # ``_forward_chunk_mega_impl`` before the chunk/padding pipeline runs - # so the backend can see raw ``all_rank_num_tokens``. Any isinstance - # check here would be dead code. - - # ========== Step 1: EPLB - Start wait GPU stage ========== - self._load_balancer_start_wait_gpu_stage(is_first_call) - - # ========== Step 2: Apply routing (only if backend supports load balancer) ========== - - if self.backend._supports_load_balancer(): - # Separated routing: ConfigurableMoE calls routing_method - token_selected_experts, token_final_scales = self.routing_method.apply(router_logits) - - # Convert to standard dtypes for consistency with other MoE implementations - token_selected_experts = token_selected_experts.to(torch.int32) - - assert token_selected_experts.shape[1] == self.routing_method.experts_per_token - assert token_selected_experts.shape == token_final_scales.shape - # CutlassFusedMoE and DenseGEMMFusedMoE expect float32, while TRTLLMGenFusedMoE uses bfloat16 - if isinstance(self.backend, (CutlassFusedMoE, DenseGEMMFusedMoE)): - assert token_final_scales.dtype == torch.float32 - assert token_selected_experts.dtype == torch.int32 - - # Convert token_final_scales to bfloat16 if needed (TRTLLMGen backend requires it) - if token_final_scales is not None and isinstance(self.backend, TRTLLMGenFusedMoE): - token_final_scales = token_final_scales.to(torch.bfloat16) - - # Apply router weight on input if enabled - 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 - # Check if using DeepEP strategies (they don't support token_final_scales=None) - if isinstance(self.comm, (DeepEP, DeepEPLowLatency)): - # DeepEP doesn't support token_final_scales is None - token_final_scales = torch.ones_like(token_final_scales) - else: - token_final_scales = None - - else: - # Fused routing: Backend handles routing internally - # EPLB must NOT be enabled for fused routing backends - assert not self._using_load_balancer(), ( - f"EPLB is enabled but backend {self.backend.__class__.__name__} " - f"has fused routing (does not support routing separation)" - ) - - # For fused routing, we don't have token_selected_experts yet - # Will be handled by backend.run_moe_with_routing() later - token_selected_experts = None - token_final_scales = None - - # ========== Step 3: EPLB - Update statistics and route ========== - # Only executed if backend supports routing separation AND EPLB is enabled - if self.layer_load_balancer and token_selected_experts is not None: - self._load_balancer_done_wait_gpu_stage(is_first_call) - - # Update EPLB statistics (method depends on communication strategy) - # Use base class method: ignore_allreduce=True for NVLINK two-sided/one-sided (uses local stats only) - ignore_allreduce = ( - self._is_using_nvlink_two_sided() or self._is_using_nvlink_one_sided() - ) - self._load_balancer_update_statistic( - token_selected_experts, - is_first_call, - is_last_call, - ignore_allreduce=ignore_allreduce, - ) - - # EPLB routing: expert IDs -> slot IDs - token_selected_slots = self._load_balancer_route(token_selected_experts, self.use_dp) - else: - token_selected_slots = token_selected_experts - - if token_selected_slots is not None: - 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 - ) - - # ========== Step 3.5: Communication Prepare Phase (BEFORE quantization) ========== - # NVLINK two-sided has a prepare phase to gather EPLB statistics - - local_statistic_tensor_for_dispatch = None - eplb_dispatch_kwargs = {} - should_update_eplb_after_dispatch = False - # Only NVLINK two-sided needs prepare_dispatch - if self._is_using_nvlink_two_sided(): - # Get local statistic info if this is the last call and EPLB is enabled - local_statistic_tensor = None - if is_last_call: - local_statistic_tensor = self._load_balancer_get_local_statistic_tensor() - - # Call prepare_dispatch (gathers statistics for NVLINK two-sided) - # prepare_dispatch stores alltoall_info in _dispatch_state and returns gathered_stats - gathered_stats = self.comm.prepare_dispatch( - token_selected_slots, all_rank_num_tokens, local_statistic_tensor - ) - - # Update EPLB with gathered statistics (if available) - if gathered_stats is not None: - gathered_stats = gathered_stats.view((self.mapping.moe_ep_size, self.num_experts)) - self._load_balancer_update_statistic_with_gathered_statistic(gathered_stats) - # TODO: The abstract does not work well as NVLinkTwoSided gathers EPLB stats in prepare_dispatch, - # while NVLinkOneSided gathers EPLB stats in dispatch. - elif self._is_using_nvlink_one_sided(): - if self.layer_load_balancer and is_last_call: - local_statistic_tensor_for_dispatch = ( - self._load_balancer_get_local_statistic_tensor() - ) - if local_statistic_tensor_for_dispatch is not None: - eplb_dispatch_kwargs["eplb_local_stats"] = local_statistic_tensor_for_dispatch - should_update_eplb_after_dispatch = True - - # ========== Step 4 & 5: Quantization and Communication Dispatch ========== - # Order depends on whether strategy supports post-quant dispatch - if self.comm is not None: - # Check if we should use post-quant dispatch - # supports_post_quant_dispatch checks strategy capability for the current quant mode - supports_post_quant = self.comm.supports_post_quant_dispatch() - - # Call dummy_allreduce before allgather for load balancing debug - if self.enable_dummy_allreduce: - self.dummy_allreduce() - - dispatch_kwargs = dict(eplb_dispatch_kwargs) - if isinstance(self.comm, DeepEP) and isinstance(self.backend, TRTLLMGenFusedMoE): - dispatch_kwargs["enable_sanitize_expert_ids"] = True - - if supports_post_quant: - # ===== Post-quant flow: Quantize → Dispatch ===== - - # Step 4a: Quantization FIRST - x, x_sf = self.backend.quantize_input(x) - - # Step 4b: Dispatch AFTER quantization - # Get pre_quant_scale for W4AFP8 if available (only DeepEPLowLatency needs it) - # Other strategies will ignore this via **kwargs, so it's safe to pass unconditionally - if hasattr(self, "quant_scales") and self.quant_scales is not None: - if hasattr(self.quant_scales, "pre_quant_scale_1"): - dispatch_kwargs["pre_quant_scale"] = self.quant_scales.pre_quant_scale_1 - x, x_sf, token_selected_slots, token_final_scales = self.comm.dispatch( - hidden_states=x, - hidden_states_sf=x_sf, - token_selected_slots=token_selected_slots, - token_final_scales=token_final_scales, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=use_dp_padding, - **dispatch_kwargs, - ) - if should_update_eplb_after_dispatch: - gathered_stats = self.comm.get_eplb_gathered_statistics() - self._load_balancer_update_statistic_with_gathered_statistic(gathered_stats) - else: - # ===== Pre-quant flow: Dispatch → Quantize ===== - - # Step 4a: Dispatch FIRST (unquantized data) - x, x_sf, token_selected_slots, token_final_scales = self.comm.dispatch( - hidden_states=x, - hidden_states_sf=None, # Not quantized yet - token_selected_slots=token_selected_slots, - token_final_scales=token_final_scales, - all_rank_num_tokens=all_rank_num_tokens, - use_dp_padding=use_dp_padding, - **dispatch_kwargs, - ) - - # Step 4b: Quantization AFTER dispatch - x, x_sf = self.backend.quantize_input(x, post_quant_comm=False) - else: - # No communication, just quantize - # (use non-post-quant-comm path for TRTLLMGenFusedMoE) - x, x_sf = self.backend.quantize_input(x, post_quant_comm=False) - - # ========== Step 6: MoE Computation ========== - - # Call unified run_moe interface with common parameters - # If EPLB is enabled, token_selected_slots represents expert slots - # Otherwise, token_selected_experts represents expert IDs - final_hidden_states = self.backend.run_moe( - x=x, - token_selected_experts=token_selected_slots, - token_final_scales=token_final_scales, - x_sf=x_sf, - **self._get_backend_kwargs( - router_logits, do_finalize, all_rank_num_tokens, output_dtype, x, workspace - ), - ) - - # ========== Step 8: EPLB - Start CPU stage ========== - self._load_balancer_start_set_cpu_stage(is_last_call) - - # ========== Step 9: Communication - Combine ========== - if self.comm is not None: - if self.enable_dummy_allreduce: - self.dummy_allreduce() - # Use unified combine interface (reads dispatch state from strategy) - all_rank_max_num_tokens = max(all_rank_num_tokens) - final_hidden_states = self.comm.combine( - final_hidden_states, - all_rank_max_num_tokens=all_rank_max_num_tokens, - ) - else: - # For non-comm case, It should be attention TP or single rank. - # only check if allreduce is needed - if self.parallel_size > 1 and self.reduce_results: - final_hidden_states = self.all_reduce(final_hidden_states) - # ========== Step 10: EPLB - Done CPU stage ========== - self._load_balancer_done_set_cpu_stage(is_last_call) - - return final_hidden_states - - def _prepare_workspaces_for_chunk( - self, - all_rank_num_tokens_list: List[Optional[List[int]]], - chunk_size_list: List[int], - use_multi_stream: bool, - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - """ - Prepare workspaces for chunked execution with DeepGemmFusedMoE backend. - This will also be used for alltoall communication in the future. - - Args: - all_rank_num_tokens_list: List of token counts per rank for each chunk (None if not using DP) - chunk_size_list: List of chunk sizes - use_multi_stream: Whether to use multi-stream execution (requires workspace_1) - - Returns: - Tuple of (workspace_0, workspace_1), where workspace_1 is None if not using multi-stream - """ - workspace_0 = None - workspace_1 = None - - if not isinstance(self.backend, DeepGemmFusedMoE): - return workspace_0, workspace_1 - - # Always need at least workspace_0 - chunk_size_0 = ( - self.mapping.moe_ep_size * max(all_rank_num_tokens_list[0]) - if self.use_dp and all_rank_num_tokens_list[0] is not None - else chunk_size_list[0] - ) - workspace_chunk_sizes = [chunk_size_0] - - # Add workspace_1 if using multi-stream for alternating between streams - # Reuse chunk_size_0 since it's always >= chunk_size_1 (first chunk is largest) - if use_multi_stream: - workspace_chunk_sizes.append(chunk_size_0) - - workspaces = self.backend.get_workspaces(workspace_chunk_sizes) - workspace_0 = workspaces[0] - if use_multi_stream: - workspace_1 = workspaces[1] - - return workspace_0, workspace_1 - - def _forward_multiple_chunks( - self, - x: Union[torch.Tensor, Fp4QuantizedTensor], - router_logits: torch.Tensor, - num_chunks: int, - output_dtype: Optional[torch.dtype], - all_rank_num_tokens: List[int], - use_dp_padding: Optional[bool], - do_finalize: bool = True, - ) -> torch.Tensor: - """ - Multiple chunks execution path with auxiliary stream for overlapping - - Same as original implementation - chunking logic is backend-agnostic - - """ - # ========== Chunk preparation ========== - if self.use_dp: - # When using DP: need all ranks' token counts for reducescatter - all_rank_chunk_size_list = [ - self.split_chunk(val, num_chunks) for val in all_rank_num_tokens - ] - all_rank_num_tokens_list = [ - [val[idx_chunk] for val in all_rank_chunk_size_list] - for idx_chunk in range(num_chunks) - ] - chunk_size_list = all_rank_chunk_size_list[self.rank] - - # For alltoall, replace 0 with 1 (avoid empty tensor) - if self.enable_alltoall: - 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 - ] - else: - # When not using DP: only need current rank's input size - all_rank_num_tokens_list = [None] * num_chunks - chunk_size_list = self.split_chunk(x.shape[0], num_chunks) - - x_list = x.split(chunk_size_list) - router_logits_list = router_logits.split(chunk_size_list) - - # Determine if we need multiple streams for overlapped execution - use_multi_stream = not self.enable_alltoall and self.aux_stream is not None - - # ========== Setup auxiliary stream ========== - if use_multi_stream: - self.event_dict[EventType.Main].record() - with torch.cuda.stream(self.aux_stream): - self.event_dict[EventType.Main].wait() - - # ========== Create workspace for DeepGemmFusedMoE ========== - workspace_0, workspace_1 = self._prepare_workspaces_for_chunk( - all_rank_num_tokens_list, chunk_size_list, use_multi_stream - ) + # DWDP: record compute and trigger next prefetch (per-layer, not per-chunk). + # Owned at the wrapper because schedulers must not run it twice (external-comm + # might enter via single- or multi-chunk paths). + if self.enable_dwdp: + self.dwdp_manager.record_compute_and_prefetch_next(self.layer_idx) - # ========== Padding empty chunk ========== - chunked_used = torch.ones(num_chunks, dtype=torch.bool) - if self.use_dp: - # For empty chunk, will use chunk 0 instead. The current split heuristic - # ensures that if an empty chunk exists, Chunk 0 contains exactly one token. - assert x_list[0].numel() != 0, "chunk 0 shouldn't be empty" - x_list = list(x_list) - router_logits_list = list(router_logits_list) - for idx_chunk in range(num_chunks): - _x = x_list[idx_chunk] - if _x.numel() == 0: - chunked_used[idx_chunk] = False - x_list[idx_chunk] = x_list[0] - router_logits_list[idx_chunk] = router_logits_list[0] - all_rank_num_tokens_list[idx_chunk][self.mapping.tp_rank] = ( - all_rank_num_tokens_list[0][self.mapping.tp_rank] - ) - x_list = tuple(x_list) - router_logits_list = tuple(router_logits_list) - - # ========== Execute chunking with overlap ========== - outputs_list = [] - for idx_chunk, (x_chunk, router_logits_chunk) in enumerate(zip(x_list, router_logits_list)): - # Calculate EPLB's first/last call - 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 use_multi_stream: - # Alternate between main stream and auxiliary stream - # Each stream processes complete chunks (forward + reducescatter) - if idx_chunk % 2 == 0: - # Even chunk: execute on auxiliary stream - with torch.cuda.stream(self.aux_stream): - outputs = self._forward_chunk_impl( - x_chunk, - router_logits_chunk, - output_dtype, - all_rank_num_tokens_list[idx_chunk], - use_dp_padding, - is_first_call, - is_last_call, - do_finalize, - workspace=workspace_0, - ) - else: - # Odd chunk: execute on main stream - outputs = self._forward_chunk_impl( - x_chunk, - router_logits_chunk, - output_dtype, - all_rank_num_tokens_list[idx_chunk], - use_dp_padding, - is_first_call, - is_last_call, - do_finalize, - workspace=workspace_1, - ) - else: - # No overlap - outputs = self._forward_chunk_impl( - x_chunk, - router_logits_chunk, - output_dtype, - all_rank_num_tokens_list[idx_chunk], - use_dp_padding, - is_first_call, - is_last_call, - do_finalize, - workspace=workspace_0, - ) - - if chunked_used[idx_chunk]: - outputs_list.append(outputs) - - # ========== Wait for auxiliary stream to complete ========== - if use_multi_stream: - # Wait for auxiliary stream to complete all its chunks - with torch.cuda.stream(self.aux_stream): - self.event_dict[EventType.MoeChunkingOverlap].record() - self.event_dict[EventType.MoeChunkingOverlap].wait() - - # ========== Concatenate outputs from all chunks ========== - outputs = torch.cat(outputs_list) + # EPLB repeat counter: advance once per forward, regardless of chunk count. + # Schedulers are forbidden from rotating ``repeat_idx`` themselves. + self.repeat_idx = (self.repeat_idx + 1) % self.repeat_count return outputs # ========== Backend Validation ========== def validate_backend(self, backend: MoE): - """ - Validate MOE backend. + """Validate MoE backend compatibility with this ConfigurableMoE. - It validates that: - 1. Backend is not None - 2. If EPLB is enabled, backend must support routing separation - - Args: - backend: MoEBackend instance to set + Generic checks (always run): + 1. ``backend`` is not None. + 2. If EPLB is enabled, the backend must support routing + separation (``backend._supports_load_balancer()``). - Raises: - ValueError: If backend is incompatible with current configuration - - Note: EPLB initialization is done in __init__, not in setter. - Setter only validates compatibility. + Backend-specific checks are delegated to + ``backend.validate_configurable_moe(self)``; backends with extra + constraints (e.g. fused-comm backends rejecting dynamic + EPLB) override that hook. EPLB / num_slots / ep_size are already + populated on ``self`` by ``MoE.__init__`` -> ``_init_load_balancer`` + before this is called, so backends may inspect them directly. """ if backend is None: raise ValueError("Backend cannot be None") - # Validate EPLB compatibility if self._using_load_balancer() and not backend._supports_load_balancer(): raise ValueError( f"EPLB is enabled but backend {backend.__class__.__name__} " @@ -1183,179 +583,7 @@ def validate_backend(self, backend: MoE): f"Either disable EPLB or use a backend that supports load balancer." ) - # ========== Helper Methods ========== - - def _is_using_nvlink_two_sided(self) -> bool: - """Check if using NVLinkTwoSided communication strategy""" - return isinstance(self.comm, NVLinkTwoSided) - - def _is_using_nvlink_one_sided(self) -> bool: - """Check if using NVLinkOneSided communication strategy""" - return isinstance(self.comm, NVLinkOneSided) - - def _get_nvlink_onesided_moe_output( - self, - all_rank_num_tokens: Optional[List[int]], - output_dtype: Optional[torch.dtype], - ) -> Optional[torch.Tensor]: - """ - Get workspace output buffer for NVLinkOneSided communication backend. - - This method handles moe_output allocation for both CutlassFusedMoE and TRTLLMGenFusedMoE - when using NVLinkOneSided communication strategy. - - Args: - all_rank_num_tokens: Token counts per rank - output_dtype: Output data type - - Returns: - moe_output tensor if NVLinkOneSided is used and backend supports it, None otherwise - """ - if not isinstance(self.comm, NVLinkOneSided): - return None - - if not self.backend.supports_moe_output_in_alltoall_workspace(): - # Ensure payload_in_workspace is False if backend doesn't support it - self.comm.payload_in_workspace = False - return None - - # Determine workspace dtype and whether backend supports workspace output - workspace_dtype = output_dtype - if isinstance(self.backend, TRTLLMGenFusedMoE): - # TRTLLMGen specific configuration - self.comm.invalid_token_expert_id = -1 - workspace_dtype = torch.bfloat16 - - # Calculate runtime max tokens per rank - assert all_rank_num_tokens is not None, ( - "all_rank_num_tokens must be provided for NVLinkOneSided backend" - ) - runtime_max_tokens_per_rank = max(all_rank_num_tokens) - - # Get workspace-backed output tensor - moe_output = self.comm.get_combine_payload_tensor_in_workspace( - runtime_max_tokens_per_rank, self.hidden_size, workspace_dtype - ) - - # Dynamically enable payload_in_workspace for this forward pass - self.comm.payload_in_workspace = True - return moe_output - - def _get_backend_kwargs( - self, - router_logits: Optional[torch.Tensor] = None, - do_finalize: bool = True, - all_rank_num_tokens: Optional[List[int]] = None, - output_dtype: Optional[torch.dtype] = None, - x: Optional[torch.Tensor] = None, - workspace: Optional[dict] = None, - ) -> Dict: - """ - Get backend-specific keyword arguments for run_moe - - Returns backend-specific parameters that are not part of the common run_moe interface. - Different backends need different parameters - this method provides them via kwargs. - - TODO: This is not finalized, will be updated later. - Common kwargs (multiple backends): - - cluster_size, cluster_rank: Cutlass, DeepGemm - - min_latency_mode: Cutlass, WideEP, DeepGemm - - use_fused_finalize: Cutlass, WideEP - - tuner_num_tokens, tuner_top_k: Cutlass, WideEP - - Backend-specific kwargs: - - Cutlass: swizzled_input_sf, enable_alltoall, output_tensor - - WideEP: swizzled_input_sf (fixed False), use_all_to_all - - DeepGemm: workspace, permutation tensors - - TRTLLMGen: router_logits, do_finalize, moe_output - - Args: - router_logits: Router logits tensor (for TRTLLMGen backend) - do_finalize: Whether to finalize output (for TRTLLMGen backend) - all_rank_num_tokens: Token counts per rank (for TRTLLMGen backend moe_output) - output_dtype: Output data type - x: Input tensor (for calculating tuner_num_tokens in Cutlass) - - Returns: - Dict: Backend-specific keyword arguments - """ - kwargs = {} - - # Common parameters for Cutlass and DeepGemm - if self.backend.__class__ in ( - CutlassFusedMoE, - DeepGemmFusedMoE, - CuteDslFusedMoE, - DenseGEMMFusedMoE, - ): - pass - - # Cutlass-specific parameters - if self.backend.__class__ == CutlassFusedMoE: - # Determine if scaling factors are swizzled based on communication flow - # In post-quant communication (quantize -> dispatch), scaling factors are not swizzled - # In pre-quant communication (dispatch -> quantize), scaling factors are swizzled - supports_post_quant = self.comm is not None and self.comm.supports_post_quant_dispatch() - kwargs["is_sf_swizzled"] = not supports_post_quant - kwargs["output_dtype"] = output_dtype - - # 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. - kwargs["enable_alltoall"] = self.enable_alltoall - if self.enable_alltoall: - if all_rank_num_tokens is not None: - kwargs["tuner_num_tokens"] = sum(all_rank_num_tokens) - else: - kwargs["tuner_num_tokens"] = ( - x.shape[0] * self.mapping.tp_size if x is not None else None - ) - kwargs["tuner_top_k"] = self.routing_method.top_k - - # Get moe_output for NVLinkOneSided backend - kwargs["moe_output"] = self._get_nvlink_onesided_moe_output( - all_rank_num_tokens=all_rank_num_tokens, output_dtype=output_dtype - ) - - # CuteDSL-specific parameters - elif self.backend.__class__ == CuteDslFusedMoE: - kwargs["enable_alltoall"] = self.enable_alltoall - - # Get moe_output for NVLinkOneSided backend - kwargs["moe_output"] = self._get_nvlink_onesided_moe_output( - all_rank_num_tokens=all_rank_num_tokens, output_dtype=output_dtype - ) - - if self.enable_dwdp: - kwargs["dwdp_weight_view"] = self.dwdp_manager.build_weight_view( - self.layer_idx, self.backend - ) - - # DeepGemm-specific parameters - elif self.backend.__class__ == DeepGemmFusedMoE: - if workspace is not None: - kwargs["workspace"] = workspace - - # TRTLLMGen-specific parameters - elif self.backend.__class__ == TRTLLMGenFusedMoE: - # Determine router_logits based on whether routing has been done - # If backend doesn't support load balancer, routing is done before communication - # In that case, router_logits should be None (routing already done) - router_logits_arg = None - if not self.backend._supports_load_balancer(): - # For fused routing backends, router_logits is only needed if routing hasn't been done yet - router_logits_arg = router_logits - - kwargs["router_logits"] = router_logits_arg - kwargs["do_finalize"] = do_finalize - - # Get moe_output for NVLinkOneSided backend - kwargs["moe_output"] = self._get_nvlink_onesided_moe_output( - all_rank_num_tokens=all_rank_num_tokens, output_dtype=output_dtype - ) - - return kwargs + backend.validate_configurable_moe(self) def create_weights(self): """ diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 68b1f6a0a8ff..9ba65b35ea8d 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -1,3 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 import os from typing import Dict, Optional, Type @@ -18,6 +20,7 @@ from .fused_moe_vanilla import VanillaMoE from .fused_moe_wide_ep import WideEPMoE from .interface import MoE, MoEWeightLoadingMode +from .mega_moe import MegaMoEDeepGemm from .moe_load_balancer import get_moe_load_balancer from .routing import BaseMoeRoutingMethod @@ -92,7 +95,7 @@ def get_moe_cls( if quant_config is None or not quant_config.quant_mode.has_w4a8_mxfp4_mxfp8( ): logger.warning( - "MegaMoEDeepGemmFusedMoE only supports W4A8_MXFP4_MXFP8. " + "MegaMoEDeepGemm only supports W4A8_MXFP4_MXFP8. " f"Check out details in quant_config: {quant_config}. Using CutlassFusedMoE instead." ) return CutlassFusedMoE @@ -100,29 +103,32 @@ def get_moe_cls( # surface. ``can_implement`` already does this full check; call it # with ``swiglu_gptoss_style=False`` (MegaMoE rejects that anyway, # and the create path doesn't know the model's SwiGLU flavor yet). - from .mega_moe import MegaMoEDeepGemmFusedMoE + # Use the same dtype / intermediate size as create_moe will use when + # instantiating the backend (prefer moe_intermediate_size for MoE). pretrained = model_config.pretrained_config - # Resolve dtype + intermediate_size from pretrained_config so the - # capability check matches the values create_moe will use to - # actually instantiate the backend (mirrors create_moe's logic - # below: prefer moe_intermediate_size for MoE models). - pretrained_dtype = getattr(pretrained, "torch_dtype", torch.bfloat16) - pretrained_inter = getattr(pretrained, "moe_intermediate_size", None) - if pretrained_inter is None: - pretrained_inter = getattr(pretrained, "intermediate_size", None) - ok, reason = MegaMoEDeepGemmFusedMoE.can_implement( + pretrained_dtype = (getattr(pretrained, "torch_dtype", torch.bfloat16) + if pretrained is not None else torch.bfloat16) + pretrained_inter = None + if pretrained is not None: + pretrained_inter = getattr(pretrained, "moe_intermediate_size", + None) + if pretrained_inter is None: + pretrained_inter = getattr(pretrained, "intermediate_size", + None) + ok, reason = MegaMoEDeepGemm.can_implement( QuantAlgo.W4A8_MXFP4_MXFP8, dtype_activation=pretrained_dtype, swiglu_gptoss_style=False, - hidden_size=getattr(pretrained, "hidden_size", None), + hidden_size=getattr(pretrained, "hidden_size", None) + if pretrained is not None else None, intermediate_size=pretrained_inter, ) if not ok: logger.warning( - f"MegaMoEDeepGemmFusedMoE rejected current environment: {reason}. " + f"MegaMoEDeepGemm rejected current environment: {reason}. " "Falling back to CutlassFusedMoE.") return CutlassFusedMoE - return MegaMoEDeepGemmFusedMoE + return MegaMoEDeepGemm else: raise ValueError(f"Unsupported moe backend: {moe_backend}") @@ -196,10 +202,19 @@ def create_moe_backend( moe_load_balancer = get_moe_load_balancer() if moe_load_balancer is not None: - assert moe_cls in [ - WideEPMoE, CutlassFusedMoE, TRTLLMGenFusedMoE, CuteDslFusedMoE, - DeepGemmFusedMoE, DenseGEMMFusedMoE - ], "MoE Load Balance is only supported in WideEPMoE, CutlassFusedMoE, TRTLLMGenFusedMoE, CuteDslFusedMoE, DeepGemmFusedMoE, and DenseGEMMFusedMoE." + supported_load_balancer_backends = ( + WideEPMoE, + CutlassFusedMoE, + TRTLLMGenFusedMoE, + CuteDslFusedMoE, + DeepGemmFusedMoE, + DenseGEMMFusedMoE, + MegaMoEDeepGemm, + ) + assert moe_cls in supported_load_balancer_backends, ( + "MoE Load Balance is only supported in " + f"{', '.join(cls.__name__ for cls in supported_load_balancer_backends)}." + ) if bias: assert moe_cls in [CutlassFusedMoE, TritonFusedMoE, TRTLLMGenFusedMoE @@ -357,10 +372,10 @@ def create_moe_backend( ) else: # Mega MoE fall-through: new backend not in the hard-coded chain. - # Import lazily to avoid pulling DG at module import time on boxes + # ``mega_moe_deepgemm`` lazily resolves DG via ``_import_deep_gemm`` + # at runtime, so a top-level import here doesn't pull DG on boxes # that don't use this backend. - from .mega_moe import MegaMoEDeepGemmFusedMoE - if moe_cls is MegaMoEDeepGemmFusedMoE: + if moe_cls is MegaMoEDeepGemm: return moe_cls( routing_method=routing_method, num_experts=num_experts, @@ -447,17 +462,9 @@ def create_moe( enable_configurable_moe = os.environ.get("ENABLE_CONFIGURABLE_MOE", "1") == "1" - # Build the ConfigurableMoE-supported set lazily so non-MegaMoE - # callers don't have to import DeepGEMM at module load time. - configurable_supported = (DeepGemmFusedMoE, TRTLLMGenFusedMoE, - CuteDslFusedMoE, CutlassFusedMoE, - DenseGEMMFusedMoE) - if model_config.moe_backend.upper() == "MEGAMOE_DEEPGEMM": - from .mega_moe import MegaMoEDeepGemmFusedMoE - configurable_supported = configurable_supported + ( - MegaMoEDeepGemmFusedMoE, ) if enable_configurable_moe or moe_cls == CuteDslFusedMoE: - if moe_cls in configurable_supported: + if moe_cls in (DeepGemmFusedMoE, TRTLLMGenFusedMoE, CuteDslFusedMoE, + CutlassFusedMoE, DenseGEMMFusedMoE, MegaMoEDeepGemm): return ConfigurableMoE( routing_method=routing_method, num_experts=num_experts, diff --git a/tensorrt_llm/_torch/modules/fused_moe/interface.py b/tensorrt_llm/_torch/modules/fused_moe/interface.py index 9dd902b83977..f5e3701a0b17 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/interface.py +++ b/tensorrt_llm/_torch/modules/fused_moe/interface.py @@ -79,6 +79,30 @@ class AlltoallMethodType(IntEnum): DeepEPLowLatency = 4 +class MoESchedulerKind(Enum): + """Selects which forward-execution scheduler ConfigurableMoE picks for a backend. + + Backends declare this via the ``scheduler_kind`` class attribute on + ``MoE``. ``ConfigurableMoE`` reads it once at init time to construct the + matching scheduler and to gate communication-strategy creation. + + The axis is whether the cross-rank EP exchange is fused into the MoE + kernel or is a separate host-orchestrated step: + + - ``EXTERNAL_COMM``: comm lives outside the MoE kernel boundary; the + scheduler issues ``Communication.dispatch`` / ``Communication.combine`` + from the host with per-chunk EPLB hooks and optional multi-stream + chunk overlap (Cutlass, DeepGemm, CuteDSL, DenseGEMM, TRTLLMGen). + - ``FUSED_COMM``: comm is fused into the backend's fused kernel via + NVLink SymmBuffer (DeepGEMM ``fp8_fp4_mega_moe``). No host comm; + lockstep chunk launches; EPLB statistic update with + ``ignore_allreduce=False``. + """ + + EXTERNAL_COMM = "external_comm" + FUSED_COMM = "fused_comm" + + def extract_extra_attrs(layer_idx: str): extra_attrs = get_model_extra_attrs() assert extra_attrs is not None, "Model extra attrs are not set" @@ -170,6 +194,11 @@ class MoE(nn.Module): aux_stream_dict (Optional[Dict[AuxStreamType, torch.cuda.Stream]]): Auxiliary CUDA streams for overlapping. """ + # Default scheduler kind for ConfigurableMoE forward dispatch. Backends + # whose fused kernel owns cross-rank exchange (e.g. MegaMoE-style) + # override this to ``MoESchedulerKind.FUSED_COMM``. + scheduler_kind: MoESchedulerKind = MoESchedulerKind.EXTERNAL_COMM + @classmethod @abstractmethod def can_implement( @@ -486,6 +515,18 @@ def _supports_load_balancer(self) -> bool: """ return False + def validate_configurable_moe(self, moe: "nn.Module") -> None: + """Backend-specific validation hook called by ``ConfigurableMoE``. + + ``ConfigurableMoE.validate_backend`` invokes this AFTER the generic + EPLB/load-balancer compatibility check, so backends may inspect + ``moe.num_slots``, ``moe.ep_size``, ``moe._using_load_balancer()``, + ``moe._using_dynamic_load_balancer()``. Default is a no-op; backends + with extra constraints (e.g. fused-comm backends rejecting + dynamic EPLB) override this. + """ + del moe + def _using_load_balancer(self) -> bool: """Check if this MoE is using load balancer.""" return self.layer_load_balancer is not None diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/__init__.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/__init__.py index 33aabdf1eaac..3a3cbdc3d26a 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/__init__.py +++ b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/__init__.py @@ -15,11 +15,11 @@ """MegaMoE — DeepGEMM ``fp8_fp4_mega_moe`` as a first-class MoE backend. Targets the W4A8_MXFP4_MXFP8 quant configuration already supported by -``TRTLLMGenFusedMoE``. Shares the ``{expert_id}.w*.weight`` / -``{expert_id}.w*.weight_scale`` loader keys so that identical MXFP4 bytes -fed to both backends produce numerically-aligned outputs. +``TRTLLMGenFusedMoE``. ``W4A8MXFP4MXFP8MegaMoEDeepGemmMethod`` owns the +DG-native weight tensors, scale conversion, and DeepGEMM weight transform. """ -from .backend import MegaMoEDeepGemmFusedMoE +from ..quantization import W4A8MXFP4MXFP8MegaMoEDeepGemmMethod +from .mega_moe_deepgemm import MegaMoEDeepGemm -__all__ = ["MegaMoEDeepGemmFusedMoE"] +__all__ = ["MegaMoEDeepGemm", "W4A8MXFP4MXFP8MegaMoEDeepGemmMethod"] diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/backend.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/backend.py deleted file mode 100644 index 8bb2763dcc92..000000000000 --- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/backend.py +++ /dev/null @@ -1,809 +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. -"""MegaMoEDeepGemmFusedMoE — DeepGEMM fp8_fp4_mega_moe as a first-class MoE backend. - -Quantization scheme -------------------- -W4A8_MXFP4_MXFP8: packed MXFP4 weights (4 bits, UE8M0 block scale per 32 -input elements) × FP8 E4M3 activations quantized per-token with UE8M0 -block scales. Matches ``TRTLLMGenFusedMoE``'s W4A8_MXFP4_MXFP8 math; the -difference is the kernel implementation — DG fuses EP dispatch + GEMM1 + -SwiGLU + GEMM2 + EP combine into a single launch. - -Weight storage (DG-native, EP-only Phase 1) ------------------------------------------- -* ``w3_w1_weight`` : uint8 [E_local, 2*I, H // 2] (MXFP4 nibbles) -* ``w3_w1_weight_scale`` : uint8 [E_local, 2*I, H // 32] (UE8M0) -* ``w2_weight`` : uint8 [E_local, H, I // 2] -* ``w2_weight_scale`` : uint8 [E_local, H, I // 32] - -The SAME raw bytes that ``torch.ops.trtllm.fp4_quantize`` emits, so the -``VANILLA`` / ``FUSED_GATE_UP_PROJ`` loader key schemas used by TRT-LLM -models work directly. On ``post_load_weights`` the uint8 scales are -lifted to DG's fp32-UE8M0 format and both weight and scale tensors are -passed through ``deep_gemm.transform_weights_for_mega_moe``. - -Forward (hot path) ------------------- -Per-forward on ``num_tokens_real`` unpadded tokens per rank: - 1. Routing (external via ``routing_method.apply``). - 2. Pre-quant hidden states BF16 → FP8 E4M3 with packed UE8M0 block - scales (``deep_gemm.utils.per_token_cast_to_fp8``). - 3. Copy FP8 hidden / SF / topk_idx / topk_weights into the shared DG - ``SymmBuffer``. - 4. Fire ``deep_gemm.fp8_fp4_mega_moe(y, t_l1, t_l2, buf, ...)``. - 5. Return ``y`` (bf16 [num_tokens_real, H], combined across EP). - -The ``num_tokens_real`` count comes from ``all_rank_num_tokens`` when -provided (attention-DP padding support), matching the shape contract -of ``MoE.forward_fake`` and ``ConfigurableMoE._forward_chunk_impl``. -""" - -from __future__ import annotations - -from typing import Dict, List, Optional, Tuple, Union - -import torch -import torch.distributed as dist -from torch import nn - -from tensorrt_llm._utils import get_sm_version -from tensorrt_llm.logger import logger -from tensorrt_llm.models.modeling_utils import QuantAlgo - -from ....model_config import ModelConfig -from ....utils import ActivationType, AuxStreamType, Fp4QuantizedTensor -from ..interface import MoE, MoEWeightLoadingMode -from ..routing import BaseMoeRoutingMethod - -__all__ = ["MegaMoEDeepGemmFusedMoE"] - - -# Process-level cache: one DG SymmBuffer shared across all MegaMoE layers. -# Same rationale as the legacy MnnvlMoe workspace singleton — MoE layers -# run serially per forward, so overlapping workspaces are unnecessary and -# per-layer allocation would need ~``num_layers × buffer_size`` symm -# memory (OOM on DSV3-scale). Keyed on (pg_id, config) so different EP -# subsets or shapes get separate buffers. -_SYMM_BUFFER_CACHE: dict = {} - - -def _import_deep_gemm(): - """Return the bundled ``tensorrt_llm.deep_gemm`` module. - - Strict: no fall-back to a standalone ``deep_gemm`` wheel. TRT-LLM - guarantees the bundled binding is the compatible-by-construction - source; silently switching to whatever external ``deep_gemm`` is - installed on the box would make capability depend on pip state and - break the guarantee. - - Raises ``_MegaMoEUnavailable`` when the bundled module is missing or - too old (mega_moe symbols absent, or ``per_token_cast_to_fp8`` lacks - ``use_packed_ue8m0``). ``can_implement`` catches this and returns - ``(False, reason)`` cleanly. - """ - import inspect - - try: - from tensorrt_llm import deep_gemm as _dg - except ImportError as e: - raise _MegaMoEUnavailable(f"tensorrt_llm.deep_gemm not importable: {e}") from e - - missing = [ - n - for n in ( - "fp8_fp4_mega_moe", - "get_symm_buffer_for_mega_moe", - "transform_weights_for_mega_moe", - ) - if not hasattr(_dg, n) - ] - if missing: - raise _MegaMoEUnavailable( - f"tensorrt_llm.deep_gemm missing mega_moe symbols {missing}; " - f"upgrade the TRT-LLM bundled DeepGEMM to a release that " - f"includes fp8_fp4_mega_moe." - ) - - p_fp8 = getattr(_dg, "per_token_cast_to_fp8", None) - if p_fp8 is None or "use_packed_ue8m0" not in inspect.signature(p_fp8).parameters: - raise _MegaMoEUnavailable( - "tensorrt_llm.deep_gemm.per_token_cast_to_fp8 does not accept " - "use_packed_ue8m0=; upgrade the bundled DeepGEMM." - ) - return _dg - - -def _import_dg_fp8_cast(): - """Return ``per_token_cast_to_fp8`` from the same bundled module the kernel lives in. - - Only this cast is used on the hot path; we do NOT require - ``per_token_cast_to_fp4`` since the backend consumes MXFP4 weights - that are already pre-quantized by the caller (see ``load_weights``) - and the transform runs on those raw bytes. - """ - dg = _import_deep_gemm() - return dg.per_token_cast_to_fp8 - - -class _MegaMoEUnavailable(RuntimeError): - """Signals that the bundled DeepGEMM doesn't expose the full mega_moe API. - - ``can_implement`` converts this into a clean ``(False, reason)`` - instead of a hard import error. - """ - - -def _ue8m0_uint8_to_fp32(sf_uint8: torch.Tensor) -> torch.Tensor: - """Convert UE8M0 stored as uint8 → fp32 with matching numeric value. - - UE8M0 is an 8-bit unsigned exponent (no sign, no mantissa). The - numerically-equivalent fp32 has the same 8 exponent bits with sign=0 - and mantissa=0 — achieved by shifting 23 bits left. - """ - assert sf_uint8.dtype == torch.uint8 - return (sf_uint8.to(torch.int32) << 23).contiguous().view(torch.float32) - - -# ---- Fused MXFP8 per-token quant backends -------------------------------- -# We want: BF16 (m, H) → FP8 E4M3 (m, H) + packed-UE8M0 SF (m, H/32/4) int32. -# Three candidates, in preference order: -# -# 1. ``torch.ops.trtllm.mxfp8_quantize(x, False, alignment=32)`` — TRT-LLM -# C++ CUDA kernel. Roundtrip-verified byte-identical to DG's Python -# helper (fp8 bytes + SF after u8→int32 reshape). Fastest by 5-25× -# vs torch.compile, one kernel launch (~11 us regardless of seq). -# Requires ``libth_common.so`` to be loaded; ``ConfigurableMoE`` pulls -# this in on construction so it's always registered by the time -# ``backend.quantize_input`` runs. -# -# 2. ``torch.compile(dg.per_token_cast_to_fp8, dynamic=True)`` — fallback -# when the TRT-LLM op isn't registered (e.g. slim builds, standalone -# DG tests). Inductor fuses the ~8 elementwise kernels into 1-2 -# Triton kernels but still pays one launch per seq boundary. -# -# ``_FUSED_PER_TOKEN_CAST`` caches the fallback so we don't re-compile on -# every module creation. -_FUSED_PER_TOKEN_CAST = None - - -def _trtllm_mxfp8_quantize_available() -> bool: - return hasattr(torch.ops, "trtllm") and hasattr(torch.ops.trtllm, "mxfp8_quantize") - - -def _quantize_bf16_to_fp8_ue8m0(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """Return (x_fp8, x_sf) in DG mega_moe's expected layout (packed int32).""" - m, n = x.shape - # Packed-UE8M0 stores 4 u8 scales per int32 over a 32-element block, - # so n must be a multiple of 128 for the int32 view below to land on - # an integer last-dim. Misaligned shapes would otherwise fail with a - # cryptic reshape/view error; surface a clear contract here instead. - if n % 128 != 0: - raise ValueError( - f"_quantize_bf16_to_fp8_ue8m0 requires hidden_size % 128 == 0 " - f"(packed-UE8M0 int32 SF stride); got hidden_size={n}" - ) - if _trtllm_mxfp8_quantize_available(): - # ``is_sf_swizzled_layout=False`` → flat row-major uint8 SF, one - # byte per 32-element group. ``alignment=32`` → MXFP8 block size. - x_fp8, x_sf_u8 = torch.ops.trtllm.mxfp8_quantize(x, False, alignment=32) - # DG wants (m, n/32/4) int32 with 4 u8 UE8M0 packed per int32. - # TRT-LLM emits (m*n/32,) uint8 in the same byte order, so a - # reshape + view is a zero-copy reinterpret. - return x_fp8, x_sf_u8.view(m, n // 32).view(torch.int32) - - global _FUSED_PER_TOKEN_CAST - if _FUSED_PER_TOKEN_CAST is None: - dg = _import_deep_gemm() - base = dg.per_token_cast_to_fp8 - - def _call(t: torch.Tensor): - return base(t, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) - - _FUSED_PER_TOKEN_CAST = torch.compile(_call, dynamic=True, fullgraph=False) - return _FUSED_PER_TOKEN_CAST(x) - - -class MegaMoEDeepGemmFusedMoE(MoE): - """MoE backend wrapping DeepGEMM's fused ``fp8_fp4_mega_moe`` kernel.""" - - _SUPPORTED_ACTIVATION_DTYPES = frozenset({torch.bfloat16}) - - # ------------------------------------------------------------------ - # Capability gating - # ------------------------------------------------------------------ - @classmethod - def can_implement( - cls, - quant_algo: Optional[QuantAlgo], - dtype_activation: torch.dtype = torch.bfloat16, - swiglu_gptoss_style: bool = False, - hidden_size: Optional[int] = None, - intermediate_size: Optional[int] = None, - ) -> Tuple[bool, Optional[str]]: - sm = get_sm_version() - if sm != 100: - return False, ( - f"MegaMoEDeepGemmFusedMoE requires SM100 (only arch with " - f"sm100_fp8_fp4_mega_moe.cuh in DeepGEMM); got SM{sm}" - ) - if dtype_activation not in cls._SUPPORTED_ACTIVATION_DTYPES: - return False, ( - f"MegaMoEDeepGemmFusedMoE supports activations in " - f"{cls._SUPPORTED_ACTIVATION_DTYPES}, got {dtype_activation}" - ) - if swiglu_gptoss_style: - return False, "MegaMoEDeepGemmFusedMoE does not support swiglu_gptoss_style" - if quant_algo != QuantAlgo.W4A8_MXFP4_MXFP8: - return False, ( - f"MegaMoEDeepGemmFusedMoE supports W4A8_MXFP4_MXFP8 only, got {quant_algo}" - ) - # Packed-UE8M0 per-token SF layout: 4 u8 scales reinterpreted as - # int32 per 128-element row stride, so hidden/intermediate must be - # divisible by 128. Divisible-by-32 shapes like ``hidden=2880`` - # would quantize cleanly but fail the int32 reshape at first - # forward — reject at ``can_implement`` time so the factory can - # fall back cleanly. - if hidden_size is not None and hidden_size % 128 != 0: - return False, ( - f"MegaMoEDeepGemmFusedMoE requires hidden_size % 128 == 0 " - f"(packed-UE8M0 int32 SF stride); got hidden_size={hidden_size}" - ) - if intermediate_size is not None and intermediate_size % 128 != 0: - return False, ( - f"MegaMoEDeepGemmFusedMoE requires intermediate_size % 128 == 0 " - f"(packed-UE8M0 int32 SF stride); got intermediate_size=" - f"{intermediate_size}" - ) - try: - _import_deep_gemm() - except _MegaMoEUnavailable as e: - return False, str(e) - return True, None - - # ------------------------------------------------------------------ - # Init - # ------------------------------------------------------------------ - 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, - activation_type: ActivationType = ActivationType.Swiglu, - init_load_balancer: bool = True, - without_comm: bool = False, - # DG tunables. - activation: str = "swiglu", - activation_clamp: Optional[float] = None, - fast_math: bool = True, - **kwargs, - ) -> None: - 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, - layer_idx=layer_idx, - activation_type=activation_type, - init_load_balancer=init_load_balancer, - ) - - # Phase 1 — assert supported topologies early. - assert self.tp_size == 1, ( - f"MegaMoEDeepGemmFusedMoE Phase 1 is EP-only (moe_tp_size=1); got tp_size={self.tp_size}" - ) - assert self.cluster_size == 1, ( - f"MegaMoEDeepGemmFusedMoE Phase 1 assumes cluster_size=1; got cluster_size={self.cluster_size}" - ) - assert num_experts % max(self.ep_size, 1) == 0 - - # ADP semantics: DG's fp8_fp4_mega_moe subsumes cross-rank token - # dispatch into its internal symm_mem exchange. When EP spans - # *all* ranks that may carry tokens (i.e. ``ep_size == - # parallel_size``), no outer allgather / reducescatter is needed: - # every token's origin rank is inside the EP group and DG returns - # results to that origin. If EP is a strict subset of - # parallel_size (e.g. attention-DP > moe_ep_size), some tokens - # live on ranks that the DG kernel cannot reach — that topology - # is not yet supported. - if self.use_dp and self.parallel_size > 1: - assert self.ep_size == self.parallel_size, ( - f"MegaMoEDeepGemmFusedMoE with enable_attention_dp=True requires " - f"ep_size == parallel_size (got ep_size={self.ep_size}, " - f"parallel_size={self.parallel_size}). Configurations " - f"with ADP > EP are not yet supported; add the standard " - f"allgather(pre) + reducescatter(post) wrapper before " - f"calling fp8_fp4_mega_moe to support them." - ) - - # apply_router_weight_on_input pre-multiplies routing weights - # onto x before the MoE compute (used by some top-1 models). DG's - # fused kernel applies the weights on the MoE output instead; - # mixing the two produces wrong math. Reject loudly — a silent - # fallback would break llama-min-latency-style paths that set - # this flag to True and assume top-1 semantics. - assert not apply_router_weight_on_input, ( - "MegaMoEDeepGemmFusedMoE does not support apply_router_weight_on_input. " - "DG's fp8_fp4_mega_moe applies routing weights on the MoE " - "output, not by pre-scaling the input — the two paths are " - "not equivalent. Use a different MoE backend for models that " - "require pre-scaling, or extend the kernel call." - ) - self.apply_router_weight_on_input = apply_router_weight_on_input - self.activation = activation - self.activation_clamp = activation_clamp - self.fast_math = fast_math - - # Buffer sizing. MoE layers execute serially per forward; a single - # process-level pool sized to worst-case per-rank tokens serves all. - self.max_num_tokens = int( - getattr(model_config, "moe_max_num_tokens", 0) - or getattr(model_config, "max_num_tokens", 0) - or 4096 - ) - - # Resolve the EP ProcessGroup at module construction — creating a - # group at forward time would be collective on a non-synchronous - # call stack and deadlock under PP / layer-skip. Construction is - # globally synchronous across ranks during model build. - self._ep_pg = self._resolve_ep_pg() - - # Deferred: weight transform + SymmBuffer allocation happen on - # ``post_load_weights`` which is also a global sync point. - self._symm_buffer = None - self._t_l1 = None - self._t_l2 = None - self._weights_loaded = False - - self._create_mega_weights() - - def _supports_load_balancer(self) -> bool: - # Phase 1: EPLB off. Follow-up via token_selected_slots. - return False - - # ------------------------------------------------------------------ - # EP process-group resolution (no collective at forward time) - # ------------------------------------------------------------------ - def _resolve_ep_pg(self): - """Return the torch.distributed ProcessGroup for the EP sub-world. - - Prefers ``mapping.moe_ep_group_pg`` (DeviceMeshTopology, Ray path) - because it was built once at Mapping init. Falls back to - ``dist.group.WORLD`` only when ``ep_size == world_size`` (single - EP subset covers all ranks). - - Does NOT call ``dist.new_group`` — that's collective and unsafe to - invoke from any path that may skip ranks (e.g. PP-isolated layer - forwards). When the mapping cannot provide a PG and EP is a - strict subset of world, we raise with a clear message pointing - at ``mpi_disabled=1`` / Ray as the supported path. - """ - if not dist.is_initialized(): - raise RuntimeError( - "MegaMoEDeepGemmFusedMoE requires torch.distributed to be " - "initialized before module construction (mpirun or Ray)." - ) - # Preferred: reuse the existing PG from the mapping (Ray / DeviceMesh). - try: - pg = self.mapping.moe_ep_group_pg - logger.info( - f"[MegaMoE] layer={self.layer_idx} using mapping.moe_ep_group_pg (DeviceMesh path)" - ) - return pg - except (NotImplementedError, AttributeError): - pass - # Fallback: degenerate to WORLD when EP spans all ranks. - world_size = dist.get_world_size() - if self.ep_size == world_size: - logger.info( - f"[MegaMoE] layer={self.layer_idx} using dist.group.WORLD " - f"(EP == world_size == {world_size})" - ) - return dist.group.WORLD - raise RuntimeError( - f"MegaMoEDeepGemmFusedMoE: cannot resolve EP ProcessGroup. The current " - f"mapping does not expose ``moe_ep_group_pg`` and EP " - f"({self.ep_size}) is a strict subset of world " - f"({world_size}). Use DeviceMeshTopology (TLLM_DISABLE_MPI=1) " - f"so the EP PG is constructed once at Mapping init, or set " - f"ep_size == world_size." - ) - - # ------------------------------------------------------------------ - # Weight lifecycle - # ------------------------------------------------------------------ - def _create_mega_weights(self) -> None: - E = self.expert_size_per_partition - H = self.hidden_size - inter = self.intermediate_size - # Divisible-by-128 (not 32) — packed-UE8M0 SF in - # ``_quantize_bf16_to_fp8_ue8m0`` reinterprets ``H/32`` bytes as - # ``H/128`` int32 values per row, so H/32 must be a multiple of 4. - # ``can_implement`` rejects before we get here under the factory - # path, but keep the asserts as defensive dev checks for direct - # MegaMoEDeepGemmFusedMoE construction. - assert H % 128 == 0, f"hidden {H} must be divisible by 128" - assert inter % 128 == 0, f"intermediate {inter} must be divisible by 128" - - self.register_parameter( - "w3_w1_weight", - nn.Parameter(torch.empty(E, inter * 2, H // 2, dtype=torch.uint8), requires_grad=False), - ) - self.register_parameter( - "w3_w1_weight_scale", - nn.Parameter( - torch.empty(E, inter * 2, H // 32, dtype=torch.uint8), requires_grad=False - ), - ) - self.register_parameter( - "w2_weight", - nn.Parameter(torch.empty(E, H, inter // 2, dtype=torch.uint8), requires_grad=False), - ) - self.register_parameter( - "w2_weight_scale", - nn.Parameter(torch.empty(E, H, inter // 32, dtype=torch.uint8), requires_grad=False), - ) - - def create_weights(self): - # No-op: allocated in __init__. Provided for MoE-contract symmetry. - return - - # ----- Per-loading-mode weight unpacking helpers ------------------- - def _iter_vanilla_expert_weights(self, w: Dict, expert_id: int): - """Return (w1, w3, w2, w1_sf, w3_sf, w2_sf) as CPU uint8 tensors. - - Used for VANILLA / W4A8_CUSTOM key schema - ``{eid}.w*.weight[_scale]``. - """ - return ( - w[f"{expert_id}.w1.weight"], - w[f"{expert_id}.w3.weight"], - w[f"{expert_id}.w2.weight"], - w[f"{expert_id}.w1.weight_scale"], - w[f"{expert_id}.w3.weight_scale"], - w[f"{expert_id}.w2.weight_scale"], - ) - - def _iter_fused_gate_up_expert_weights(self, w: Dict, expert_id: int): - """FUSED_GATE_UP_PROJ schema (gpt-oss / llama). - - Mirrors ``MoEWeightLoader.load_expert_weights`` in quantization.py: - ``gate_up_proj[expert_id]`` is transposed then chunked along dim 0 - into (w1, w3). ``down_proj[expert_id]`` is transposed to get w2. - """ - w1w3 = w["gate_up_proj"][expert_id].transpose(0, 1).contiguous() - w1, w3 = w1w3.chunk(2, dim=0) - w2 = w["down_proj"][expert_id].transpose(0, 1).contiguous() - - w1w3_sf = w["gate_up_proj_weight_scale"][expert_id].transpose(0, 1).contiguous() - w1_sf, w3_sf = w1w3_sf.chunk(2, dim=0) - w2_sf = w["down_proj_weight_scale"][expert_id].transpose(0, 1).contiguous() - return w1, w3, w2, w1_sf, w3_sf, w2_sf - - def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False) -> None: - """Load MXFP4 weights + UE8M0 block scales for this rank's experts. - - Supports VANILLA / W4A8_CUSTOM (per-expert ``{eid}.w*.*`` keys) - and FUSED_GATE_UP_PROJ (stacked ``gate_up_proj`` / ``down_proj``) - loading modes, matching ``MoEWeightLoader`` conventions. - """ - assert len(weights) == 1, f"MegaMoE expects one weight dict, got {len(weights)}" - w = weights[0] - - mode = self.weight_loading_mode - if mode in (MoEWeightLoadingMode.VANILLA, MoEWeightLoadingMode.W4A8_CUSTOM): - get_expert = self._iter_vanilla_expert_weights - elif mode == MoEWeightLoadingMode.FUSED_GATE_UP_PROJ: - get_expert = self._iter_fused_gate_up_expert_weights - else: - raise NotImplementedError( - f"MegaMoE load_weights unsupported weight_loading_mode={mode}" - ) - - def _to_u8(t: torch.Tensor) -> torch.Tensor: - return t.cuda().view(torch.uint8) - - local_ids = list(self.initial_local_expert_ids) - for slot_id, expert_id in enumerate(local_ids): - w1, w3, w2, w1_sf, w3_sf, w2_sf = get_expert(w, expert_id) - - # Stack [w1 | w3] along intermediate dim. DG's - # ``_interleave_l1_weights`` (deep_gemm/mega/__init__.py:78) - # interprets the first half of the L1 weight as gate and the - # second half as up. TRT-LLM's MoE convention (consistent - # across ``modeling_gpt_oss.py`` and the ``FUSED_GATE_UP_PROJ`` - # loader at ``quantization.py:362-365``) maps - # ``w1 = gate_proj``, ``w3 = up_proj`` (HF ``gate_up_proj`` - # is ``[gate | up]`` along out_dim, ``chunk(2)[0]`` -> w1, - # ``chunk(2)[1]`` -> w3). Therefore the right cat order is - # ``[w1 | w3]`` so DG sees ``[gate | up]`` and computes - # ``silu(gate) * up`` correctly. The earlier ``[w3, w1]`` - # order silently swapped which side the silu was applied to, - # producing ``silu(up) * gate`` and ~94% mismatch vs reference. - self.w3_w1_weight.data[slot_id].copy_( - torch.cat([_to_u8(w1), _to_u8(w3)], dim=0), non_blocking=True - ) - self.w3_w1_weight_scale.data[slot_id].copy_( - torch.cat([_to_u8(w1_sf), _to_u8(w3_sf)], dim=0), non_blocking=True - ) - self.w2_weight.data[slot_id].copy_(_to_u8(w2), non_blocking=True) - self.w2_weight_scale.data[slot_id].copy_(_to_u8(w2_sf), non_blocking=True) - - self._weights_loaded = True - - def post_load_weights(self) -> None: - """Finalize loaded weights for the MegaMoE hot path. - - Allocates the DG SymmBuffer (collective rendezvous) and runs - ``transform_weights_for_mega_moe`` on this rank's weights. - - Both operations happen here because (a) rendezvous is collective - and must run at a globally-synchronous time (post-load is), and - (b) the transform needs the loaded weight bytes. Both are - idempotent via internal guards. - - Phase 1 reload limitations (tracked as follow-up, consciously - skipped): - * ``allow_partial_loading`` is ignored in ``load_weights`` — - a partial load will overwrite an incomplete weight subset. - * This method returns early once ``_t_l1`` is set, so repeated - ``post_load_weights`` calls after a weight reload will keep - the **stale** transformed tensors. Invalidate ``_t_l1`` / - ``_t_l2`` manually if you reload weights at runtime. - * ``_SYMM_BUFFER_CACHE`` is a process-global cache keyed on - (pg_id, shape, experts-per-token) and is never evicted; - long-running processes that construct many MegaMoE layers - of the same shape share one buffer. OOM at buffer allocation - is the only signal that invalidation is needed. - EPLB migration (Phase 2) will require addressing all three. - """ - assert self._weights_loaded, "post_load_weights before load_weights" - dg = _import_deep_gemm() - - if self._symm_buffer is None: - key = ( - id(self._ep_pg), - self.num_experts, - self.max_num_tokens, - self.routing_method.experts_per_token, - self.hidden_size, - self.intermediate_size, - self.activation, - ) - cached = _SYMM_BUFFER_CACHE.get(key) - if cached is None: - cached = dg.get_symm_buffer_for_mega_moe( - self._ep_pg, - self.num_experts, - self.max_num_tokens, - self.routing_method.experts_per_token, - self.hidden_size, - self.intermediate_size, - True, # use_fp8_dispatch - self.activation, - ) - _SYMM_BUFFER_CACHE[key] = cached - logger.info( - f"[MegaMoE] layer={self.layer_idx} allocated DG " - f"SymmBuffer: {cached.buffer.nbytes / 2**30:.2f} GiB" - ) - self._symm_buffer = cached - - if self._t_l1 is not None: - return - - E = self.expert_size_per_partition - H = self.hidden_size - inter = self.intermediate_size - - l1_sf_fp32 = _ue8m0_uint8_to_fp32(self.w3_w1_weight_scale) - # Bundled ``tensorrt_llm.deep_gemm.transform_sf_into_required_layout`` - # expects a 3-tuple ``(gm_sfa, gm_sfb, gk)`` recipe; for a 3-tuple - # recipe the C++ side also requires ``is_sfa`` to be set so it - # picks between the SFA / SFB granularity. These weight scales are - # the SFB (B-operand) side of the grouped GEMM so pass - # ``is_sfa=False``. ``num_groups=E`` marks this as a per-expert - # grouped SF tensor. - l1_sf = dg.transform_sf_into_required_layout( - l1_sf_fp32, mn=inter * 2, k=H, recipe=(1, 1, 32), num_groups=E, is_sfa=False - ) - - l2_sf_fp32 = _ue8m0_uint8_to_fp32(self.w2_weight_scale) - l2_sf = dg.transform_sf_into_required_layout( - l2_sf_fp32, mn=H, k=inter, recipe=(1, 1, 32), num_groups=E, is_sfa=False - ) - - l1_w = self.w3_w1_weight.view(torch.int8) - l2_w = self.w2_weight.view(torch.int8) - - self._t_l1, self._t_l2 = dg.transform_weights_for_mega_moe((l1_w, l1_sf), (l2_w, l2_sf)) - logger.info( - f"[MegaMoE] layer={self.layer_idx} weight transform done " - f"t_l1=(w {tuple(self._t_l1[0].shape)}/{self._t_l1[0].dtype}, " - f"sf {tuple(self._t_l1[1].shape)}/{self._t_l1[1].dtype})" - ) - - # ------------------------------------------------------------------ - # Abstract MoE-contract methods (not used in this backend) - # ------------------------------------------------------------------ - def quantize_input(self, x, *, post_quant_comm: bool = False, **kwargs): - """BF16 → FP8-E4M3 + packed-UE8M0 per-token SF (gran_k=32). - - Delegates to ``_quantize_bf16_to_fp8_ue8m0`` which picks the - fastest available backend (TRT-LLM C++ op ~11 us at any seq, - or ``torch.compile`` fallback ~60-260 us). Byte-identical - output across all paths so DG's ``fp8_fp4_mega_moe`` consumes - it unchanged. - """ - del post_quant_comm # MegaMoE runs pre-quant comm via DG SymmBuffer - x_bf16 = x.to(torch.bfloat16).contiguous() - return _quantize_bf16_to_fp8_ue8m0(x_bf16) - - def run_moe(self, *args, **kwargs): - raise NotImplementedError( - "MegaMoE's fused kernel replaces run_moe — call ``forward_impl`` " - "or ``run_with_prequant`` with pre-computed FP8+SF+topk." - ) - - # ------------------------------------------------------------------ - # Fast path: accept already-quantized inputs from the outer pipeline. - # ------------------------------------------------------------------ - def run_with_prequant( - self, - x_fp8: torch.Tensor, - x_sf: torch.Tensor, - topk_idx: torch.Tensor, - topk_weights: torch.Tensor, - num_tokens: int, - output_dtype: torch.dtype, - ) -> torch.Tensor: - """Kernel-only path: 4 x ``buf.copy_()`` + empty-alloc + kernel launch. - - Matches DG's own ``run_fused`` shape contract so the GPU work - here is exactly what DG benchmarks report. - - Caller is responsible for: - * slicing ``x_real`` / ``router_logits_real`` to ``num_tokens`` - * running the routing method to produce ``topk_idx`` (int64) and - ``topk_weights`` (float32) - * BF16 → FP8 per-token quant (``quantize_input`` above) - * sizing ``num_tokens`` appropriately vs ``max_num_tokens`` - """ - dg = _import_deep_gemm() - buf = self._symm_buffer - assert buf is not None, "MegaMoE SymmBuffer not allocated — post_load_weights missing?" - assert num_tokens <= self.max_num_tokens, ( - f"MegaMoE got {num_tokens} tokens but buffer is sized for " - f"{self.max_num_tokens}. Raise model_config.moe_max_num_tokens." - ) - - if num_tokens > 0: - buf.x[:num_tokens].copy_(x_fp8) - buf.x_sf[:num_tokens].copy_(x_sf) - buf.topk_idx[:num_tokens].copy_(topk_idx) - buf.topk_weights[:num_tokens].copy_(topk_weights) - - y = torch.empty((num_tokens, self.hidden_size), dtype=torch.bfloat16, device=buf.x.device) - dg.fp8_fp4_mega_moe( - y, - self._t_l1, - self._t_l2, - buf, - activation=self.activation, - activation_clamp=self.activation_clamp, - fast_math=self.fast_math, - ) - return y.to(output_dtype) - - # ------------------------------------------------------------------ - # Hot path - # ------------------------------------------------------------------ - 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, - **kwargs, - ) -> torch.Tensor: - if isinstance(x, Fp4QuantizedTensor): - raise NotImplementedError( - "MegaMoE Phase 1 expects BF16 activation; kernel does its own FP8 quant internally." - ) - dg = _import_deep_gemm() - - assert do_finalize, "MegaMoE always finalizes inside the fused kernel" - if output_dtype is None: - output_dtype = x.dtype - - # ----- Resolve real (unpadded) token count ----------------------- - # MoE.forward_fake contract: return shape [num_tokens_real, H] - # where num_tokens_real = all_rank_num_tokens[moe_ep_rank] when - # provided (attention-DP padding case). x.shape[0] may be larger - # than num_tokens_real under ``use_dp_padding=True``. - # Phase 1 asserts ``ep_size == parallel_size`` so the per-EP-rank - # entry is also the per-DP-rank entry. - if all_rank_num_tokens is not None: - num_tokens = int(all_rank_num_tokens[self.mapping.moe_ep_rank]) - else: - num_tokens = x.shape[0] - assert num_tokens <= x.shape[0], f"num_tokens ({num_tokens}) > x.shape[0] ({x.shape[0]})" - assert num_tokens <= self.max_num_tokens, ( - f"MegaMoE got {num_tokens} tokens but buffer is sized for " - f"{self.max_num_tokens}. Raise model_config.moe_max_num_tokens." - ) - - # Note: DO NOT short-circuit when num_tokens == 0. DG's - # fp8_fp4_mega_moe is a collective on the EP symm-mem group — - # skipping the kernel on a zero-token rank would hang peers - # whose tokens route to experts on this rank (they block waiting - # for this rank's kernel entry). DG's own test exercises this - # with uneven per-rank token counts; we mirror that contract. - - buf = self._symm_buffer - assert buf is not None, "MegaMoE SymmBuffer not allocated — post_load_weights missing?" - - if num_tokens > 0: - # Slice to real tokens (skip DP-padded rows, if any). - x_real = x[:num_tokens] - router_logits_real = router_logits[:num_tokens] - - # ----- Routing ---------------------------------------------- - # Upstream ``BaseMoeRoutingMethod.apply`` takes only - # ``router_logits``. ``input_ids`` is accepted by this method - # for forward-compat but ignored at this layer. - topk_idx, topk_weights = self.routing_method.apply(router_logits_real) - topk_idx = topk_idx.to(torch.int64) - topk_weights = topk_weights.to(torch.float32) - - # ----- Pre-quant activations via the fused Inductor path ---- - x_fp8, x_sf = self.quantize_input(x_real) - - # ----- Write into symm buffer ------------------------------- - buf.x[:num_tokens].copy_(x_fp8) - buf.x_sf[:num_tokens].copy_(x_sf) - buf.topk_idx[:num_tokens].copy_(topk_idx) - buf.topk_weights[:num_tokens].copy_(topk_weights) - - # ----- Kernel launch (always, even when num_tokens == 0) -------- - y = torch.empty((num_tokens, self.hidden_size), dtype=torch.bfloat16, device=x.device) - dg.fp8_fp4_mega_moe( - y, - self._t_l1, - self._t_l2, - buf, - activation=self.activation, - activation_clamp=self.activation_clamp, - fast_math=self.fast_math, - ) - return y.to(output_dtype) diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py new file mode 100644 index 000000000000..5c8f37a5caa8 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py @@ -0,0 +1,591 @@ +# 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. +"""MegaMoE — DeepGEMM ``fp8_fp4_mega_moe`` as a first-class MoE backend. + +This backend owns capability checks, routing/activation quantization, and the +fused kernel entry point. ``W4A8MXFP4MXFP8MegaMoEDeepGemmMethod`` owns the +DG-native weight tensors, checkpoint loading, scale conversion, SymmBuffer +allocation, and DeepGEMM weight transform. +""" + +from __future__ import annotations + +from typing import Dict, List, Optional, Tuple + +import torch +import torch.distributed as dist + +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger +from tensorrt_llm.models.modeling_utils import QuantAlgo + +from ....model_config import ModelConfig +from ....utils import ActivationType, AuxStreamType +from ..interface import MoE, MoESchedulerKind, MoEWeightLoadingMode +from ..quantization import ( + W4A8MXFP4MXFP8MegaMoEDeepGemmMethod, + _import_deep_gemm, + _MegaMoEUnavailable, +) +from ..routing import BaseMoeRoutingMethod + +__all__ = ["MegaMoEDeepGemm"] + +# Process-global DG SymmBuffer cache. The cached object is mutable +# forward-time activation workspace (input ``x`` / routing slots / +# L1+L2 GEMM intermediates), not immutable weight state. Reuse relies +# on the current TRT-LLM execution contract that MegaMoE layers run +# serially within a forward pass; concurrent MegaMoE forwards sharing +# a key would race on the same scratch buffers. +_MEGA_MOE_SYMM_BUFFER_CACHE: Dict[tuple, object] = {} + +# ---- Fused MXFP8 per-token quant backends -------------------------------- +# We want: BF16 (m, H) → FP8 E4M3 (m, H) + packed-UE8M0 SF (m, H/32/4) int32. +# Three candidates, in preference order: +# +# 1. ``torch.ops.trtllm.mxfp8_quantize(x, False, alignment=32)`` — TRT-LLM +# C++ CUDA kernel. Roundtrip-verified byte-identical to DG's Python +# helper (fp8 bytes + SF after u8→int32 reshape). Fastest by 5-25× +# vs torch.compile, one kernel launch (~11 us regardless of seq). +# Requires ``libth_common.so`` to be loaded; ``ConfigurableMoE`` pulls +# this in on construction so it's always registered by the time +# ``backend.quantize_input`` runs. +# +# 2. ``torch.compile(dg.per_token_cast_to_fp8, dynamic=True)`` — fallback +# when the TRT-LLM op isn't registered (e.g. slim builds, standalone +# DG tests). Inductor fuses the ~8 elementwise kernels into 1-2 +# Triton kernels but still pays one launch per seq boundary. +# +# ``_FUSED_PER_TOKEN_CAST`` caches the fallback so we don't re-compile on +# every module creation. +_FUSED_PER_TOKEN_CAST = None + + +def _trtllm_mxfp8_quantize_available() -> bool: + return hasattr(torch.ops, "trtllm") and hasattr(torch.ops.trtllm, "mxfp8_quantize") + + +def _quantize_bf16_to_fp8_ue8m0(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Return (x_fp8, x_sf) in DG mega_moe's expected layout (packed int32).""" + m, n = x.shape + # Packed-UE8M0 stores 4 u8 scales per int32 over a 32-element block, + # so n must be a multiple of 128 for the int32 view below to land on + # an integer last-dim. Misaligned shapes would otherwise fail with a + # cryptic reshape/view error; surface a clear contract here instead. + if n % 128 != 0: + raise ValueError( + f"_quantize_bf16_to_fp8_ue8m0 requires hidden_size % 128 == 0 " + f"(packed-UE8M0 int32 SF stride); got hidden_size={n}" + ) + if _trtllm_mxfp8_quantize_available(): + # ``is_sf_swizzled_layout=False`` → flat row-major uint8 SF, one + # byte per 32-element group. ``alignment=32`` → MXFP8 block size. + x_fp8, x_sf_u8 = torch.ops.trtllm.mxfp8_quantize(x, False, alignment=32) + # DG wants (m, n/32/4) int32 with 4 u8 UE8M0 packed per int32. + # TRT-LLM emits (m*n/32,) uint8 in the same byte order, so a + # reshape + view is a zero-copy reinterpret. + return x_fp8, x_sf_u8.view(m, n // 32).view(torch.int32) + + global _FUSED_PER_TOKEN_CAST + if _FUSED_PER_TOKEN_CAST is None: + dg = _import_deep_gemm() + base = dg.per_token_cast_to_fp8 + + def _call(t: torch.Tensor): + return base(t, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + + _FUSED_PER_TOKEN_CAST = torch.compile(_call, dynamic=True, fullgraph=False) + return _FUSED_PER_TOKEN_CAST(x) + + +class MegaMoEDeepGemm(MoE): + """MoE backend wrapping DeepGEMM's fused ``fp8_fp4_mega_moe`` kernel.""" + + _SUPPORTED_ACTIVATION_DTYPES = frozenset({torch.bfloat16}) + + # Kernel owns dispatch + GEMM1 + SwiGLU + GEMM2 + combine via NVLink + # SymmBuffer; ConfigurableMoE must NOT layer host-side comm on top. + scheduler_kind = MoESchedulerKind.FUSED_COMM + + # ------------------------------------------------------------------ + # Capability gating + # ------------------------------------------------------------------ + @classmethod + def can_implement( + cls, + quant_algo: Optional[QuantAlgo], + dtype_activation: torch.dtype = torch.bfloat16, + swiglu_gptoss_style: bool = False, + hidden_size: Optional[int] = None, + intermediate_size: Optional[int] = None, + ) -> Tuple[bool, Optional[str]]: + # Note: we intentionally do NOT probe ``torch.distributed`` state here. + # ``can_implement`` is a static capability query (SM / dtype / quant / + # shape). Whether a live EP ProcessGroup exists is a runtime concern, + # not a capability one, and ``__init__``'s ``_resolve_ep_pg`` will + # surface a clear error if dist is not initialized by the launcher. + sm = get_sm_version() + if sm != 100: + return False, ( + f"MegaMoEDeepGemm requires SM100 (only arch with " + f"sm100_fp8_fp4_mega_moe.cuh in DeepGEMM); got SM{sm}" + ) + if dtype_activation not in cls._SUPPORTED_ACTIVATION_DTYPES: + return False, ( + f"MegaMoEDeepGemm supports activations in " + f"{cls._SUPPORTED_ACTIVATION_DTYPES}, got {dtype_activation}" + ) + if swiglu_gptoss_style: + return False, "MegaMoEDeepGemm does not support swiglu_gptoss_style" + if quant_algo != QuantAlgo.W4A8_MXFP4_MXFP8: + return False, (f"MegaMoEDeepGemm supports W4A8_MXFP4_MXFP8 only, got {quant_algo}") + # Packed-UE8M0 per-token SF layout has two constraints. First, + # the quantizer reinterprets 4 u8 scales as one int32, so K must + # be divisible by 128. Second, DeepGEMM MegaMoE feeds SF buffers + # through TMA; one u8 scale is stored per 32 K elements and the + # per-token SF row must be 16B aligned. The TMA constraint is + # stricter: (K / 32) % 16 == 0, so K must be divisible by 512. + # Enforce the backend constraint here so the factory can fall + # back cleanly before DG SymmBuffer allocation. + if hidden_size is not None and hidden_size % 512 != 0: + return False, ( + f"MegaMoEDeepGemm requires hidden_size % 512 == 0 " + f"(DeepGEMM TMA-aligned packed-UE8M0 SF row); " + f"got hidden_size={hidden_size}" + ) + if intermediate_size is not None and intermediate_size % 512 != 0: + return False, ( + f"MegaMoEDeepGemm requires intermediate_size % 512 == 0 " + f"(DeepGEMM TMA-aligned packed-UE8M0 SF row); " + f"got intermediate_size={intermediate_size}" + ) + try: + _import_deep_gemm() + except _MegaMoEUnavailable as e: + return False, str(e) + return True, None + + # ------------------------------------------------------------------ + # Init + # ------------------------------------------------------------------ + 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, + activation_type: ActivationType = ActivationType.Swiglu, + init_load_balancer: bool = True, + without_comm: bool = False, + # DG tunables. + activation: str = "swiglu", + activation_clamp: Optional[float] = None, + fast_math: bool = True, + **kwargs, + ) -> None: + 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, + layer_idx=layer_idx, + activation_type=activation_type, + init_load_balancer=init_load_balancer, + ) + + # Assert supported topologies early so unsupported configurations + # fall back via ``can_implement`` rather than crashing later in DG. + assert self.tp_size == 1, ( + f"MegaMoEDeepGemm is EP-only (moe_tp_size=1); got tp_size={self.tp_size}" + ) + assert self.cluster_size == 1, ( + f"MegaMoEDeepGemm assumes cluster_size=1; got cluster_size={self.cluster_size}" + ) + # The DG SymmBuffer is sized to ``num_slots`` and sharded evenly over + # EP ranks. Without EPLB, ``num_slots == num_experts`` so the two + # constraints collapse; with EPLB ``num_slots`` may exceed + # ``num_experts`` and ``num_experts % ep_size == 0`` is too strict. + if self.num_slots % max(self.ep_size, 1) != 0: + raise ValueError( + f"MegaMoEDeepGemm requires num_slots ({self.num_slots}) " + f"divisible by ep_size ({self.ep_size})." + ) + + # ADP semantics: DG's fp8_fp4_mega_moe subsumes cross-rank token + # dispatch into its internal symm_mem exchange. When EP spans + # *all* ranks that may carry tokens (i.e. ``ep_size == + # parallel_size``), no outer allgather / reducescatter is needed: + # every token's origin rank is inside the EP group and DG returns + # results to that origin. If EP is a strict subset of + # parallel_size (e.g. attention-DP > moe_ep_size), some tokens + # live on ranks that the DG kernel cannot reach — that topology + # is not yet supported. + if self.use_dp and self.parallel_size > 1: + assert self.ep_size == self.parallel_size, ( + f"MegaMoEDeepGemm with enable_attention_dp=True requires " + f"ep_size == parallel_size (got ep_size={self.ep_size}, " + f"parallel_size={self.parallel_size}). Configurations " + f"with ADP > EP are not yet supported; add the standard " + f"allgather(pre) + reducescatter(post) wrapper before " + f"calling fp8_fp4_mega_moe to support them." + ) + + # apply_router_weight_on_input pre-multiplies routing weights + # onto x before the MoE compute (used by some top-1 models). DG's + # fused kernel applies the weights on the MoE output instead; + # mixing the two produces wrong math. Reject loudly — a silent + # fallback would break llama-min-latency-style paths that set + # this flag to True and assume top-1 semantics. + assert not apply_router_weight_on_input, ( + "MegaMoEDeepGemm does not support apply_router_weight_on_input. " + "DG's fp8_fp4_mega_moe applies routing weights on the MoE " + "output, not by pre-scaling the input — the two paths are " + "not equivalent. Use a different MoE backend for models that " + "require pre-scaling, or extend the kernel call." + ) + # DG's fp8_fp4_mega_moe currently only ships a fused SwiGLU + # activation path. Reject other ActivationType values explicitly so + # ``create_moe_backend`` callers do not silently get SwiGLU when + # they asked for GELU / etc. + if activation_type != ActivationType.Swiglu: + raise ValueError( + f"MegaMoEDeepGemm only supports ActivationType.Swiglu (got {activation_type})." + ) + self.apply_router_weight_on_input = apply_router_weight_on_input + self.activation = activation + self.activation_clamp = activation_clamp + self.fast_math = fast_math + + # Buffer sizing. MoE layers execute serially per forward; a single + # process-level pool sized to worst-case per-rank tokens serves all. + self.max_num_tokens = int( + getattr(model_config, "moe_max_num_tokens", 0) + or getattr(model_config, "max_num_tokens", 0) + or 4096 + ) + + # Resolve the EP ProcessGroup at module construction — creating a + # group at forward time would be collective on a non-synchronous + # call stack and deadlock under PP / layer-skip. Construction is + # globally synchronous across ranks during model build. + self._ep_pg = self._resolve_ep_pg() + + # Cache the bundled DeepGEMM module once at construction. ``_import_deep_gemm`` + # does a fresh ``hasattr`` / ``inspect.signature`` check on every call; + # paying that on every forward (``run_moe`` path) shows up in host-side + # CPU overhead even though the underlying ``import`` is cached by Python. + self._dg = _import_deep_gemm() + + # NVLink SymmBuffer activation workspace. Allocation is a + # model-build-period collective (``symm_mem.rendezvous`` over the + # EP group); allocating from ``run_moe`` would deadlock under + # PP / layer-skip paths where some ranks may not enter this + # layer in lockstep, and would also fail under CUDA graph + # capture because rendezvous is a host-side IPC operation. + # ``_resolve_ep_pg`` above relies on the same lockstep window. + # + # The actual allocation is deferred to ``create_weights`` so that + # ConfigurableMoE has a chance to overwrite EPLB-derived + # attributes (``num_slots``, ``expert_size_per_partition``, ...) + # via ``_BACKEND_SYNC_ATTRS`` before we size the buffer. When the + # backend is constructed with ``init_load_balancer=False`` (the + # ConfigurableMoE path), ``MoE.__init__`` only seeds + # ``num_slots = num_experts`` as a placeholder; under EPLB the + # real slot count is larger, and sizing the SymmBuffer here would + # produce ``sym_buffer.num_experts != num_experts_per_rank * + # num_ranks`` at forward time. ``create_weights`` is also + # collective across ranks (called by ConfigurableMoE on all ranks + # right after the sync, or by the backend itself when used + # standalone with ``init_load_balancer=True``), so deferring keeps + # the rendezvous lockstep guarantee intact. + # See ``_alloc_symm_buffer`` for the cache contract. + self._symm_buffer = None + + # Weight tensors and DG transforms are owned by the quant method. + self._t_l1 = None + self._t_l2 = None + self._weights_loaded = False + self._weights_created = False + self.quant_method = None + if not model_config.skip_create_weights_in_init: + self.create_weights() + + def _supports_load_balancer(self) -> bool: + # The DeepGEMM mega kernel routes by `topk_idx` interpreted as slot id + # (range [0, num_slots)) once the SymmBuffer is sized to num_slots. + # Dynamic EPLB migrates the transformed DG tensors registered by the + # quantization method, not the raw checkpoint-layout weights. + return True + + def validate_configurable_moe(self, moe) -> None: + """Assert ``num_slots % ep_size == 0`` for the DG global slot count. + + ``moe`` is the owning ``ConfigurableMoE``; its ``num_slots`` / + ``ep_size`` / load-balancer flags are populated by ``MoE.__init__`` + before ``validate_backend`` runs, so they're stable here. + """ + # SymmBuffer.num_experts (= num_slots in the DG kernel) must divide + # evenly across EP ranks because each rank's weight shard is + # ``num_slots // ep_size`` slots. + if moe.num_slots % moe.ep_size != 0: + raise ValueError( + f"MegaMoEDeepGemm requires num_slots ({moe.num_slots}) " + f"divisible by ep_size ({moe.ep_size}). Adjust the EPLB " + f"replication factor or ep_size." + ) + + # ------------------------------------------------------------------ + # EP process-group resolution (no collective at forward time) + # ------------------------------------------------------------------ + def _resolve_ep_pg(self): + """Return the torch.distributed ProcessGroup for the EP sub-world. + + Prefers ``mapping.moe_ep_group_pg`` (DeviceMeshTopology, Ray path) + because it was built once at Mapping init. Falls back to + ``dist.group.WORLD`` only when ``ep_size == world_size`` (single + EP subset covers all ranks). + + Does NOT call ``dist.new_group`` — that's collective and unsafe to + invoke from any path that may skip ranks (e.g. PP-isolated layer + forwards). When the mapping cannot provide a PG and EP is a + strict subset of world, we raise with a clear message pointing + at ``mpi_disabled=1`` / Ray as the supported path. + """ + if not dist.is_initialized(): + raise RuntimeError( + "MegaMoEDeepGemm requires torch.distributed to be " + "initialized before module construction (mpirun or Ray)." + ) + # Preferred: reuse the existing PG from the mapping (Ray / DeviceMesh). + # Log at info() only on layer 0 so deep models do not spam N copies of + # the same message; deeper layers log at debug() for triage. + try: + pg = self.mapping.moe_ep_group_pg + log_fn = logger.info if self.layer_idx == 0 else logger.debug + log_fn( + f"[MegaMoE] layer={self.layer_idx} using mapping.moe_ep_group_pg (DeviceMesh path)" + ) + return pg + except (NotImplementedError, AttributeError): + pass + # Fallback: degenerate to WORLD when EP spans all ranks. + world_size = dist.get_world_size() + if self.ep_size == world_size: + log_fn = logger.info if self.layer_idx == 0 else logger.debug + log_fn( + f"[MegaMoE] layer={self.layer_idx} using dist.group.WORLD " + f"(EP == world_size == {world_size})" + ) + return dist.group.WORLD + raise RuntimeError( + f"MegaMoEDeepGemm: cannot resolve EP ProcessGroup. The current " + f"mapping does not expose ``moe_ep_group_pg`` and EP " + f"({self.ep_size}) is a strict subset of world " + f"({world_size}). Use DeviceMeshTopology (TLLM_DISABLE_MPI=1) " + f"so the EP PG is constructed once at Mapping init, or set " + f"ep_size == world_size." + ) + + # ------------------------------------------------------------------ + # SymmBuffer activation workspace (collective resource) + # ------------------------------------------------------------------ + def _alloc_symm_buffer(self) -> None: + """Allocate (or fetch from cache) the DG NVLink SymmBuffer. + + The SymmBuffer is forward-time activation workspace + (input ``x`` / ``x_sf``, ``topk_idx``/``topk_weights``, L1/L2 + GEMM intermediates) backed by NVLink symmetric memory. Allocation + runs ``symm_mem.rendezvous`` over the EP group plus a barrier and + ``cuda.synchronize`` (see DeepGEMM ``mega/__init__.py``); this is + a build-time collective and must not run on ``run_moe``: a + non-lockstep rank would deadlock the rendezvous, and CUDA graph + capture would fail on the host-side IPC handle exchange. + + Buffers are shared across layers via ``_MEGA_MOE_SYMM_BUFFER_CACHE`` + keyed on the (EP-PG, slot/expert/topk/shape/activation) tuple. + Sharing is safe only while MegaMoE layer forwards are issued + serially within a forward pass; concurrent MegaMoE forwards + sharing a key would race on the same scratch buffers. + + Both ``num_slots`` and ``num_experts`` participate in the cache + key because two layers with the same ``num_experts`` but + different EPLB replication factors must not collide on the same + cached buffer. + + Invariant: the SymmBuffer's ``num_experts`` parameter is the + GLOBAL slot count (``kNumExperts`` in the DG kernel). With EPLB + this equals ``num_slots`` (``>= num_experts``); without EPLB + ``ConfigurableMoE`` syncs ``num_slots == num_experts`` so the + contract holds in both cases. See ``CHUNKING_DESIGN.md §5.3.2`` + for the local-vs-global axis split. + """ + if self._symm_buffer is not None: + return + key = ( + id(self._ep_pg), + self.num_experts, + self.num_slots, + self.max_num_tokens, + self.routing_method.experts_per_token, + self.hidden_size, + self.intermediate_size, + self.activation, + ) + cached = _MEGA_MOE_SYMM_BUFFER_CACHE.get(key) + if cached is None: + cached = self._dg.get_symm_buffer_for_mega_moe( + self._ep_pg, + self.num_slots, + self.max_num_tokens, + self.routing_method.experts_per_token, + self.hidden_size, + self.intermediate_size, + True, + self.activation, + ) + _MEGA_MOE_SYMM_BUFFER_CACHE[key] = cached + # Log only on the first layer; deeper layers reuse the cache + # and would otherwise spam N copies of an identical line. + log_fn = logger.info if self.layer_idx == 0 else logger.debug + log_fn( + f"[MegaMoE] layer={self.layer_idx} allocated DG " + f"SymmBuffer: {cached.buffer.nbytes / 2**30:.2f} GiB" + ) + self._symm_buffer = cached + + # ------------------------------------------------------------------ + # Weight lifecycle + # ------------------------------------------------------------------ + def _get_quant_method(self): + if ( + self.quant_config is None + or not self.quant_config.layer_quant_mode.has_w4a8_mxfp4_mxfp8() + ): + raise NotImplementedError("MegaMoEDeepGemm supports W4A8_MXFP4_MXFP8 quantization only") + return W4A8MXFP4MXFP8MegaMoEDeepGemmMethod() + + def create_weights(self): + if self._weights_created: + return + # Allocate the DG NVLink SymmBuffer here (lazily) rather than from + # ``__init__`` because ConfigurableMoE only syncs the EPLB-derived + # attributes (``num_slots``, ``expert_size_per_partition``, ...) + # onto the backend AFTER backend ``__init__`` returns, just before + # calling ``backend.create_weights()``. Sizing the SymmBuffer in + # ``__init__`` would therefore use the placeholder + # ``num_slots = num_experts`` and break EPLB at forward time + # (DeepGEMM asserts ``num_experts == num_experts_per_rank * + # num_ranks`` in ``mega.hpp``). Both call sites (the + # ConfigurableMoE-driven path and the standalone + # ``init_load_balancer=True`` path that runs ``create_weights`` + # from ``__init__``) reach this point on every EP rank in + # lockstep, preserving the rendezvous safety invariant. + self._alloc_symm_buffer() + self.quant_method = self._get_quant_method() + self.quant_method.create_weights(self) + self._weights_created = True + + def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False) -> None: + if self.quant_method is None: + self.create_weights() + self.quant_method.load_weights(self, weights, allow_partial_loading) + + def post_load_weights(self) -> None: + if self.quant_method is None: + self.create_weights() + self.quant_method.post_load_weights(self) + + # ------------------------------------------------------------------ + # MoE-contract methods + # ------------------------------------------------------------------ + def quantize_input(self, x, *, post_quant_comm: bool = False, **kwargs): + """BF16 → FP8-E4M3 + packed-UE8M0 per-token SF (gran_k=32). + + Delegates to ``_quantize_bf16_to_fp8_ue8m0`` which picks the + fastest available backend (TRT-LLM C++ op ~11 us at any seq, + or ``torch.compile`` fallback ~60-260 us). Byte-identical + output across all paths so DG's ``fp8_fp4_mega_moe`` consumes + it unchanged. + """ + del post_quant_comm # MegaMoE runs pre-quant comm via DG SymmBuffer + x_bf16 = x.to(torch.bfloat16).contiguous() + return _quantize_bf16_to_fp8_ue8m0(x_bf16) + + def run_moe( + self, + x: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + x_sf: Optional[torch.Tensor] = None, + *, + output_dtype: Optional[torch.dtype] = None, + **unused_kwargs, + ) -> torch.Tensor: + """Run the fused kernel with pre-quantized activations. + + ConfigurableMoE computes routing and calls ``quantize_input`` before + invoking this method, so the backend receives the same FP8+SF+topk + contract at this unified backend entry point. + """ + assert not unused_kwargs, ( + f"MegaMoEDeepGemm.run_moe got unexpected kwargs: {sorted(unused_kwargs)}" + ) + if output_dtype is None: + output_dtype = self.dtype or torch.bfloat16 + if x_sf is None: + raise ValueError("MegaMoEDeepGemm requires x_sf from quantize_input") + dg = self._dg + buf = self._symm_buffer + assert buf is not None, ( + "MegaMoE SymmBuffer not allocated — _alloc_symm_buffer should " + "run unconditionally in __init__; check for a subclass that " + "skipped the parent constructor." + ) + num_tokens = x.shape[0] + assert num_tokens <= self.max_num_tokens, ( + f"MegaMoE got {num_tokens} tokens but buffer is sized for " + f"{self.max_num_tokens}. Raise model_config.moe_max_num_tokens." + ) + + if num_tokens > 0: + buf.x[:num_tokens].copy_(x) + buf.x_sf[:num_tokens].copy_(x_sf) + buf.topk_idx[:num_tokens].copy_(token_selected_experts.to(torch.int64)) + buf.topk_weights[:num_tokens].copy_(token_final_scales.to(torch.float32)) + + y = torch.empty((num_tokens, self.hidden_size), dtype=torch.bfloat16, device=buf.x.device) + dg.fp8_fp4_mega_moe( + y, + self._t_l1, + self._t_l2, + buf, + activation=self.activation, + activation_clamp=self.activation_clamp, + fast_math=self.fast_math, + ) + return y.to(output_dtype) diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py new file mode 100644 index 000000000000..fb9f408a8f89 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py @@ -0,0 +1,1113 @@ +# 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. + +"""MoE forward-execution schedulers. + +ConfigurableMoE owns module lifecycle (backend creation, attribute sync, +weight loading delegation, comm strategy lifetime, EPLB init, repeat_idx +advancement, DWDP record). Schedulers own forward-time decisions: padding, +chunking, communication ordering, EPLB hook ordering, and backend +``run_moe`` invocation. + +Schedulers are read-mostly with respect to ``ConfigurableMoE``: they may +call ``moe.X`` helpers and read ``moe.``, but must NOT write +``moe.repeat_idx`` (advanced by the wrapper) and must only mutate +``moe.comm`` through ``moe.determine_communication_method`` (the documented +AllToAll -> AllGather fallback). See MOE_SCHEDULER_DESIGN.md for the full +contract. + +Two schedulers exist today, distinguished by where the cross-rank EP +exchange runs: + +- ``ExternalCommMoEScheduler``: comm lives outside the MoE kernel; the + scheduler issues ``Communication.dispatch`` / ``Communication.combine`` + from the host with per-chunk EPLB hooks and optional multi-stream + chunk overlap. +- ``FusedCommMoEScheduler``: comm is fused into the backend's fused + kernel (DeepGEMM ``fp8_fp4_mega_moe``-style "MegaMoE") via NVLink + SymmBuffer; no host comm, lockstep chunk launches. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union + +import torch + +from tensorrt_llm._torch.expert_statistic import ExpertStatistic +from tensorrt_llm._torch.utils import EventType, Fp4QuantizedTensor +from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator + +from .communication import DeepEP, DeepEPLowLatency, NVLinkOneSided, NVLinkTwoSided +from .fused_moe_cute_dsl import CuteDslFusedMoE +from .fused_moe_cutlass import CutlassFusedMoE +from .fused_moe_deepgemm import DeepGemmFusedMoE +from .fused_moe_densegemm import DenseGEMMFusedMoE +from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE +from .interface import MoESchedulerKind + +__all__ = [ + "MoEScheduler", + "ExternalCommMoEScheduler", + "FusedCommMoEScheduler", + "create_moe_scheduler", +] + +if TYPE_CHECKING: + from .configurable_moe import ConfigurableMoE + + +class MoEScheduler(ABC): + """Forward-execution strategy for ConfigurableMoE. + + Stateless w.r.t. model configuration. Holds a back-reference to the + owning ``ConfigurableMoE`` and reads (but does not write) wrapper + state. See module docstring for the contract. + """ + + def __init__(self, moe: "ConfigurableMoE") -> None: + self.moe = moe + + @abstractmethod + def forward( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + *, + do_finalize: bool, + output_dtype: Optional[torch.dtype], + all_rank_num_tokens: Optional[List[int]], + use_dp_padding: Optional[bool], + ) -> torch.Tensor: ... + + +# ============================================================================ +# External-comm scheduler +# ============================================================================ + + +class ExternalCommMoEScheduler(MoEScheduler): + """External-comm forward path: host-side dispatch/combine + per-chunk EPLB hooks. + + Steps: + + 1. Fill ``all_rank_num_tokens`` with local token count when missing. + 2. Apply DP padding metadata when requested. + 3. Compute ``num_chunks`` via ``moe.calculate_num_chunks``. + 4. Validate / fallback comm strategy via + ``moe.determine_communication_method``. + 5. Dispatch to single- or multi-chunk implementation. + 6. Truncate DP padding from outputs. + + ``repeat_idx`` advancement and DWDP record are owned by + ``ConfigurableMoE.forward_impl`` after the scheduler returns. + + ``TRTLLM_ENABLE_DUMMY_ALLREDUCE`` is a performance-debug knob that + injects symmetric synchronization around dispatch/combine. It helps + separate MoE communication timing from rank skew or load-imbalance + artifacts when analyzing traces. + """ + + def forward( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + *, + do_finalize: bool, + output_dtype: Optional[torch.dtype], + all_rank_num_tokens: Optional[List[int]], + use_dp_padding: Optional[bool], + ) -> torch.Tensor: + moe = self.moe + + # ========== Step 1: Handle padding ========== + if all_rank_num_tokens is None: + all_rank_num_tokens = [x.shape[0]] + + 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 + + # ========== Step 2: Determine communication method ========== + num_chunks = moe.calculate_num_chunks(all_rank_num_tokens_padded) + + # May fall back AllToAll -> AllGather; this is the only sanctioned + # mutation of ``moe.comm`` from a scheduler. + moe.determine_communication_method(all_rank_num_tokens_padded, num_chunks) + + # ========== Step 3: Execute MoE computation ========== + if num_chunks == 1: + outputs = self._forward_single_chunk( + x, + router_logits, + output_dtype, + all_rank_num_tokens_padded, + use_dp_padding, + do_finalize, + ) + else: + outputs = self._forward_multiple_chunks( + x, + router_logits, + num_chunks, + output_dtype, + all_rank_num_tokens_padded, + use_dp_padding, + do_finalize, + ) + + # ========== Step 4: Truncate DP padding ========== + if moe.use_dp and moe.parallel_size > 1: + outputs = outputs[: all_rank_num_tokens[moe.mapping.tp_rank]] + + return outputs + + # ------------------------------------------------------------------ + # Communication-strategy probes (used by _forward_chunk_impl to gate + # NVLink-specific EPLB stat-gather paths) + # ------------------------------------------------------------------ + def _is_using_nvlink_two_sided(self) -> bool: + return isinstance(self.moe.comm, NVLinkTwoSided) + + def _is_using_nvlink_one_sided(self) -> bool: + return isinstance(self.moe.comm, NVLinkOneSided) + + # ------------------------------------------------------------------ + # DeepGemm workspace allocation + # ------------------------------------------------------------------ + def _prepare_workspace_deepgemm( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + all_rank_num_tokens: List[int], + ) -> Optional[torch.Tensor]: + """Single-chunk workspace for DeepGemmFusedMoE; otherwise ``None``. + + Multi-chunk execution uses ``_prepare_workspaces_for_chunk`` instead. + """ + moe = self.moe + if not isinstance(moe.backend, DeepGemmFusedMoE): + return None + + num_rows = x.shape[0] + if moe.use_dp and moe.comm is not None: + # Communication path padding: dispatch outputs are + # ``[ep_size * max_tokens_per_rank, ...]`` (or expert-major for + # DeepEPLowLatency). Workspace must cover that footprint. + if isinstance(moe.comm, DeepEPLowLatency): + num_rows = moe.num_slots * max(all_rank_num_tokens) + else: + num_rows = moe.mapping.moe_ep_size * max(all_rank_num_tokens) + + workspaces = moe.backend.get_workspaces([num_rows]) + return workspaces[0] + + def _prepare_workspaces_for_chunk( + self, + all_rank_num_tokens_list: List[Optional[List[int]]], + chunk_size_list: List[int], + use_multi_stream: bool, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Multi-chunk workspaces for DeepGemmFusedMoE; ``(None, None)`` otherwise. + + Single-chunk execution uses ``_prepare_workspace_deepgemm`` instead. + """ + moe = self.moe + workspace_0 = None + workspace_1 = None + + if not isinstance(moe.backend, DeepGemmFusedMoE): + return workspace_0, workspace_1 + + # Always need at least workspace_0; reuse chunk_0 size for workspace_1 + # since chunk 0 is always >= subsequent chunks under split_chunk. + # Mirror ``_prepare_workspace_deepgemm``: DeepEPLowLatency dispatches + # expert-major outputs sized ``num_slots * max_tokens_per_rank`` per + # rank (one shard per slot), while other comms produce + # ``ep_size * max_tokens_per_rank``. Using the wrong formula + # under-allocates the workspace for DeepEPLowLatency multi-chunk + # runs and is caught by ``DeepGemmFusedMoE.run_moe``. + if moe.use_dp and all_rank_num_tokens_list[0] is not None: + max_tokens = max(all_rank_num_tokens_list[0]) + if isinstance(moe.comm, DeepEPLowLatency): + chunk_size_0 = moe.num_slots * max_tokens + else: + chunk_size_0 = moe.mapping.moe_ep_size * max_tokens + else: + chunk_size_0 = chunk_size_list[0] + workspace_chunk_sizes = [chunk_size_0] + + if use_multi_stream: + workspace_chunk_sizes.append(chunk_size_0) + + workspaces = moe.backend.get_workspaces(workspace_chunk_sizes) + workspace_0 = workspaces[0] + if use_multi_stream: + workspace_1 = workspaces[1] + + return workspace_0, workspace_1 + + # ------------------------------------------------------------------ + # Single / multi chunk dispatch + # ------------------------------------------------------------------ + def _forward_single_chunk( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + output_dtype: Optional[torch.dtype], + all_rank_num_tokens: List[int], + use_dp_padding: Optional[bool], + do_finalize: bool = True, + ) -> torch.Tensor: + moe = self.moe + is_first_call = moe.repeat_idx == 0 + is_last_call = moe.repeat_idx == moe.repeat_count - 1 + + workspace = self._prepare_workspace_deepgemm(x, all_rank_num_tokens) + + return self._forward_chunk_impl( + x, + router_logits, + output_dtype, + all_rank_num_tokens, + use_dp_padding, + is_first_call, + is_last_call, + do_finalize, + workspace=workspace, + ) + + def _forward_chunk_impl( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + output_dtype: Optional[torch.dtype], + all_rank_num_tokens: List[int], + use_dp_padding: bool, + is_first_call: bool, + is_last_call: bool, + do_finalize: bool = True, + workspace: Optional[dict] = None, + ) -> torch.Tensor: + """Unified per-chunk execution flow for all external-comm backends. + + Flow: + 1. EPLB - Start wait GPU stage (first call only, dynamic only) + 2. Apply routing (only if backend supports routing separation) + 3. EPLB - Update statistics and route (only if EPLB enabled) + 4. Communication prepare phase (NVLINK two-sided only) + 5. Quantization + dispatch (pre/post-quant adaptive ordering) + 6. backend.run_moe + 7. EPLB - Start CPU stage (last call only, dynamic only) + 8. Communication combine + 9. EPLB - Done CPU stage (last call only, dynamic only) + """ + moe = self.moe + + # ========== Step 1: EPLB - Start wait GPU stage ========== + moe._load_balancer_start_wait_gpu_stage(is_first_call) + + # ========== Step 2: Apply routing (only if backend supports load balancer) ========== + if moe.backend._supports_load_balancer(): + # Separated routing: ConfigurableMoE calls routing_method + token_selected_experts, token_final_scales = moe.routing_method.apply(router_logits) + + token_selected_experts = token_selected_experts.to(torch.int32) + + assert token_selected_experts.shape[1] == moe.routing_method.experts_per_token + assert token_selected_experts.shape == token_final_scales.shape + # CutlassFusedMoE and DenseGEMMFusedMoE expect float32; TRTLLMGen expects bfloat16 + if isinstance(moe.backend, (CutlassFusedMoE, DenseGEMMFusedMoE)): + assert token_final_scales.dtype == torch.float32 + assert token_selected_experts.dtype == torch.int32 + + if token_final_scales is not None and isinstance(moe.backend, TRTLLMGenFusedMoE): + token_final_scales = token_final_scales.to(torch.bfloat16) + + # apply_router_weight_on_input: fuse top-k weight onto x + if moe.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) + # DeepEP variants need a non-None token_final_scales tensor + # (they don't tolerate None), so feed all-ones; other strategies + # accept None and skip the multiply. + if isinstance(moe.comm, (DeepEP, DeepEPLowLatency)): + token_final_scales = torch.ones_like(token_final_scales) + else: + token_final_scales = None + + else: + # Fused routing: backend handles routing internally; EPLB must be off. + assert not moe._using_load_balancer(), ( + f"EPLB is enabled but backend {moe.backend.__class__.__name__} " + f"has fused routing (does not support routing separation)" + ) + token_selected_experts = None + token_final_scales = None + + # ========== Step 3: EPLB - Update statistics and route ========== + if moe.layer_load_balancer and token_selected_experts is not None: + moe._load_balancer_done_wait_gpu_stage(is_first_call) + + # NVLink two-sided / one-sided gather EPLB stats themselves; skip the + # base helper's own AllReduce in that case (ignore_allreduce=True). + ignore_allreduce = ( + self._is_using_nvlink_two_sided() or self._is_using_nvlink_one_sided() + ) + moe._load_balancer_update_statistic( + token_selected_experts, + is_first_call, + is_last_call, + ignore_allreduce=ignore_allreduce, + ) + + token_selected_slots = moe._load_balancer_route(token_selected_experts, moe.use_dp) + else: + token_selected_slots = token_selected_experts + + if token_selected_slots is not None: + ExpertStatistic.set_layer(moe.layer_idx) + ExpertStatistic.maybe_add_info(moe.num_slots, token_selected_slots) + token_selected_slots = get_calibrator().maybe_collect_or_replay_slots( + moe.num_slots, token_selected_slots + ) + + # ========== Step 4: Communication prepare phase (NVLINK two-sided only) ========== + local_statistic_tensor_for_dispatch = None + eplb_dispatch_kwargs = {} + should_update_eplb_after_dispatch = False + if self._is_using_nvlink_two_sided(): + local_statistic_tensor = None + if is_last_call: + local_statistic_tensor = moe._load_balancer_get_local_statistic_tensor() + + # prepare_dispatch stores alltoall_info in _dispatch_state and returns gathered_stats + gathered_stats = moe.comm.prepare_dispatch( + token_selected_slots, all_rank_num_tokens, local_statistic_tensor + ) + + if gathered_stats is not None: + gathered_stats = gathered_stats.view((moe.mapping.moe_ep_size, moe.num_experts)) + moe._load_balancer_update_statistic_with_gathered_statistic(gathered_stats) + # NVLinkOneSided gathers EPLB stats inside dispatch, not prepare_dispatch + elif self._is_using_nvlink_one_sided(): + if moe.layer_load_balancer and is_last_call: + local_statistic_tensor_for_dispatch = ( + moe._load_balancer_get_local_statistic_tensor() + ) + if local_statistic_tensor_for_dispatch is not None: + eplb_dispatch_kwargs["eplb_local_stats"] = local_statistic_tensor_for_dispatch + should_update_eplb_after_dispatch = True + + # ========== Step 5: Quantization + dispatch (pre/post-quant adaptive ordering) ========== + if moe.comm is not None: + supports_post_quant = moe.comm.supports_post_quant_dispatch() + + # Debug: optional dummy AllReduce to break load-balancing artifacts + if moe.enable_dummy_allreduce: + moe.dummy_allreduce() + + dispatch_kwargs = dict(eplb_dispatch_kwargs) + if isinstance(moe.comm, DeepEP) and isinstance(moe.backend, TRTLLMGenFusedMoE): + dispatch_kwargs["enable_sanitize_expert_ids"] = True + + if supports_post_quant: + # Quantize -> Dispatch + x, x_sf = moe.backend.quantize_input(x) + + # W4AFP8 + DeepEPLowLatency needs pre_quant_scale_1; other strategies + # absorb the kwarg via **kwargs so unconditional passing is safe. + if hasattr(moe, "quant_scales") and moe.quant_scales is not None: + if hasattr(moe.quant_scales, "pre_quant_scale_1"): + dispatch_kwargs["pre_quant_scale"] = moe.quant_scales.pre_quant_scale_1 + x, x_sf, token_selected_slots, token_final_scales = moe.comm.dispatch( + hidden_states=x, + hidden_states_sf=x_sf, + token_selected_slots=token_selected_slots, + token_final_scales=token_final_scales, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding, + **dispatch_kwargs, + ) + if should_update_eplb_after_dispatch: + gathered_stats = moe.comm.get_eplb_gathered_statistics() + moe._load_balancer_update_statistic_with_gathered_statistic(gathered_stats) + else: + # Dispatch -> Quantize + x, x_sf, token_selected_slots, token_final_scales = moe.comm.dispatch( + hidden_states=x, + hidden_states_sf=None, # not quantized yet + token_selected_slots=token_selected_slots, + token_final_scales=token_final_scales, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding, + **dispatch_kwargs, + ) + x, x_sf = moe.backend.quantize_input(x, post_quant_comm=False) + else: + # No comm: just quantize + x, x_sf = moe.backend.quantize_input(x, post_quant_comm=False) + + # ========== Step 6: MoE computation ========== + # If EPLB is enabled, token_selected_slots is slot ids; otherwise expert ids. + final_hidden_states = moe.backend.run_moe( + x=x, + token_selected_experts=token_selected_slots, + token_final_scales=token_final_scales, + x_sf=x_sf, + **self._get_backend_kwargs( + router_logits, do_finalize, all_rank_num_tokens, output_dtype, x, workspace + ), + ) + + # ========== Step 7: EPLB - Start CPU stage ========== + moe._load_balancer_start_set_cpu_stage(is_last_call) + + # ========== Step 8: Communication combine ========== + if moe.comm is not None: + if moe.enable_dummy_allreduce: + moe.dummy_allreduce() + all_rank_max_num_tokens = max(all_rank_num_tokens) + final_hidden_states = moe.comm.combine( + final_hidden_states, + all_rank_max_num_tokens=all_rank_max_num_tokens, + ) + else: + # Non-comm path: attention TP or single rank; only AllReduce if reduce_results + if moe.parallel_size > 1 and moe.reduce_results: + final_hidden_states = moe.all_reduce(final_hidden_states) + + # ========== Step 9: EPLB - Done CPU stage ========== + moe._load_balancer_done_set_cpu_stage(is_last_call) + + return final_hidden_states + + def _forward_multiple_chunks( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + num_chunks: int, + output_dtype: Optional[torch.dtype], + all_rank_num_tokens: List[int], + use_dp_padding: Optional[bool], + do_finalize: bool = True, + ) -> torch.Tensor: + """Multiple-chunk path with optional aux-stream overlap.""" + moe = self.moe + + # ========== Chunk preparation ========== + if moe.use_dp: + # DP: need all ranks' token counts for reducescatter + all_rank_chunk_size_list = [ + moe.split_chunk(val, num_chunks) for val in all_rank_num_tokens + ] + all_rank_num_tokens_list = [ + [val[idx_chunk] for val in all_rank_chunk_size_list] + for idx_chunk in range(num_chunks) + ] + chunk_size_list = all_rank_chunk_size_list[moe.rank] + + # AllToAll cannot consume an all-zero rank; substitute 1 token. + if moe.enable_alltoall: + 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 + ] + else: + all_rank_num_tokens_list = [None] * num_chunks + chunk_size_list = moe.split_chunk(x.shape[0], num_chunks) + + x_list = x.split(chunk_size_list) + router_logits_list = router_logits.split(chunk_size_list) + + use_multi_stream = not moe.enable_alltoall and moe.aux_stream is not None + + # ========== Setup auxiliary stream ========== + if use_multi_stream: + moe.event_dict[EventType.Main].record() + with torch.cuda.stream(moe.aux_stream): + moe.event_dict[EventType.Main].wait() + + # ========== DeepGemm workspaces ========== + workspace_0, workspace_1 = self._prepare_workspaces_for_chunk( + all_rank_num_tokens_list, chunk_size_list, use_multi_stream + ) + + # ========== Empty-chunk substitution (DP only) ========== + chunked_used = torch.ones(num_chunks, dtype=torch.bool) + if moe.use_dp: + # The split heuristic guarantees chunk 0 has >= 1 token, so it can + # stand in for any empty chunk on this rank. Without substitution, + # the per-chunk dispatch would launch with 0-token shape and the + # peers would see a barrier mismatch. + assert x_list[0].numel() != 0, "chunk 0 shouldn't be empty" + x_list = list(x_list) + router_logits_list = list(router_logits_list) + for idx_chunk in range(num_chunks): + _x = x_list[idx_chunk] + if _x.numel() == 0: + chunked_used[idx_chunk] = False + x_list[idx_chunk] = x_list[0] + router_logits_list[idx_chunk] = router_logits_list[0] + all_rank_num_tokens_list[idx_chunk][moe.mapping.tp_rank] = ( + all_rank_num_tokens_list[0][moe.mapping.tp_rank] + ) + x_list = tuple(x_list) + router_logits_list = tuple(router_logits_list) + + # ========== Execute chunking with overlap ========== + outputs_list = [] + for idx_chunk, (x_chunk, router_logits_chunk) in enumerate(zip(x_list, router_logits_list)): + is_first_call = idx_chunk == 0 and moe.repeat_idx == 0 + is_last_call = idx_chunk == num_chunks - 1 and moe.repeat_idx == moe.repeat_count - 1 + + if use_multi_stream: + # Alternate streams; each chunk fully owns its (forward + reducescatter). + # Even chunks use aux_stream so chunk 0 is isolated from outer main-stream traffic. + if idx_chunk % 2 == 0: + with torch.cuda.stream(moe.aux_stream): + outputs = self._forward_chunk_impl( + x_chunk, + router_logits_chunk, + output_dtype, + all_rank_num_tokens_list[idx_chunk], + use_dp_padding, + is_first_call, + is_last_call, + do_finalize, + workspace=workspace_0, + ) + else: + outputs = self._forward_chunk_impl( + x_chunk, + router_logits_chunk, + output_dtype, + all_rank_num_tokens_list[idx_chunk], + use_dp_padding, + is_first_call, + is_last_call, + do_finalize, + workspace=workspace_1, + ) + else: + outputs = self._forward_chunk_impl( + x_chunk, + router_logits_chunk, + output_dtype, + all_rank_num_tokens_list[idx_chunk], + use_dp_padding, + is_first_call, + is_last_call, + do_finalize, + workspace=workspace_0, + ) + + if chunked_used[idx_chunk]: + outputs_list.append(outputs) + + # ========== Wait for auxiliary stream to complete ========== + if use_multi_stream: + with torch.cuda.stream(moe.aux_stream): + moe.event_dict[EventType.MoeChunkingOverlap].record() + moe.event_dict[EventType.MoeChunkingOverlap].wait() + + outputs = torch.cat(outputs_list) + return outputs + + # ------------------------------------------------------------------ + # Backend run_moe kwargs builder (external-comm only) + # ------------------------------------------------------------------ + def _get_nvlink_onesided_moe_output( + self, + all_rank_num_tokens: Optional[List[int]], + output_dtype: Optional[torch.dtype], + ) -> Optional[torch.Tensor]: + """Workspace-backed output buffer for NVLinkOneSided combine, or None. + + Only meaningful when ``moe.comm`` is NVLinkOneSided AND the backend + supports payload-in-workspace combine. Returns None for all other + comm strategies; callers should always set the resulting kwarg + unconditionally and let backends ignore None. + """ + moe = self.moe + if not isinstance(moe.comm, NVLinkOneSided): + return None + + if not moe.backend.supports_moe_output_in_alltoall_workspace(): + # Backend opts out: keep payload off the workspace path. + moe.comm.payload_in_workspace = False + return None + + workspace_dtype = output_dtype + if isinstance(moe.backend, TRTLLMGenFusedMoE): + # TRTLLMGen sentinel for unfilled rows; bf16 workspace is the + # combine reduction precision used by the kernel. + moe.comm.invalid_token_expert_id = -1 + workspace_dtype = torch.bfloat16 + + assert all_rank_num_tokens is not None, ( + "all_rank_num_tokens must be provided for NVLinkOneSided backend" + ) + runtime_max_tokens_per_rank = max(all_rank_num_tokens) + + moe_output = moe.comm.get_combine_payload_tensor_in_workspace( + runtime_max_tokens_per_rank, moe.hidden_size, workspace_dtype + ) + + # Toggle on for this forward; combine() reads this flag to decide + # whether to emit into the workspace tensor. + moe.comm.payload_in_workspace = True + return moe_output + + def _get_backend_kwargs( + self, + router_logits: Optional[torch.Tensor] = None, + do_finalize: bool = True, + all_rank_num_tokens: Optional[List[int]] = None, + output_dtype: Optional[torch.dtype] = None, + x: Optional[torch.Tensor] = None, + workspace: Optional[dict] = None, + ) -> Dict: + """Backend-specific kwargs for ``backend.run_moe`` (external-comm only). + + ``FusedCommMoEScheduler`` constructs its own kwargs and never + calls this helper, so all branches here are EXTERNAL_COMM backends. + + Backend-specific kwargs: + - Cutlass: is_sf_swizzled, enable_alltoall, tuner_*, moe_output + - CuteDSL: enable_alltoall, moe_output, dwdp_weight_view + - DeepGemm: workspace + - TRTLLMGen: router_logits, do_finalize, moe_output + """ + moe = self.moe + kwargs: Dict = {} + + if moe.backend.__class__ == CutlassFusedMoE: + # Pre-quant dispatch: SFs arrive swizzled; post-quant dispatch: + # SFs arrive unswizzled. Backend uses this to skip a re-swizzle. + supports_post_quant = moe.comm is not None and moe.comm.supports_post_quant_dispatch() + kwargs["is_sf_swizzled"] = not supports_post_quant + kwargs["output_dtype"] = output_dtype + + # Tuner sees pre-alltoall token shapes so cached tactics from the + # warmup (no-alltoall) phase still apply at runtime. + kwargs["enable_alltoall"] = moe.enable_alltoall + if moe.enable_alltoall: + if all_rank_num_tokens is not None: + kwargs["tuner_num_tokens"] = sum(all_rank_num_tokens) + else: + kwargs["tuner_num_tokens"] = ( + x.shape[0] * moe.mapping.tp_size if x is not None else None + ) + kwargs["tuner_top_k"] = moe.routing_method.top_k + + kwargs["moe_output"] = self._get_nvlink_onesided_moe_output( + all_rank_num_tokens=all_rank_num_tokens, output_dtype=output_dtype + ) + + elif moe.backend.__class__ == CuteDslFusedMoE: + kwargs["enable_alltoall"] = moe.enable_alltoall + kwargs["moe_output"] = self._get_nvlink_onesided_moe_output( + all_rank_num_tokens=all_rank_num_tokens, output_dtype=output_dtype + ) + + if moe.enable_dwdp: + kwargs["dwdp_weight_view"] = moe.dwdp_manager.build_weight_view( + moe.layer_idx, moe.backend + ) + + elif moe.backend.__class__ == DeepGemmFusedMoE: + if workspace is not None: + kwargs["workspace"] = workspace + + elif moe.backend.__class__ == TRTLLMGenFusedMoE: + # When the scheduler precomputes top-k for DP/load-balancer paths, + # the backend must not route again. Single-rank TRTLLMGen paths do + # not get precomputed top-k, so they still need router_logits. + router_logits_arg = None if moe.backend._supports_load_balancer() else router_logits + kwargs["router_logits"] = router_logits_arg + kwargs["do_finalize"] = do_finalize + kwargs["moe_output"] = self._get_nvlink_onesided_moe_output( + all_rank_num_tokens=all_rank_num_tokens, output_dtype=output_dtype + ) + + return kwargs + + +# ============================================================================ +# Fused-comm scheduler (MegaMoE-style) +# ============================================================================ + + +class FusedCommMoEScheduler(MoEScheduler): + """Fused-comm scheduler: backend's fused kernel owns the EP exchange. + + Invariants (see MOE_SCHEDULER_DESIGN.md / mega_moe/CHUNKING_DESIGN.md): + + 1. Reject ``Fp4QuantizedTensor`` activation; backend.quantize_input + owns the BF16 -> FP8 conversion. + 2. Ignore ``use_dp_padding`` (no host-side cross-rank shape alignment). + 3. Use ``mapping.moe_ep_rank`` for local token count, not global rank. + 4. Strip ADP padding before splitting tensors. + 5. ``had_meta=False`` -> pass ``None`` per-chunk so inner falls back to + ``num_tokens=x.shape[0]`` (avoids IndexError on moe_ep_rank>0). + 6. ``num_chunks = max(real_all_rank_num_tokens)`` (not the generic + ``calculate_num_chunks``; that one falls back to ``sum()`` for + ``comm is None`` and would diverge per rank). + 7. Launch every chunk on every EP rank, including zero-token chunks, + so peers can cross the in-kernel NVLink barrier. + 8. No external Communication.dispatch / Communication.combine. + 9. No multi-stream chunk overlap. + + ``repeat_idx`` advancement is done by ``ConfigurableMoE.forward_impl`` + after this scheduler returns. The scheduler must not rotate + ``moe.repeat_idx``. + """ + + def forward( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + *, + do_finalize: bool, + output_dtype: Optional[torch.dtype], + all_rank_num_tokens: Optional[List[int]], + use_dp_padding: Optional[bool], + ) -> torch.Tensor: + """Sequential multi-chunk path for MegaMoE-style backends. + + Single-chunk case is just ``num_chunks == 1`` -- no separate fast + path. Invariants enforced here (see class docstring): identical + ``num_chunks`` per rank computed from ``max()``, ADP padding + stripped before splitting, zero-token chunks still launch the + kernel for the cross-rank barrier. + """ + del use_dp_padding # MegaMoE has no host-side cross-rank shape alignment. + + if isinstance(x, Fp4QuantizedTensor): + raise NotImplementedError( + "Fused-comm MoE expects BF16 activation; " + "quantization happens in backend.quantize_input." + ) + + x_real, rl_real, real_all_rank_num_tokens, ep_rank, had_meta = self._strip_adp_padding( + x, router_logits, all_rank_num_tokens + ) + num_chunks, x_chunks, rl_chunks, all_rank_chunk_size_list = self._compute_chunk_layout( + x_real, rl_real, real_all_rank_num_tokens, ep_rank + ) + outputs = self._run_chunks( + x_chunks, + rl_chunks, + num_chunks=num_chunks, + x_real=x_real, + rl_real=rl_real, + all_rank_chunk_size_list=all_rank_chunk_size_list, + had_meta=had_meta, + output_dtype=output_dtype, + do_finalize=do_finalize, + ) + if not outputs: + cast_dtype = output_dtype if output_dtype is not None else x.dtype + return x.new_empty((0, x.shape[1]), dtype=cast_dtype) + return torch.cat(outputs, dim=0) + + def _strip_adp_padding( + self, + x: torch.Tensor, + router_logits: torch.Tensor, + all_rank_num_tokens: Optional[List[int]], + ) -> Tuple[torch.Tensor, torch.Tensor, List[int], int, bool]: + """Slice ADP padding off ``x`` / ``router_logits`` using moe_ep_rank. + + SymmBuffer exchange is EP-scoped, so we index the per-rank token count + via ``moe.mapping.moe_ep_rank``, not ``self.rank``. ``had_meta`` lets + the per-chunk impl fall back to ``num_tokens=x.shape[0]`` (avoids + ``[len-1 list][moe_ep_rank>0]`` IndexError when no metadata is + provided, e.g. dummy / single-rank forwards). + """ + moe = self.moe + had_meta = all_rank_num_tokens is not None + if had_meta: + # Force plain Python int: downstream torch.Tensor.split and range() + # reject torch 0-d tensor / numpy scalar elements, and the public + # ``Optional[List[int]]`` type hint is not runtime-enforced. + real_all_rank_num_tokens = [int(v) for v in all_rank_num_tokens] + ep_rank = moe.mapping.moe_ep_rank + else: + real_all_rank_num_tokens = [int(x.shape[0])] + ep_rank = 0 + real_local = real_all_rank_num_tokens[ep_rank] + assert real_local <= x.shape[0], ( + f"real_local ({real_local}) > x.shape[0] ({x.shape[0]}); " + "all_rank_num_tokens may not be indexed correctly." + ) + # ADP padding stripped before split, else trailing rows silently + # drift into chunk-0 or torch.split shape-errors. + x_real = x[:real_local] + rl_real = router_logits[:real_local] + return x_real, rl_real, real_all_rank_num_tokens, ep_rank, had_meta + + def _compute_chunk_layout( + self, + x_real: torch.Tensor, + rl_real: torch.Tensor, + real_all_rank_num_tokens: List[int], + ep_rank: int, + ) -> Tuple[int, List[torch.Tensor], List[torch.Tensor], List[List[int]]]: + """Compute per-rank/per-chunk shape and the actual tensor splits. + + ``num_chunks`` uses ``max()``, not ``moe.calculate_num_chunks``: the + latter falls back to ``sum()`` when ``comm is None`` and would diverge + per rank, breaking the in-kernel cross-rank barrier (class invariant 6). + ``... else 0`` defends against an empty meta list (caller passing + ``[]`` instead of ``None``); ``max([])`` would otherwise raise. + + ``all_rank_chunk_size_list[r][c]`` = tokens rank r contributes to + chunk c. ``split_chunk`` evenly partitions ``v`` into exactly + ``num_chunks`` pieces (zero-padded when v < num_chunks, including + v == 0), so every row has the same length and ``chunk_size_list`` + below is this rank's row. + """ + moe = self.moe + real_local = real_all_rank_num_tokens[ep_rank] + + max_real = max(real_all_rank_num_tokens) if real_all_rank_num_tokens else 0 + num_chunks = max( + 1, + (max_real + moe.moe_max_num_tokens - 1) // moe.moe_max_num_tokens, + ) + + all_rank_chunk_size_list = [ + moe.split_chunk(v, num_chunks) for v in real_all_rank_num_tokens + ] + chunk_size_list = all_rank_chunk_size_list[ep_rank] + # ``else []`` shortcut for real_local == 0: equivalent to + # x_real.split([0]*num_chunks) but skips the no-op torch call. The + # zero-token fallback in ``_run_chunks`` then fires for every chunk. + x_chunks = list(x_real.split(chunk_size_list)) if real_local > 0 else [] + rl_chunks = list(rl_real.split(chunk_size_list)) if real_local > 0 else [] + return num_chunks, x_chunks, rl_chunks, all_rank_chunk_size_list + + def _run_chunks( + self, + x_chunks: List[torch.Tensor], + rl_chunks: List[torch.Tensor], + *, + num_chunks: int, + x_real: torch.Tensor, + rl_real: torch.Tensor, + all_rank_chunk_size_list: List[List[int]], + had_meta: bool, + output_dtype: Optional[torch.dtype], + do_finalize: bool, + ) -> List[torch.Tensor]: + """Drive the per-chunk kernel launches, padding zero-token chunks. + + Stage hooks + AllReduce only fire at the (first|last) chunk of the + (first|last) repeat, matching the external-comm path. The + ``idx_chunk >= len(x_chunks)`` branch only triggers when + ``real_local == 0`` (this rank has no tokens but peers do): class + invariant 7 says launch every chunk on every EP rank so the in-kernel + NVLink barrier (SymmBuffer collective) can synchronize. + """ + moe = self.moe + outputs: List[torch.Tensor] = [] + for idx_chunk in range(num_chunks): + is_first_call = idx_chunk == 0 and moe.repeat_idx == 0 + is_last_call = idx_chunk == num_chunks - 1 and moe.repeat_idx == moe.repeat_count - 1 + + if idx_chunk < len(x_chunks): + x_chunk = x_chunks[idx_chunk] + rl_chunk = rl_chunks[idx_chunk] + else: + # Shape ``(0, hidden_size)`` keeps dtype/device/column-width + # intact so routing / quantize / run_moe execute as no-ops + # without shape errors before reaching the barrier. + x_chunk = x_real.new_empty((0, x_real.shape[1])) + rl_chunk = rl_real.new_empty((0, rl_real.shape[1])) + + per_chunk_all_rank = ( + [lst[idx_chunk] for lst in all_rank_chunk_size_list] if had_meta else None + ) + + out_chunk = self._forward_chunk( + x_chunk, + rl_chunk, + output_dtype=output_dtype, + all_rank_num_tokens=per_chunk_all_rank, + do_finalize=do_finalize, + is_first_call=is_first_call, + is_last_call=is_last_call, + ) + outputs.append(out_chunk) + return outputs + + def _forward_chunk( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + router_logits: torch.Tensor, + *, + output_dtype: Optional[torch.dtype], + all_rank_num_tokens: Optional[List[int]], + do_finalize: bool, + is_first_call: bool = True, + is_last_call: bool = True, + ) -> torch.Tensor: + """Run a single chunk through the fused-comm backend. + + Inputs are already ADP-stripped by the caller; ``x.shape[0]`` is + the true unpadded per-rank token count for this chunk. + ``x.shape[0] == 0`` is valid: the kernel still launches so peers + can cross ``nvlink_barrier``. + + EPLB hook ordering (matches ``ExternalCommMoEScheduler._forward_chunk_impl``): + ``start_wait_gpu_stage`` -> routing -> ``done_wait_gpu_stage`` -> + ``update_statistic(ignore_allreduce=False)`` -> ``route`` -> + quantize -> ``run_moe`` -> ``start_set_cpu_stage`` -> + ``done_set_cpu_stage``. ``start/done_set_cpu_stage`` are placed + AFTER ``run_moe``; otherwise dynamic-EPLB weight migration would + race with the fused kernel using those weights. + """ + moe = self.moe + assert not moe.apply_router_weight_on_input, ( + "Fused-comm MoE does not support apply_router_weight_on_input" + ) + assert do_finalize, "Fused-comm MoE always finalizes inside the fused kernel" + + if isinstance(x, Fp4QuantizedTensor): + raise NotImplementedError( + "Fused-comm MoE expects BF16 activation; " + "quantization happens in backend.quantize_input." + ) + if output_dtype is None: + output_dtype = x.dtype + + # Index per moe_ep_rank, not self.rank: SymmBuffer exchange is EP-scoped. + if all_rank_num_tokens is not None: + num_tokens = int(all_rank_num_tokens[moe.mapping.moe_ep_rank]) + else: + num_tokens = x.shape[0] + assert num_tokens <= x.shape[0], f"num_tokens ({num_tokens}) > x.shape[0] ({x.shape[0]})" + + x_chunk_real = x[:num_tokens] + router_logits_chunk_real = router_logits[:num_tokens] + + # ----- EPLB: drain previous CPU rebalance ----- + # Static EPLB early-returns inside the helper; only the dynamic + # balancer actually waits. + moe._load_balancer_start_wait_gpu_stage(is_first_call) + + # ----- routing ----- + # int32 matches the EPLB stats kernel contract used by the external-comm + # path; the fused-comm backend casts to int64 internally. + if num_tokens > 0: + token_selected_experts, token_final_scales = moe.routing_method.apply( + router_logits_chunk_real + ) + token_selected_experts = token_selected_experts.to(torch.int32) + token_final_scales = token_final_scales.to(torch.float32) + else: + device = x.device + token_selected_experts = torch.empty( + (0, moe.routing_method.experts_per_token), + dtype=torch.int32, + device=device, + ) + token_final_scales = torch.empty( + (0, moe.routing_method.experts_per_token), + dtype=torch.float32, + device=device, + ) + + # ----- EPLB: update stats + remap expert ids -> slot ids ----- + if moe.layer_load_balancer: + moe._load_balancer_done_wait_gpu_stage(is_first_call) + # ignore_allreduce=False: the fused kernel has no side channel + # for an external stats gather. The base helper runs its own + # EP-wide AllReduce, gated to is_last_call=True. + moe._load_balancer_update_statistic( + token_selected_experts, + is_first_call, + is_last_call, + ignore_allreduce=False, + ) + token_selected_slots = moe._load_balancer_route(token_selected_experts, moe.use_dp) + else: + token_selected_slots = token_selected_experts + + if token_selected_slots is not None: + ExpertStatistic.set_layer(moe.layer_idx) + ExpertStatistic.maybe_add_info(moe.num_slots, token_selected_slots) + token_selected_slots = get_calibrator().maybe_collect_or_replay_slots( + moe.num_slots, token_selected_slots + ) + + # ----- quantize ----- + if num_tokens > 0: + x_fp8, x_sf = moe.backend.quantize_input(x_chunk_real) + else: + device = x.device + x_fp8 = torch.empty((0, moe.hidden_size), dtype=torch.float8_e4m3fn, device=device) + # Packed-UE8M0 int32 SF: one int32 per 128 input elements per row, + # same stride contract as the non-empty runs for run_moe. + x_sf = torch.empty((0, moe.hidden_size // 128), dtype=torch.int32, device=device) + + # ----- MoE compute ----- + # ``token_selected_slots`` is in [0, num_slots), matching the kernel's + # ``num_experts`` template parameter (SymmBuffer / weights sized to + # num_slots in quantization.py). + out = moe.backend.run_moe( + x=x_fp8, + token_selected_experts=token_selected_slots, + token_final_scales=token_final_scales, + x_sf=x_sf, + output_dtype=output_dtype, + ) + + # ----- EPLB: start/done CPU rebalance, AFTER run_moe ----- + # The external-comm path overlaps CPU stage with ``comm.combine``; + # fused-comm has no external combine, so start_set fires + # immediately after the fused kernel and done_set drains it. Placing + # start_set before run_moe would let dynamic-EPLB migration race the + # kernel. + moe._load_balancer_start_set_cpu_stage(is_last_call) + moe._load_balancer_done_set_cpu_stage(is_last_call) + + return out + + +# ============================================================================ +# Factory +# ============================================================================ + + +def create_moe_scheduler(moe: "ConfigurableMoE") -> MoEScheduler: + """Pick the scheduler matching ``moe.backend.scheduler_kind``.""" + kind = moe.backend.scheduler_kind + if kind == MoESchedulerKind.FUSED_COMM: + return FusedCommMoEScheduler(moe) + if kind == MoESchedulerKind.EXTERNAL_COMM: + return ExternalCommMoEScheduler(moe) + raise ValueError( + f"Unknown MoE scheduler kind {kind!r} on backend " + f"{type(moe.backend).__name__}. Set ``scheduler_kind`` to one of " + f"{[k.name for k in MoESchedulerKind]}." + ) diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 0feb706ffa8b..3a5183a54b38 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -4528,3 +4528,459 @@ def create_weights(self, module: torch.nn.Module): def load_quant_scales(self, module: torch.nn.Module, weights: Dict): # Load weight block scales. super().load_quant_scales(module, weights) + + +class _MegaMoEUnavailable(RuntimeError): + """Bundled DeepGEMM does not expose the full MegaMoE API.""" + + +def _import_deep_gemm(): + """Return the bundled ``tensorrt_llm.deep_gemm`` module.""" + try: + from tensorrt_llm import deep_gemm as _dg + except ImportError as e: + raise _MegaMoEUnavailable( + f"tensorrt_llm.deep_gemm not importable: {e}") from e + + missing = [ + name for name in ( + "fp8_fp4_mega_moe", + "get_symm_buffer_for_mega_moe", + "transform_sf_into_required_layout", + "transform_weights_for_mega_moe", + ) if not hasattr(_dg, name) + ] + if missing: + raise _MegaMoEUnavailable( + f"tensorrt_llm.deep_gemm missing mega_moe symbols {missing}; " + f"upgrade the TRT-LLM bundled DeepGEMM to a release that " + f"includes fp8_fp4_mega_moe.") + + p_fp8 = getattr(_dg, "per_token_cast_to_fp8", None) + if p_fp8 is None or "use_packed_ue8m0" not in inspect.signature( + p_fp8).parameters: + raise _MegaMoEUnavailable( + "tensorrt_llm.deep_gemm.per_token_cast_to_fp8 does not accept " + "use_packed_ue8m0=; upgrade the bundled DeepGEMM.") + return _dg + + +def _ue8m0_uint8_to_fp32(sf_uint8: torch.Tensor) -> torch.Tensor: + """Convert UE8M0 stored as uint8 to fp32 with matching numeric value. + + Shifting left by 23 places each uint8 scale into the IEEE-754 fp32 + exponent field; the final view reinterprets those bits as fp32. + """ + assert sf_uint8.dtype == torch.uint8 + return (sf_uint8.to(torch.int32) << 23).contiguous().view(torch.float32) + + +class W4A8MXFP4MXFP8MegaMoEDeepGemmMethod(FusedMoEMethodBase): + """Weight lifecycle for DeepGEMM MegaMoE W4A8_MXFP4_MXFP8 weights. + + The NVLink SymmBuffer (forward-time activation workspace, not + weight storage) is owned by ``MegaMoEDeepGemm`` itself and + allocated in its ``__init__`` because the allocation is a build-time + EP collective; this class only handles weight tensors and DG + weight transforms. + """ + + eplb_support_status = EplbSupportStatus.SUPPORTED + weight_dtype = torch.uint8 + block_scales_dtype = torch.uint8 + weight_alignment = 128 + input_hidden_alignment = 128 + + def create_weights(self, module: torch.nn.Module) -> None: + expert_count = module.expert_size_per_partition + hidden_size = module.hidden_size + intermediate_size = module.intermediate_size + + # Packed UE8M0 SF reinterprets H / 32 bytes as H / 128 int32 values, + # so H / 32 must be divisible by 4. + assert hidden_size % self.input_hidden_alignment == 0, ( + f"hidden {hidden_size} must be divisible by " + f"{self.input_hidden_alignment}") + assert intermediate_size % self.weight_alignment == 0, ( + f"intermediate {intermediate_size} must be divisible by " + f"{self.weight_alignment}") + + module.register_parameter( + "w3_w1_weight", + nn.Parameter(torch.empty(expert_count, + intermediate_size * 2, + hidden_size // 2, + dtype=self.weight_dtype), + requires_grad=False), + ) + module.register_parameter( + "w3_w1_weight_scale", + nn.Parameter(torch.empty(expert_count, + intermediate_size * 2, + hidden_size // 32, + dtype=self.block_scales_dtype), + requires_grad=False), + ) + module.register_parameter( + "w2_weight", + nn.Parameter(torch.empty(expert_count, + hidden_size, + intermediate_size // 2, + dtype=self.weight_dtype), + requires_grad=False), + ) + module.register_parameter( + "w2_weight_scale", + nn.Parameter(torch.empty(expert_count, + hidden_size, + intermediate_size // 32, + dtype=self.block_scales_dtype), + requires_grad=False), + ) + # Downstream reload/EPLB metadata path; populated lazily when parameter + # replacement records tensors that need rebuilding before reload. + module.rebuild_tensor_metadata = {} + self.setup_quant_scales(module) + + def setup_quant_scales(self, module: torch.nn.Module): + module.quant_scales = tuple() + + def _iter_vanilla_expert_weights(self, weights: Dict, expert_id: int): + return ( + weights[f"{expert_id}.w1.weight"], + weights[f"{expert_id}.w3.weight"], + weights[f"{expert_id}.w2.weight"], + weights[f"{expert_id}.w1.weight_scale"], + weights[f"{expert_id}.w3.weight_scale"], + weights[f"{expert_id}.w2.weight_scale"], + ) + + def _iter_fused_gate_up_expert_weights(self, weights: Dict, expert_id: int): + w1_w3 = weights["gate_up_proj"][expert_id].transpose(0, 1).contiguous() + w1_weight, w3_weight = w1_w3.chunk(2, dim=0) + w2_weight = weights["down_proj"][expert_id].transpose(0, 1).contiguous() + + w1_w3_scale = weights["gate_up_proj_weight_scale"][expert_id].transpose( + 0, 1).contiguous() + w1_scale, w3_scale = w1_w3_scale.chunk(2, dim=0) + w2_scale = weights["down_proj_weight_scale"][expert_id].transpose( + 0, 1).contiguous() + return w1_weight, w3_weight, w2_weight, w1_scale, w3_scale, w2_scale + + def _to_weight_device_uint8(self, tensor: torch.Tensor, + dst: torch.Tensor) -> torch.Tensor: + return tensor.to(device=dst.device, non_blocking=True).view(torch.uint8) + + def _load_expert_weights_to_dst( + self, + module: torch.nn.Module, + weight_dict: Dict, + load_expert_ids: List[int], + dst_w3_w1_weight: torch.Tensor, + dst_w3_w1_weight_scale: torch.Tensor, + dst_w2_weight: torch.Tensor, + dst_w2_weight_scale: torch.Tensor, + ) -> None: + mode = module.weight_loading_mode + if mode in (MoEWeightLoadingMode.VANILLA, + MoEWeightLoadingMode.W4A8_CUSTOM): + get_expert = self._iter_vanilla_expert_weights + elif mode == MoEWeightLoadingMode.FUSED_GATE_UP_PROJ: + get_expert = self._iter_fused_gate_up_expert_weights + else: + raise NotImplementedError( + f"MegaMoEDeepGemm load_weights unsupported " + f"weight_loading_mode={mode}") + + for slot_id, expert_id in enumerate(load_expert_ids): + w1, w3, w2, w1_scale, w3_scale, w2_scale = get_expert( + weight_dict, expert_id) + + # DeepGEMM expects L1 in [gate | up] order before + # transform_weights_for_mega_moe interleaves gate/up rows. + # TRT-LLM checkpoints map gate_proj -> w1 and up_proj -> w3. + dst_w3_w1_weight[slot_id].copy_( + torch.cat([ + self._to_weight_device_uint8(w1, dst_w3_w1_weight), + self._to_weight_device_uint8(w3, dst_w3_w1_weight), + ], + dim=0), + non_blocking=True, + ) + dst_w3_w1_weight_scale[slot_id].copy_( + torch.cat([ + self._to_weight_device_uint8(w1_scale, + dst_w3_w1_weight_scale), + self._to_weight_device_uint8(w3_scale, + dst_w3_w1_weight_scale), + ], + dim=0), + non_blocking=True, + ) + dst_w2_weight[slot_id].copy_(self._to_weight_device_uint8( + w2, dst_w2_weight), + non_blocking=True) + dst_w2_weight_scale[slot_id].copy_(self._to_weight_device_uint8( + w2_scale, dst_w2_weight_scale), + non_blocking=True) + + def load_weights( + self, + module: torch.nn.Module, + weights: List[Dict], + allow_partial_loading: bool = False, + ) -> None: + if allow_partial_loading: + raise NotImplementedError("Partial loading is not supported for " + f"{type(self).__name__}") + assert len(weights) == 1, ( + f"MegaMoEDeepGemm expects one weight dict, got {len(weights)}") + weight_dict = weights[0] + + self._load_expert_weights_to_dst( + module, + weight_dict, + module.initial_local_expert_ids, + module.w3_w1_weight.data, + module.w3_w1_weight_scale.data, + module.w2_weight.data, + module.w2_weight_scale.data, + ) + + # ----- EPLB shared-weights migration buffers ----- + # When dynamic EPLB is on, the load balancer migrates experts across + # ranks at runtime by pulling host-side copies into device-side slots. + # ``layer_load_balancer.get_load_expert_ids()`` returns the EXTRA + # experts this rank must hold a CPU copy of (beyond + # ``initial_local_expert_ids`` which already populate the device + # weight tensors above). We allocate matching CPU tensors with the + # same per-expert shape/dtype as the device weights and load the + # same MXFP4 byte layout into them. ``post_load_weights`` will later + # transform these into DG-required form and register them with the + # host_tensor_sharer so peer ranks can read them during migration. + # The CPU staging is required because EPLB's host_tensor_sharer + # exchanges weights via host-pinned memory. + if self.need_load_shared_weights(module): + local_shared_load_expert_ids = module.layer_load_balancer.get_load_expert_ids( + ) + module.local_shared_w3_w1_tensors = torch.empty( + (len(local_shared_load_expert_ids), ) + + module.w3_w1_weight.data.shape[1:], + dtype=module.w3_w1_weight.data.dtype, + device='cpu') + module.local_shared_w3_w1_scale_tensors = torch.empty( + (len(local_shared_load_expert_ids), ) + + module.w3_w1_weight_scale.data.shape[1:], + dtype=module.w3_w1_weight_scale.data.dtype, + device='cpu') + module.local_shared_w2_tensors = torch.empty( + (len(local_shared_load_expert_ids), ) + + module.w2_weight.data.shape[1:], + dtype=module.w2_weight.data.dtype, + device='cpu') + module.local_shared_w2_scale_tensors = torch.empty( + (len(local_shared_load_expert_ids), ) + + module.w2_weight_scale.data.shape[1:], + dtype=module.w2_weight_scale.data.dtype, + device='cpu') + self._load_expert_weights_to_dst( + module, + weight_dict, + local_shared_load_expert_ids, + module.local_shared_w3_w1_tensors, + module.local_shared_w3_w1_scale_tensors, + module.local_shared_w2_tensors, + module.local_shared_w2_scale_tensors, + ) + + module._weights_loaded = True + + def _transform_weights_for_mega_moe( + self, + module: torch.nn.Module, + w3_w1_weight: torch.Tensor, + w3_w1_weight_scale: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale: torch.Tensor, + *, + device: torch.device, + ): + # ``module._dg`` is the DeepGEMM module cached by ``MegaMoEDeepGemm.__init__``. + # Calling ``_import_deep_gemm()`` here would re-run the ``hasattr`` / + # ``inspect.signature`` API checks for every layer. + dg = module._dg + expert_count = w3_w1_weight.shape[0] + hidden_size = module.hidden_size + intermediate_size = module.intermediate_size + + w3_w1_weight = w3_w1_weight.to(device=device, + non_blocking=True).contiguous() + w3_w1_weight_scale = w3_w1_weight_scale.to( + device=device, non_blocking=True).contiguous() + w2_weight = w2_weight.to(device=device, non_blocking=True).contiguous() + w2_weight_scale = w2_weight_scale.to(device=device, + non_blocking=True).contiguous() + + l1_sf_fp32 = _ue8m0_uint8_to_fp32(w3_w1_weight_scale) + l1_sf = dg.transform_sf_into_required_layout( + l1_sf_fp32, + mn=intermediate_size * 2, + k=hidden_size, + recipe=(1, 1, 32), + num_groups=expert_count, + is_sfa=False, + ) + + l2_sf_fp32 = _ue8m0_uint8_to_fp32(w2_weight_scale) + l2_sf = dg.transform_sf_into_required_layout( + l2_sf_fp32, + mn=hidden_size, + k=intermediate_size, + recipe=(1, 1, 32), + num_groups=expert_count, + is_sfa=False, + ) + + l1_weight = w3_w1_weight.view(torch.int8) + l2_weight = w2_weight.view(torch.int8) + return dg.transform_weights_for_mega_moe((l1_weight, l1_sf), + (l2_weight, l2_sf)) + + def post_load_weights(self, module: torch.nn.Module) -> None: + """Transform loaded MXFP4 weights into DG-native form. + + Pipeline (each step is independent and idempotent on its own guard): + 1. ``_transform_main_weights`` - DG-form L1/L2 + EPLB-friendly slot views + 2. ``_setup_shared_weights_for_eplb`` - host-side shared copies for dynamic EPLB + 3. ``_attach_initial_weight_assignments`` - tell load_balancer the initial layout + + The NVLink SymmBuffer (forward-time activation workspace, not + weight storage) is allocated by ``MegaMoEDeepGemm.__init__`` itself + because that is the build-time lockstep window where the + ``symm_mem.rendezvous`` collective is safe; see + ``MegaMoEDeepGemm._alloc_symm_buffer``. + """ + assert module._weights_loaded, "post_load_weights before load_weights" + self._transform_main_weights(module) + self._setup_shared_weights_for_eplb(module) + self._attach_initial_weight_assignments(module) + + def _transform_main_weights(self, module: torch.nn.Module) -> None: + """Build DG-form ``_t_l1`` / ``_t_l2`` and EPLB-friendly slot views. + + Invariant: the scale tensors returned by ``_transform_weights_for_mega_moe`` + are produced by DeepGEMM's ``get_mn_major_tma_aligned_packed_ue8m0_tensor`` + helper, which allocates a non-contiguous storage with shape + ``(num_groups, mn, packed_sf_k)`` and strides + ``(packed_sf_k * tma_aligned_mn, 1, tma_aligned_mn)``. The MN dimension + is the fast-changing axis (stride 1) so that DeepGEMM kernels can issue + MN-major TMA loads on packed UE8M0 SF rows. + + ``MegaMoEDeepGemm.can_implement`` enforces ``hidden_size % 512 == 0`` + and ``intermediate_size % 512 == 0``, which guarantees ``mn`` + (= ``hidden_size`` for L2 / ``2 * intermediate_size`` for L1) is + already aligned to the int32 TMA boundary, so ``tma_aligned_mn == mn``. + Substituting that back, the actual storage stride is + ``(packed_sf_k * mn, 1, mn)``. + + Swapping the last two axes via ``transpose(-2, -1)`` produces: + shape = (num_groups, packed_sf_k, mn) + stride = (packed_sf_k * mn, mn, 1) + which is exactly the row-major contiguous stride for that shape. The + transposed view therefore satisfies ``is_contiguous() == True`` without + any data copy, and EPLB slot registration (which requires contiguous + storage for migration buffers) can reuse the same memory while the + DeepGEMM-facing tuple ``module._t_l*`` keeps the original MN-major view. + + The asserts below are contract guards: if a future change relaxes the + 512-alignment constraint or alters DeepGEMM's TMA-aligned SF layout, + ``tma_aligned_mn != mn`` will break the contiguity property and these + asserts will fire instead of silently corrupting EPLB weight migration. + """ + if module._t_l1 is not None: + return + device = module.w3_w1_weight.device + module._t_l1, module._t_l2 = self._transform_weights_for_mega_moe( + module, + module.w3_w1_weight, + module.w3_w1_weight_scale, + module.w2_weight, + module.w2_weight_scale, + device=device, + ) + module._t_l1_weight, module._t_l1_scale = module._t_l1 + module._t_l2_weight, module._t_l2_scale = module._t_l2 + module._t_l1_scale_slot = module._t_l1_scale.transpose(-2, -1) + module._t_l2_scale_slot = module._t_l2_scale.transpose(-2, -1) + assert module._t_l1_scale_slot.is_contiguous() + assert module._t_l2_scale_slot.is_contiguous() + log_fn = logger.info if module.layer_idx == 0 else logger.debug + log_fn(f"[MegaMoE] layer={module.layer_idx} weight transform done " + f"t_l1=(w {tuple(module._t_l1[0].shape)}/" + f"{module._t_l1[0].dtype}, " + f"sf {tuple(module._t_l1[1].shape)}/" + f"{module._t_l1[1].dtype})") + + def _setup_shared_weights_for_eplb(self, module: torch.nn.Module) -> None: + """Transform & register host-side shared weights for dynamic EPLB. + + Background: ``load_weights`` already populated CPU staging tensors + (``module.local_shared_w*_tensors`` and matching ``*_scale_tensors``) + with the raw MXFP4 layout for the extra experts this rank has been + asked to keep host copies of (see the EPLB shared-weights migration + block in ``load_weights``). Here we DG-transform those into the same + layout as ``_t_l1`` / ``_t_l2``, hand them to + ``register_all_parameter_slot_and_to_fix_weight_fns`` so the load + balancer can copy them into device slots during runtime migration, + and register a fix-up callback that re-derives the MN-major DG views + from the slot-major storage after each migration. + + Finally we drop the staging tensors so they don't keep CPU memory + pinned for the rest of the run. + """ + if not self.need_load_shared_weights(module): + return + device = module.w3_w1_weight.device + shared_t_l1, shared_t_l2 = self._transform_weights_for_mega_moe( + module, + module.local_shared_w3_w1_tensors, + module.local_shared_w3_w1_scale_tensors, + module.local_shared_w2_tensors, + module.local_shared_w2_scale_tensors, + device=device, + ) + module.register_all_parameter_slot_and_to_fix_weight_fns({ + '_t_l1_weight': + shared_t_l1[0].cpu().contiguous(), + '_t_l1_scale_slot': + shared_t_l1[1].transpose(-2, -1).cpu().contiguous(), + '_t_l2_weight': + shared_t_l2[0].cpu().contiguous(), + '_t_l2_scale_slot': + shared_t_l2[1].transpose(-2, -1).cpu().contiguous(), + }) + + def refresh_deepgemm_scale_views(): + module._t_l1_scale = module._t_l1_scale_slot.transpose(-2, -1) + module._t_l2_scale = module._t_l2_scale_slot.transpose(-2, -1) + module._t_l1 = (module._t_l1_weight, module._t_l1_scale) + module._t_l2 = (module._t_l2_weight, module._t_l2_scale) + + module.layer_load_balancer.add_to_migrate_weight_fn( + refresh_deepgemm_scale_views, ()) + for attr in ( + 'local_shared_w3_w1_tensors', + 'local_shared_w3_w1_scale_tensors', + 'local_shared_w2_tensors', + 'local_shared_w2_scale_tensors', + ): + delattr(module, attr) + module.layer_load_balancer.host_tensor_sharer.finalize_layer_weights() + + @staticmethod + def _attach_initial_weight_assignments(module: torch.nn.Module) -> None: + """Hand the initial expert->slot assignments to the load balancer.""" + if hasattr(module, + "layer_load_balancer") and module.layer_load_balancer: + module.layer_load_balancer.set_initial_weight_assignments( + module.initial_global_assignments) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 1391d1fdab7d..8b7dc720824d 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -77,7 +77,10 @@ l0_dgx_b200: # --- CUTEDSL (NVFP4 only) --- - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTEDSL" # --- DEEPGEMM (FP8_BLOCK_SCALES only) --- - - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "DEEPGEMM" + - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "DEEPGEMM and not MEGAMOE_DEEPGEMM" + # --- MEGAMOE_DEEPGEMM (W4A8_MXFP4_MXFP8 only) --- + - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "MEGAMOE_DEEPGEMM" + - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "MEGAMOE_DEEPGEMM" # ------------- MoE: test_multi_gpu_eplb --------------- - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb - condition: diff --git a/tests/integration/test_lists/test-db/l0_dgx_b300.yml b/tests/integration/test_lists/test-db/l0_dgx_b300.yml index 5e9305a38132..575365b0f68b 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b300.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b300.yml @@ -39,10 +39,14 @@ l0_dgx_b300: - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=CUTEDSL-quant=NVFP4-routing=Renormalize] # DEEPGEMM backend: FP8_BLOCK_SCALES - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=DEEPEP-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=DEEPGEMM-quant=FP8_BLOCK_SCALES-routing=Renormalize] + # MEGAMOE_DEEPGEMM backend: W4A8_MXFP4_MXFP8 + - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu[parallel=DEP-comm=IGNORE-e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=MEGAMOE_DEEPGEMM-quant=W4A8_MXFP4_MXFP8-routing=DeepSeekV3] + - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e256_k8_h7168_i2048-seq=1-dtype=torch.bfloat16-backend=MEGAMOE_DEEPGEMM-quant=W4A8_MXFP4_MXFP8-routing=DeepSeekV3] # ------------- MoE: EPLB (Expert Load Balancing) tests --------------- - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb[parallel=DEP-comm=NVLINK_ONE_SIDED-e8_k2_h512_i512-slots=16-dtype=torch.bfloat16-backend=TRTLLM-quant=W4A16_MXFP4-routing=Renormalize] + - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb -k "MEGAMOE_DEEPGEMM" - unittest/_torch/modeling -k "modeling_llama" - unittest/_torch/modeling -k "modeling_mixtral" - unittest/_torch/modeling -k "modeling_gpt_oss" diff --git a/tests/unittest/_torch/modules/moe/moe_test_utils.py b/tests/unittest/_torch/modules/moe/moe_test_utils.py index 5badee4e51fe..22176fb21955 100644 --- a/tests/unittest/_torch/modules/moe/moe_test_utils.py +++ b/tests/unittest/_torch/modules/moe/moe_test_utils.py @@ -47,6 +47,7 @@ from tensorrt_llm._torch.modules.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE from tensorrt_llm._torch.modules.fused_moe.fused_moe_densegemm import DenseGEMMFusedMoE from tensorrt_llm._torch.modules.fused_moe.interface import MoE +from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoEDeepGemm from tensorrt_llm._torch.utils import ActivationType, is_gated_activation from tensorrt_llm.models.modeling_utils import QuantAlgo @@ -64,6 +65,7 @@ class MoeBackendType(str, Enum): CUTEDSL = "CUTEDSL" DEEPGEMM = "DEEPGEMM" DENSEGEMM = "DENSEGEMM" + MEGAMOE = "MEGAMOE_DEEPGEMM" def get_backend_class(backend_type: MoeBackendType) -> Type[MoE]: @@ -74,6 +76,7 @@ def get_backend_class(backend_type: MoeBackendType) -> Type[MoE]: MoeBackendType.CUTEDSL: CuteDslFusedMoE, MoeBackendType.DEEPGEMM: DeepGemmFusedMoE, MoeBackendType.DENSEGEMM: DenseGEMMFusedMoE, + MoeBackendType.MEGAMOE: MegaMoEDeepGemm, } return backend_class_map[backend_type] @@ -614,6 +617,54 @@ def should_skip_densegemm( return None +def should_skip_megamoe( + backend_type: MoeBackendType, + quant_algo: Optional[QuantAlgo] = None, + dtype: Optional[torch.dtype] = None, + model_config: "MoeModelConfig" = None, + comm_method: Optional[str] = None, + moe_tp_size: int = 1, + parallel_mode: Optional[str] = None, + swiglu_gptoss_style: bool = False, +) -> Optional[str]: + """Check MegaMoE-specific constraints for the generic MoE test matrix.""" + if backend_type != MoeBackendType.MEGAMOE: + return None + + if not torch.cuda.is_available(): + return "MegaMoEDeepGemm requires CUDA" + + if comm_method is not None or parallel_mode is not None: + return ( + "MegaMoEDeepGemm generic multi-GPU coverage requires a torch.distributed " + "EP ProcessGroup-aware launcher; keep it out of this MPIPool-based matrix." + ) + + if quant_algo != QuantAlgo.W4A8_MXFP4_MXFP8: + return f"MegaMoEDeepGemm only supports W4A8_MXFP4_MXFP8 (got quant_algo={quant_algo})" + + if dtype is not None and dtype != torch.bfloat16: + return f"MegaMoEDeepGemm only supports bfloat16 activations (got dtype={dtype})" + + if swiglu_gptoss_style: + return "MegaMoEDeepGemm does not support swiglu_gptoss_style" + + if moe_tp_size != 1: + return f"MegaMoEDeepGemm Phase 1 is EP-only (got moe_tp_size={moe_tp_size})" + + if model_config is not None: + hidden_size = model_config.hidden_size + intermediate_size = model_config.intermediate_size + if hidden_size % 512 != 0 or intermediate_size % 512 != 0: + return ( + f"MegaMoEDeepGemm requires 512-aligned hidden/intermediate sizes " + f"for DeepGEMM TMA-packed SF rows " + f"(got h={hidden_size}, i={intermediate_size})" + ) + + return None + + def should_skip_multi_gpu( parallel_mode: str, model_config: "MoeModelConfig", @@ -720,8 +771,8 @@ def supports_autotuner_capture( Returns: True if autotuner capture/replay is supported, False otherwise """ - # DEEPGEMM does not support autotuner capture - if backend_type == MoeBackendType.DEEPGEMM: + # DEEPGEMM and MEGAMOE do not support autotuner capture + if backend_type in (MoeBackendType.DEEPGEMM, MoeBackendType.MEGAMOE): return False if use_flashinfer: @@ -763,6 +814,9 @@ def get_quick_skip_reason( can_impl_kwargs = {"dtype_activation": dtype} if swiglu_gptoss_style: can_impl_kwargs["swiglu_gptoss_style"] = swiglu_gptoss_style + if backend_type == MoeBackendType.MEGAMOE and model_config is not None: + can_impl_kwargs["hidden_size"] = model_config.hidden_size + can_impl_kwargs["intermediate_size"] = model_config.intermediate_size can_impl, skip_reason = backend_cls.can_implement(quant_algo, **can_impl_kwargs) if not can_impl: return skip_reason @@ -787,6 +841,13 @@ def get_quick_skip_reason( lambda: should_skip_densegemm( backend_type, quant_algo=quant_algo, model_config=model_config ), + lambda: should_skip_megamoe( + backend_type, + quant_algo=quant_algo, + dtype=dtype, + model_config=model_config, + swiglu_gptoss_style=swiglu_gptoss_style, + ), ] for check in skip_checks: skip_reason = check() diff --git a/tests/unittest/_torch/modules/moe/quantize_utils.py b/tests/unittest/_torch/modules/moe/quantize_utils.py index a5c7a6de9fa8..9d76824b871b 100644 --- a/tests/unittest/_torch/modules/moe/quantize_utils.py +++ b/tests/unittest/_torch/modules/moe/quantize_utils.py @@ -168,6 +168,8 @@ def get_test_quant_params(quant_algo, x, backend_type=None): # CUTLASS and others use weight_alignment for both quant_kwargs["weight_alignment"] = 128 quant_kwargs["input_hidden_alignment"] = 128 + elif backend_name == "MEGAMOE_DEEPGEMM": + quant_kwargs["ref_cls"] = MXFP4MXFP8RefMegaMoEDeepGemm elif quant_algo == QuantAlgo.W4A16_MXFP4: quantize_util_cls = WFP4A16QuantizeUtil quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_MXFP4) @@ -1344,6 +1346,41 @@ def check_accuracy(self, output, ref_output): check_accuracy(output, ref_output, rtol=0.10, atol=0.2, percent=0.85) +class MXFP4MXFP8RefMegaMoEDeepGemm(MXFP4MXFP8RefGatedMLPFusedMoE): + """Reference matching DeepGEMM MegaMoE's pre-L2 routing-weight placement.""" + + def forward(self, hidden_states: torch.Tensor, router_logits: torch.Tensor) -> torch.Tensor: + if self.hidden_size_unpadded < self.hidden_size: + pad_size = self.hidden_size - self.hidden_size_unpadded + hidden_states = torch.nn.functional.pad(hidden_states, (0, pad_size)) + + assert hidden_states.shape[-1] == self.hidden_size + hidden_states = hidden_states.view(-1, self.hidden_size) + selected_experts, routing_weights = self.routing_method.apply(router_logits) + final_hidden_states = torch.zeros( + hidden_states.shape, dtype=hidden_states.dtype, device=hidden_states.device + ) + + for expert_id in range(self.num_experts): + if not torch.any(selected_experts == expert_id): + continue + batch_idx, nth_expert = torch.where(selected_experts == expert_id) + expert_inputs = hidden_states[batch_idx] + expert = self.experts[expert_id] + l1_output = expert.gate_up_proj(expert_inputs) + act_output = expert._apply_activation(l1_output) + act_output = act_output * routing_weights[batch_idx, nth_expert, None].to( + act_output.dtype + ) + output = expert.down_proj(act_output) + final_hidden_states[batch_idx] += output.float() + + final_hidden_states = final_hidden_states.reshape(hidden_states.shape) + if self.hidden_size_unpadded < self.hidden_size: + final_hidden_states = final_hidden_states[:, : self.hidden_size_unpadded] + return final_hidden_states + + class MXFP4MXFP8QuantizeUtil(BaseQuantizeUtil): """ MXFP4MXFP8QuantizeUtil inherits from BaseQuantizeUtil to support correctness testing @@ -1364,23 +1401,29 @@ def prepare_weights_from_backend(self, backend, **quant_kwargs): Returns: (backend_weights, ref_weights, ref_module_kwargs) """ - # Get actual shapes from backend + # Get actual shapes from backend. MoE TP stores per-rank + # intermediate shards, but the checkpoint-style weights passed to + # load_weights are global and are sharded by the loader. num_elts_per_dtype = torch.iinfo(backend.quant_method.weight_dtype).bits // 4 hidden_size_in = backend.w3_w1_weight.shape[-1] * num_elts_per_dtype # hidden_size_out_padded is used for weight creation (padded value) hidden_size_out_padded = backend.w2_weight.shape[-2] - inter_size = backend.w2_weight.shape[-1] * num_elts_per_dtype + local_inter_size = backend.w2_weight.shape[-1] * num_elts_per_dtype + tp_size = getattr(backend, "tp_size", 1) + inter_size = local_inter_size * tp_size weight_align = backend.quant_method.weight_alignment input_hidden_align = getattr(backend.quant_method, "input_hidden_alignment", weight_align) - # Backend weights: contamination padding + # Backend weights: TP pads global checkpoint weights before sharding. + # The padded intermediate extent participates in CUTLASS kernels, so + # zero TP padding to preserve the original unpadded model semantics. backend_kwargs = dict( quant_kwargs, hidden_size_in=hidden_size_in, hidden_size_out=hidden_size_out_padded, intermediate_size=inter_size, input_hidden_alignment=input_hidden_align, - pad_zero_or_val=False, + pad_zero_or_val=tp_size > 1, bias=self.bias, # Pass bias from self to create bias weights ) backend_weights = self.create_weights(**backend_kwargs) diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index aa8563731d1d..5dfc0ab3c9f6 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -28,10 +28,12 @@ import itertools import logging +import os from typing import List, Optional import pytest import torch +import torch.distributed as dist from _torch.modules.moe.moe_test_utils import ( IS_CI_MODE, MoeBackendType, @@ -52,6 +54,8 @@ from tensorrt_llm._torch.modules.fused_moe import RenormalizeMoeRoutingMethod from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend from tensorrt_llm._torch.modules.fused_moe.interface import MoE, MoEWeightLoadingMode +from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoEDeepGemm +from tensorrt_llm._torch.modules.fused_moe.quantization import W4A8MXFP4MXFP8MegaMoEDeepGemmMethod from tensorrt_llm._torch.utils import ActivationType, is_gated_activation from tensorrt_llm._utils import mpi_rank from tensorrt_llm.mapping import Mapping @@ -60,6 +64,23 @@ logger = logging.getLogger(__name__) +def _ensure_single_proc_dist_for_megamoe(backend_type: MoeBackendType, rank: int) -> None: + """MegaMoE resolves an EP ProcessGroup at construction time.""" + if backend_type != MoeBackendType.MEGAMOE: + return + if not torch.cuda.is_available(): + pytest.skip("CUDA required for MegaMoE tests") + if dist.is_initialized(): + return + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29561") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", str(rank)) + torch.cuda.set_device(rank) + dist.init_process_group(backend="nccl", rank=0, world_size=1) + + def should_skip_gptoss( backend_type: MoeBackendType, quant_algo: Optional[QuantAlgo], @@ -144,6 +165,52 @@ def create_test_backend( ) +def test_megamoe_init_rejects_uneven_num_slots_with_value_error(): + routing_method = RenormalizeMoeRoutingMethod(top_k=1) + model_config = ModelConfig( + mapping=Mapping( + world_size=4, + rank=0, + tp_size=4, + moe_tp_size=1, + moe_ep_size=4, + ), + moe_backend=MoeBackendType.MEGAMOE.value, + ) + + with pytest.raises( + ValueError, + match=r"MegaMoEDeepGemm requires num_slots \(10\) divisible by ep_size \(4\)", + ): + MegaMoEDeepGemm( + routing_method=routing_method, + num_experts=10, + hidden_size=512, + intermediate_size=512, + dtype=torch.bfloat16, + model_config=model_config, + init_load_balancer=False, + ) + + +def test_megamoe_post_load_rejects_uneven_num_slots_with_value_error(monkeypatch): + import tensorrt_llm._torch.modules.fused_moe.quantization as quantization_module + + class DummyModule: + _weights_loaded = True + num_slots = 10 + ep_size = 4 + + monkeypatch.setattr(quantization_module, "_import_deep_gemm", lambda: object()) + method = W4A8MXFP4MXFP8MegaMoEDeepGemmMethod() + + with pytest.raises( + ValueError, + match=r"MegaMoEDeepGemm requires num_slots \(10\) divisible by ep_size \(4\)", + ): + method.post_load_weights(DummyModule()) + + def run_backend_moe( backend: MoE, backend_type: MoeBackendType, @@ -163,6 +230,7 @@ def run_backend_moe( - TRTLLM: token_final_scales=bfloat16, optionally router_logits - CUTEDSL: token_final_scales=float32 - DEEPGEMM: workspace, token_final_scales=float32 + - MEGAMOE_DEEPGEMM: token_selected_experts=int64, output_dtype Args: trtllm_use_router_logits: If True, TRTLLM backend uses router_logits for routing. @@ -193,6 +261,9 @@ def run_backend_moe( m_max = fp8_utils.align(x_quantized.shape[0], 128) args["workspace"] = backend.get_workspace(m_max, 128) + elif backend_type == MoeBackendType.MEGAMOE: + args["token_selected_experts"] = token_selected_experts.to(torch.int64) + args["output_dtype"] = dtype return backend.run_moe(**args) @@ -221,6 +292,7 @@ def run_backend_moe( MoeBackendType.CUTEDSL, MoeBackendType.DEEPGEMM, MoeBackendType.DENSEGEMM, + MoeBackendType.MEGAMOE, ] # Data types to test @@ -247,6 +319,7 @@ def run_backend_moe( LOCAL_MOE_MODEL_CONFIGS = CI_MOE_MODEL_CONFIGS + [ MoeModelConfig(256, 8, 7168, 2048), # DeepSeek-V3 + MoeModelConfig(256, 6, 4096, 2048), # DeepSeek-V4-Flash MoeModelConfig(8, 2, 4096, 14336), # Mixtral-8x7B MoeModelConfig(64, 6, 2048, 1408), # DeepSeek-MoE-16B / DeepSeek-V2-Lite MoeModelConfig(8, 2, 6144, 32768), # Grok-1 @@ -504,6 +577,7 @@ def test_moe_backend( # Create mapping mapping = Mapping() mapping.rank = mpi_rank() + _ensure_single_proc_dist_for_megamoe(backend_type, mapping.rank) with torch.device(f"cuda:{mapping.rank}"): torch.manual_seed(0) @@ -574,8 +648,8 @@ def test_moe_backend( activation_type=activation_type, ) - # W4A8_MXFP4_MXFP8 requires different weights for backend and reference - # due to different padding/alignment requirements + # W4A8_MXFP4_MXFP8 requires backend-layout-aware weights. CUTLASS and + # MegaMoE use 128 hidden alignment; TRTLLMGen pads FC1 input to 512. ref_cls = quant_kwargs.pop("ref_cls", None) ref_module_kwargs = {} if quant_algo == QuantAlgo.W4A8_MXFP4_MXFP8: diff --git a/tests/unittest/_torch/modules/moe/test_moe_module.py b/tests/unittest/_torch/modules/moe/test_moe_module.py index 02bf8f3539f1..eff8dafd7cea 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_module.py +++ b/tests/unittest/_torch/modules/moe/test_moe_module.py @@ -30,6 +30,7 @@ import logging import os import pickle +import socket import sys import tempfile import traceback @@ -40,6 +41,7 @@ import cloudpickle import pytest import torch +import torch.distributed as dist from _torch.modules.moe.moe_test_utils import ( IS_CI_MODE, MoeBackendType, @@ -52,6 +54,7 @@ should_skip_cutlass, should_skip_deepgemm, should_skip_densegemm, + should_skip_megamoe, should_skip_multi_gpu, should_skip_to_accelerate_ci, should_skip_trtllm, @@ -115,6 +118,30 @@ ) +def _get_free_tcp_port() -> int: + """Return a local TCP port for MPI-worker torch.distributed rendezvous.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _ensure_dist_for_megamoe(moe_backend: str, rank: int, world_size: int) -> None: + """MegaMoE resolves an EP ProcessGroup at construction time.""" + if moe_backend != MoeBackendType.MEGAMOE.value: + return + if not torch.cuda.is_available(): + pytest.skip("CUDA required for MegaMoE tests") + if dist.is_initialized(): + return + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29561") + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["LOCAL_RANK"] = str(rank) + torch.cuda.set_device(rank) + dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) + + def _create_mapping_for_parallel_mode(world_size, parallel_mode): """Create Mapping for different parallelism strategies. @@ -479,6 +506,7 @@ def _test_moe_worker_impl( mapping.rank = mpi_rank() all_rank_num_tokens = [seq_len] * mapping.world_size torch.cuda.set_device(mapping.rank) + _ensure_dist_for_megamoe(moe_backend, mapping.rank, mapping.world_size) with torch.device(f"cuda:{mapping.rank}"): torch.manual_seed(0) @@ -525,14 +553,8 @@ def _test_moe_worker_impl( swiglu_limit=swiglu_limit if swiglu_gptoss_style else None, num_local_experts=num_local_experts, ) - weights = quantize_util.create_weights(**quant_kwargs) - - # For EPLB, keep weights on CPU - if enable_eplb: - for key in weights: - if isinstance(weights[key], torch.Tensor): - weights[key] = weights[key].to("cpu") - ref_weights = copy.deepcopy(weights) if enable_eplb else weights + ref_cls = quant_kwargs.pop("ref_cls", None) + ref_module_kwargs = {} # Use a small max_num_tokens for unit tests to avoid NVSHMEM buffer # allocation failures. DeepEP low-latency buffers are sized by @@ -581,6 +603,28 @@ def _test_moe_worker_impl( weight_loading_mode=weight_loading_mode, ) as fused_moe, ): + # W4A8_MXFP4_MXFP8 needs backend-layout-aware weights. In + # particular, MegaMoEDeepGemm and TRTLLMGen can have different + # padded backend/ref tensor layouts, so create weights after the + # backend exposes its quant_method shapes. + if quant_algo == QuantAlgo.W4A8_MXFP4_MXFP8: + ( + weights, + ref_weights, + ref_module_kwargs, + ) = quantize_util.prepare_weights_from_backend(fused_moe, **quant_kwargs) + else: + weights = quantize_util.create_weights(**quant_kwargs) + ref_weights = weights + + # For EPLB, keep backend weights on CPU. + if enable_eplb: + for key, value in weights.items(): + if isinstance(value, torch.Tensor): + weights[key] = value.to("cpu") + if ref_weights is weights: + ref_weights = copy.deepcopy(weights) + fused_moe.load_weights([weights]) fused_moe.post_load_weights() fused_moe.cuda(f"cuda:{mapping.rank}") @@ -599,7 +643,12 @@ def _test_moe_worker_impl( G_LOGGER.info(f"[EPLB Debug] Initial expert_ids (after init): {initial_expert_ids}") # Create reference module - ref_fused_moe = quantize_util.create_ref_module(routing_method) + if ref_cls is not None: + ref_fused_moe = quantize_util.create_ref_module( + routing_method, ref_cls=ref_cls, **ref_module_kwargs + ) + else: + ref_fused_moe = quantize_util.create_ref_module(routing_method, **ref_module_kwargs) ref_fused_moe.moe_tp_size = mapping.moe_tp_size ref_fused_moe.load_weights([ref_weights]) ref_fused_moe.cuda(f"cuda:{mapping.rank}") @@ -689,19 +738,26 @@ def _test_moe_multi_gpu( swiglu_limit: SwiGLU limit parameter (default=inf, non-gptoss) """ - def init_worker(custom_paths, comm_method_type): + def init_worker(custom_paths, comm_method_type, master_port): # Update the sys.path to align with main process for submodule import for custom_path in custom_paths: if custom_path.endswith("tests/unittest") and custom_path not in sys.path: sys.path.append(custom_path) - # Set comm method - os.environ["TRTLLM_FORCE_COMM_METHOD"] = comm_method_type + if comm_method_type == MEGAMOE_DEEPGEMM_IGNORE_COMM_METHOD: + os.environ.pop("TRTLLM_FORCE_COMM_METHOD", None) + else: + os.environ["TRTLLM_FORCE_COMM_METHOD"] = comm_method_type + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ["MASTER_PORT"] = str(master_port) mapping = _create_mapping_for_parallel_mode(world_size, parallel_mode) + master_port = _get_free_tcp_port() with MPIPoolExecutor( - initializer=init_worker, initargs=(sys.path, comm_method_type), max_workers=world_size + initializer=init_worker, + initargs=(sys.path, comm_method_type, master_port), + max_workers=world_size, ) as executor: results = executor.map( _test_moe_worker, @@ -756,6 +812,7 @@ def init_worker(custom_paths, comm_method_type): MoeBackendType.CUTEDSL, MoeBackendType.DEEPGEMM, MoeBackendType.DENSEGEMM, + MoeBackendType.MEGAMOE, ] # Data types to test @@ -778,6 +835,7 @@ def init_worker(custom_paths, comm_method_type): LOCAL_MOE_MODEL_CONFIGS = CI_MOE_MODEL_CONFIGS + [ MoeModelConfig(64, 6, 2048, 1408), # DeepSeek-MoE-16B / DeepSeek-V2-Lite + MoeModelConfig(256, 6, 4096, 2048), # DeepSeek-V4-Flash MoeModelConfig(384, 8, 7168, 2048), # Kimi-K2 # === Boundary Tests: num_experts / top_k === MoeModelConfig(4, 4, 512, 512), # top_k=num_experts, all experts activated @@ -830,6 +888,9 @@ def init_worker(custom_paths, comm_method_type): "DEEPEPLOWLATENCY", ] +MEGAMOE_DEEPGEMM_IGNORE_COMM_METHOD = "IGNORE" +MEGAMOE_DEEPGEMM_COMM_METHODS = [MEGAMOE_DEEPGEMM_IGNORE_COMM_METHOD] +MEGAMOE_DEEPGEMM_PARALLEL_MODES = ["DEP"] if IS_CI_MODE else ["DEP", "TEP"] # SwiGLU parameters for swiglu_gptoss_style testing SWIGLU_ALPHAS = [1, 1.702] # default, GPT-OSS (modeling_gpt_oss.py) SWIGLU_BETAS = [0, 1.0] # default, GPT-OSS @@ -907,6 +968,53 @@ def _get_comm_method_skip_reason( return None +def should_skip_MegaMoEDeepGemm( + parallel_mode: str, + comm_method: str, + backend_type: MoeBackendType, + quant_algo: Optional[QuantAlgo], + dtype: torch.dtype, + model_config: MoeModelConfig, + routing_method_cls, + swiglu_gptoss_style: bool, +) -> Optional[str]: + """Check MegaMoEDeepGemm constraints for module-level multi-GPU tests.""" + if backend_type != MoeBackendType.MEGAMOE: + return None + + if comm_method != MEGAMOE_DEEPGEMM_IGNORE_COMM_METHOD: + return ( + "MegaMoEDeepGemm uses DeepGEMM internal EP communication; " + f"use comm={MEGAMOE_DEEPGEMM_IGNORE_COMM_METHOD} instead of " + f"forcing {comm_method}." + ) + + if parallel_mode not in ("DEP", "TEP"): + return f"MegaMoEDeepGemm Phase 1 is MoE-EP only (got {parallel_mode})" + + base_reason = should_skip_megamoe( + backend_type, + quant_algo=quant_algo, + dtype=dtype, + model_config=model_config, + moe_tp_size=1, + swiglu_gptoss_style=swiglu_gptoss_style, + ) + if base_reason: + return base_reason + + # The DeepGEMM mega kernel consumes precomputed top-k expert ids and + # routing weights. Routing itself runs before the kernel, so module tests + # can cover the routing methods already used by the multi-GPU matrix. + if routing_method_cls not in (RenormalizeMoeRoutingMethod, DeepSeekV3MoeRoutingMethod): + return ( + "MegaMoEDeepGemm module multi-GPU coverage is limited to " + "Renormalize and DeepSeekV3 routing methods" + ) + + return None + + def generate_multi_gpu_test_params( parallel_modes, comm_methods, @@ -1001,6 +1109,18 @@ def generate_multi_gpu_test_params( moe_tp_size=moe_tp_size, parallel_mode=parallel_mode, ), + should_skip_megamoe( + backend_type, + quant_algo=quant_algo, + dtype=dtype, + model_config=model_config, + comm_method=comm_method, + moe_tp_size=moe_tp_size, + parallel_mode=parallel_mode, + swiglu_gptoss_style=swiglu_alpha != 1 + or swiglu_beta != 0 + or swiglu_limit != float("inf"), + ), should_skip_multi_gpu( parallel_mode, model_config, world_size=4, comm_method=comm_method ), @@ -1031,6 +1151,74 @@ def generate_multi_gpu_test_params( return params +def generate_megamoe_deepgemm_multi_gpu_test_params() -> List: + """Generate focused MegaMoEDeepGemm module multi-GPU coverage.""" + params: List = [] + seq_lens = [8] if IS_CI_MODE else SEQ_LENS + + for parallel_mode, comm_method in product( + MEGAMOE_DEEPGEMM_PARALLEL_MODES, MEGAMOE_DEEPGEMM_COMM_METHODS + ): + for ( + swiglu_alpha, + swiglu_beta, + swiglu_limit, + model_config, + seq_len, + dtype, + backend_type, + quant_algo, + routing_method_cls, + skip_reason, + base_test_id, + ) in iter_base_test_configs( + [(1, 0, float("inf"))], + MOE_MODEL_CONFIGS, + seq_lens, + [torch.bfloat16], + [MoeBackendType.MEGAMOE], + [QuantAlgo.W4A8_MXFP4_MXFP8], + MULTI_GPU_ROUTING_METHODS, + ): + if not skip_reason: + skip_reason = should_skip_MegaMoEDeepGemm( + parallel_mode, + comm_method, + backend_type, + quant_algo, + dtype, + model_config, + routing_method_cls, + swiglu_alpha != 1 or swiglu_beta != 0 or swiglu_limit != float("inf"), + ) + + if not skip_reason: + skip_reason = should_skip_multi_gpu( + parallel_mode, model_config, world_size=4, comm_method=comm_method + ) + + if skip_reason: + continue + + test_id = f"parallel={parallel_mode}-comm={comm_method}-{base_test_id}" + param_values = ( + parallel_mode, + comm_method, + dtype, + backend_type.value, + quant_algo, + seq_len, + model_config, + routing_method_cls, + swiglu_alpha, + swiglu_beta, + swiglu_limit, + ) + params.append(create_test_param(param_values, test_id)) + + return params + + def generate_base_test_params( swiglu_combos, model_configs, seq_lens, dtypes, backend_types, quant_algos, routing_methods ) -> List: @@ -1287,6 +1475,7 @@ def test_trtllm_gen_fp32_routing_bias(routing_method_cls, moe_model_config, quan quant_algos=QUANT_ALGOS, routing_methods=MULTI_GPU_ROUTING_METHODS, ) +MULTI_GPU_TEST_PARAMS += generate_megamoe_deepgemm_multi_gpu_test_params() @pytest.mark.skipif(torch.cuda.device_count() < 4, reason="needs 4 GPUs to run this test") @@ -1553,16 +1742,94 @@ def generate_eplb_test_params( return params +def generate_megamoe_deepgemm_eplb_test_params() -> List: + """Generate focused dynamic-EPLB params for MegaMoEDeepGemm.""" + params: List = [] + ep_size = 4 + + for parallel_mode, comm_method, num_slots in product( + EPLB_PARALLEL_MODES, MEGAMOE_DEEPGEMM_COMM_METHODS, EPLB_NUM_SLOTS_LIST + ): + for ( + swiglu_alpha, + swiglu_beta, + swiglu_limit, + model_config, + _seq_len, + dtype, + backend_type, + quant_algo, + routing_method_cls, + skip_reason, + base_test_id, + ) in iter_base_test_configs( + [(1, 0, float("inf"))], + EPLB_MODEL_CONFIGS, + [8], + [torch.bfloat16], + [MoeBackendType.MEGAMOE], + [QuantAlgo.W4A8_MXFP4_MXFP8], + EPLB_ROUTING_METHODS, + ): + if not skip_reason: + skip_reason = should_skip_MegaMoEDeepGemm( + parallel_mode, + comm_method, + backend_type, + quant_algo, + dtype, + model_config, + routing_method_cls, + swiglu_alpha != 1 or swiglu_beta != 0 or swiglu_limit != float("inf"), + ) + + if not skip_reason and num_slots <= model_config.num_experts: + skip_reason = ( + f"EPLB requires num_slots ({num_slots}) > " + f"num_experts ({model_config.num_experts})" + ) + + if not skip_reason and num_slots % ep_size != 0: + skip_reason = ( + f"MegaMoEDeepGemm requires num_slots ({num_slots}) " + f"divisible by ep_size ({ep_size})." + ) + + if skip_reason: + continue + + test_id = ( + f"parallel={parallel_mode}-comm={comm_method}-{base_test_id}-" + f"slots={num_slots}-eplb=dynamic" + ) + param_values = ( + parallel_mode, + comm_method, + dtype, + backend_type.value, + quant_algo, + model_config, + num_slots, + routing_method_cls, + ) + params.append(create_test_param(param_values, test_id)) + + return params + + # Pre-generate EPLB test parameters at module load time -EPLB_TEST_PARAMS = generate_eplb_test_params( - parallel_modes=EPLB_PARALLEL_MODES, - comm_methods=EPLB_COMM_METHODS, - model_configs=EPLB_MODEL_CONFIGS, - num_slots_list=EPLB_NUM_SLOTS_LIST, - dtypes=DTYPES, - backend_types=BACKEND_TYPES, - quant_algos=QUANT_ALGOS, - routing_methods=EPLB_ROUTING_METHODS, +EPLB_TEST_PARAMS = ( + generate_eplb_test_params( + parallel_modes=EPLB_PARALLEL_MODES, + comm_methods=EPLB_COMM_METHODS, + model_configs=EPLB_MODEL_CONFIGS, + num_slots_list=EPLB_NUM_SLOTS_LIST, + dtypes=DTYPES, + backend_types=[b for b in BACKEND_TYPES if b != MoeBackendType.MEGAMOE], + quant_algos=QUANT_ALGOS, + routing_methods=EPLB_ROUTING_METHODS, + ) + + generate_megamoe_deepgemm_eplb_test_params() )