diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index fc9925c6a749..43158e23f931 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -360,6 +360,30 @@ def resolve_moe_backend(moe_backend: str, quant_config: Optional[QuantConfig] = None) -> str: """Resolve AUTO moe_backend to a specific backend based on model architecture. + **Not the implementation-selection entry point.** That is + ``moe_resolution.resolve_moe_impl``, and the two run in different + phases on different questions. This one turns the literal ``AUTO`` into + a concrete backend name while the checkpoint is being read; the other + turns a concrete backend name into an impl class while a layer is being + built, by asking each candidate's ``can_implement``. + + The phases cannot be merged, and the reason is a genuine cycle rather + than an accident of layering: several quant formats pick their + ``quant_algo`` from the backend name (see ``get_mxfp4_quant_algo`` and + ``load_hf_quant_config``), so a backend name is needed to finish + building ``quant_config`` -- while ``resolve_moe_impl`` needs a + finished ``quant_config`` to state the problem at all. Hence the + deliberate two-step in ``from_pretrained``: an architecture-only hint + first, then a quant-aware resolution once ``quant_config`` exists. + + What this must therefore never grow is capability knowledge. Every + rule here is a *preference* ("on Blackwell we would rather run + TRTLLM-Gen"), and preferences that turn out to be unservable are caught + downstream, where ``resolve_moe_impl`` records the substitution in a + ``MoEResolutionReport``. A "can it run" test added here would be a + second copy of a gate that already exists in a ``can_implement``, and + the two copies would drift. + Args: moe_backend: The configured moe_backend (may be "AUTO") architecture: The model architecture name (e.g., "GptOssForCausalLM") diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 7bbcd0ce9186..da7c281238c2 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -73,7 +73,7 @@ TritonFusedMoE, TRTLLMGenFusedMoE, create_moe, - get_moe_cls, + resolve_moe_cls, ) from ..modules.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE from ..modules.fused_moe.fused_moe_wide_ep import WideEPMoE @@ -1504,10 +1504,22 @@ def __init__( moe_swiglu_limit = None if swiglu_limit is not None: # `create_moe` only accepts swiglu_limit for these MoE classes; - # resolve via get_moe_cls so backend-string fallbacks (e.g. - # TRTLLM/CUTEDSL/DENSEGEMM dropping back to CutlassFusedMoE on - # unsupported quant) are handled correctly. - moe_cls = get_moe_cls(model_config, override_quant_config=experts_quant_config) + # ask the resolver rather than the backend string so that a + # degradation (e.g. TRTLLM/CUTEDSL/DENSEGEMM dropping back to + # CutlassFusedMoE on unsupported quant) is accounted for here too. + moe_cls = resolve_moe_cls( + model_config, + override_quant_config=experts_quant_config, + dtype=dtype, + # Same routing object as create_moe below. + routing=self.gate.routing_method, + # create_moe below passes no bias and no swiglu alpha/beta, so + # it resolves with the plain SwiGLU package. Say so here too: + # leaving this unknown lets gates abstain that create_moe + # rejects, and the two calls would pick different backends. + swiglu_gptoss_style=False, + layer_idx=layer_idx, + ) supports_swiglu_limit = moe_cls in ( CutlassFusedMoE, TritonFusedMoE, diff --git a/tensorrt_llm/_torch/models/modeling_laguna.py b/tensorrt_llm/_torch/models/modeling_laguna.py index ce170899f899..03257d22cedf 100644 --- a/tensorrt_llm/_torch/models/modeling_laguna.py +++ b/tensorrt_llm/_torch/models/modeling_laguna.py @@ -34,7 +34,12 @@ from ..modules.attention import _helix_cp_allgather_input, _helix_cp_output_projection from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding -from ..modules.fused_moe import MiniMaxM2MoeRoutingMethod, create_moe, get_moe_cls +from ..modules.fused_moe import ( + MiniMaxM2MoeRoutingMethod, + RoutingMethodType, + create_moe, + resolve_moe_cls, +) from ..modules.fused_moe.interface import MoE, MoEWeightLoadingMode from ..modules.fused_moe.interface import MoE as MoEInterface from ..modules.gated_mlp import GatedMLP @@ -127,7 +132,15 @@ def __init__(self, model_config, layer_idx, aux_stream_dict): num_experts=self.num_experts, top_k=self.top_k, dtype=config.torch_dtype, - moe_backend_cls=get_moe_cls(model_config), + moe_backend_cls=resolve_moe_cls( + model_config, + routing=RoutingMethodType.MiniMax2, + # Match the create_moe call below, which passes no bias and no + # swiglu alpha/beta. Left unknown, gates that create_moe + # rejects abstain here and the gate would name a backend the + # layer does not run. + swiglu_gptoss_style=False, + ), ) self.experts = create_moe( diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py index 571e3fe503c0..32f38cc7ec31 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py @@ -17,7 +17,7 @@ RenormalizeMoeRoutingMethod, RenormalizeNaiveMoeRoutingMethod, RoutingMethodType, TRTLLMGenFusedMoE, - create_moe, get_moe_cls) + create_moe, resolve_moe_cls) from ..modules.fused_moe.interface import MoE, MoEWeightLoadingMode from ..modules.linear import TensorParallelMode from ..modules.rms_norm import RMSNorm @@ -111,7 +111,16 @@ def __init__( dtype=config.torch_dtype, apply_routing=False, routing_method_type=RoutingMethodType.Renormalize, - moe_backend_cls=get_moe_cls(model_config, layer_idx=layer_idx), + moe_backend_cls=resolve_moe_cls( + model_config, + routing=RoutingMethodType.Renormalize, + # Match the create_moe call below, which passes no bias and no + # swiglu alpha/beta. Left unknown, gates that create_moe + # rejects abstain here and the gate would name a backend the + # layer does not run. + swiglu_gptoss_style=False, + layer_idx=layer_idx, + ), ) self.weight_loading_mode = MoEWeightLoadingMode.FUSED_GATE_UP_PROJ if config.model_type == "qwen3_vl_moe_text" else MoEWeightLoadingMode.VANILLA 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 7ca9a88c813c..86b9e97adf99 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -100,7 +100,7 @@ No external `Communication.dispatch` / `.combine`. Zero-token chunks still launc 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** — `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 +5. **Backend declares capabilities** — `can_implement(p, d)` is the single source of truth for what a backend supports, and it is a **pure function** of its two arguments (see [Backend Selection](#backend-selection)); the resolver holds no capability knowledge of its own 6. **Backend declares scheduler** — `scheduler_kind` class attribute selects the forward path; lifecycle code stays generic, forward path stays specialized ## Architecture Transition (IMPORTANT) @@ -132,7 +132,11 @@ Still on old path (standalone, with embedded communication): |------|------| | `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` | +| `create_moe.py` | Factory — builds the layer once `moe_resolution` has named the class | +| `moe_resolution.py` | **The one selection entry point** (`resolve_moe_impl`) — orders candidates, asks each one's `can_implement`, returns a `MoEResolutionReport` | +| `impl_contract.py` | Selection vocabulary — `MoEProblem`, `MoEDeployment`, `MoEEnvironment`, `MoEEligibility`, `MoERejectReason`, `MoEResolutionReport` | +| `impl_environment.py` | The only place that probes the machine (SM, optional wheels, env flags) and freezes the result | +| `impl_identity.py` | `MoEImplId` / `MoEImplDescriptor` / registry — the stable one-id-per-leaf-class mechanism used after an implementation migrates | | `interface.py` | Base class `MoE` and enums (`MoEWeightLoadingMode`, `MoESchedulerKind`, `AlltoallMethodType`) | | `quantization.py` | Quantization method implementations (`FusedMoEMethod` subclasses: weight creation, loading, quant/dequant ops per quant mode) | | `routing.py` | Routing methods (`TopKRouting`, etc.) | @@ -144,16 +148,16 @@ Still on old path (standalone, with embedded communication): | 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_trtllm_gen.py` | `TRTLLMGenFusedMoE` | SM100/SM103 | Min-latency and high-throughput on Blackwell; also serves unquantized BF16 through FlashInfer's `trtllm_bf16_moe` (gated on `MoEDep.FLASHINFER_BF16_MOE`, not on a quant algo) | `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` | -| `fused_moe_cute_dsl_b12x.py` | `CuteDslB12xFusedMoE` | SM120/SM121 | NVFP4 hybrid CUTLASS-prefill / FlashInfer NVFP4 MoE decode — best perf on RTX PRO 6000 (SM120) and DGX Spark (SM121); select via the `CUTEDSL` backend path (auto-promoted when flashinfer is importable) | `EXTERNAL_COMM` | +| `fused_moe_cute_dsl_b12x.py` | `CuteDslB12xFusedMoE` | SM120/SM121 | NVFP4 hybrid CUTLASS-prefill / FlashInfer NVFP4 MoE decode — best perf on RTX PRO 6000 (SM120) and DGX Spark (SM121); select via the `CUTEDSL` backend path (it heads that family's candidate list, so it wins on SM120/121 when flashinfer is present and yields to `CuteDslFusedMoE` otherwise); single-GPU-shaped topology only — it rejects both `ep_size > 1` and attention-DP, because it has no dispatch/combine kernel and has never been exercised behind a DP allgather | `EXTERNAL_COMM` | | `mega_moe/mega_moe_deepgemm.py` | `MegaMoEDeepGemm` | SM100/SM103 | W4A8_MXFP4_MXFP8 via DeepGEMM `fp8_fp4_mega_moe` fused dispatch+GEMM+act+GEMM+combine kernel; requires `hidden_size % 512 == 0` | `FUSED_COMM` | | `mega_moe/mega_moe_cute_dsl.py` | `MegaMoECuteDsl` | SM100/SM103 | NVFP4 via ported CuteDSL `Sm100MegaMoEKernel` fused dispatch+FC1+act+FC2+combine kernel; requires CUDA 13 Cutlass DSL runtime (PR #14354) and NVSHMEM provider (hard gate); threads per-expert `fc31_alpha`/`fc2_alpha`/`fc1_norm_const` through the kernel ABI and supports SwiGLU clamp via `swiglu_limit`; default deepgemm graph (topk score folded before fc1-out quant, host `combine_output.sum(dim=1)`) | `FUSED_COMM` | -| `fused_moe_marlin.py` | `MarlinFusedMoE` | SM89-SM99 | W4A16 NVFP4 on Ada/Hopper (BF16 activations + FP4 weights, fused single-launch `marlin_nvfp4_moe_gemm` kernel); supports attention-DP + EP via external comm (scheduler precomputes routing; dispatch payload is plain BF16, no activation scales); non-NVFP4 layers (e.g. unquantized MTP draft layers) fall back to Cutlass in `get_moe_cls`; no dynamic EPLB | `EXTERNAL_COMM` | +| `fused_moe_marlin.py` | `MarlinFusedMoE` | SM89-SM99 | W4A16 NVFP4 on Ada/Hopper (BF16 activations + FP4 weights, fused single-launch `marlin_nvfp4_moe_gemm` kernel); supports attention-DP + EP via external comm (scheduler precomputes routing; dispatch payload is plain BF16, no activation scales); non-NVFP4 layers (e.g. unquantized MTP draft layers) degrade to Cutlass in `resolve_moe_impl`, recorded in the layer's `MoEResolutionReport`; no dynamic EPLB | `EXTERNAL_COMM` | | `fused_moe_triton.py` | `TritonFusedMoE` | SM90 only | GPT-OSS on Hopper (requires `swiglu_gptoss_style=True`) | (legacy path) | -| `fused_moe_wide_ep.py` | `WideEPMoE` | All GPUs | Deprecated — `create_moe.py` rejects the `WIDEEP` backend. Wide EP and EPLB are available on the other backends: use `DEEPGEMM` for FP8 block-scale checkpoints, or `TRTLLM` / `CUTEDSL` / `CUTLASS` otherwise. Class kept for reference only | (legacy path) | +| `fused_moe_wide_ep.py` | `WideEPMoE` | All GPUs | Deprecated — `moe_resolution.py` raises on the `WIDEEP` backend literal. Wide EP and EPLB are available on the other backends: use `DEEPGEMM` for FP8 block-scale checkpoints, or `TRTLLM` / `CUTEDSL` / `CUTLASS` otherwise. Class kept for reference only | (legacy path) | | `fused_moe_vanilla.py` | `VanillaMoE` | All devices | Reference / debugging only | (legacy path) | ### Communication (`fused_moe/communication/`) @@ -193,18 +197,126 @@ is available. | `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 Selection + +One function decides which implementation runs: `moe_resolution.resolve_moe_impl`. +It owns no capability knowledge. It orders candidates, asks each one, and returns +the first that accepts — so "which impl will run" has a single answer no matter who +asks, and the factory can no longer admit something a backend rejects. + +```python +report = resolve_moe_impl(model_config, layer_idx=layer_idx) +impl_cls = impl_class_for(report) # raises, with the full trail, if nothing fits +``` + +### `can_implement` must be pure + +```python +@classmethod +def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + if d.env.sm not in (100, 103): + return _reject(MoERejectReason.SM_UNSUPPORTED, f"... got SM{d.env.sm}") + if not d.env.has_dep(MoEDep.FLASHINFER): + return _reject(MoERejectReason.DEP_MISSING, "... requires flashinfer") + return MoEEligibility.ok() +``` + +Read `p` and `d`, nothing else. Specifically: **no** `get_sm_version()`, **no** +`import` to test whether a wheel exists, **no** `os.environ`. Those probes happen +once in `impl_environment.collect_moe_environment()` and arrive frozen as +`d.env`. A gate that probes the host instead answers a different question on +every machine, which is exactly the irreproducibility the frozen environment +exists to remove. + +The same rule is why a gate reads `d.eplb_enabled` and `d.moe_lora_enabled` +instead of `self._supports_load_balancer()` — a predicate on `self` can only be +consulted after the object it might have to reject already exists. + +Unknown is not false. A shape field can be `None` when the caller does not know +it yet, and a gate that reads `None` must abstain rather than reject, or a +missing `pretrained_config` attribute turns into a backend downgrade. + +Abstaining has no state of its own: `MoEEligibility` is two-valued, so a gate +abstains by skipping its check and returning `MoEEligibility.ok()`. Read that +literally — `resolve_moe_impl` counts the candidate as eligible and may pick it +without the shape constraint ever having been proven. That is the accepted cost: +rejecting on absent information downgrades a backend that would have been +perfectly legal, and a caller who wants the constraint checked has to supply the +shape. + +### Adding a new probe + +Add a member to `MoEDep` or `MoEEnvFlag` and a probe function to the table in +`impl_environment.py`. Both enums are closed on purpose: a name not declared +there cannot be read during selection, which is what keeps the environment an +explicit input instead of a growing set of implicit ones. + +### Degradation is allowed, silence is not + +`moe_backend` is a **preference, not a pin**. A requested backend that cannot +serve the layer is turned down and a substitute runs — production depends on +this, because an unquantized MTP draft layer in a MIXED_PRECISION checkpoint +must not take down a model whose other layers are NVFP4. + +What is not allowed is doing it quietly. Every resolution returns a +`MoEResolutionReport` naming the winner, every rejected candidate, its +`MoERejectReason`, and the environment fingerprint; a degradation additionally +logs a warning once per layer. `report.degraded_from` is the answer to "why did +my `moe_backend` not take effect". + +Until the one-class-per-implementation migration is complete, the report is +diagnostic rather than pinnable: it records the legacy backend class name, +`problem.quant`, and `deployment.env.env_flags`. Legacy classes deliberately do +not synthesize a `MoEImplId`, because one class still spans several quantization +formats. A canonical ID is attached only when a leaf class owns one fixed +`MoEImplDescriptor.identity`. + +Two things fail hard, and they fail in different places. An unknown or retired +backend literal raises `ValueError` before any candidate is considered, so there +is no report to inspect — a misspelled backend is a config error to fix, not +something to route around. Nothing being able to serve the layer is the opposite: +`resolve_moe_impl` still returns a full report, with `winner is None` and every +rejection recorded, and `impl_class_for` is what raises. Catch the first, read +the second. + +`NO_FALLBACK_BACKENDS` is the one exception to "degradation is allowed", and +`VANILLA` is its only member. Vanilla exists to produce reference numerics, so a +caller comparing a kernel against it is not helped by getting Cutlass back with a +warning — that silently answers a different question than the one asked. A +`VANILLA` request whose gates reject therefore takes the `winner is None` path and +raises with the whole rejection trail. + +### Cross-rank agreement + +Every rank resolves independently, so a wheel installed on some nodes and not +others makes ranks pick different impls. The symptom is a hang, not an error: +the ranks allocate differently shaped expert weights and then wait for each +other in a collective that no longer matches. + +Selection does not police this itself. `MoEEnvironment.fingerprint()` is recorded +in every `MoEResolutionReport`, so comparing two ranks' reports names the +divergence immediately — but that is a diagnosis after the fact, not a guard. +An automatic check has to be a collective, and selection is the wrong place to +start one: `resolve_moe_impl` runs per MoE layer, so its participants are +"the ranks that happen to build a MoE layer", which under pipeline parallelism +is not every rank (a stage holding only dense layers never calls it). A +collective entered by that set on the world group deadlocks. If such a check is +added later it belongs at an initialization point every rank reaches +unconditionally, not 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. +Each backend's `can_implement(p, d)` classmethod declares what it supports. Source of truth: the `can_implement` classmethod in each backend file. -| Quantization | Cutlass | TRTLLMGen | DeepGemm | DenseGEMM | CuteDSL | MegaMoE-DG | MegaMoE-CuteDSL | Triton | Marlin | WideEP (deprecated) | Vanilla | +| Quantization | Cutlass | TRTLLMGen | DeepGemm | DenseGEMM | CuteDSL | MegaMoE-DG | MegaMoE-CuteDSL | Triton | Marlin | WideEP (retired)† | Vanilla | |---|---|---|---|---|---|---|---|---|---|---|---| -| Unquantized (BF16/FP16) | Y (SM80+) | N | N | N | N | N | N | Y (SM90, BF16) | N | Y | Y | +| Unquantized (BF16/FP16) | Y (SM80+) | Y (SM100/103, BF16, needs FlashInfer `trtllm_bf16_moe`)§ | N | N | N | N | N | Y (SM90, BF16) | N | Y | Y | | FP8 QDQ | Y (SM89+) | N | N | N | N | N | N | Y (SM90) | N | Y | Y | -| FP8 Block Scales | Y (SM90, SM120) | Y (SM100/103) | Y (SM100/103) | N | Y (SM100/103) | N | N | N | N | Y | Y | -| NVFP4 | Y (SM100/103/120/121) | Y (SM100/103) | N | Y (SM100/103) | Y (SM100/103/120/121) | N | Y (SM100/103, cu13 cutlass-dsl + NVSHMEM provider; per-expert alpha/norm_const + SwiGLU clamp) | N | Y (SM89-SM99, W4A16) | Y | Y | +| FP8 Block Scales | Y (SM90, SM120) | Y (SM100/103) | Y (SM100/103) | N | N‡ | N | N | N | N | Y | Y | +| NVFP4 | Y (SM100/103/120/121) | Y (SM100/103) | N | Y (SM100/103) | Y (SM100/103/120/121) | N | Y (SM100/103, cu13 cutlass-dsl + NVSHMEM provider; per-expert alpha/norm_const + SwiGLU clamp) | N | Y (SM89-SM99) | Y | Y | +| W4A16 NVFP4 | Y (SM80+, dequant-on-the-fly) | N | N | N | Y (SM120/121 via `CuteDslB12xFusedMoE`, needs flashinfer) | N | N | N | Y (SM89-SM99, BF16) | N | Y | | W4A8 NVFP4 FP8 | N | Y (SM100/103) | N | N | N | N | N | N | N | N | N | | W4A16 MXFP4 | Y (SM90) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | N | | W4A8 MXFP4 FP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | N | @@ -214,6 +326,93 @@ Each backend's `can_implement(quant_algo, dtype_activation, swiglu_gptoss_style, | W8A16 | Y (SM80+) | N | N | N | N | N | N | N | N | N | N | | INT4 WoQ (W4AFP8) | N | N | N | N | N | N | N | N | N | Y | N | +† The `WideEP` column is history, not an option. `resolve_moe_impl` raises on the +`WIDEEP` literal, so nothing in that column can be selected — read a `Y` there as +"the retired class implemented this", never as "you may request this". For +`INT4 WoQ (W4AFP8)` that leaves the row without a selectable backend. + +§ The unquantized `TRTLLMGenFusedMoE` path is not a TRTLLM-Gen kernel at all: it +calls FlashInfer's `trtllm_bf16_moe` / `trtllm_bf16_routed_moe`, which is why it +is gated on `MoEDep.FLASHINFER_BF16_MOE` rather than on a quant algo, and why +`TRTLLMOpBackend` raises `NotImplementedError` for it. The row reads `Y` because +`can_implement` really can select it; without the FlashInfer symbols the layer +degrades to Cutlass with `DEP_MISSING` recorded in `degraded_from`, where the +pre-resolver code raised `RuntimeError` instead. The same path also requires +`intermediate_size_per_partition % 128 == 0` (`Bf16MoeLauncher::check_moe`); +a non-aligned shard is `SHAPE_UNALIGNED` and falls back to Cutlass. + +Cutlass covers `W4A16 NVFP4` on a wider SM range than plain `NVFP4` because the +two run different kernels: `W4A16NVFP4CutlassFusedMoEMethod` dequantizes the FP4 +weights into the activation dtype each forward and then calls the unquantized +kernel, so it inherits that path's `SM80+` floor instead of needing NVFP4 +tensor cores. This is what makes Cutlass the landing spot when a W4A16 NVFP4 +layer finds no specialized backend — `CuteDslB12xFusedMoE` without flashinfer, +or any SM outside Marlin's 89-99 and B12x's 120/121. + +‡ `CuteDslFusedMoE` has FP8-block-scale *plumbing* but no FP8-block-scale kernel: +`run_moe_fp8_block_scales` ends in `cute_dsl_fp8_group_blockwise_gemm_ref`, a +local pure-PyTorch helper that upcasts to fp32, materializes the expanded scales, +and loops `torch.einsum` per expert. The only CuteDSL runners the file imports +are the `Sm100BlockScaledContiguous*` NVFP4 ones. `can_implement` therefore +declines FP8 block scales rather than claiming a reference path as a backend, and +the algorithm's real owners are `DeepGemmFusedMoE` / `TRTLLMGenFusedMoE` on +SM100/103 and Cutlass on SM90/SM120. Consequence worth knowing before changing +this: because a `CUTEDSL` request only ever considers the CuteDSL family plus the +Cutlass fallback, and Cutlass's FP8-block kernel stops at SM90/SM120, an explicit +`moe_backend="CUTEDSL"` on an FP8-block checkpoint at SM100 now fails resolution +instead of silently running the reference GEMM. `test_cute_dsl_fp8_block_scales` +and `test_cute_dsl_fp8_block_scales_4gpus` in +`tests/integration/defs/accuracy/test_llm_api_pytorch.py` are exactly that +configuration and are `pytest.mark.skip`-ed for this reason; they were never +scheduled in any `tests/integration/test_lists/` entry, so the skip changes no CI +stage. Point them at `DEEPGEMM` / `TRTLLM` if this checkpoint needs coverage +again. + +### Activation Support + +The matrix above is quantization only; activation style is a separate axis. The +gpt-oss SwiGLU package (per-expert bias plus `swiglu_alpha` / `swiglu_beta` / +`swiglu_limit`, surfaced as `MoEProblem.swiglu_gptoss_style`) is rejected by +every specialized backend — `CuteDslFusedMoE`, `CuteDslB12xFusedMoE`, +`DeepGemmFusedMoE`, `DenseGEMMFusedMoE`, `MarlinFusedMoE` — while +`TRTLLMGenFusedMoE` accepts only the algorithms in its `_GPTOSS_SUPPORTED_ALGOS`. + +Cutlass gates gpt-oss / MiniMax SwiGLU on unquantized, NVFP4, and the MXFP4 +family (`CutlassFusedMoE._GPTOSS_SUPPORTED_ALGOS` = `None`, `NVFP4`, +`W4A16_MXFP4`, `W4A8_MXFP4_FP8`, `W4A8_MXFP4_MXFP8`). The CUDA kernel is not +the constraint — `torch.ops.trtllm.fused_moe` takes `swiglu_alpha` / +`swiglu_beta` / `swiglu_limit` on the same call for every path, including +NVFP4 (`CutlassMoeFCRunner<__nv_fp4_e2m1, __nv_fp4_e2m1>`), and TMA-WS GEMM1 +applies `SwigluBiasAdaptor` in `doActivation`. NVFP4 is eligible only when +there is no expert bias (`MoEProblem.bias is not True`): MiniMax-M3 NVFP4 +passes `ActivationType.SwigluBias` + alpha/beta/limit with `bias=False`. +gpt-oss 1-D bias still goes through `NVFP4CutlassFusedMoEMethod`'s 2-D +weight pad and is rejected at selection. Unquantized and the MXFP4 family +can load that 1-D bias. W8A16 / W4A8_AWQ stay rejected because they inherit +the base `w3_w1_weight_shape[:2]` default (wrong for transposed layouts). +Widening the set without that distinction converts a selection-time +rejection into a weight-loading crash. + +Three things make this easy to get wrong in either direction. First, the SM +asymmetry: `ModelConfig.get_mxfp4_quant_algo` maps a gpt-oss checkpoint to +`W4A16_MXFP4` below SM100 and to the `W4A8_MXFP4_*` pair at SM100+, so a gate +keyed on `W4A8_MXFP4_MXFP8` alone excludes Hopper entirely and — because Cutlass +is `FALLBACK_IMPL` and every other backend abstains — leaves gpt-oss unservable +there. Second, dropping the gate altogether is equally wrong: it un-skips the +`test_configurable_moe_single_gpu` gpt-oss × CUTLASS matrix, which then fails +inside weight loading rather than being rejected up front. Third, omitting +`None` rejects dummy / unquantized gpt-oss (`test_gpt_oss_trtllmgen[CUTLASS]`) +even though the kernel path is valid. Fourth, treating MiniMax SwigluBias as +"gpt-oss bias load" and excluding NVFP4 rejects +`TestMiniMaxM3::test_nvfp4` (`MoeConfig(backend="CUTLASS")`): MiniMax has no +expert bias, and the NVFP4 TMA-WS runner already applies `SwigluBiasAdaptor`. + +The unquantized `TRTLLMGenFusedMoE` FlashInfer path has a separate shape +gate: `Bf16MoeLauncher::check_moe` requires +`intermediate_size % 128 == 0`, and the wrapper passes +`intermediate_size_per_partition`. Qwen3.5-35B BF16 TP8 shards 512 → 64 and +must degrade to Cutlass instead of dying in the kernel launcher. + ### Scheduler / EPLB Constraints - `FUSED_COMM` backends (`MegaMoEDeepGemm`, `MegaMoECuteDsl`) **must not** layer host-side `Communication.dispatch` / `.combine` on top of the fused kernel — `ConfigurableMoE._create_comm_strategy_auto` returns `None` for them. @@ -231,7 +430,7 @@ When adding new components, use these reference implementations: | Task | Reference | Key methods to implement | |------|-----------|--------------------------| -| New `EXTERNAL_COMM` Backend | `fused_moe_cutlass.py` (`CutlassFusedMoE`) | `can_implement`, `run_moe`, `create_weights`, `load_weights` | +| New `EXTERNAL_COMM` Backend | `fused_moe_cutlass.py` (`CutlassFusedMoE`) | `capabilities`, `can_implement`, `run_moe`, `create_weights`, `load_weights`; then add the class to `moe_resolution.BACKEND_CANDIDATES`. Add a fixed `descriptor.identity` only for a one-implementation leaf class | | New `FUSED_COMM` Backend | `mega_moe/mega_moe_deepgemm.py` (`MegaMoEDeepGemm`), `mega_moe/mega_moe_cute_dsl.py` (`MegaMoECuteDsl`) | Same as above + override `scheduler_kind = MoESchedulerKind.FUSED_COMM` and `validate_configurable_moe` for backend-specific constraints. For NVFP4 CuteDSL specifically, mirror the `MegaMoECuteDsl` pattern: capability probe for the CUDA 13 Cutlass DSL runtime, JSON-friendly tactic dict, lazy kernel import via `cute_dsl_kernels/mega_moe_nvfp4/import_kernel()`, and `quantize_input` that short-circuits zero-token input. | | 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` | @@ -247,7 +446,10 @@ When adding new components, use these reference implementations: - **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 skip `can_implement()` checks** — Every backend must declare what it supports; an unsupported combination returns `MoEEligibility.no(MoERejectReason., detail)`, never a bare `False` and never a free-form string a test would have to pattern-match +- **Do NOT probe the machine inside `can_implement()`** — No `get_sm_version()`, no `import` as a presence test, no `os.environ`. Read `d.env`; add the probe to `impl_environment.py` if it does not exist yet +- **Do NOT add a second selection entry point** — `resolve_moe_impl` is the only one. A helper that picks a class on the side is how `get_moe_cls` and the old `resolve_moe_cls` drifted apart in the first place +- **Do NOT substitute a backend without recording it** — A degradation must be visible in the `MoEResolutionReport`, not only in a log line - **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 - **Do NOT allocate symmetric memory from `run_moe` in `FUSED_COMM` backends** — Symmetric-memory rendezvous is a build-time collective and is unsafe under PP / layer-skip or CUDA graph capture; allocate from `create_weights()` after `ConfigurableMoE` has synchronized EPLB-derived attributes. See `mega_moe/mega_moe_deepgemm.py` for the DG pattern and `mega_moe/mega_moe_cute_dsl.py:_alloc_symm_provider` for the NVSHMEM-equivalent provider. diff --git a/tensorrt_llm/_torch/modules/fused_moe/__init__.py b/tensorrt_llm/_torch/modules/fused_moe/__init__.py index f578feb247bf..1e99f8e21b2f 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/__init__.py +++ b/tensorrt_llm/_torch/modules/fused_moe/__init__.py @@ -1,5 +1,5 @@ from .configurable_moe import ConfigurableMoE -from .create_moe import create_moe, get_moe_cls +from .create_moe import create_moe, resolve_moe_cls, resolve_moe_impl from .fused_moe_cute_dsl import CuteDslFusedMoE from .fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE from .fused_moe_cutlass import CutlassFusedMoE @@ -36,7 +36,8 @@ "DeepSeekV3MoeRoutingMethod", "DefaultMoeRoutingMethod", "FusedMoEQuantScalesFP8", - "get_moe_cls", + "resolve_moe_cls", + "resolve_moe_impl", "Llama4RenormalizeMoeRoutingMethod", "MarlinFusedMoE", "LoadBalancedMoeRoutingMethod", diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index 17b209db1099..72e504b48f45 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -29,12 +29,18 @@ import copy from contextlib import contextmanager -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional, Type, Union import torch from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe.interface import MoE, MoESchedulerKind +from tensorrt_llm._torch.modules.fused_moe.impl_contract import ( + MoEDeployment, + MoEEligibility, + MoEProblem, + MoERejectReason, +) +from tensorrt_llm._torch.modules.fused_moe.interface import MoE, MoESchedulerKind, _reject 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 ( @@ -81,7 +87,8 @@ class ConfigurableMoE(MoE): This class orchestrates the MoE execution flow by composing: - moe_backend: Existing FusedMoE implementation used as a pluggable backend. - Currently supported backends (see ``create_moe.get_moe_cls``): + Currently supported backends (see + ``moe_resolution.IMPL_PRIORITY``): CutlassFusedMoE, TRTLLMGenFusedMoE, DeepGemmFusedMoE, CuteDslFusedMoE, DenseGEMMFusedMoE, MegaMoEDeepGemm. Note: Current FusedMoE implementations are used as backends (transitional). @@ -111,40 +118,29 @@ class ConfigurableMoE(MoE): Auto-Detection: - EPLB: Enabled if get_moe_load_balancer() is not None - - 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. + - Backend: Resolved from ``model_config.moe_backend`` by + ``moe_resolution.resolve_moe_impl``, which degrades to + CutlassFusedMoE when the requested backend cannot serve the + layer and records the reason in the returned report. + ``create_moe`` passes the resolved class in as ``moe_cls``; + constructing this wrapper directly resolves on demand. - Communication: Auto-selected based on hardware (NVLINK > DeepEP > AllGather); skipped entirely for FUSED_COMM backends (e.g. MegaMoEDeepGemm). """ @classmethod - def can_implement( - cls, - quant_algo, - dtype_activation: torch.dtype = torch.bfloat16, - swiglu_gptoss_style: bool = False, - ): - """ - ConfigurableMoE is a wrapper class that delegates to specific backends. - - To check capability, query the specific backend class directly: - - CutlassFusedMoE.can_implement(quant_algo, dtype_activation, swiglu_gptoss_style) - - TRTLLMGenFusedMoE.can_implement(quant_algo, dtype_activation, swiglu_gptoss_style) - - etc. - - Args: - quant_algo: The quantization algorithm to check (None for unquantized) - dtype_activation: The activation data type - swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + """Always ineligible: ConfigurableMoE delegates, it does not compute. - Returns: - Tuple[bool, Optional[str]]: Always returns (False, reason) + Answering ``False`` rather than raising is what lets a registry walk + include this class harmlessly. Query the backend it would delegate to + (``CutlassFusedMoE``, ``TRTLLMGenFusedMoE``, ...) instead. """ - del quant_algo, dtype_activation, swiglu_gptoss_style # Unused - wrapper class - return False, ( + del p, d # a wrapper's answer cannot depend on the question + return _reject( + MoERejectReason.NOT_AN_IMPL, "ConfigurableMoE is a wrapper class. " - "Query the specific backend (CutlassFusedMoE, TRTLLMGenFusedMoE, etc.) directly." + "Query the specific backend (CutlassFusedMoE, TRTLLMGenFusedMoE, etc.) directly.", ) def __init__( @@ -162,6 +158,7 @@ def __init__( apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, override_quant_config: Optional["QuantConfig"] = None, + moe_cls: Optional[Type] = None, activation: Optional[str] = None, situ_beta: Optional[float] = None, situ_linear_beta: Optional[float] = None, @@ -200,6 +197,7 @@ def __init__( model_config=model_config, routing_method=routing_method, override_quant_config=override_quant_config, + moe_cls=moe_cls, activation=activation, situ_beta=situ_beta, situ_linear_beta=situ_linear_beta, @@ -289,12 +287,13 @@ def _create_and_sync_backend( model_config: ModelConfig, routing_method: BaseMoeRoutingMethod, override_quant_config: Optional["QuantConfig"], - activation: Optional[str], - situ_beta: Optional[float], - situ_linear_beta: Optional[float], - trtllm_gen_activation_type: Optional[ActType_TrtllmGen], - trtllm_gen_activation_alpha: Optional[float], - trtllm_gen_activation_beta: Optional[float], + moe_cls: Optional[Type] = None, + activation: Optional[str] = None, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, + trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None, + trtllm_gen_activation_alpha: Optional[float] = None, + trtllm_gen_activation_beta: Optional[float] = None, **kwargs, ) -> None: """Build the MoE backend, mirror EPLB attrs, then create weights. @@ -313,15 +312,30 @@ def _create_and_sync_backend( """ from tensorrt_llm._torch.modules.fused_moe.create_moe import ( create_moe_backend, + infer_swiglu_gptoss_style, resolve_moe_cls, ) - moe_cls = resolve_moe_cls( - model_config, - routing_method, - self.dtype, - override_quant_config=override_quant_config, - ) + # create_moe already resolved; direct constructors resolve here. + if moe_cls is None: + moe_cls = resolve_moe_cls( + model_config, + override_quant_config=override_quant_config, + dtype=self.dtype, + num_experts=self.num_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + swiglu_gptoss_style=infer_swiglu_gptoss_style( + bias=kwargs.get("bias", False), + swiglu_alpha=kwargs.get("swiglu_alpha"), + swiglu_beta=kwargs.get("swiglu_beta"), + activation_type=self.activation_type, + ), + bias=kwargs.get("bias", False), + activation_type=self.activation_type, + routing=self.routing_method, + layer_idx=self.layer_idx, + ) backend_model_config = model_config if override_quant_config is not None: diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 8b1fd45725cd..011f77500322 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -4,11 +4,9 @@ import torch -from tensorrt_llm.logger import logger -from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig +from tensorrt_llm.models.modeling_utils import QuantConfig from ...model_config import ModelConfig -from ...peft.lora.validation import check_moe_lora_supported from ...utils import ActivationType, ActType_TrtllmGen, AuxStreamType from .configurable_moe import ConfigurableMoE from .fused_moe_cute_dsl import CuteDslFusedMoE @@ -24,247 +22,19 @@ from .interface import MoE, MoEWeightLoadingMode from .mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm from .moe_load_balancer import get_moe_load_balancer +from .moe_resolution import (WIDEEP_DEPRECATION_MESSAGE, + derive_moe_layer_shapes, infer_swiglu_gptoss_style, + resolve_moe_cls, resolve_moe_impl) from .routing import BaseMoeRoutingMethod -WIDEEP_DEPRECATION_MESSAGE = ( - "The WIDEEP MoE backend is deprecated and can no longer be selected. Wide " - "expert parallelism and EPLB are supported by the other backends: use " - "DEEPGEMM for FP8 block-scale checkpoints, or TRTLLM / CUTEDSL / CUTLASS " - "otherwise.") - - -def _get_pretrained_megamoe_capability_args( - model_config: ModelConfig) -> Dict[str, Optional[object]]: - """Extract dtype / hidden / intermediate kwargs for MegaMoE - ``can_implement`` from ``model_config.pretrained_config``. - - Both MegaMoE backends (``MEGAMOE_DEEPGEMM`` and ``MEGAMOE_CUTEDSL``) - perform the same pretrained-config probe before instantiating the - backend; centralising it keeps the probe and fallback logic - consistent across backends. - """ - pretrained = model_config.pretrained_config - 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) - pretrained_hidden = (getattr(pretrained, "hidden_size", None) - if pretrained is not None else None) - return dict(dtype_activation=pretrained_dtype, - hidden_size=pretrained_hidden, - intermediate_size=pretrained_inter) - - -def get_moe_cls( - model_config: ModelConfig, - override_quant_config: Optional[QuantConfig] = None, - layer_idx: Optional[int] = None, -) -> Type[MoE]: - moe_backend = model_config.moe_backend - quant_config = model_config.quant_config - if override_quant_config is not None: - quant_config = override_quant_config - layer_prefix = f"[layer_idx={layer_idx}] " if layer_idx is not None else "" - if moe_backend.upper() == "MARLIN": - # Marlin MoE is an Ada/Hopper NVFP4 W4A16 backend. Layers without - # NVFP4 quantization (e.g. deliberately-unquantized MTP draft layers in - # MIXED_PRECISION checkpoints) fall back to CutlassFusedMoE, matching - # the CUTEDSL / DENSEGEMM / MEGAMOE_* fallback behavior below. - if quant_config is None or not quant_config.quant_mode.has_nvfp4(): - logger.warning(f"{layer_prefix}MarlinFusedMoE only supports NVFP4 " - "quantization. Check out details in quant_config: " - f"{quant_config}. Using CutlassFusedMoE instead.") - return CutlassFusedMoE - return MarlinFusedMoE - if moe_backend.upper() == "CUTLASS": - return CutlassFusedMoE - elif moe_backend.upper() == "VANILLA": - return VanillaMoE - elif moe_backend.upper() == "CUTEDSL": - has_w4a16_nvfp4 = (quant_config is not None - and quant_config.quant_algo == QuantAlgo.W4A16_NVFP4) - if quant_config is not None and ( - quant_config.quant_mode.has_fp8_block_scales() - or quant_config.quant_mode.has_nvfp4() or has_w4a16_nvfp4): - # On SM120 / SM121 + NVFP4/W4A16_NVFP4 the cuteDSL family member is the - # hybrid CUTLASS-prefill / FlashInfer NVFP4 MoE decode backend - # (CuteDslB12xFusedMoE). Prefer it when flashinfer is importable; - # otherwise fall through to CuteDslFusedMoE for SM100 / SM103. - has_nvfp4 = (quant_config.quant_mode.has_nvfp4() - and not has_w4a16_nvfp4) - if has_nvfp4 or has_w4a16_nvfp4: - from tensorrt_llm._utils import get_sm_version - sm_version = get_sm_version() - if sm_version in CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS: - mapping = model_config.mapping - if mapping.moe_ep_size > 1 or mapping.dp_size > 1: - logger.warning( - "CuteDslB12xFusedMoE does not support expert " - "parallelism or attention-DP/all-to-all; selecting " - "CutlassFusedMoE.") - return CutlassFusedMoE - try: - import flashinfer # noqa: F401 - logger.info( - "Selecting CuteDslB12xFusedMoE for hybrid " - "CUTLASS-prefill / FlashInfer NVFP4 MoE decode " - "(SM%d + NVFP4).", - sm_version, - ) - return CuteDslB12xFusedMoE - except ImportError: - logger.warning( - "CuteDslB12xFusedMoE eligible (SM%d + NVFP4) " - "but flashinfer is not importable; using %s.", - sm_version, - "CutlassFusedMoE" - if has_w4a16_nvfp4 else "CuteDslFusedMoE", - ) - if has_w4a16_nvfp4: - return CutlassFusedMoE - elif has_w4a16_nvfp4: - logger.warning( - "CuteDslB12xFusedMoE requires SM120/121 for W4A16_NVFP4 " - "(got SM%d). Using CutlassFusedMoE.", - sm_version, - ) - return CutlassFusedMoE - return CuteDslFusedMoE - else: - logger.warning( - f"{layer_prefix}CuteDslFusedMoE only supports fp8_block_scales, nvfp4, and w4a16_nvfp4. " - f"Check out details in quant_config: {quant_config}. Using CutlassFusedMoE instead." - ) - return CutlassFusedMoE - elif moe_backend.upper() == "DEEPGEMM": - return DeepGemmFusedMoE - elif moe_backend.upper() == "DENSEGEMM": - if quant_config is None or not quant_config.quant_mode.has_nvfp4(): - logger.warning( - f"{layer_prefix}DenseGEMMFusedMoE only supports nvfp4. " - f"Check out details in quant_config: {quant_config}. Using CutlassFusedMoE instead." - ) - return CutlassFusedMoE - # DenseGEMM CuTe DSL kernels only support SM100/SM103. - from tensorrt_llm._utils import get_sm_version - sm_version = get_sm_version() - if sm_version not in DenseGEMMFusedMoE._SUPPORTED_SM_VERSIONS: - logger.warning( - f"{layer_prefix}DenseGEMMFusedMoE only supports SM {DenseGEMMFusedMoE._SUPPORTED_SM_VERSIONS} " - f"(got SM {sm_version}). Using CutlassFusedMoE instead.") - return CutlassFusedMoE - return DenseGEMMFusedMoE - elif moe_backend.upper() == "TRTLLM": - has_quant = quant_config is not None and quant_config.quant_mode.has_any_quant( - exclude_kv_cache=True) - if has_quant and (quant_config.quant_mode.has_fp8_block_scales() - or quant_config.quant_mode.has_nvfp4() - or quant_config.quant_mode.has_w4a16_mxfp4() - or quant_config.quant_mode.has_w4a8_nvfp4_fp8() - or quant_config.quant_mode.has_w4a8_mxfp4_fp8() - or quant_config.quant_mode.has_w4a8_mxfp4_mxfp8()): - return TRTLLMGenFusedMoE - if not has_quant and model_config.pretrained_config is not None and getattr( - model_config.pretrained_config, "torch_dtype", - None) == torch.bfloat16: - if TRTLLMGenFusedMoE._is_flashinfer_fused_moe_available(): - return TRTLLMGenFusedMoE - raise RuntimeError( - "TRTLLMGenFusedMoE BF16 path requires FlashInfer fused MoE with " - "trtllm_bf16_moe support, but it is not available.") - else: - logger.warning( - f"{layer_prefix}TRTLLMGenFusedMoE only supports fp8_block_scales, nvfp4, w4a16_mxfp4, w4a8_nvfp4_fp8, w4a8_mxfp4_fp8, and w4a8_mxfp4_mxfp8. " - f"Check out details in quant_config: {quant_config}. Using CutlassFusedMoE instead." - ) - return CutlassFusedMoE - elif moe_backend.upper() == "WIDEEP": - raise ValueError(WIDEEP_DEPRECATION_MESSAGE) - elif moe_backend.upper() == "TRITON": - return TritonFusedMoE - elif moe_backend.upper() == "MEGAMOE_DEEPGEMM": - # MegaMoE (DeepGEMM): DeepGEMM fp8_fp4_mega_moe fused kernel for - # W4A8_MXFP4_MXFP8 weights. Falls back to CutlassFusedMoE whenever - # the env cannot serve the backend (wrong quant / SM family / - # missing DG symbols) so we never allocate MegaMoE-specific weight - # tensors we cannot use. - if quant_config is None or not quant_config.quant_mode.has_w4a8_mxfp4_mxfp8( - ): - logger.warning( - "MegaMoEDeepGemm only supports W4A8_MXFP4_MXFP8. " - f"Check out details in quant_config: {quant_config}. Using CutlassFusedMoE instead." - ) - return CutlassFusedMoE - ok, reason = MegaMoEDeepGemm.can_implement( - QuantAlgo.W4A8_MXFP4_MXFP8, - swiglu_gptoss_style=False, - **_get_pretrained_megamoe_capability_args(model_config), - ) - if not ok: - logger.warning( - f"MegaMoEDeepGemm rejected current environment: {reason}. " - "Falling back to CutlassFusedMoE.") - return CutlassFusedMoE - return MegaMoEDeepGemm - elif moe_backend.upper() == "MEGAMOE_CUTEDSL": - # MegaMoE (CuteDSL): ported Sm100MegaMoEKernel fused - # dispatch+GEMM+activation+GEMM+combine kernel for NVFP4 weights on - # SM100-family GPUs. Same fall-back pattern as MEGAMOE_DEEPGEMM - # when the env cannot serve the backend. - if quant_config is None or not quant_config.quant_mode.has_nvfp4(): - logger.warning( - "MegaMoECuteDsl only supports NVFP4. " - f"Check out details in quant_config: {quant_config}. " - "Using CutlassFusedMoE instead.") - return CutlassFusedMoE - ok, reason = MegaMoECuteDsl.can_implement( - QuantAlgo.NVFP4, - swiglu_gptoss_style=False, - **_get_pretrained_megamoe_capability_args(model_config), - ) - if not ok: - logger.warning( - f"MegaMoECuteDsl rejected current environment: {reason}. " - "Falling back to CutlassFusedMoE.") - return CutlassFusedMoE - return MegaMoECuteDsl - else: - raise ValueError(f"Unsupported moe backend: {moe_backend}") - - -def resolve_moe_cls( - model_config: ModelConfig, - routing_method: BaseMoeRoutingMethod, - dtype: Optional[torch.dtype], - override_quant_config: Optional[QuantConfig] = None, - layer_idx: Optional[int] = None, -) -> Type[MoE]: - moe_cls = get_moe_cls(model_config, override_quant_config, layer_idx) - - effective_quant_config = override_quant_config or model_config.quant_config - has_quant = (effective_quant_config is not None - and effective_quant_config.layer_quant_mode.has_any_quant( - exclude_kv_cache=True)) - if (moe_cls == TRTLLMGenFusedMoE and not has_quant): - moe_cls = CutlassFusedMoE - - # Routed-expert LoRA is supported only on CutlassFusedMoE with unquantized - # base weights. Fail loudly here rather than at runtime if the user-selected - # backend cannot serve the LoRA request. Use the resolved class name, so a - # fallback to CutlassFusedMoE keeps LoRA supportable even when the user - # requested TRTLLM/CUTEDSL. - resolved_backend = "CUTLASS" if moe_cls is CutlassFusedMoE else model_config.moe_backend - check_moe_lora_supported( - moe_backend_name=resolved_backend, - lora_config=getattr(model_config, "lora_config", None), - quant_config=effective_quant_config, - layer_idx=layer_idx, - ) - - return moe_cls +__all__ = [ + "create_moe", + "create_moe_backend", + "infer_swiglu_gptoss_style", + "resolve_moe_cls", + "resolve_moe_impl", + "WIDEEP_DEPRECATION_MESSAGE", +] def create_moe_backend( @@ -331,24 +101,27 @@ def create_moe_backend( if moe_cls is WideEPMoE: raise ValueError(WIDEEP_DEPRECATION_MESSAGE) - # Get parameters from pretrained_config if not explicitly provided - pretrained_config = model_config.pretrained_config - if num_experts is None: - assert pretrained_config is not None, "num_experts must be provided or model_config.pretrained_config must be set" - num_experts = pretrained_config.num_experts - if hidden_size is None: - assert pretrained_config is not None, "hidden_size must be provided or model_config.pretrained_config must be set" - hidden_size = pretrained_config.hidden_size - if intermediate_size is None: - assert pretrained_config is not None, "intermediate_size must be provided or model_config.pretrained_config must be set" - # For MoE models, prefer moe_intermediate_size if available - if hasattr(pretrained_config, 'moe_intermediate_size'): - intermediate_size = pretrained_config.moe_intermediate_size - else: - intermediate_size = pretrained_config.intermediate_size - if dtype is None and pretrained_config is not None and hasattr( - pretrained_config, 'torch_dtype'): - dtype = pretrained_config.torch_dtype + shapes = derive_moe_layer_shapes( + model_config, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + routing=routing_method, + ) + num_experts = shapes.num_experts + hidden_size = shapes.hidden_size + intermediate_size = shapes.intermediate_size + dtype = shapes.dtype + assert num_experts is not None, ( + "num_experts must be provided or model_config.pretrained_config must " + "expose num_experts / n_routed_experts / num_local_experts") + assert hidden_size is not None, ( + "hidden_size must be provided or model_config.pretrained_config must be set" + ) + assert intermediate_size is not None, ( + "intermediate_size must be provided or model_config.pretrained_config " + "must expose moe_intermediate_size / intermediate_size") moe_load_balancer = get_moe_load_balancer() if moe_load_balancer is not None: @@ -637,27 +410,47 @@ def create_moe( Returns: MoE: MoE instance """ - # Get parameters from pretrained_config if not explicitly provided - pretrained_config = model_config.pretrained_config - if num_experts is None: - assert pretrained_config is not None, "num_experts must be provided or model_config.pretrained_config must be set" - num_experts = pretrained_config.num_experts - if hidden_size is None: - assert pretrained_config is not None, "hidden_size must be provided or model_config.pretrained_config must be set" - hidden_size = pretrained_config.hidden_size - if intermediate_size is None: - assert pretrained_config is not None, "intermediate_size must be provided or model_config.pretrained_config must be set" - # For MoE models, prefer moe_intermediate_size if available - if hasattr(pretrained_config, 'moe_intermediate_size'): - intermediate_size = pretrained_config.moe_intermediate_size - else: - intermediate_size = pretrained_config.intermediate_size - if dtype is None and pretrained_config is not None and hasattr( - pretrained_config, 'torch_dtype'): - dtype = pretrained_config.torch_dtype - - moe_cls = resolve_moe_cls(model_config, routing_method, dtype, - override_quant_config, layer_idx) + shapes = derive_moe_layer_shapes( + model_config, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + routing=routing_method, + ) + num_experts = shapes.num_experts + hidden_size = shapes.hidden_size + intermediate_size = shapes.intermediate_size + dtype = shapes.dtype + assert num_experts is not None, ( + "num_experts must be provided or model_config.pretrained_config must " + "expose num_experts / n_routed_experts / num_local_experts") + assert hidden_size is not None, ( + "hidden_size must be provided or model_config.pretrained_config must be set" + ) + assert intermediate_size is not None, ( + "intermediate_size must be provided or model_config.pretrained_config " + "must expose moe_intermediate_size / intermediate_size") + + # Pass the same shapes / activation package the layer will be built with. + moe_cls = resolve_moe_cls( + model_config, + override_quant_config=override_quant_config, + dtype=dtype, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + swiglu_gptoss_style=infer_swiglu_gptoss_style( + bias=bias, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + activation_type=activation_type, + ), + bias=bias, + activation_type=activation_type, + routing=routing_method, + layer_idx=layer_idx, + ) if (any(value is not None for value in (activation, situ_beta, situ_linear_beta)) and moe_cls is not MegaMoEDeepGemm): @@ -678,6 +471,7 @@ def create_moe( CuteDslB12xFusedMoE, CutlassFusedMoE, DenseGEMMFusedMoE, MegaMoEDeepGemm, MegaMoECuteDsl, MarlinFusedMoE): return ConfigurableMoE( + moe_cls=moe_cls, routing_method=routing_method, num_experts=num_experts, hidden_size=hidden_size, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py index 4e88d11848bd..d2eb8a592c98 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py @@ -20,7 +20,7 @@ import torch import torch.nn.functional as F -from tensorrt_llm._utils import get_sm_version, is_sm_100f +from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.models.modeling_utils import QuantAlgo from ...autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec, @@ -37,7 +37,10 @@ get_last_power_of_2_num_tokens_buckets, last_positive_power_of_2) from .fused_moe_cutlass import CutlassFusedMoE -from .impl_contract import MoERunContext, MoEStaticCapability, require_comm_plan +from .impl_contract import (MoEDeployment, MoEEligibility, MoEProblem, + MoERejectReason, MoERunContext, MoEStaticCapability, + require_comm_plan) +from .interface import _reject from .quantization import MoEWeightLoadingMode, NVFP4CuteDslFusedMoEMethod from .routing import BaseMoeRoutingMethod @@ -363,66 +366,53 @@ class CuteDslFusedMoE(CutlassFusedMoE): supports_dwdp=True) @classmethod - def can_implement( - cls, - quant_algo: Optional[QuantAlgo], - dtype_activation: torch.dtype = torch.bfloat16, - swiglu_gptoss_style: bool = False, - ) -> Tuple[bool, Optional[str]]: - """ - Check if CuteDslFusedMoE can implement the given quantization algorithm. - - CuteDslFusedMoE supports: - - NVFP4: SM in {100, 103} - - Does NOT support unquantized mode. Output dtype is hardcoded to bfloat16. - Does NOT support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit). - - Args: - quant_algo: The quantization algorithm to check (None for unquantized) - dtype_activation: The activation input data type. Only bfloat16 is supported - because output dtype is hardcoded to bfloat16 (input/output dtype must match). - swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled. - CuteDslFusedMoE does NOT support swiglu_gptoss_style. - - Returns: - Tuple[bool, Optional[str]]: (can_implement, skip_reason) - """ - from .interface import _warn_and_return - - sm_version = get_sm_version() + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + """CuteDSL grouped GEMM: NVFP4 on SM100/SM103, bfloat16 activations.""" + sm_version = d.env.sm + quant_algo = p.quant_algo # CuteDslFusedMoE requires at least SM90 if sm_version < 90: - return _warn_and_return( + return _reject( + MoERejectReason.SM_UNSUPPORTED, f"CuteDslFusedMoE requires SM >= 90, got SM{sm_version}") - # Check dtype_activation: output is hardcoded to bfloat16, so input must also be bfloat16 - # to maintain input/output dtype consistency - if dtype_activation != torch.bfloat16: - return _warn_and_return( + # Output is hardcoded to bfloat16, so input must also be bfloat16 to + # maintain input/output dtype consistency. + if p.dtype_act != torch.bfloat16: + return _reject( + MoERejectReason.DTYPE_UNSUPPORTED, f"CuteDslFusedMoE only supports bfloat16 activation (output is hardcoded to bfloat16), " - f"got {dtype_activation}") + f"got {p.dtype_act}") # CuteDslFusedMoE does NOT support unquantized mode if quant_algo is None: - return _warn_and_return( - "CuteDslFusedMoE does not support unquantized mode") + return _reject(MoERejectReason.QUANT_UNSUPPORTED, + "CuteDslFusedMoE does not support unquantized mode") # CuteDslFusedMoE does NOT support swiglu_gptoss_style - if swiglu_gptoss_style: - return _warn_and_return( + if p.swiglu_gptoss_style: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, "CuteDslFusedMoE does not support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit)" ) # NVFP4 - SM in {100, 103} if quant_algo == QuantAlgo.NVFP4: if sm_version not in {100, 103}: - return _warn_and_return( + return _reject( + MoERejectReason.SM_UNSUPPORTED, f"NVFP4 requires SM100 or SM103, got SM{sm_version}") - return True, None - - return _warn_and_return( + return MoEEligibility.ok() + + # FP8_BLOCK_SCALES lands here on purpose. ``run_moe_fp8_block_scales`` + # exists, but its GEMM is ``cute_dsl_fp8_group_blockwise_gemm_ref`` -- + # an fp32 einsum-per-expert reference, not a CuteDSL kernel -- so + # claiming the algorithm here would advertise a reference path as a + # backend. DeepGemm / TRTLLMGen own it on SM100/103, Cutlass on + # SM90/SM120. See the FP8-block note in MOE_DEVELOPER_GUIDE.md. + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, f"CuteDslFusedMoE does not support quant_algo={quant_algo}") def __init__( diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py index 1d54d27fefac..16ac8c7d7be6 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py @@ -18,14 +18,23 @@ import torch -from tensorrt_llm._utils import get_sm_version, nvtx_range +from tensorrt_llm._utils import nvtx_range from tensorrt_llm.models.modeling_utils import QuantAlgo from ...utils import ActivationType, Fp4QuantizedTensor from .fused_moe_cute_dsl import CuteDslFusedMoE from .fused_moe_cutlass import CutlassFusedMoE -from .impl_contract import MoERunContext, MoEStaticCapability, require_comm_plan -from .interface import _warn_and_return +from .impl_contract import ( + MoEDeployment, + MoEEligibility, + MoEProblem, + MoERejectReason, + MoERunContext, + MoEStaticCapability, + require_comm_plan, +) +from .impl_environment import MoEDep +from .interface import _reject # Shared MoE output buffer pool, keyed by (max_num_tokens, hidden_size, dtype, # device). ``B12xMoEWrapper.__init__`` allocates a private @@ -48,43 +57,7 @@ class CuteDslB12xFusedMoE(CuteDslFusedMoE): """B12x NVFP4 fused-MoE backend for SM120 / SM121. - Member of the cuteDSL backend family: the decode kernel - (``flashinfer.B12xMoEWrapper.run``) is JIT-compiled CuTe DSL, so the - backend slots in next to :class:`CuteDslFusedMoE` (which targets SM100 / - SM103). Plain NVFP4 prefill can route through the C++ CUTLASS NVFP4 - GroupGEMM via explicit :class:`CutlassFusedMoE` method calls; the parent - class on the MRO does not change which kernels execute, only where the - b12x backend sits in the family. - - Composition (see ``MOE_DEVELOPER_GUIDE.md`` for the full explainer): - - - **NVFP4 prefill (``m >= _PREFILL_VIA_CUTLASS_THRESHOLD``)** explicitly - invokes :class:`CutlassFusedMoE` NVFP4 GroupGEMM. The b12x kernel's - 12-CTA-per-token MMA pattern is suboptimal at large ``m``. - - **Decode (``m < _PREFILL_VIA_CUTLASS_THRESHOLD``)** dispatches to - FlashInfer's ``B12xMoEWrapper.run`` — a kernel purpose-built for - ``m=1`` / small routed-row counts. - - **W4A16_NVFP4** stays on the b12x path for both prefill and decode. - - NVFP4 weights are loaded via :class:`NVFP4CuteDslB12xFusedMoEMethod` - (an :class:`NVFP4CutlassFusedMoEMethod` subclass returned by - ``_get_quant_method``). The inherited CUTLASS NVFP4 layout is finalised - by the base class, and the b12x-shaped tensors (un-normalised FP8 SF, - ``convert_sf_to_mma_layout`` reshape, ``B12xMoEWrapper`` instance) are - materialised on top by the quant method's ``transform_weights``. Both - layouts coexist in memory and the dispatcher picks per call based on - ``x.shape[0]``. - - CUDA graph capture only covers decode, so captured graphs always replay - the b12x path; eager prefill always runs CUTLASS — there is no graph - capture conflict. - - The backend hard-rejects EP (b12x has no dispatch / combine kernel), - MoE alltoall, ``Fp4QuantizedTensor`` input, ``swiglu_gptoss_style`` - biased SwiGLU, and activations outside ``{Relu2, Swiglu}``. It is - selected on the ``CUTEDSL`` MoE path when SM120 / SM121 + NVFP4 or - W4A16_NVFP4 + flashinfer-importable gates pass (see - ``create_moe.get_moe_cls``). + Large prefill chunks use CUTLASS; decode uses FlashInfer's b12x kernel. """ # Restated rather than inherited: the LoRA gate this replaces compared the @@ -104,29 +77,62 @@ class on the MRO does not change which kernels execute, only where the _PREFILL_VIA_CUTLASS_THRESHOLD = 64 @classmethod - def can_implement( - cls, - quant_algo: Optional[QuantAlgo], - dtype_activation: torch.dtype = torch.bfloat16, - swiglu_gptoss_style: bool = False, - ) -> Tuple[bool, Optional[str]]: - sm_version = get_sm_version() + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + sm_version = d.env.sm if sm_version not in cls._SUPPORTED_SM_VERSIONS: sm_list = "/".join(f"SM{v}" for v in sorted(cls._SUPPORTED_SM_VERSIONS)) - return _warn_and_return(f"CuteDslB12xFusedMoE requires {sm_list}, got SM{sm_version}") - if quant_algo not in {QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4}: - return _warn_and_return( + return _reject( + MoERejectReason.SM_UNSUPPORTED, + f"CuteDslB12xFusedMoE requires {sm_list}, got SM{sm_version}", + ) + if p.quant_algo not in {QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4}: + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, f"CuteDslB12xFusedMoE only supports NVFP4 or W4A16_NVFP4 quantization " - f"(got quant_algo={quant_algo})" + f"(got quant_algo={p.quant_algo})", ) - if dtype_activation not in {torch.float16, torch.bfloat16}: - return _warn_and_return( + if p.dtype_act not in {torch.float16, torch.bfloat16}: + return _reject( + MoERejectReason.DTYPE_UNSUPPORTED, f"CuteDslB12xFusedMoE NVFP4 requires float16 or bfloat16 " - f"activation dtype (got {dtype_activation})" + f"activation dtype (got {p.dtype_act})", + ) + if p.swiglu_gptoss_style: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + "CuteDslB12xFusedMoE does not support swiglu_gptoss_style", + ) + if p.activation_type not in _ACTIVATION_MAP: + supported = ", ".join(a.name for a in _ACTIVATION_MAP) + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"CuteDslB12xFusedMoE does not support activation " + f"{p.activation}; supported: {supported}", ) - if swiglu_gptoss_style: - return _warn_and_return("CuteDslB12xFusedMoE does not support swiglu_gptoss_style") - return True, None + # The decode kernel ships in the FlashInfer wheel. + if not d.env.has_dep(MoEDep.FLASHINFER): + return _reject( + MoERejectReason.DEP_MISSING, + "CuteDslB12xFusedMoE requires the flashinfer package", + ) + # No expert-parallel dispatch/combine kernel: EP must stay at 1. + if d.ep_size != 1: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"CuteDslB12xFusedMoE requires ep_size == 1 (got {d.ep_size})", + ) + # Attention-DP is a separate axis from EP: with moe_tp == tp the layer + # can have ep_size == 1 and still sit behind a DP allgather / + # reducescatter that the b12x wrapper has never been exercised under. + # ``use_dp and parallel_size > 1`` is exactly ``mapping.dp_size > 1`` + # (``Mapping.dp_size`` is ``tp_size`` when attention-DP is on). + if d.use_dp and d.parallel_size > 1: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"CuteDslB12xFusedMoE does not support attention-DP " + f"(parallel_size={d.parallel_size})", + ) + return MoEEligibility.ok() def __init__(self, *args, **kwargs): # ``ModelConfig`` is consumed by the inherited ``__init__`` for cache @@ -138,22 +144,11 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - # b12x has no expert-parallel dispatch/combine kernel, so EP must be - # disabled. dp_size > 1 implies the alltoall path which b12x can't run. - if self.ep_size != 1: - raise ValueError( - f"CuteDslB12xFusedMoE requires ep_size == 1 " - f"(got ep_size={self.ep_size}); use --moe_backend CUTLASS for EP." - ) + # Eligibility (SM / quant / activation / EP) is owned by + # ``can_implement``. ``enable_alltoall`` is a run-time flag not yet on + # ``MoEDeployment``, so keep the construction guard here. if self.enable_alltoall: raise ValueError("CuteDslB12xFusedMoE does not support MoE alltoall communication.") - if self.activation_type not in _ACTIVATION_MAP: - supported = ", ".join(a.name for a in _ACTIVATION_MAP) - raise ValueError( - f"CuteDslB12xFusedMoE does not support activation " - f"{ActivationType(self.activation_type).name}; " - f"supported: {supported}." - ) self._b12x_weights: Optional[dict] = None self.b12x_wrapper = None diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index 70daa39d561c..203a4c4fbaa1 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -28,9 +28,10 @@ from ...peft.lora.validation import has_moe_lora_targets from ...utils import (ActivationType, AuxStreamType, EventType, Fp4QuantizedTensor) -from .impl_contract import (MoEInputRequirement, MoERunContext, +from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement, + MoEProblem, MoERejectReason, MoERunContext, MoEStaticCapability, require_comm_plan) -from .interface import MoE +from .interface import MoE, _reject from .quantization import UnquantizedFusedMoEMethod # isort: off @@ -121,6 +122,13 @@ class CutlassFusedMoE(MoE): "sm_constraint": ("in", {100, 103, 120, 121}), "dtypes": {torch.float16, torch.bfloat16, torch.float8_e4m3fn}, }, + # W4A16_NVFP4: weights stay NVFP4 but are dequantized to the activation + # dtype every forward, so what finally runs is the unquantized kernel -- + # this entry tracks that path's limits, not NVFP4 tensor-core support. + QuantAlgo.W4A16_NVFP4: { + "sm_constraint": ("min", 80), + "dtypes": {torch.float16, torch.bfloat16}, + }, # W4A8_AWQ: SM in {89, 90} only QuantAlgo.W4A8_AWQ: { "sm_constraint": ("in", {89, 90}), @@ -154,62 +162,58 @@ class CutlassFusedMoE(MoE): }, } - _GPTOSS_SUPPORTED_ALGOS = {QuantAlgo.W4A8_MXFP4_MXFP8} - """set[QuantAlgo]: Quantization algorithms that support swiglu_gptoss_style.""" + _GPTOSS_SUPPORTED_ALGOS: frozenset[Optional[QuantAlgo]] = frozenset({ + None, + QuantAlgo.NVFP4, + QuantAlgo.W4A16_MXFP4, + QuantAlgo.W4A8_MXFP4_FP8, + QuantAlgo.W4A8_MXFP4_MXFP8, + }) + """Algorithms whose weight methods can serve gpt-oss / MiniMax SwiGLU. + + Unquantized and the MXFP4 family can load a 1-D gpt-oss expert bias. + NVFP4 is included for MiniMax-style SwigluBias without expert bias. + ``can_implement`` still rejects NVFP4 when ``p.bias is True`` because the + NVFP4 weight pad only accepts 2-D tensors. + """ @classmethod - def can_implement( - cls, - quant_algo: Optional[QuantAlgo], - dtype_activation: torch.dtype = torch.bfloat16, - swiglu_gptoss_style: bool = False, - ) -> Tuple[bool, Optional[str]]: - """ - Check if CutlassFusedMoE can implement the given quantization algorithm. - - CutlassFusedMoE supports: - - Unquantized (FP16/BF16): SM >= 80 - - FP8 per-tensor (QDQ): SM >= 89 - - FP8_BLOCK_SCALES: SM in {90, 120} - - NVFP4: SM in {100, 103, 120, 121} - - W4A8_AWQ: SM in {89, 90} only - - W8A16: SM >= 80 - - W4A16_MXFP4: SM == 90 only - - W4A8_MXFP4_FP8: SM in {100, 103} - - W4A8_MXFP4_MXFP8: SM in {100, 103, 120, 121} + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + """Cutlass grouped-GEMM MoE: the widest quant and SM coverage there is. - Args: - quant_algo: The quantization algorithm to check (None for unquantized) - dtype_activation: The activation input data type (before quantization). - Supported dtypes vary by quantization mode: - - Unquantized: float16, bfloat16 - - FP8/FP8_BLOCK_SCALES/W4A8_MXFP4_FP8: float16, bfloat16, float32 - - NVFP4: float16, bfloat16, float8_e4m3fn - - W4A16_MXFP4/W4A8_AWQ/W8A16/W4A8_MXFP4_MXFP8: float16, bfloat16 - swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled. - CutlassFusedMoE only supports swiglu_gptoss_style for W4A8_MXFP4_MXFP8 quantization. - - Returns: - Tuple[bool, Optional[str]]: (can_implement, skip_reason) + Per-algorithm SM and dtype support lives in ``_QUANT_SUPPORT_TABLE``; + this method is only the interpreter for it. """ - from .interface import _warn_and_return - - sm_version = get_sm_version() + sm_version = d.env.sm + quant_algo = p.quant_algo # Check minimum SM version for Cutlass backend if sm_version < 80: - return _warn_and_return( + return _reject( + MoERejectReason.SM_UNSUPPORTED, f"CutlassFusedMoE requires SM >= 80, got SM{sm_version}") - # Check swiglu_gptoss_style support - if swiglu_gptoss_style and quant_algo not in cls._GPTOSS_SUPPORTED_ALGOS: - return _warn_and_return( - f"CutlassFusedMoE swiglu_gptoss_style only supports W4A8_MXFP4_MXFP8 " - f"(got quant_algo={quant_algo})") + if p.swiglu_gptoss_style and quant_algo not in cls._GPTOSS_SUPPORTED_ALGOS: + supported = sorted("unquantized" if a is None else a.name + for a in cls._GPTOSS_SUPPORTED_ALGOS) + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"CutlassFusedMoE cannot load a gpt-oss bias for " + f"quant_algo={quant_algo}; supported: {supported}") + + # NVFP4 can run SwigluBias, but cannot pad a 1-D gpt-oss expert bias. + if (p.swiglu_gptoss_style and quant_algo == QuantAlgo.NVFP4 + and p.bias is True): + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + "CutlassFusedMoE NVFP4 cannot load a 1-D gpt-oss expert bias " + "(weight-pad assert is 2-D); MiniMax-style SwigluBias without " + "bias is eligible") # Check if quant_algo is supported if quant_algo not in cls._QUANT_SUPPORT_TABLE: - return _warn_and_return( + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, f"CutlassFusedMoE does not support quant_algo={quant_algo}") support_info = cls._QUANT_SUPPORT_TABLE[quant_algo] @@ -220,30 +224,42 @@ def can_implement( if constraint_type == "min": if sm_version < constraint_value: - return _warn_and_return( + return _reject( + MoERejectReason.SM_UNSUPPORTED, f"CutlassFusedMoE {algo_name} requires SM >= {constraint_value}, " f"got SM{sm_version}") elif constraint_type == "exact": if sm_version != constraint_value: - return _warn_and_return( + return _reject( + MoERejectReason.SM_UNSUPPORTED, f"CutlassFusedMoE {algo_name} only supports SM{constraint_value}, " f"got SM{sm_version}") elif constraint_type == "in": if sm_version not in constraint_value: sm_list = "/".join(f"SM{v}" for v in sorted(constraint_value)) - return _warn_and_return( + return _reject( + MoERejectReason.SM_UNSUPPORTED, f"CutlassFusedMoE {algo_name} only supports {sm_list}, " f"got SM{sm_version}") - # Check dtype_activation + # Check activation dtype supported_dtypes = support_info["dtypes"] - if dtype_activation not in supported_dtypes: - dtype_list = ", ".join(str(d) for d in supported_dtypes) - return _warn_and_return( + if p.dtype_act not in supported_dtypes: + dtype_list = ", ".join(str(dtype) for dtype in supported_dtypes) + return _reject( + MoERejectReason.DTYPE_UNSUPPORTED, f"CutlassFusedMoE {algo_name} requires {dtype_list}, " - f"got {dtype_activation}") + f"got {p.dtype_act}") + + # Routed-expert MoE LoRA supports unquantized fp16/bf16 or per-tensor FP8 only. + if d.moe_lora_enabled and quant_algo not in (None, QuantAlgo.FP8): + return _reject( + MoERejectReason.LORA_UNSUPPORTED, + "CutlassFusedMoE MoE LoRA only supports unquantized " + f"fp16/bf16 or per-tensor FP8 (qdq); got quant_algo={quant_algo}" + ) - return True, None + return MoEEligibility.ok() def __init__( self, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py index 7d3951762326..d5861f0798e7 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, Optional, Tuple, Union +from typing import Dict, Optional, Union import torch import triton @@ -21,15 +21,17 @@ import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils from tensorrt_llm import deep_gemm -from tensorrt_llm._utils import get_sm_version, nvtx_range +from tensorrt_llm._utils import nvtx_range from tensorrt_llm.models.modeling_utils import QuantAlgo from ...memory_buffer_utils import get_memory_buffers from ...model_config import ModelConfig from ...utils import AuxStreamType, Fp4QuantizedTensor from .fused_moe_cutlass import CutlassFusedMoE -from .impl_contract import (MoEInputRequirement, MoERunContext, +from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement, + MoEProblem, MoERejectReason, MoERunContext, MoEStaticCapability) +from .interface import _reject from .quantization import (DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm, MoEWeightLoadingMode, UnquantizedFusedMoEMethod) from .routing import BaseMoeRoutingMethod @@ -737,64 +739,42 @@ def supports_moe_output_in_alltoall_workspace(self): return False @classmethod - def can_implement( - cls, - quant_algo: Optional[QuantAlgo], - dtype_activation: torch.dtype = torch.bfloat16, - swiglu_gptoss_style: bool = False, - ) -> Tuple[bool, Optional[str]]: - """ - Check if DeepGemmFusedMoE can implement the given quantization algorithm. - - DeepGemmFusedMoE supports: - - FP8_BLOCK_SCALES: SM in {100, 103} - - Does NOT support unquantized mode. Output dtype is hardcoded to bfloat16. - Does NOT support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit). - - Args: - quant_algo: The quantization algorithm to check (None for unquantized) - dtype_activation: The activation input data type. Supported types are - float32, bfloat16, and float16 (required by moe_permute_op kernel). - Note: Output dtype is always bfloat16 regardless of input dtype. - swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled. - DeepGemmFusedMoE does NOT support swiglu_gptoss_style. - - Returns: - Tuple[bool, Optional[str]]: (can_implement, skip_reason) - """ - from .interface import _warn_and_return - - sm_version = get_sm_version() + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + """DeepGEMM grouped GEMM: FP8 block scales on SM100/SM103.""" + sm_version = d.env.sm + quant_algo = p.quant_algo if sm_version not in {100, 103}: - return _warn_and_return( + return _reject( + MoERejectReason.SM_UNSUPPORTED, f"DeepGemmFusedMoE requires SM100 or SM103, got SM{sm_version}") - # Check dtype_activation: moe_permute_op only supports float32, bfloat16, float16 - if dtype_activation not in { - torch.float32, torch.bfloat16, torch.float16 - }: - return _warn_and_return( + # moe_permute_op only supports float32, bfloat16, float16 + if p.dtype_act not in {torch.float32, torch.bfloat16, torch.float16}: + return _reject( + MoERejectReason.DTYPE_UNSUPPORTED, f"DeepGemmFusedMoE requires float32, bfloat16, or float16 activation, " - f"got {dtype_activation}") + f"got {p.dtype_act}") # DeepGemmFusedMoE does NOT support unquantized mode if quant_algo is None: - return _warn_and_return( + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, "DeepGemmFusedMoE does not support unquantized mode") # DeepGemmFusedMoE does NOT support swiglu_gptoss_style - if swiglu_gptoss_style: - return _warn_and_return( + if p.swiglu_gptoss_style: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, "DeepGemmFusedMoE does not support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit)" ) # Only FP8_BLOCK_SCALES is supported if quant_algo == QuantAlgo.FP8_BLOCK_SCALES: - return True, None + return MoEEligibility.ok() - return _warn_and_return( + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, f"DeepGemmFusedMoE does not support quant_algo={quant_algo}") # To reuse pytorch memory segments allocated during graph capture. diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py index 1dad6dab2ec5..e5ae59f83345 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py @@ -12,9 +12,24 @@ from ...memory_buffer_utils import get_memory_buffers from ...model_config import ModelConfig -from ...utils import AuxStreamType, EventType, Fp4QuantizedTensor, swizzle_sf, unswizzle_sf -from .impl_contract import MoEInputRequirement, MoERunContext, require_comm_plan -from .interface import MoE, MoEWeightLoadingMode +from ...utils import ( + ActivationType, + AuxStreamType, + EventType, + Fp4QuantizedTensor, + swizzle_sf, + unswizzle_sf, +) +from .impl_contract import ( + MoEDeployment, + MoEEligibility, + MoEInputRequirement, + MoEProblem, + MoERejectReason, + MoERunContext, + require_comm_plan, +) +from .interface import MoE, MoEWeightLoadingMode, _reject from .quantization import NVFP4CuteDslFusedMoEMethod from .routing import BaseMoeRoutingMethod @@ -80,6 +95,11 @@ def gen_fc2_alpha_fused( return fc2_alpha.scatter_(1, token_selected_experts.long(), scaled_values) +# MMA tile size the FC2 DenseGEMM kernel tiles the K dimension with. Module +# level so that the selection gate and its rejection message read one number. +_FC2_MMA_TILE_K = 256 + + class DenseGEMMFusedMoE(MoE): """CuteDSL DenseGEMM flow of fused mixture of experts (MoE) Layer. @@ -111,38 +131,53 @@ class DenseGEMMFusedMoE(MoE): _SUPPORTED_SM_VERSIONS = (100, 103) @classmethod - def can_implement( - cls, - quant_algo: Optional[QuantAlgo], - dtype_activation: torch.dtype = torch.bfloat16, - swiglu_gptoss_style: bool = False, - ) -> tuple: - """Check if DenseGEMMFusedMoE can implement the given configuration. - - DenseGEMMFusedMoE supports: - - NVFP4 quantization only - - SM100/SM103 (Blackwell) only - - SwiGLU activation only (swiglu_gptoss_style not supported) - """ - from tensorrt_llm._utils import get_sm_version + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + """DenseGEMM CuTe DSL kernels: NVFP4 on SM100/SM103, SwiGLU only.""" + sm_version = d.env.sm + if sm_version not in cls._SUPPORTED_SM_VERSIONS: + return _reject( + MoERejectReason.SM_UNSUPPORTED, + f"DenseGEMMFusedMoE requires SM {cls._SUPPORTED_SM_VERSIONS}, got SM{sm_version}", + ) - from .interface import _warn_and_return + if p.quant_algo != QuantAlgo.NVFP4: + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, + f"DenseGEMMFusedMoE only supports NVFP4 quantization (got quant_algo={p.quant_algo})", + ) - sm_version = get_sm_version() - if sm_version not in cls._SUPPORTED_SM_VERSIONS: - return _warn_and_return( - f"DenseGEMMFusedMoE requires SM {cls._SUPPORTED_SM_VERSIONS}, got SM{sm_version}" + if p.swiglu_gptoss_style: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + "DenseGEMMFusedMoE does not support swiglu_gptoss_style", ) - if quant_algo != QuantAlgo.NVFP4: - return _warn_and_return( - f"DenseGEMMFusedMoE only supports NVFP4 quantization (got quant_algo={quant_algo})" + if p.activation_type != ActivationType.Swiglu: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"DenseGEMMFusedMoE fuses SwiGLU into the FC1 GEMM epilogue " + f"and serves no other activation (got {p.activation})", ) - if swiglu_gptoss_style: - return _warn_and_return("DenseGEMMFusedMoE does not support swiglu_gptoss_style") + if d.ep_size != 1: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"DenseGEMMFusedMoE is TP-only; expert parallelism would need " + f"an alltoall this backend does not implement (got " + f"ep_size={d.ep_size})", + ) + + # The FC2 kernel tiles K by an MMA tile of 256 and splits alpha_scale + # at expert boundaries, so a weight_per_expert that is not tile-aligned + # silently applies the wrong scale rather than failing. + if p.intermediate_size is not None and p.intermediate_size % _FC2_MMA_TILE_K != 0: + return _reject( + MoERejectReason.SHAPE_UNALIGNED, + f"DenseGEMMFusedMoE requires intermediate_size divisible by " + f"{_FC2_MMA_TILE_K} (FC2 MMA tile-K); got {p.intermediate_size}", + ) - return (True, None) + return MoEEligibility.ok() def __init__( self, @@ -161,45 +196,10 @@ def __init__( init_load_balancer: bool = True, activation_type=None, ): - # DenseGEMM CuTe DSL kernels only support SM100 and SM103. - from tensorrt_llm._utils import get_sm_version - - from ...utils import ActivationType - - sm_version = get_sm_version() - assert sm_version in self._SUPPORTED_SM_VERSIONS, ( - f"DenseGEMMFusedMoE only supports SM {self._SUPPORTED_SM_VERSIONS} " - f"(got SM {sm_version}). The CuTe DSL kernels require Blackwell architecture." - ) - - # DenseGEMM kernel hardcodes SwiGLU fusion — reject other activation types - # before calling super().__init__() to fail fast with a clear message. + # Eligibility (SM / quant / SwiGLU / EP / intermediate alignment) is + # owned by ``can_implement``; do not re-assert it here. if activation_type is None: activation_type = ActivationType.Swiglu - assert activation_type == ActivationType.Swiglu, ( - f"DenseGEMMFusedMoE only supports SwiGLU activation " - f"(got activation_type={activation_type}). " - f"The FC1 kernel fuses SwiGLU into the GEMM epilogue." - ) - - # FC2 DenseGEMM kernel tiles K dimension with MMA tile size 256. - # weight_per_expert (= intermediate_size) must be 256-aligned so that - # expert boundaries align with MMA tile boundaries. - _MMA_TILE_K = 256 - assert intermediate_size % _MMA_TILE_K == 0, ( - f"DenseGEMMFusedMoE requires intermediate_size to be a multiple of " - f"{_MMA_TILE_K} (got intermediate_size={intermediate_size}). " - f"FC2 kernel cannot correctly split alpha_scale at expert boundaries " - f"when weight_per_expert is not MMA tile-K aligned." - ) - - # DenseGEMM only supports TP; EP requires alltoall communication not implemented here. - ep_size = model_config.mapping.moe_ep_size - assert ep_size == 1, ( - f"DenseGEMMFusedMoE does not support Expert Parallelism " - f"(got ep_size={ep_size}). Use a different MoE backend (e.g. CutlassFusedMoE) " - f"when EP is enabled." - ) # Call MoE base class directly (not CutlassFusedMoE). # Note: `apply_router_weight_on_input` is accepted for API diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py index f9a7d1ad90dc..5826dfe27aa8 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py @@ -26,13 +26,19 @@ import torch.nn.functional as F from tensorrt_llm._torch.utils import Fp4QuantizedTensor, is_nvfp4_marlin_supported_sm -from tensorrt_llm._utils import get_sm_version from tensorrt_llm.models.modeling_utils import QuantAlgo from ...utils import ActivationType, is_gated_activation, relu2 from .fused_moe_cutlass import CutlassFusedMoE -from .impl_contract import MoERunContext, MoEStaticCapability -from .interface import _warn_and_return +from .impl_contract import ( + MoEDeployment, + MoEEligibility, + MoEProblem, + MoERejectReason, + MoERunContext, + MoEStaticCapability, +) +from .interface import _reject from .quantization import NVFP4MarlinFusedMoEMethod # Block size for moe_align_block_size — must match TILE_M in the kernel @@ -69,33 +75,43 @@ class MarlinFusedMoE(CutlassFusedMoE): } @classmethod - def can_implement( - cls, - quant_algo: Optional[QuantAlgo], - dtype_activation: torch.dtype = torch.bfloat16, - swiglu_gptoss_style: bool = False, - ) -> Tuple[bool, Optional[str]]: - sm_version = get_sm_version() - - if quant_algo not in cls._QUANT_SUPPORT_TABLE: - return _warn_and_return( - f"MarlinFusedMoE only supports NVFP4 or W4A16_NVFP4 (got quant_algo={quant_algo})" + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + if p.quant_algo not in cls._QUANT_SUPPORT_TABLE: + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, + f"MarlinFusedMoE only supports NVFP4 or W4A16_NVFP4 " + f"(got quant_algo={p.quant_algo})", ) - if not is_nvfp4_marlin_supported_sm(sm_version): - return _warn_and_return( - f"MarlinFusedMoE only supports SM89-SM99 (Ada/Hopper), got SM{sm_version}" + if not is_nvfp4_marlin_supported_sm(d.env.sm): + return _reject( + MoERejectReason.SM_UNSUPPORTED, + f"MarlinFusedMoE only supports SM89-SM99 (Ada/Hopper), got SM{d.env.sm}", ) - if swiglu_gptoss_style: - return _warn_and_return("MarlinFusedMoE does not support swiglu_gptoss_style") + if p.swiglu_gptoss_style: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + "MarlinFusedMoE does not support swiglu_gptoss_style", + ) + + if p.dtype_act != torch.bfloat16: + return _reject( + MoERejectReason.DTYPE_UNSUPPORTED, + f"MarlinFusedMoE W4A16 requires bfloat16 activations, got {p.dtype_act}", + ) - if dtype_activation != torch.bfloat16: - return _warn_and_return( - f"MarlinFusedMoE W4A16 requires bfloat16 activations, got {dtype_activation}" + # Sorted-token dispatch has no EPLB slot layout, so a layer that + # registered a load balancer cannot run here. Answered from ``d`` + # rather than from ``self._supports_load_balancer()``, which could + # only be consulted after the object existed. + if d.eplb_enabled: + return _reject( + MoERejectReason.EPLB_UNSUPPORTED, + "MarlinFusedMoE has no EPLB slot layout", ) - return True, None + return MoEEligibility.ok() def quantize_input( self, x: torch.Tensor | Fp4QuantizedTensor, post_quant_comm: bool = True, **kwargs diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py index c6164fc7b54f..00157e09eadd 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py @@ -16,7 +16,7 @@ from __future__ import annotations import os -from typing import Dict, List, NamedTuple, Optional, Tuple +from typing import Dict, List, NamedTuple, Optional import torch import torch.nn as nn @@ -33,13 +33,18 @@ from triton_kernels.tensor import convert_layout, wrap_torch_tensor from triton_kernels.tensor_details import layout +from tensorrt_llm.models.modeling_utils import QuantAlgo + from ...model_config import ModelConfig from ..linear import TensorParallelMode, load_weight_shard -from .interface import MoE +from .impl_contract import (MoEDeployment, MoEEligibility, MoEProblem, + MoERejectReason) +from .interface import MoE, _reject from .quantization import (FusedMoEMethodBase, MoEWeightLoadingMode, load_activation_scales_fp8_qdq, requantize_expert_w3_w1_weight_fp8_qdq) -from .routing import BaseMoeRoutingMethod, RenormalizeMoeRoutingMethod +from .routing import (ROUTING_METHOD_TYPE_TO_CLASS, BaseMoeRoutingMethod, + RenormalizeMoeRoutingMethod) # Triton kernels has hardcoded beta = 1, so we use this implementation when beta is not 1 @@ -1468,70 +1473,73 @@ def transform_weights(self, module: torch.nn.Module) -> None: class TritonFusedMoE(MoE): @classmethod - def can_implement( - cls, - quant_algo: Optional["QuantAlgo"], - dtype_activation: torch.dtype = torch.bfloat16, - swiglu_gptoss_style: bool = False, - ) -> Tuple[bool, Optional[str]]: - """ - Check if TritonFusedMoE can implement the given quantization algorithm. - - TritonFusedMoE supports (SM90 only, swiglu_gptoss_style=True only): - - Unquantized (BF16 only) - - FP8 per-tensor (QDQ) - - W4A8_MXFP4_FP8 - - W4A16_MXFP4 - - Args: - quant_algo: The quantization algorithm to check (None for unquantized) - dtype_activation: The activation data type. In unquantized mode, activation, - weight, and output dtypes must all match (only bfloat16 supported). - swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled. - TritonFusedMoE ONLY supports swiglu_gptoss_style=True. - - Returns: - Tuple[bool, Optional[str]]: (can_implement, skip_reason) - """ - from tensorrt_llm._utils import get_sm_version - from tensorrt_llm.models.modeling_utils import QuantAlgo - - from .interface import _warn_and_return + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + """Triton MoE: SM90 only, and only the gpt-oss style swiglu. - sm_version = get_sm_version() + Supports unquantized BF16, FP8 per-tensor QDQ, W4A8_MXFP4_FP8 and + W4A16_MXFP4. + """ + sm_version = d.env.sm + quant_algo = p.quant_algo # TritonFusedMoE only supports SM90 if sm_version != 90: - return _warn_and_return( + return _reject( + MoERejectReason.SM_UNSUPPORTED, f"TritonFusedMoE only supports SM90, got SM{sm_version}") - # TritonFusedMoE ONLY supports swiglu_gptoss_style=True - if not swiglu_gptoss_style: - return _warn_and_return( + if d.eplb_enabled: + return _reject( + MoERejectReason.EPLB_UNSUPPORTED, + "TritonFusedMoE does not implement the EPLB slot hooks") + + # Require gpt-oss SwiGLU; abstain when the style is unknown. + if p.swiglu_gptoss_style is False: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, "TritonFusedMoE only supports swiglu_gptoss_style=True") + if d.smart_router: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"TritonFusedMoE has no smart-router path (moe_cluster_size=" + f"{d.cluster_size})") + + # Require renormalize-family routing; abstain when routing is unknown. + routing_type = p.routing_method_type + if routing_type is not None and not issubclass( + ROUTING_METHOD_TYPE_TO_CLASS[routing_type], + RenormalizeMoeRoutingMethod): + return _reject( + MoERejectReason.ROUTING_UNSUPPORTED, + f"TritonFusedMoE fuses renormalize routing only, got {p.routing}" + ) + # Unquantized mode - only bfloat16 is supported if quant_algo is None: - if dtype_activation != torch.bfloat16: - return _warn_and_return( - f"TritonFusedMoE unquantized mode only supports bfloat16, got {dtype_activation}" + if p.dtype_act != torch.bfloat16: + return _reject( + MoERejectReason.DTYPE_UNSUPPORTED, + f"TritonFusedMoE unquantized mode only supports bfloat16, got {p.dtype_act}" ) - return True, None + return MoEEligibility.ok() - # FP8 per-tensor (QDQ) and W4A8_MXFP4_FP8 - no dtype_activation restriction + # FP8 per-tensor (QDQ) and W4A8_MXFP4_FP8 - no activation dtype restriction if quant_algo in {QuantAlgo.FP8, QuantAlgo.W4A8_MXFP4_FP8}: - return True, None + return MoEEligibility.ok() # W4A16_MXFP4 - only bfloat16 and float16 are supported if quant_algo == QuantAlgo.W4A16_MXFP4: - if dtype_activation not in {torch.bfloat16, torch.float16}: - return _warn_and_return( + if p.dtype_act not in {torch.bfloat16, torch.float16}: + return _reject( + MoERejectReason.DTYPE_UNSUPPORTED, f"TritonFusedMoE W4A16_MXFP4 only supports bfloat16 or float16, " - f"got {dtype_activation}") - return True, None + f"got {p.dtype_act}") + return MoEEligibility.ok() # Unsupported quantization algorithm - return _warn_and_return( + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, f"TritonFusedMoE does not support quant_algo={quant_algo}") def __init__( @@ -1563,13 +1571,8 @@ def __init__( weight_loading_mode=weight_loading_mode, layer_idx=layer_idx, ) - if torch.cuda.get_device_capability()[0] != 9 and self.ep_size > 1: - raise NotImplementedError( - "TritonFusedMoE is only supported on Hopper with EP size > 1.") - - assert isinstance(self.routing_method, RenormalizeMoeRoutingMethod), \ - "routing_method must be an instance of RenormalizeMoeRoutingMethod for TritonFusedMoE" - assert not self.smart_router, "Smart router is not supported in TritonFusedMoE." + # Eligibility (SM / routing / smart_router / quant) is owned by + # ``can_implement``; do not re-assert it here. self.num_slots = self.num_experts self.expert_size_per_partition = self.num_experts // self.ep_size diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index ba01e7ad536a..2d1e1fefac43 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -16,7 +16,7 @@ import inspect import os from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Union import torch from torch import nn @@ -31,8 +31,12 @@ from ...utils import (ActivationType, ActType_TrtllmGen, AuxStreamType, Fp4QuantizedTensor) from ..gated_mlp import GatedMLP -from .impl_contract import MoEInputRequirement, MoERunContext, require_comm_plan -from .interface import FORCE_SEPARATED_ROUTING, MoE, MoEWeightLoadingMode +from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement, + MoEProblem, MoERejectReason, MoERunContext, + require_comm_plan) +from .impl_environment import MoEDep +from .interface import (FORCE_SEPARATED_ROUTING, MoE, MoEWeightLoadingMode, + _reject) from .moe_op_backend import MoEOpBackend, TRTLLMOpBackend, get_op_backend # isort: off @@ -126,75 +130,93 @@ class TRTLLMGenFusedMoE(MoE): } @classmethod - def can_implement( - cls, - quant_algo: Optional[QuantAlgo], - dtype_activation: torch.dtype = torch.bfloat16, - swiglu_gptoss_style: bool = False, - ) -> Tuple[bool, Optional[str]]: - """ - Check if TRTLLMGenFusedMoE can implement the given quantization algorithm. - - TRTLLMGenFusedMoE only supports SM in {100, 103} and the following quantizations: - - NVFP4 - - FP8_BLOCK_SCALES - - W4A8_NVFP4_FP8 - - W4A16_MXFP4 - - W4A8_MXFP4_FP8 - - W4A8_MXFP4_MXFP8 + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + """TRTLLM-Gen kernels: SM100/SM103, bfloat16 activations. - Unquantized BF16 path is supported only with FlashInfer fused MoE backend. - Output dtype is hardcoded to bfloat16. - - Args: - quant_algo: The quantization algorithm to check (None for unquantized) - dtype_activation: The activation input data type. Only bfloat16 is supported. - See: forward_impl() assert x.dtype == torch.bfloat16 (line 722). - swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled. - Only supported for nvfp4 and mxfp4 variants. - - Returns: - Tuple[bool, Optional[str]]: (can_implement, skip_reason) + Quantized coverage is ``_SUPPORTED_QUANT_ALGOS``. The unquantized BF16 + path is served by a FlashInfer kernel, so it is available only where + that wheel exposes ``trtllm_bf16_moe``. """ - from .interface import _warn_and_return - - sm_version = get_sm_version() + sm_version = d.env.sm + quant_algo = p.quant_algo # TRTLLMGenFusedMoE requires SM in {100, 103} if sm_version not in {100, 103}: - return _warn_and_return( + return _reject( + MoERejectReason.SM_UNSUPPORTED, f"TRTLLMGenFusedMoE requires SM100 or SM103, got SM{sm_version}" ) - # Check dtype_activation: only bfloat16 is supported - if dtype_activation != torch.bfloat16: - return _warn_and_return( - f"TRTLLMGenFusedMoE only supports bfloat16 activation, got {dtype_activation}" + # forward_impl asserts x.dtype == torch.bfloat16 + if p.dtype_act != torch.bfloat16: + return _reject( + MoERejectReason.DTYPE_UNSUPPORTED, + f"TRTLLMGenFusedMoE only supports bfloat16 activation, got {p.dtype_act}" ) + if d.smart_router: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"TRTLLMGenFusedMoE has no smart-router path (moe_cluster_size=" + f"{d.cluster_size})") + if quant_algo is None: - if swiglu_gptoss_style: - return _warn_and_return( + if p.swiglu_gptoss_style: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, "TRTLLMGenFusedMoE BF16 path does not support bias/swiglu custom parameters." ) - if not cls._is_flashinfer_fused_moe_available(): - return _warn_and_return( + # Same set _check_configs asserts on, so the verdict and the + # constructor agree instead of failing later at create_weights. + if p.activation_type not in cls._BF16_SUPPORTED_ACTIVATIONS: + supported = ", ".join( + sorted(activation.name + for activation in cls._BF16_SUPPORTED_ACTIVATIONS)) + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"TRTLLMGenFusedMoE BF16 path only supports {supported} " + f"activations, got {p.activation}") + if not d.env.has_dep(MoEDep.FLASHINFER_BF16_MOE): + return _reject( + MoERejectReason.DEP_MISSING, "TRTLLMGenFusedMoE unquantized BF16 path requires FlashInfer fused MoE " "with trtllm_bf16_moe support.") - return True, None + # FlashInfer BF16 kernels require the per-rank intermediate size + # to be a multiple of 128. + if p.intermediate_size is not None: + inter = p.intermediate_size + if d.tp_size > 1: + if inter % d.tp_size != 0: + return _reject( + MoERejectReason.SHAPE_UNALIGNED, + "TRTLLMGenFusedMoE BF16 FlashInfer path requires " + f"intermediate_size ({inter}) divisible by " + f"moe_tp_size ({d.tp_size})") + inter = inter // d.tp_size + if inter % 128 != 0: + return _reject( + MoERejectReason.SHAPE_UNALIGNED, + "TRTLLMGenFusedMoE BF16 FlashInfer path requires " + "intermediate_size_per_partition % 128 == 0; " + f"got {inter} " + f"(full intermediate_size={p.intermediate_size}, " + f"moe_tp_size={d.tp_size})") + return MoEEligibility.ok() # Check if quant_algo is supported if quant_algo not in cls._SUPPORTED_QUANT_ALGOS: - return _warn_and_return( + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, f"TRTLLMGenFusedMoE does not support quant_algo={quant_algo}") - # Check swiglu_gptoss_style support: only supported for nvfp4 and mxfp4 variants - if swiglu_gptoss_style and quant_algo not in cls._GPTOSS_SUPPORTED_ALGOS: - return _warn_and_return( + # swiglu_gptoss_style is only supported for nvfp4 and mxfp4 variants + if p.swiglu_gptoss_style and quant_algo not in cls._GPTOSS_SUPPORTED_ALGOS: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, f"TRTLLMGenFusedMoE supports swiglu_gptoss_style (bias/swiglu) only for nvfp4 and mxfp4 variants, " f"got quant_algo={quant_algo}") - return True, None + return MoEEligibility.ok() def __init__( self, @@ -253,21 +275,9 @@ def __init__( # tune_max_num_tokens to the MoE op). self.max_num_tokens = model_config.max_num_tokens - sm_version = get_sm_version() - if sm_version >= 120: - raise NotImplementedError( - "TRTLLMGenFusedMoE does not support SM120 and above.") - - assert not self.smart_router, "Smart router is not supported in TRTLLMGenFusedMoE." - + # Eligibility (SM / smart_router / BF16 FlashInfer dep) is owned by + # ``can_implement``. Keep only the provider selection for the op. self.use_flashinfer = self._check_flashinfer_backend_support() - if (self.quant_config is None - or not self.quant_config.layer_quant_mode.has_any_quant( - exclude_kv_cache=True)) and not self.use_flashinfer: - raise NotImplementedError( - "TRTLLMGenFusedMoE BF16 path requires FlashInfer fused MoE. " - "Please install a FlashInfer build with trtllm_bf16_moe support." - ) backend_name = "flashinfer" if self.use_flashinfer else "trtllm" self.op_backend: MoEOpBackend = get_op_backend(backend_name) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py index e93d67d20398..0bf5e9b7ff85 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py @@ -13,12 +13,76 @@ from ...utils import ActivationType, is_gated_activation, relu2 from ..gated_mlp import GatedMLP from ..mlp import MLP -from .interface import MoEWeightLoadingMode +from .impl_contract import (MoEDeployment, MoEEligibility, MoEProblem, + MoERejectReason, MoEStaticCapability) +from .interface import MoEWeightLoadingMode, _reject from .routing import BaseMoeRoutingMethod class VanillaMoE(nn.ModuleList): + #: Declared explicitly because the resolver may return this ModuleList. + capabilities = MoEStaticCapability() + + #: Quantization labels supported by the Linear dispatcher. + _SUPPORTED_QUANT_LABELS = frozenset({ + "FP8", + "FP8_PER_CHANNEL_PER_TOKEN", + "FP8_BLOCK_SCALES", + "NVFP4", + "W4A16_NVFP4", + "W4A8_NVFP4_FP8", + "W4A8_MXFP4_FP8", + "W4A8_MXFP4_MXFP8", + "MXFP8", + "W8A16", + "W4A16", + "W4A16_AWQ", + "W4A8_AWQ", + }) + + @classmethod + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + """Check whether the PyTorch reference path has the required plumbing.""" + if p.quant is not None and p.quant not in cls._SUPPORTED_QUANT_LABELS: + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, + f"VanillaMoE dequantizes through Linear, which has no " + f"quant method for {p.quant}") + if p.swiglu_gptoss_style: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + "VanillaMoE has no bias / swiglu alpha-beta-limit parameters") + activation = p.activation_type + if (not is_gated_activation(activation) + and activation != ActivationType.Relu2): + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"VanillaMoE builds non-gated experts as MLP, whose only " + f"non-gated activation is Relu2 (got {p.activation})") + if d.eplb_enabled: + return _reject( + MoERejectReason.EPLB_UNSUPPORTED, + "VanillaMoE holds expert weights as plain submodules and " + "cannot expose them as migratable EPLB slots") + if d.smart_router: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"VanillaMoE has no smart-router path (moe_cluster_size=" + f"{d.cluster_size})") + # Uniform partitioning only. ``MoE._compute_ep_partition`` would hand a + # non-divisible count a ceil/floor split that this backend's local + # expert range does not implement, so it would produce wrong ranges + # rather than fail. + if (p.num_experts is not None and d.ep_size > 0 + and p.num_experts % d.ep_size != 0): + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"VanillaMoE partitions experts uniformly, so num_experts " + f"({p.num_experts}) must be divisible by ep_size " + f"({d.ep_size})") + return MoEEligibility.ok() + def __init__( self, *, @@ -49,16 +113,13 @@ def __init__( self.activation_type = activation_type self.is_gated_activation = is_gated_activation(activation_type) - # Limit support for VanillaMoE to non-gated activations for now. - if not self.is_gated_activation: - if pack_weights: - raise ValueError( - "pack_weights must be False for non-gated activations. Otherwise please update `create_weights`." - ) - if self.activation_type != ActivationType.Relu2: - raise ValueError( - f"Unsupported activation type: {self.activation_type} for non-gated activations. Only Relu2 is supported." - ) + # Activation eligibility (gated vs Relu2-only non-gated) is owned by + # ``can_implement``. ``pack_weights`` is a construction option that is + # not part of the problem/deployment question, so keep it here. + if not self.is_gated_activation and pack_weights: + raise ValueError( + "pack_weights must be False for non-gated activations. Otherwise please update `create_weights`." + ) self.dtype = dtype self.reduce_results = reduce_results @@ -69,9 +130,8 @@ def __init__( self.cluster_rank = model_config.mapping.moe_cluster_rank self.cluster_size = model_config.mapping.moe_cluster_size self.smart_router = True if self.cluster_size > 1 else False - assert not self.smart_router, ( - "Smart router is not supported in vanilla MoE, " - "please set moe_cluster_size to 1.") + # Smart-router / non-divisible EP eligibility is owned by + # ``can_implement``. self.rank = model_config.mapping.rank @@ -92,16 +152,6 @@ def __init__( self.intermediate_size_per_partition = intermediate_size // self.tp_size - # VanillaMoE uses uniform expert partitioning; non-divisible EP would require - # ceil/floor distribution (see MoE._compute_ep_partition in interface.py). - # Reject explicitly rather than silently producing wrong local-expert ranges. - if num_experts % self.ep_size != 0: - raise ValueError( - f"VanillaMoE does not support non-divisible EP: " - f"num_experts ({num_experts}) must be divisible by ep_size ({self.ep_size}). " - f"Use CutlassFusedMoE / TRTLLMGenFusedMoE with NVLINK_ONE_SIDED comm for non-divisible EP." - ) - self.expert_size_per_partition = num_experts // self.ep_size self.expert_start = self.ep_rank * self.expert_size_per_partition self.expert_end = min( diff --git a/tensorrt_llm/_torch/modules/fused_moe/impl_contract.py b/tensorrt_llm/_torch/modules/fused_moe/impl_contract.py index 4c2d76d3b5b1..258ead6e42a9 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/modules/fused_moe/impl_contract.py @@ -12,34 +12,21 @@ # 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. -"""Declaration, selection and execution contracts for MoE implementations. - -Every type here is a frozen, GPU-free dataclass: it can be constructed and -asserted on a machine with no device, which is what lets selection be unit -tested and lets offline tuning enumerate candidates ahead of deployment. - -The types split along three axes, and putting a field on the wrong one is the -mistake this file is shaped to prevent: - -- What an impl *can do* regardless of input -> MoEStaticCapability -- What an impl *demands of its caller* -> MoEInputRequirement -- Whether an impl fits *one concrete question* -> MoEEligibility - -A capability is a class-level declaration. An eligibility is a verdict about a -single (problem, deployment) pair. An input requirement is neither: it never -disqualifies an impl, it just tells the scheduler what to prepare. -""" +"""Contracts for declaring, selecting, and executing MoE implementations.""" import hashlib from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, List, Optional, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch if TYPE_CHECKING: - from .impl_identity import MoEImplId + from tensorrt_llm._torch.utils import ActivationType + from tensorrt_llm.models.modeling_utils import QuantAlgo + from .moe_load_balancer import SingleLayerMoeLoadBalancer + from .routing import BaseMoeRoutingMethod, RoutingMethodType # --------------------------------------------------------------------------- # Declaration @@ -48,33 +35,17 @@ @dataclass(frozen=True) class MoEStaticCapability: - """What this impl CAN do, independent of any particular input. - - Feeds SELECTION. Every field defaults to the conservative answer, so an - impl that forgets to declare gets excluded rather than silently accepted. - - Two rules keep this class from absorbing everything: a condition that - depends on the actual problem shape belongs in ``can_implement``, and a - condition the caller can simply satisfy by preparing its input belongs in - :class:`MoEInputRequirement`. - """ + """Static abilities used during selection. Defaults are conservative.""" # Legacy gate: ``moe.backend.__class__ == CutlassFusedMoE`` in MoEScheduler. supports_moe_lora: bool = False - # Legacy gate: the CuteDslFusedMoE isinstance check in - # ``ConfigurableMoE._should_enable_dwdp``. + # Legacy gate: CuteDslFusedMoE isinstance check in ConfigurableMoE DWDP. supports_dwdp: bool = False @dataclass(frozen=True) class MoEInputRequirement: - """What this impl REQUIRES THE CALLER to hand it. - - Read by MoEScheduler while assembling :class:`MoERunContext`, and by the - comm strategy while building :class:`MoECommPlan`. Deliberately NOT part of - selection: an unmet input requirement is the scheduler's job to satisfy, - never a reason to pick a different impl. - """ + """Caller-side inputs the scheduler must prepare; not a selection filter.""" # Legacy: the ``token_final_scales`` bfloat16 / float32 casts in # MoEScheduler. @@ -88,10 +59,7 @@ class MoEInputRequirement: # buffer follows the model output dtype. onesided_workspace_dtype: Optional[torch.dtype] = None - # There is deliberately no sentinel field here. Every Communication sets - # ``invalid_token_expert_id = -1`` in its own __init__, and TRTLLM-Gen - # kernels accept nothing else, so the value is a comm-side invariant rather - # than something an impl needs the caller to supply. + # No sentinel for invalid_token_expert_id: every Communication uses -1. # There is deliberately no ``requires_router_logits`` field either, though # the design sketched one to replace the scheduler's router-logits filter. @@ -108,27 +76,96 @@ class MoEInputRequirement: @dataclass(frozen=True) class MoEProblem: - """The part of the question that is REUSABLE ACROSS DEPLOYMENTS. - - Normative rule for adding a field here: if changing the field must - invalidate an already-persisted tuning result, it belongs in MoEProblem; if - it only changes WHICH impls are eligible, it belongs in - :class:`MoEDeployment`. This is also the tuning key, so field order and - types are part of the on-disk format. - - ``quant`` / ``dtype_act`` / ``swiglu_gptoss_style`` are exactly today's - three arguments of today's ``MoE.can_implement``. The shape fields are - new: the tuning winner depends on them, yet today they are only checked - later, inside ``__init__`` / ``validate``. + """Reusable tuning-key inputs for a MoE layer. + + Eligibility gates abstain when an optional field is unknown. """ - quant: Optional[str] # canonical quant name, or None for bf16 + quant: Optional[str] # QuantAlgo value, or None for unquantized dtype_act: torch.dtype # activation dtype BEFORE quantization - hidden_size: int - intermediate_size: int - num_experts: int - top_k: int - swiglu_gptoss_style: bool = False + hidden_size: Optional[int] = None + intermediate_size: Optional[int] = None + num_experts: Optional[int] = None + top_k: Optional[int] = None + #: Tri-state because some call sites cannot distinguish gpt-oss SwiGLU. + swiglu_gptoss_style: Optional[bool] = None + #: Expert FC bias, distinct from ``swiglu_gptoss_style``. MiniMax sets + #: SwigluBias + alpha/beta/limit with ``bias=False``; gpt-oss sets both. + bias: Optional[bool] = None + #: ``ActivationType`` member name; omitted values canonicalize to SwiGLU. + activation: str = "Swiglu" + #: ``RoutingMethodType`` member name; None means the call site did not say. + routing: Optional[str] = None + + @property + def routing_method_type(self) -> Optional["RoutingMethodType"]: + """``routing`` as the enum member, or ``None`` when unknown.""" + from .routing import RoutingMethodType + + if self.routing is None: + return None + return RoutingMethodType[self.routing] + + @property + def activation_type(self) -> "ActivationType": + """Return ``activation`` as an enum member.""" + from tensorrt_llm._torch.utils import ActivationType + + return ActivationType[self.activation] + + @property + def quant_algo(self) -> Optional["QuantAlgo"]: + """Return ``quant`` as an enum member.""" + if self.quant is None: + return None + from tensorrt_llm.models.modeling_utils import QuantAlgo + + return QuantAlgo(self.quant) + + @property + def is_fully_specified(self) -> bool: + """Whether this problem can key a persisted tuning result.""" + return None not in (self.hidden_size, self.intermediate_size, self.num_experts, self.top_k) + + +def canonical_quant(quant_algo: Optional["QuantAlgo"]) -> Optional[str]: + """Canonicalize a quantization algorithm for the tuning key.""" + if quant_algo is None: + return None + from tensorrt_llm.models.modeling_utils import QuantAlgo + + aliases = { + # Calibration recipes; the weights and the kernel are plain NVFP4. + QuantAlgo.NVFP4_AWQ: QuantAlgo.NVFP4, + QuantAlgo.NVFP4_ARC: QuantAlgo.NVFP4, + # MIXED_PRECISION is a model-level marker, not a layer format. + QuantAlgo.MIXED_PRECISION: None, + QuantAlgo.NO_QUANT: None, + } + resolved = aliases.get(quant_algo, quant_algo) + return None if resolved is None else str(resolved.value) + + +def canonical_activation(activation_type: Optional["ActivationType"]) -> str: + """Canonicalize an activation for the tuning key.""" + from tensorrt_llm._torch.utils import ActivationType + + if activation_type is None: + return ActivationType.Swiglu.name + return ActivationType(activation_type).name + + +def canonical_routing( + routing: Optional["BaseMoeRoutingMethod | RoutingMethodType"], +) -> Optional[str]: + """Canonicalize a routing method or method type for the tuning key.""" + from .routing import RoutingMethodType + + if routing is None: + return None + if not isinstance(routing, RoutingMethodType): + routing = routing.routing_method_type + return RoutingMethodType(routing).name @dataclass(frozen=True) @@ -149,35 +186,33 @@ def has_dep(self, name: str) -> bool: return name in self.available_deps def fingerprint(self) -> str: - """Provenance stamp for a tuning result. - - Records WHICH machine state produced a given winner, so that replaying - under a different environment is detectable instead of silently - selecting someone else. - """ + """Return a stable fingerprint for the selection environment.""" payload = repr((self.sm, sorted(self.available_deps), self.env_flags)) return hashlib.sha256(payload.encode()).hexdigest()[:16] @dataclass(frozen=True) class MoEDeployment: - """Topology and slot layout. - - Changing these changes WHICH impls are eligible, but never invalidates a - tuning result. - """ + """Topology and slot layout used for eligibility.""" ep_size: int tp_size: int use_dp: bool num_slots: int env: MoEEnvironment + # Whole parallel-group width from mapping.tp_size. + parallel_size: int + # mapping.moe_cluster_size; values above one enable the smart router. + cluster_size: int = 1 + # True only when an EPLB load balancer is registered. + eplb_enabled: bool = False + # True only for routed-expert LoRA targets. + moe_lora_enabled: bool = False @property - def parallel_size(self) -> int: - # Matches today's ``self.use_dp and self.parallel_size > 1`` test in - # ``TRTLLMGenFusedMoE._supports_load_balancer``. - return self.ep_size * self.tp_size + def smart_router(self) -> bool: + """Mirrors ``MoE.smart_router`` (``interface.py``), its only definition.""" + return self.cluster_size > 1 # --------------------------------------------------------------------------- @@ -186,7 +221,12 @@ def parallel_size(self) -> int: class MoERejectReason(str, Enum): - """Closed enum. Tests assert on these, never on log substrings.""" + """Closed enum. Tests assert on these, never on log substrings. + + Closed because the reasons are an API: a test that wants "this request was + turned down for the right cause" must be able to name the cause, and a + free-form string cannot be named without pinning the wording too. + """ QUANT_UNSUPPORTED = "quant_unsupported" DTYPE_UNSUPPORTED = "dtype_unsupported" @@ -196,6 +236,24 @@ class MoERejectReason(str, Enum): SLOTS_NOT_DIVISIBLE_BY_EP = "slots_not_divisible_by_ep" TOPOLOGY_UNSUPPORTED = "topology_unsupported" LORA_UNSUPPORTED = "lora_unsupported" + # Activation shape the impl cannot serve (today: swiglu_gptoss_style, i.e. + # bias plus custom swiglu alpha/beta/limit). + ACTIVATION_UNSUPPORTED = "activation_unsupported" + # The routing method produces scores in a form the impl's kernel cannot + # consume. Distinct from TOPOLOGY_UNSUPPORTED: nothing about the parallel + # layout is wrong, the impl just fuses one routing shape and no other. + ROUTING_UNSUPPORTED = "routing_unsupported" + # EPLB is registered for this layer and the impl cannot lay out slots for + # it. Distinct from TOPOLOGY_UNSUPPORTED: the parallel sizes are fine. + EPLB_UNSUPPORTED = "eplb_unsupported" + # Not a capability verdict: the impl could run, but the resolver refuses to + # route production traffic there. Kept separate so that "we chose not to" + # never reads as "it cannot". + PATH_NOT_ENABLED = "path_not_enabled" + # The named backend no longer exists (today: WIDEEP). + BACKEND_DEPRECATED = "backend_deprecated" + # A wrapper or aggregate that is never itself an execution unit. + NOT_AN_IMPL = "not_an_impl" @dataclass(frozen=True) @@ -204,7 +262,7 @@ class MoEEligibility: It carries the verdict and, on rejection, a closed-enum reason. It holds no execution parameters and no identity -- those live in :class:`MoERunContext` - and ``MoEImplId``. + and, after the leaf-class migration, ``MoEImplDescriptor``. """ eligible: bool @@ -218,25 +276,130 @@ def __post_init__(self) -> None: if not self.eligible and self.reject_reason is None: raise ValueError("a rejection must name a MoERejectReason") + def __bool__(self) -> bool: + return self.eligible + + @classmethod + def ok(cls) -> "MoEEligibility": + return cls(eligible=True) + + @classmethod + def no(cls, reason: MoERejectReason, detail: str) -> "MoEEligibility": + """Reject with a machine-readable cause and a human-readable detail. + + ``detail`` is required rather than optional: a reason code narrows the + cause to a category, and the operator still needs the specific value + that failed the gate. + """ + return cls(eligible=False, reject_reason=reason, detail=detail) + @dataclass(frozen=True) -class MoEResolutionReport: - """Structured answer to "who got picked, and why was everyone else out". +class MoERejection: + """One candidate that did not win, and why.""" + + legacy_backend: str + reason: MoERejectReason + detail: str = "" + + def to_dict(self) -> Dict[str, str]: + return { + "legacy_backend": self.legacy_backend, + "reason": self.reason.value, + "detail": self.detail, + } - Replaces today's read-the-logs workflow. Two consumers: offline tuning - (this IS the candidate set) and tests (assert on ``reject_reason``). - """ + +@dataclass(frozen=True) +class MoEResolutionReport: + """Selected implementation, eligible alternatives, and rejection reasons.""" problem: MoEProblem deployment: MoEDeployment - winner: Optional["MoEImplId"] # None => hard failure - rejected: Tuple[Tuple["MoEImplId", MoERejectReason], ...] = () - selected_by: str = "auto" # "auto" | "pin" + winner: Optional[str] # legacy backend class name; None => hard failure + # Selection mode: pinned, heuristic fallback, or failed. + selected_by: str + rejected: Tuple[MoERejection, ...] = () + # Eligible candidates in priority order; eligible[0] is the winner. + eligible: Tuple[str, ...] = () + requested: Optional[str] = None # backend literal, as written env_fingerprint: str = "" @property - def eligible(self) -> Tuple["MoEImplId", ...]: - return (self.winner,) if self.winner is not None else () + def alternatives(self) -> Tuple[str, ...]: + """Eligible impls that lost to the winner on priority alone. + + Not rejections: nothing is wrong with these, they simply ranked lower. + Keeping the two lists apart is the point -- "could not run" and "ran + second" call for opposite responses from whoever reads the report. + """ + return self.eligible[1:] + + @property + def degraded(self) -> bool: + """Whether the caller got something other than what it asked for.""" + return self.selected_by == "heuristic" + + @property + def degraded_from(self) -> Optional[MoERejection]: + """The rejection that caused the substitution, if there was one.""" + if not self.degraded or not self.rejected: + return None + # The last family rejection best explains the fallback. + return self.rejected[-1] + + def to_dict(self) -> Dict[str, object]: + """Serializable form. Field names are part of the artifact format.""" + return { + "winner": self.winner, + "requested": self.requested, + "selected_by": self.selected_by, + "env_fingerprint": self.env_fingerprint, + "eligible": list(self.eligible), + "rejected": [rejection.to_dict() for rejection in self.rejected], + "problem": { + "quant": self.problem.quant, + "dtype_act": str(self.problem.dtype_act), + "hidden_size": self.problem.hidden_size, + "intermediate_size": self.problem.intermediate_size, + "num_experts": self.problem.num_experts, + "top_k": self.problem.top_k, + "swiglu_gptoss_style": self.problem.swiglu_gptoss_style, + "bias": self.problem.bias, + "activation": self.problem.activation, + "routing": self.problem.routing, + }, + "deployment": { + "ep_size": self.deployment.ep_size, + "tp_size": self.deployment.tp_size, + "parallel_size": self.deployment.parallel_size, + "cluster_size": self.deployment.cluster_size, + "use_dp": self.deployment.use_dp, + "num_slots": self.deployment.num_slots, + "eplb_enabled": self.deployment.eplb_enabled, + "moe_lora_enabled": self.deployment.moe_lora_enabled, + "sm": self.deployment.env.sm, + "env_flags": dict(self.deployment.env.env_flags), + }, + } + + def describe(self) -> str: + """One line for the log. Reads as a sentence, not as a dict dump.""" + winner = "none" if self.winner is None else self.winner + head = f"MoE resolution: {winner} (via {self.selected_by}" + if self.requested is not None: + head += f", requested {self.requested}" + head += f", env {self.env_fingerprint})" + if self.alternatives: + # Named in the same line as the winner, because this is the list an + # operator retries one by one when the default is not fast enough. + head += f"; also eligible: {', '.join(self.alternatives)}" + if not self.rejected: + return head + turned_down = ", ".join( + f"{rejection.legacy_backend}={rejection.reason.value}" for rejection in self.rejected + ) + return f"{head}; turned down: {turned_down}" # --------------------------------------------------------------------------- @@ -296,42 +459,18 @@ class MoERunContext: def require_comm_plan(impl: object, ctx: MoERunContext) -> MoECommPlan: - """The plan for this forward, for impls that cannot run without one. - - ``comm_plan`` is optional on the context because a fused-comm impl owns the - EP exchange itself, so nothing outside its kernel decided anything about the - forward. Every external-comm impl is the opposite case: it is only reachable - through ``ExternalCommMoEScheduler``, which builds a plan on every path. - - Substituting defaults instead of failing is what this guards against. A - wrong ``moe_output`` or ``enable_alltoall`` surfaces as a shape or kernel - error, but a wrong ``input_sf_swizzled`` does not: the kernel reads scale - factors at the stride it was told, so a plan-less default of "swizzled" - against unswizzled input returns silently wrong numbers. - """ - if ctx.comm_plan is None: - # Not an assert: silently-wrong output is the failure mode this exists - # to prevent, so the check must survive ``python -O``. - raise ValueError( - f"{type(impl).__name__}.run_moe needs ctx.comm_plan, and the scheduler " - "that drives it always supplies one. A missing plan means run_moe was " - "called without going through ExternalCommMoEScheduler." - ) + """Return the required external-communication plan for this forward.""" + assert ctx.comm_plan is not None, ( + f"{type(impl).__name__}.run_moe needs ctx.comm_plan, and the scheduler " + "that drives it always supplies one. A missing plan means run_moe was " + "called without going through ExternalCommMoEScheduler." + ) return ctx.comm_plan @dataclass(frozen=True) class MoEEplbBinding: - """Everything an impl needs to lay out and load its expert weights. - - Computed once by whoever owns the load-balancer registration, then passed as - an explicit constructor argument -- never ``setattr``'d after construction. - That is the entire point: weight shapes depend on these values, so they must - be known BEFORE ``create_weights()``, not patched in afterwards. - - Excludes ``repeat_idx`` / ``repeat_count`` on purpose: those are - forward-time scheduling state owned by the wrapper. - """ + """EPLB expert layout required before weight creation.""" layer_idx: int num_slots: int diff --git a/tensorrt_llm/_torch/modules/fused_moe/impl_environment.py b/tensorrt_llm/_torch/modules/fused_moe/impl_environment.py new file mode 100644 index 000000000000..4dbbe1f19bfe --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/impl_environment.py @@ -0,0 +1,178 @@ +# 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. +"""Collect and freeze machine inputs used for MoE selection.""" + +import os +from contextlib import contextmanager +from enum import Enum +from typing import Callable, Dict, Optional, Tuple + +from tensorrt_llm.logger import logger + +from .impl_contract import MoEEnvironment + + +class MoEDep(str, Enum): + """Optional dependencies that affect MoE selection.""" + + #: ``import flashinfer`` succeeds. Gates the SM120/SM121 NVFP4 decode + #: backend (``CuteDslB12xFusedMoE``). + FLASHINFER = "flashinfer" + #: FlashInfer additionally exposes ``trtllm_bf16_moe`` / + #: ``trtllm_bf16_routed_moe``. Strictly stronger than :attr:`FLASHINFER` + #: and gates the TRTLLM-Gen unquantized BF16 path. + FLASHINFER_BF16_MOE = "flashinfer_bf16_moe" + #: The bundled DeepGEMM build exposes the ``fp8_fp4_mega_moe`` kernel. + DEEPGEMM_MEGAMOE = "deepgemm_megamoe" + #: ``nvidia-cutlass-dsl[cu13]`` is new enough for the MegaMoE CuteDSL ABI. + MEGAMOE_CUTEDSL_RUNTIME = "megamoe_cutedsl_runtime" + #: The ``trtllm::cute_dsl_megamoe_nvfp4_*`` custom ops are registered. + MEGAMOE_CUTEDSL_OP = "megamoe_cutedsl_op" + + +class MoEEnvFlag(str, Enum): + """Environment variables that MoE selection is allowed to read.""" + + #: Opt-in to the FlashInfer provider for quantized TRTLLM-Gen. Changes the + #: routing split, which is why load-balancer eligibility depends on it. + TRTLLM_GEN_USE_FLASHINFER = "TRTLLM_GEN_FUSED_MOE_USE_FLASHINFER" + + +# Probe details are logged but excluded from the stable fingerprint. +DepProbe = Callable[[], Tuple[bool, str]] + + +def _probe_flashinfer() -> Tuple[bool, str]: + try: + import flashinfer # noqa: F401 + except Exception as exc: # noqa: BLE001 - any import failure means absent + return False, f"import flashinfer failed: {exc}" + return True, "" + + +def _probe_flashinfer_bf16_moe() -> Tuple[bool, str]: + try: + from flashinfer.fused_moe import core as _core + except Exception as exc: # noqa: BLE001 - any import failure means absent + return False, f"import flashinfer.fused_moe.core failed: {exc}" + missing = [ + symbol + for symbol in ("trtllm_bf16_moe", "trtllm_bf16_routed_moe") + if not hasattr(_core, symbol) + ] + if missing: + return False, f"flashinfer.fused_moe.core lacks {', '.join(missing)}" + return True, "" + + +def _probe_deepgemm_megamoe() -> Tuple[bool, str]: + from .quantization import _import_deep_gemm, _MegaMoEUnavailable + + try: + _import_deep_gemm() + except _MegaMoEUnavailable as exc: + return False, str(exc) + return True, "" + + +def _probe_megamoe_cutedsl_runtime() -> Tuple[bool, str]: + from .mega_moe.mega_moe_cute_dsl import is_megamoe_cute_dsl_runtime_available + + available, reason = is_megamoe_cute_dsl_runtime_available() + return bool(available), "" if available else str(reason) + + +def _probe_megamoe_cutedsl_op() -> Tuple[bool, str]: + # Read the module because registration updates this flag after import. + from ...custom_ops import cute_dsl_megamoe_custom_op as megamoe_op + + if megamoe_op.IS_MEGAMOE_OP_AVAILABLE: + return True, "" + return False, str(megamoe_op.MEGAMOE_OP_UNAVAILABLE_REASON) + + +_DEP_PROBES: Dict[MoEDep, DepProbe] = { + MoEDep.FLASHINFER: _probe_flashinfer, + MoEDep.FLASHINFER_BF16_MOE: _probe_flashinfer_bf16_moe, + MoEDep.DEEPGEMM_MEGAMOE: _probe_deepgemm_megamoe, + MoEDep.MEGAMOE_CUTEDSL_RUNTIME: _probe_megamoe_cutedsl_runtime, + MoEDep.MEGAMOE_CUTEDSL_OP: _probe_megamoe_cutedsl_op, +} + +# Preserve prior defaults when environment variables are unset. +_ENV_FLAG_DEFAULTS: Dict[MoEEnvFlag, str] = { + MoEEnvFlag.TRTLLM_GEN_USE_FLASHINFER: "0", +} + +_CACHED_ENVIRONMENT: Optional[MoEEnvironment] = None +_OVERRIDE_ENVIRONMENT: Optional[MoEEnvironment] = None + + +def _run_probe(dep: MoEDep, probe: DepProbe) -> bool: + try: + available, detail = probe() + except Exception as exc: # noqa: BLE001 - a broken probe means "absent" + # Treat broken probes as unavailable while keeping the failure visible. + logger.warning(f"MoE dependency probe {dep.value} raised {type(exc).__name__}: {exc}") + return False + if not available: + logger.debug(f"MoE dependency {dep.value} unavailable: {detail}") + return available + + +def collect_moe_environment(force: bool = False) -> MoEEnvironment: + """Collect and cache the frozen MoE selection environment.""" + global _CACHED_ENVIRONMENT + if _OVERRIDE_ENVIRONMENT is not None: + return _OVERRIDE_ENVIRONMENT + if _CACHED_ENVIRONMENT is not None and not force: + return _CACHED_ENVIRONMENT + + from tensorrt_llm._utils import get_sm_version + + available = tuple( + sorted(dep.value for dep, probe in _DEP_PROBES.items() if _run_probe(dep, probe)) + ) + env_flags = tuple( + sorted( + (flag.value, os.environ.get(flag.value, default)) + for flag, default in _ENV_FLAG_DEFAULTS.items() + ) + ) + environment = MoEEnvironment(sm=get_sm_version(), available_deps=available, env_flags=env_flags) + logger.debug( + f"collected MoE environment: sm={environment.sm} deps={available} " + f"flags={env_flags} ({environment.fingerprint()})" + ) + _CACHED_ENVIRONMENT = environment + return environment + + +def reset_moe_environment_cache() -> None: + """Drop the cached probe result. For tests that change probe outcomes.""" + global _CACHED_ENVIRONMENT + _CACHED_ENVIRONMENT = None + + +@contextmanager +def override_moe_environment(environment: MoEEnvironment): + """Temporarily override the collected MoE selection environment.""" + global _OVERRIDE_ENVIRONMENT + previous = _OVERRIDE_ENVIRONMENT + _OVERRIDE_ENVIRONMENT = environment + try: + yield environment + finally: + _OVERRIDE_ENVIRONMENT = previous diff --git a/tensorrt_llm/_torch/modules/fused_moe/impl_identity.py b/tensorrt_llm/_torch/modules/fused_moe/impl_identity.py index e59f76574571..9ca4157f1c2f 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/impl_identity.py +++ b/tensorrt_llm/_torch/modules/fused_moe/impl_identity.py @@ -12,72 +12,72 @@ # 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. -"""Stable identity and registration for MoE implementations. - -A Python class name cannot serve as the identity of an implementation: today a -single ``TRTLLMGenFusedMoE`` class covers eleven distinct quant x provider -combinations, so "which implementation ran" is not answerable from the class -alone. :class:`MoEImplId` makes that identity explicit, serializable and -round-trippable, which is what a persisted tuning result must key on. -""" +"""Stable identities, queries, and registration for leaf MoE implementations.""" import re from dataclasses import dataclass, field -from typing import Optional, TypeVar +from typing import Dict, List, Optional, Tuple, Type, TypeVar from .impl_contract import MoEInputRequirement, MoEStaticCapability from .interface import MoESchedulerKind _FIELD_RE = re.compile(r"^[a-z0-9]+(_[a-z0-9]+)*$") _SEP = "." - -# Registration is a decorator, so it must hand the CONCRETE class back. A bare -# ``type`` return would erase the decorated impl down to ``type[Any]``. T = TypeVar("T") +# Canonical written order; tokens are assigned to fields by value. +_ID_FIELDS: Tuple[str, ...] = ("provider", "technique", "kernel_name", "quant") +_FIELD_INDEX: Dict[str, int] = {name: i for i, name in enumerate(_ID_FIELDS)} + +# Explicit "any value here". Never required -- an omitted field is already +# unconstrained -- but accepted so that a query can be written out at full +# width, and so that MoEImplQuery.describe() round-trips back through +# MoEImplRegistry.parse_query(). +_WILDCARD = "*" + + +def _normalize(name: str, value: str) -> str: + """Lowercase an incoming segment and reject anything still malformed. + + Case folding happens here rather than at each call site so that a value + typed as ``Cutlass`` and one typed as ``cutlass`` cannot become two + different registry keys. The canonical form is always lowercase. + """ + if not isinstance(value, str): + raise TypeError(f"MoEImplId.{name} must be a string, got {type(value).__name__}") + folded = value.strip().lower() + if not _FIELD_RE.match(folded): + raise ValueError(f"MoEImplId.{name}={value!r} must match {_FIELD_RE.pattern}") + return folded + @dataclass(frozen=True) class MoEImplId: - """Serializable identity of ONE MoE implementation. - - Four fields joined by ``'.'``. Each field may contain ``'_'`` internally, - which is exactly why ``'.'`` is the separator: quant values such as - ``w4a8_mxfp4_mxfp8`` already carry underscores, so an all-underscore - encoding would not be parseable back into four fields. - - The four fields together must pick out exactly ONE kernel. The first three - narrow the space; ``kernel_name`` carries the final disambiguation, so it - names the specific kernel and not a category it belongs to. Naming a family - instead -- ``blockscale`` when two block-scale kernels differ in, say, - their tiling -- makes both kernels compute the same id, and - :class:`MoEImplRegistry` then rejects the second as a duplicate, so it - cannot be registered at all. When two kernels would otherwise collide, - extend ``kernel_name`` with whatever actually separates them rather than - overloading one of the other three fields. - """ + """Exact ``provider.technique.kernel_name.quant`` implementation identity.""" - provider: str # maintaining entity: trtllm_native | flashinfer | deepgemm - technique: str # impl tech: cutlass | cutedsl | cuda_cpp | trtllm_gen - quant: str # weight/act format: nvfp4 | w4a8_mxfp4_mxfp8 | bf16 - # Where uniqueness of the whole id is won or lost: anything the three - # fields above leave ambiguous has to be spelled out here, or two distinct - # kernels collide on one id and the registry refuses to accept the second. - kernel_name: str # densegemm | blockscale | blockscale_splitk + # Kernel lineage, e.g. trtllm, flashinfer, deepgemm, or marlin. + provider: str # trtllm | flashinfer | deepgemm | marlin | triton_kernels + # Implementation technology, independent of provider. + technique: str # cutlass | cutedsl | cuda_cpp | trtllm_gen | triton | torch + # Specific kernel name that makes the full identity unique. + kernel_name: str # grouped_gemm | dense_gemm | fused_moe | mega_moe | vanilla + quant: str # weight/act format; ``none`` when unquantized - def __post_init__(self) -> None: - for name in ("provider", "technique", "quant", "kernel_name"): - value = getattr(self, name) - if not _FIELD_RE.match(value): - raise ValueError(f"MoEImplId.{name}={value!r} must match {_FIELD_RE.pattern}") + def __post_init__(self): + for name in _ID_FIELDS: + # Frozen dataclass, so normalization has to go around __setattr__. + object.__setattr__(self, name, _normalize(name, getattr(self, name))) def canonical(self) -> str: - return _SEP.join((self.provider, self.technique, self.quant, self.kernel_name)) + return _SEP.join(getattr(self, name) for name in _ID_FIELDS) @classmethod def parse(cls, text: str) -> "MoEImplId": parts = text.split(_SEP) - if len(parts) != 4: - raise ValueError(f"expected 4 {_SEP!r}-separated fields, got {len(parts)}: {text!r}") + if len(parts) != len(_ID_FIELDS): + raise ValueError( + f"expected {len(_ID_FIELDS)} {_SEP!r}-separated fields, got {len(parts)}: {text!r}" + ) # Going through the constructor is what checks the segments: the count # check above says nothing about their contents, and parse() is the # untrusted door -- user YAML and persisted tuning results come in here. @@ -87,6 +87,60 @@ def __str__(self) -> str: return self.canonical() +@dataclass(frozen=True) +class MoEImplQuery: + """Partial implementation identity; None leaves a field unconstrained.""" + + provider: Optional[str] = None + technique: Optional[str] = None + kernel_name: Optional[str] = None + quant: Optional[str] = None + + def __post_init__(self): + for name in _ID_FIELDS: + value = getattr(self, name) + if value is not None: + object.__setattr__(self, name, _normalize(name, value)) + + @property + def is_empty(self) -> bool: + """No constraint at all, i.e. every registered impl matches.""" + return all(getattr(self, name) is None for name in _ID_FIELDS) + + @property + def is_exact(self) -> bool: + """Every field pinned, so this names at most one implementation. + + The distinction drives failure semantics: an exact query that matches + nothing is a hard error, while a partial one that matches several is + resolved by priority. + """ + return all(getattr(self, name) is not None for name in _ID_FIELDS) + + def as_impl_id(self) -> MoEImplId: + """The single id this query names. Only valid when :attr:`is_exact`.""" + if not self.is_exact: + raise ValueError(f"query {self.describe()} does not pin all {len(_ID_FIELDS)} fields") + return MoEImplId(**{name: getattr(self, name) for name in _ID_FIELDS}) + + def matches(self, identity: MoEImplId) -> bool: + """Whether ``identity`` satisfies every field this query does pin.""" + return all( + getattr(self, name) is None or getattr(self, name) == getattr(identity, name) + for name in _ID_FIELDS + ) + + def describe(self) -> str: + """Full-width rendering, ``*`` for the fields left open. + + Round-trips: :meth:`MoEImplRegistry.parse_query` accepts this back. + """ + return _SEP.join(getattr(self, name) or _WILDCARD for name in _ID_FIELDS) + + def __str__(self): + return self.describe() + + @dataclass(frozen=True) class MoEImplDescriptor: """Declaration-time metadata attached to one implementation class. @@ -108,10 +162,22 @@ def impl_id(self) -> str: class MoEImplRegistry: - """MoEImplId -> implementation class. Duplicate identity is a hard error.""" + """Map unique implementation identities and query tokens to classes.""" + + def __init__(self): + self._store: Dict[MoEImplId, Type] = {} + self._token_to_field: Dict[str, str] = {} - def __init__(self) -> None: - self._store: dict[MoEImplId, type] = {} + def _check_tokens_disjoint(self, identity: MoEImplId) -> None: + for name in _ID_FIELDS: + token = getattr(identity, name) + owner = self._token_to_field.get(token) + if owner is not None and owner != name: + raise ValueError( + f"cannot register {identity.canonical()}: token {token!r} is already a " + f"value of field {owner!r}, so a user writing {token!r} could mean either " + f"field. Rename one of them -- value sets must stay disjoint across fields." + ) def register(self, cls: type[T]) -> type[T]: descriptor = getattr(cls, "descriptor", None) @@ -125,13 +191,82 @@ def register(self, cls: type[T]) -> type[T]: raise ValueError( f"duplicate MoEImplId {identity.canonical()}: {previous.__name__} vs {cls.__name__}" ) + self._check_tokens_disjoint(identity) self._store[identity] = cls + for name in _ID_FIELDS: + self._token_to_field[getattr(identity, name)] = name return cls def lookup(self, identity: MoEImplId) -> Optional[type]: return self._store.get(identity) - def __len__(self) -> int: + def field_of(self, token: str) -> Optional[str]: + """Which id field a bare token belongs to, or ``None`` if unknown.""" + return self._token_to_field.get(token.strip().lower()) + + def known_tokens(self) -> Dict[str, str]: + """Copy of the token vocabulary, for diagnostics and error messages.""" + return dict(self._token_to_field) + + def parse_query(self, text: str) -> MoEImplQuery: + """Parse a partial, canonically ordered implementation query.""" + tokens = [] + for raw in text.split(_SEP): + token = raw.strip().lower() + if not token: + raise ValueError(f"empty segment in MoE impl specification {text!r}") + tokens.append(token) + + # Pass 1: assign by value. + assignment: Dict[str, str] = {} + for token in tokens: + if token == _WILDCARD: + continue + name = self._token_to_field.get(token) + if name is None: + raise ValueError( + f"unknown MoE impl token {token!r} in {text!r}. " + f"Known tokens: {sorted(self._token_to_field)}" + ) + if name in assignment: + raise ValueError( + f"MoE impl specification {text!r} sets field {name!r} twice: " + f"{assignment[name]!r} and {token!r}" + ) + assignment[name] = token + query = MoEImplQuery(**assignment) + + # Pass 2: check canonical order. + cursor = 0 + for token in tokens: + if token == _WILDCARD: + # Only a wildcard can overrun: a named token past the end is + # out of order, and saying so points at the real mistake. + if cursor >= len(_ID_FIELDS): + raise ValueError( + f"MoE impl specification {text!r} has more segments than the " + f"{len(_ID_FIELDS)} fields {_SEP.join(_ID_FIELDS)}" + ) + cursor += 1 + continue + index = _FIELD_INDEX[self._token_to_field[token]] + if index < cursor: + raise ValueError( + f"MoE impl specification {text!r} is out of order at {token!r}: " + f"segments must follow {_SEP.join(_ID_FIELDS)}. " + f"Write it as {query.describe()!r}." + ) + cursor = index + 1 + return query + + def find(self, query: MoEImplQuery) -> List[Tuple[MoEImplId, Type]]: + """Every registered impl the query matches, in registration order.""" + return [(ident, cls) for ident, cls in self._store.items() if query.matches(ident)] + + def identities(self) -> Tuple[MoEImplId, ...]: + return tuple(self._store) + + def __len__(self): return len(self._store) diff --git a/tensorrt_llm/_torch/modules/fused_moe/interface.py b/tensorrt_llm/_torch/modules/fused_moe/interface.py index 85e2689e9905..dc4041699411 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/interface.py +++ b/tensorrt_llm/_torch/modules/fused_moe/interface.py @@ -22,11 +22,9 @@ import torch from torch import nn -from tensorrt_llm.logger import logger -from tensorrt_llm.models.modeling_utils import QuantAlgo - from ...distributed.ops import reducescatter -from .impl_contract import (MoEInputRequirement, MoERunContext, +from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement, + MoEProblem, MoERejectReason, MoERunContext, MoEStaticCapability) # Route on the host (fused noaux_tc + post-topk pipeline) instead of inside @@ -42,22 +40,9 @@ "TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING", "0") == "1" -def _warn_and_return(reason: str) -> Tuple[bool, Optional[str]]: - """ - Log a warning and return (False, reason) for can_implement() checks. - - This is a common utility function used by all MoE backend implementations - to provide consistent logging and return values when a configuration - is not supported. - - Args: - reason: The reason why the configuration is not supported. - - Returns: - Tuple[bool, Optional[str]]: Always returns (False, reason) - """ - logger.warning(reason) - return False, reason +def _reject(reason: MoERejectReason, detail: str) -> MoEEligibility: + """Create a silent ``can_implement`` rejection.""" + return MoEEligibility.no(reason, detail) from ...model_config import ModelConfig @@ -241,26 +226,10 @@ class MoE(nn.Module): # override this to ``MoESchedulerKind.FUSED_COMM``. scheduler_kind: MoESchedulerKind = MoESchedulerKind.EXTERNAL_COMM - # What this backend can do, read by callers that would otherwise test its - # class. A backend deriving from another backend MUST restate every field - # rather than inherit it: the exact-class comparisons these fields replace - # answered False for subclasses, so a capability picked up through - # inheritance would silently widen behaviour. - # - # That restatement rule is transitional, not the intended end state. It only - # has to exist while backends still derive from other backends, which today - # they do: CuteDsl, DeepGemm and Marlin derive from Cutlass, B12x from - # CuteDsl, and Llama4MinLatency from Cutlass. The per-backend tickets - # TRTLLM-14960..14969 cut those inheritance edges as each impl moves its - # run_moe into its own leaf class, and once no impl derives from another the - # rule has nothing left to guard and should be deleted with it. + # Subclasses must restate capabilities to preserve exact-class behavior. capabilities: MoEStaticCapability = MoEStaticCapability() - # What this backend needs the scheduler to hand it. Unlike - # ``capabilities``, the checks these fields replace used isinstance, so - # inheriting a value is correct here. Overriding is not partial though: - # the whole object is replaced, so a subclass that sets one field must - # restate the ones it still wants from its parent. + # Scheduler-provided inputs; inherited values remain valid for subclasses. input_requirement: MoEInputRequirement = MoEInputRequirement() # Opt-in flag for non-divisible EP (num_experts % ep_size != 0). False by default @@ -271,34 +240,10 @@ class MoE(nn.Module): @classmethod @abstractmethod - def can_implement( - cls, - quant_algo: Optional[QuantAlgo], - dtype_activation: torch.dtype = torch.bfloat16, - swiglu_gptoss_style: bool = False, - ) -> Tuple[bool, Optional[str]]: - """ - Check if this MoE backend can implement the given quantization algorithm. - - NOTE: This is a TRANSITIONAL interface. In the future, this method will be moved - to the MoEBackend interface as part of the backend abstraction layer. During this - transition period, it remains in the MoE base class to maintain compatibility. - - This method checks both: - 1. Whether the backend supports the specified quantization algorithm - 2. Whether the current platform (SM version) supports the backend and quantization - - Each backend MUST override this method to provide accurate capability information. + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + """Purely evaluate ``p`` and ``d`` without probing runtime state. - Args: - quant_algo: The quantization algorithm to check (None for unquantized) - dtype_activation: The activation data type. - swiglu_gptoss_style: Whether swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit) is enabled. - - Returns: - Tuple[bool, Optional[str]]: (can_implement, skip_reason) - - can_implement: True if the backend can implement this configuration - - skip_reason: None if can_implement is True, otherwise a string explaining why not + Abstain rather than reject when a required problem field is unknown. """ raise NotImplementedError( f"{cls.__name__} must implement can_implement method") diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py index 38f057f999eb..c6f4784b8110 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -77,7 +77,7 @@ import torch import torch.distributed as dist -from tensorrt_llm._utils import get_sm_version, is_sm_100f +from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.logger import logger from tensorrt_llm.math_utils import ceil_div from tensorrt_llm.models.modeling_utils import QuantAlgo @@ -93,8 +93,15 @@ from ....cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ....model_config import ModelConfig from ....utils import ActivationType, AuxStreamType, Fp4QuantizedTensor -from ..impl_contract import MoERunContext -from ..interface import MoE, MoESchedulerKind, MoEWeightLoadingMode +from ..impl_contract import ( + MoEDeployment, + MoEEligibility, + MoEProblem, + MoERejectReason, + MoERunContext, +) +from ..impl_environment import MoEDep +from ..interface import MoE, MoESchedulerKind, MoEWeightLoadingMode, _reject from ..quantization import NVFP4MegaMoECuteDslMethod from ..routing import BaseMoeRoutingMethod @@ -348,73 +355,97 @@ class MegaMoECuteDsl(MoE): # ------------------------------------------------------------------ # 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]]: - """Static capability query: SM/dtype/quant/shape only. - - Does NOT probe checkpoint tensor values. The kernel ABI consumes - per-expert scales directly, so there is no checkpoint-value - rejection for non-1 alpha products. The SwiGLU - clamp (``swiglu_limit``) is validated for uniformity in - ``__init__`` (``_resolve_gate_up_clamp``), not here, because - ``can_implement`` is a static query that does not see per-tensor - checkpoint values. - - Multi-rank execution gate (NVSHMEM provider) is NOT in this - query either, by analogy to ``MegaMoEDeepGemm.can_implement``; - ``run_moe`` is where the provider absence becomes a hard error - for ``ep_size > 1`` topologies. - """ - sm = get_sm_version() - if not is_sm_100f(sm): - return False, (f"MegaMoECuteDsl requires SM100 family (SM100 or SM103); got SM{sm}.") - if dtype_activation not in cls._SUPPORTED_ACTIVATION_DTYPES: - return False, ( + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + """Check static eligibility; runtime providers and tensor values are validated later.""" + if not is_sm_100f(d.env.sm): + return _reject( + MoERejectReason.SM_UNSUPPORTED, + f"MegaMoECuteDsl requires SM100 family (SM100 or SM103); got SM{d.env.sm}.", + ) + if p.dtype_act not in cls._SUPPORTED_ACTIVATION_DTYPES: + return _reject( + MoERejectReason.DTYPE_UNSUPPORTED, f"MegaMoECuteDsl supports activations in " - f"{cls._SUPPORTED_ACTIVATION_DTYPES}, got {dtype_activation}." + f"{cls._SUPPORTED_ACTIVATION_DTYPES}, got {p.dtype_act}.", + ) + if p.swiglu_gptoss_style: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + "MegaMoECuteDsl does not support swiglu_gptoss_style.", + ) + if p.quant_algo != QuantAlgo.NVFP4: + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, + f"MegaMoECuteDsl supports NVFP4 only, got quant_algo={p.quant_algo}.", ) - if swiglu_gptoss_style: - return False, "MegaMoECuteDsl does not support swiglu_gptoss_style." - if quant_algo != QuantAlgo.NVFP4: - return False, (f"MegaMoECuteDsl supports NVFP4 only, got quant_algo={quant_algo}.") # ``hidden_size % 32`` covers the kernel's NVFP4 SF leg # alignment; the SF row width is padded to # ``round_up(ceil(hidden/16), 4)`` at every allocation site (see # ``megamoe_activation_sf_bytes_per_row``). - if hidden_size is not None and (hidden_size <= 0 or hidden_size % 32 != 0): - return False, ( + if p.hidden_size is not None and (p.hidden_size <= 0 or p.hidden_size % 32 != 0): + return _reject( + MoERejectReason.SHAPE_UNALIGNED, f"MegaMoECuteDsl requires positive hidden_size divisible " - f"by 32 (NVFP4 SF leg alignment); got {hidden_size}." + f"by 32 (NVFP4 SF leg alignment); got {p.hidden_size}.", ) # The kernel's expand_intermediate = 2 * intermediate must be # divisible by 2 * Fc1GateUpInterleave (32) -> intermediate % 16. - if intermediate_size is not None and ( - intermediate_size <= 0 or intermediate_size % 16 != 0 + if p.intermediate_size is not None and ( + p.intermediate_size <= 0 or p.intermediate_size % 16 != 0 ): - return False, ( + return _reject( + MoERejectReason.SHAPE_UNALIGNED, f"MegaMoECuteDsl requires positive intermediate_size " f"divisible by 16 (Fc1GateUpInterleave); got " - f"{intermediate_size}." + f"{p.intermediate_size}.", + ) + if not d.env.has_dep(MoEDep.MEGAMOE_CUTEDSL_RUNTIME): + return _reject( + MoERejectReason.DEP_MISSING, + "MegaMoECuteDsl requires nvidia-cutlass-dsl[cu13] >= 4.5.0", ) - ok, reason = is_megamoe_cute_dsl_runtime_available() - if not ok: - return False, reason # The fused path also requires the ``trtllm::cute_dsl_megamoe_nvfp4_*`` - # custom op to be registered (strict import of every kernel symbol in - # cute_dsl_megamoe_custom_op). Read the flag dynamically from the - # custom-op module so it reflects the live registration state. - from ....custom_ops import cute_dsl_megamoe_custom_op as _megamoe_op - - if not _megamoe_op.IS_MEGAMOE_OP_AVAILABLE: - return False, _megamoe_op.MEGAMOE_OP_UNAVAILABLE_REASON - return True, None + # custom ops to be registered (strict import of every kernel symbol in + # cute_dsl_megamoe_custom_op). + if p.activation_type != ActivationType.Swiglu: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"MegaMoECuteDsl only supports ActivationType.Swiglu (got {p.activation}).", + ) + if d.tp_size != 1: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"MegaMoECuteDsl is EP-only (moe_tp_size=1); got tp_size={d.tp_size}.", + ) + if d.cluster_size != 1: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"MegaMoECuteDsl assumes cluster_size=1; got cluster_size={d.cluster_size}.", + ) + if d.num_slots % max(d.ep_size, 1) != 0: + return _reject( + MoERejectReason.SLOTS_NOT_DIVISIBLE_BY_EP, + f"MegaMoECuteDsl requires num_slots ({d.num_slots}) " + f"divisible by ep_size ({d.ep_size}).", + ) + # ADP wider than EP would need an outer allgather + reducescatter + # wrapper that this backend does not have. Unlike MegaMoEDeepGemm, TEP + # itself is fine here, so only the attention-DP case is constrained. + if d.use_dp and d.parallel_size > 1 and d.ep_size != d.parallel_size: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"MegaMoECuteDsl with enable_attention_dp=True requires " + f"ep_size == parallel_size (got ep_size={d.ep_size}, " + f"parallel_size={d.parallel_size}).", + ) + if not d.env.has_dep(MoEDep.MEGAMOE_CUTEDSL_OP): + return _reject( + MoERejectReason.DEP_MISSING, + "MegaMoECuteDsl requires the trtllm::cute_dsl_megamoe_nvfp4_* custom ops", + ) + return MoEEligibility.ok() # ------------------------------------------------------------------ # Init @@ -458,40 +489,13 @@ def __init__( init_load_balancer=init_load_balancer, ) - # Constructor-time invariant checks raise ValueError so that - # Python ``-O`` (which strips ``assert``) does not silently let an - # invalid topology through. - if self.tp_size != 1: - raise ValueError( - f"MegaMoECuteDsl is EP-only (moe_tp_size=1); got tp_size={self.tp_size}." - ) - if self.cluster_size != 1: - raise ValueError( - f"MegaMoECuteDsl assumes cluster_size=1; got cluster_size={self.cluster_size}." - ) - if self.num_slots % max(self.ep_size, 1) != 0: - raise ValueError( - f"MegaMoECuteDsl requires num_slots ({self.num_slots}) " - f"divisible by ep_size ({self.ep_size})." - ) - - if self.use_dp and self.parallel_size > 1 and self.ep_size != self.parallel_size: - raise ValueError( - f"MegaMoECuteDsl with enable_attention_dp=True requires " - f"ep_size == parallel_size (got ep_size={self.ep_size}, " - f"parallel_size={self.parallel_size}). ADP > EP would " - f"require an outer allgather + reducescatter wrapper." - ) - + # Topology / activation eligibility is owned by ``can_implement``. + # Keep construction-only invariants that are not part of (p, d). if apply_router_weight_on_input: raise ValueError( "MegaMoECuteDsl does not support apply_router_weight_on_input; " "the fused kernel applies routing weights on the MoE output." ) - if activation_type != ActivationType.Swiglu: - raise ValueError( - f"MegaMoECuteDsl only supports ActivationType.Swiglu (got {activation_type})." - ) self.apply_router_weight_on_input = apply_router_weight_on_input # topk-score application point. v2 default is the deepgemm graph @@ -701,18 +705,12 @@ def _supports_load_balancer(self) -> bool: return True def validate_configurable_moe(self, moe) -> None: - """Mirrors :meth:`MegaMoEDeepGemm.validate_configurable_moe`. - - Enforces the MegaMoECuteDsl wrapper-level invariants (EP-only, - ``moe.comm is None``, ``num_slots % moe_ep_size == 0``, - ``experts_per_token <= 13``, ``moe_max_num_tokens > 0``) listed - inline below. - - ``ConfigurableMoE.__init__`` calls this at the very end (after - ``self.comm`` / ``self.moe_max_num_tokens`` and every EPLB / - num_slots / ep_size attribute are populated -- see - ``configurable_moe.py`` ``validate_backend`` docstring), so - every attribute touched below may be read directly. + """Wrapper-only checks that ``can_implement`` cannot see. + + Topology / slot / ADP eligibility is owned by ``can_implement``. + This validates ConfigurableMoE attributes populated after the backend + is constructed: host-side ``comm`` must stay None for FUSED_COMM, and + top-k / max-token bounds used to size the fused kernel. """ if moe.comm is not None: raise ValueError( @@ -720,28 +718,6 @@ def validate_configurable_moe(self, moe) -> None: f"backends must not layer host-side communication on top " f"of the fused kernel); got moe.comm={type(moe.comm).__name__}." ) - if moe.mapping.moe_tp_size != 1: - raise ValueError( - f"MegaMoECuteDsl is EP-only (moe_tp_size=1); got {moe.mapping.moe_tp_size}." - ) - # NOTE: ``mapping.tp_size`` is the *wrapper-level* TP size used by - # attention, not by the MoE layer. In DEP / TEP modes the wrapper - # sets ``tp_size = world_size`` while ``moe_tp_size = 1``; the - # MegaMoECuteDsl kernel only cares about the MoE axes - # (``moe_ep_size`` / ``moe_tp_size``) — see - # ``_create_mapping_for_parallel_mode`` in test_moe_module.py. - if moe.num_slots % moe.mapping.moe_ep_size != 0: - raise ValueError( - f"MegaMoECuteDsl requires num_slots ({moe.num_slots}) " - f"divisible by moe_ep_size ({moe.mapping.moe_ep_size})." - ) - if moe.use_dp and moe.parallel_size > 1 and moe.mapping.moe_ep_size != moe.parallel_size: - raise ValueError( - f"MegaMoECuteDsl with enable_attention_dp requires " - f"moe_ep_size == parallel_size (got " - f"moe_ep_size={moe.mapping.moe_ep_size}, " - f"parallel_size={moe.parallel_size})." - ) top_k = moe.routing_method.experts_per_token if top_k > 13: raise ValueError( 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 index 8e0a77acb0ec..753612fa4a43 100644 --- 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 @@ -29,19 +29,22 @@ import torch import torch.distributed as dist -from tensorrt_llm._utils import get_sm_version, is_sm_100f +from tensorrt_llm._utils import is_sm_100f 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 ..impl_contract import MoERunContext -from ..interface import MoE, MoESchedulerKind, MoEWeightLoadingMode -from ..quantization import ( - W4A8MXFP4MXFP8MegaMoEDeepGemmMethod, - _import_deep_gemm, - _MegaMoEUnavailable, +from ..impl_contract import ( + MoEDeployment, + MoEEligibility, + MoEProblem, + MoERejectReason, + MoERunContext, ) +from ..impl_environment import MoEDep +from ..interface import MoE, MoESchedulerKind, MoEWeightLoadingMode, _reject +from ..quantization import W4A8MXFP4MXFP8MegaMoEDeepGemmMethod, _import_deep_gemm from ..routing import BaseMoeRoutingMethod __all__ = ["MegaMoEDeepGemm"] @@ -114,6 +117,23 @@ def _call(t: torch.Tensor): return _FUSED_PER_TOKEN_CAST(x) +def _assert_num_slots_divisible_by_ep(num_slots: int, ep_size: int) -> None: + """The DG SymmBuffer is sized to the global slot count and sharded evenly. + + Each rank's weight shard is ``num_slots // ep_size`` slots, so a + non-divisible layout does not fail -- it silently produces wrong per-rank + slot ranges. That wrong-answer shape is why this is re-checked outside + ``can_implement`` (see the callers) instead of trusting the resolution path + alone: direct construction and post-``__init__`` EPLB syncs never go through + it. + """ + if num_slots % max(ep_size, 1) != 0: + raise ValueError( + f"MegaMoEDeepGemm requires num_slots ({num_slots}) divisible by " + f"ep_size ({ep_size}). Adjust the EPLB replication factor or ep_size." + ) + + class MegaMoEDeepGemm(MoE): """MoE backend wrapping DeepGEMM's fused ``fp8_fp4_mega_moe`` kernel.""" @@ -130,60 +150,92 @@ class MegaMoEDeepGemm(MoE): # ------------------------------------------------------------------ # 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 not is_sm_100f(sm): - return False, ( + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + # Process-group availability is validated during construction. + if not is_sm_100f(d.env.sm): + return _reject( + MoERejectReason.SM_UNSUPPORTED, f"MegaMoEDeepGemm requires SM100 family (SM100 or SM103) " - f"for DeepGEMM's fp8_fp4_mega_moe kernel; got SM{sm}" + f"for DeepGEMM's fp8_fp4_mega_moe kernel; got SM{d.env.sm}", ) - if dtype_activation not in cls._SUPPORTED_ACTIVATION_DTYPES: - return False, ( + if p.dtype_act not in cls._SUPPORTED_ACTIVATION_DTYPES: + return _reject( + MoERejectReason.DTYPE_UNSUPPORTED, f"MegaMoEDeepGemm supports activations in " - f"{cls._SUPPORTED_ACTIVATION_DTYPES}, got {dtype_activation}" + f"{cls._SUPPORTED_ACTIVATION_DTYPES}, got {p.dtype_act}", + ) + if p.swiglu_gptoss_style: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + "MegaMoEDeepGemm does not support swiglu_gptoss_style", + ) + if p.quant_algo != QuantAlgo.W4A8_MXFP4_MXFP8: + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, + f"MegaMoEDeepGemm supports W4A8_MXFP4_MXFP8 only, got {p.quant_algo}", ) - 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, ( + # TMA requires packed-UE8M0 scale-factor rows to be 16-byte aligned (K % 512 == 0). + if p.hidden_size is not None and p.hidden_size % 512 != 0: + return _reject( + MoERejectReason.SHAPE_UNALIGNED, f"MegaMoEDeepGemm requires hidden_size % 512 == 0 " f"(DeepGEMM TMA-aligned packed-UE8M0 SF row); " - f"got hidden_size={hidden_size}" + f"got hidden_size={p.hidden_size}", ) - if intermediate_size is not None and intermediate_size % 512 != 0: - return False, ( + if p.intermediate_size is not None and p.intermediate_size % 512 != 0: + return _reject( + MoERejectReason.SHAPE_UNALIGNED, f"MegaMoEDeepGemm requires intermediate_size % 512 == 0 " f"(DeepGEMM TMA-aligned packed-UE8M0 SF row); " - f"got intermediate_size={intermediate_size}" + f"got intermediate_size={p.intermediate_size}", ) - try: - _import_deep_gemm() - except _MegaMoEUnavailable as e: - return False, str(e) - return True, None + if p.activation_type != ActivationType.Swiglu: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"MegaMoEDeepGemm only supports ActivationType.Swiglu (got {p.activation})", + ) + if d.tp_size != 1: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"MegaMoEDeepGemm is EP-only (moe_tp_size=1); got tp_size={d.tp_size}", + ) + if d.cluster_size != 1: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"MegaMoEDeepGemm assumes cluster_size=1; got cluster_size={d.cluster_size}", + ) + if d.num_slots % max(d.ep_size, 1) != 0: + return _reject( + MoERejectReason.SLOTS_NOT_DIVISIBLE_BY_EP, + f"MegaMoEDeepGemm requires num_slots ({d.num_slots}) " + f"divisible by ep_size ({d.ep_size})", + ) + # DG's fp8_fp4_mega_moe assumes the MoE input is partitioned across + # ranks. DEP > EP leaves some tokens unreachable; TEP replicates the + # input so dispatch sees parallel_size duplicate copies, which is + # arithmetically correct but ~parallel_size times slower. + if d.parallel_size > 1: + if not d.use_dp: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"MegaMoEDeepGemm does not support TEP " + f"(enable_attention_dp=False, parallel_size={d.parallel_size})", + ) + if d.ep_size != d.parallel_size: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"MegaMoEDeepGemm with enable_attention_dp=True requires " + f"ep_size == parallel_size (got ep_size={d.ep_size}, " + f"parallel_size={d.parallel_size})", + ) + if not d.env.has_dep(MoEDep.DEEPGEMM_MEGAMOE): + return _reject( + MoERejectReason.DEP_MISSING, + "MegaMoEDeepGemm requires a DeepGEMM build exposing fp8_fp4_mega_moe", + ) + return MoEEligibility.ok() # ------------------------------------------------------------------ # Init @@ -228,64 +280,21 @@ def __init__( 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})." - ) - - # DG's fp8_fp4_mega_moe assumes the MoE input is partitioned across - # ranks (each rank's SymmBuffer holds a unique slice). Supported: - # single rank, or DEP with ep_size == parallel_size. Reject both - # DEP > EP (some tokens unreachable) and TEP (input TP-replicated - # so dispatch sees parallel_size duplicate copies — math correct, - # wall ~parallel_size× slower). - 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})." - ) - elif (not self.use_dp) and self.parallel_size > 1: - raise NotImplementedError( - f"MegaMoEDeepGemm does not support TEP " - f"(enable_attention_dp=False, parallel_size=" - f"{self.parallel_size}>1). Use moe_config.backend=TRTLLM " - f"or enable attention-DP with ep_size == parallel_size." - ) - - # 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." - ) - # ``ActivationType.Swiglu`` describes the gated FC1 tensor geometry - # shared by SwiGLU and SiTU. The DeepGEMM-specific activation selects - # the actual elementwise function below. - if activation_type != ActivationType.Swiglu: + # Topology / activation eligibility is owned by ``can_implement``. + # Keep construction-only invariants that are not part of (p, d): + # apply_router_weight_on_input is a call-site flag, not a deployment + # field, so it stays here until it is modeled on MoEProblem/Deployment. + if apply_router_weight_on_input: raise ValueError( - f"MegaMoEDeepGemm only supports ActivationType.Swiglu (got {activation_type})." + "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." ) + # Also gated in ``can_implement``, but that only covers the resolution + # path; this catches direct construction. + _assert_num_slots_divisible_by_ep(self.num_slots, self.ep_size) activation, situ_beta, situ_linear_beta = self._resolve_activation_config( model_config, activation=activation, @@ -407,21 +416,15 @@ def _supports_load_balancer(self) -> bool: return True def validate_configurable_moe(self, moe) -> None: - """Assert ``num_slots % ep_size == 0`` for the DG global slot count. + """Re-assert the DG global slot count after ``ConfigurableMoE`` wiring. - ``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. + ``can_implement`` gates the same invariant, but only on the resolution + path and only on the slot count the balancer config advertises at select + time. ``moe`` is the owning ``ConfigurableMoE``, whose ``num_slots`` / + ``ep_size`` have since been synced through ``_BACKEND_SYNC_ATTRS``, so + this is the one place that sees the layout the kernel will actually run. """ - # 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." - ) + _assert_num_slots_divisible_by_ep(moe.num_slots, moe.ep_size) @staticmethod def _maybe_init_dist_from_mpi() -> None: diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_resolution.py b/tensorrt_llm/_torch/modules/fused_moe/moe_resolution.py new file mode 100644 index 000000000000..7deb101217f0 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_resolution.py @@ -0,0 +1,424 @@ +# 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. +"""Resolve the MoE implementation from backend preference and capabilities. + +Records rejected candidates; falls back when the requested backend cannot serve. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Dict, FrozenSet, List, Optional, Tuple, Type, Union + +import torch + +from tensorrt_llm.logger import logger +from tensorrt_llm.models.modeling_utils import QuantConfig + +from ...model_config import ModelConfig +from ...peft.lora.validation import has_moe_lora_targets +from ...utils import ActivationType +from .fused_moe_cute_dsl import CuteDslFusedMoE +from .fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE +from .fused_moe_cutlass import CutlassFusedMoE +from .fused_moe_deepgemm import DeepGemmFusedMoE +from .fused_moe_densegemm import DenseGEMMFusedMoE +from .fused_moe_marlin import MarlinFusedMoE +from .fused_moe_triton import TritonFusedMoE +from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE +from .fused_moe_vanilla import VanillaMoE +from .impl_contract import ( + MoEDeployment, + MoEEnvironment, + MoEProblem, + MoERejection, + MoERejectReason, + MoEResolutionReport, + canonical_activation, + canonical_quant, + canonical_routing, +) +from .impl_environment import collect_moe_environment +from .mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm +from .moe_load_balancer import get_moe_load_balancer + +if TYPE_CHECKING: + from .routing import BaseMoeRoutingMethod, RoutingMethodType + +WIDEEP_DEPRECATION_MESSAGE = ( + "The WIDEEP MoE backend is deprecated and can no longer be selected. Wide " + "expert parallelism and EPLB are supported by the other backends: use " + "DEEPGEMM for FP8 block-scale checkpoints, or TRTLLM / CUTEDSL / CUTLASS " + "otherwise." +) + +# Global priority: specialized first, broad fallbacks last. +IMPL_PRIORITY: Tuple[Type, ...] = ( + CuteDslB12xFusedMoE, # SM120/121 NVFP4 decode only -- narrowest, so first + MegaMoEDeepGemm, # ahead of plain CuteDSL / DeepGEMM: better perf when eligible + MegaMoECuteDsl, + CuteDslFusedMoE, + TRTLLMGenFusedMoE, + DeepGemmFusedMoE, + DenseGEMMFusedMoE, + MarlinFusedMoE, + TritonFusedMoE, + CutlassFusedMoE, # widest coverage, hence the fallback + VanillaMoE, # reference implementation, never preferred +) + +# Family membership only; IMPL_PRIORITY decides try order. +BACKEND_FAMILY: Dict[str, FrozenSet[Type]] = { + "CUTLASS": frozenset({CutlassFusedMoE}), + "VANILLA": frozenset({VanillaMoE}), + "MARLIN": frozenset({MarlinFusedMoE}), + "CUTEDSL": frozenset({CuteDslB12xFusedMoE, CuteDslFusedMoE}), + "DEEPGEMM": frozenset({DeepGemmFusedMoE}), + "DENSEGEMM": frozenset({DenseGEMMFusedMoE}), + "TRTLLM": frozenset({TRTLLMGenFusedMoE}), + "TRITON": frozenset({TritonFusedMoE}), + "MEGAMOE_DEEPGEMM": frozenset({MegaMoEDeepGemm}), + "MEGAMOE_CUTEDSL": frozenset({MegaMoECuteDsl}), +} + +# Catch table drift at import time. +_UNRANKED = {cls for family in BACKEND_FAMILY.values() for cls in family} - set(IMPL_PRIORITY) +if _UNRANKED: + raise RuntimeError( + f"MoE impls named by BACKEND_FAMILY but absent from IMPL_PRIORITY: " + f"{sorted(cls.__name__ for cls in _UNRANKED)}" + ) + +# Widest coverage; default degradation target. +FALLBACK_IMPL: Type = CutlassFusedMoE + + +def _legacy_backend_name(impl_cls: Type) -> str: + """Diagnostic name used until each leaf class owns one fixed impl id.""" + return impl_cls.__name__ + + +# --------------------------------------------------------------------------- +# Building the question +# --------------------------------------------------------------------------- + + +# HF configs use different names for the same fields; ModelConfig does not unify them. +_NUM_EXPERTS_ATTRS = ("num_experts", "n_routed_experts", "num_local_experts") +_TOP_K_ATTRS = ("num_experts_per_tok", "experts_per_token") + + +@dataclass(frozen=True) +class MoELayerShapes: + """Resolved shapes for MoE construction / selection.""" + + num_experts: Optional[int] + hidden_size: Optional[int] + intermediate_size: Optional[int] + dtype: Optional[torch.dtype] + top_k: Optional[int] + + +def derive_moe_layer_shapes( + model_config: ModelConfig, + *, + num_experts: Optional[int] = None, + hidden_size: Optional[int] = None, + intermediate_size: Optional[int] = None, + dtype: Optional[torch.dtype] = None, + top_k: Optional[int] = None, + routing: Optional["BaseMoeRoutingMethod | RoutingMethodType"] = None, +) -> MoELayerShapes: + """Fill unset fields from ``pretrained_config`` (and routing for ``top_k``). + + Explicit args win. ``top_k`` order: explicit, then routing object's + ``experts_per_token``, then pretrained. A bare ``RoutingMethodType`` has + no k, so it falls through. + """ + from .routing import BaseMoeRoutingMethod + + pretrained = model_config.pretrained_config + + if dtype is None and pretrained is not None: + dtype = getattr(pretrained, "torch_dtype", None) + if hidden_size is None and pretrained is not None: + hidden_size = getattr(pretrained, "hidden_size", None) + if intermediate_size is None and pretrained is not None: + # Prefer MoE width; getattr so a present-but-None field still falls through. + intermediate_size = getattr(pretrained, "moe_intermediate_size", None) + if intermediate_size is None: + intermediate_size = getattr(pretrained, "intermediate_size", None) + if num_experts is None and pretrained is not None: + for attr in _NUM_EXPERTS_ATTRS: + value = getattr(pretrained, attr, None) + if value is not None: + num_experts = value + break + + if top_k is None and isinstance(routing, BaseMoeRoutingMethod): + top_k = routing.experts_per_token + if top_k is None and pretrained is not None: + for attr in _TOP_K_ATTRS: + value = getattr(pretrained, attr, None) + if value is not None: + top_k = value + break + + return MoELayerShapes( + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + top_k=top_k, + ) + + +def build_moe_problem( + model_config: ModelConfig, + *, + override_quant_config: Optional[QuantConfig] = None, + dtype: Optional[torch.dtype] = None, + num_experts: Optional[int] = None, + hidden_size: Optional[int] = None, + intermediate_size: Optional[int] = None, + top_k: Optional[int] = None, + swiglu_gptoss_style: Optional[bool] = None, + bias: Optional[bool] = None, + activation_type: Optional[ActivationType] = None, + routing: Optional["BaseMoeRoutingMethod | RoutingMethodType"] = None, +) -> MoEProblem: + """Assemble the problem half of a selection question. + + Explicit args win over ``pretrained_config``. Missing fields stay ``None`` + (unknown): shape gates abstain instead of rejecting on absent info. + """ + shapes = derive_moe_layer_shapes( + model_config, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + top_k=top_k, + routing=routing, + ) + quant_config = override_quant_config or model_config.quant_config + quant_algo = None if quant_config is None else quant_config.quant_algo + + return MoEProblem( + quant=canonical_quant(quant_algo), + dtype_act=shapes.dtype if shapes.dtype is not None else torch.bfloat16, + hidden_size=shapes.hidden_size, + intermediate_size=shapes.intermediate_size, + num_experts=shapes.num_experts, + top_k=shapes.top_k, + swiglu_gptoss_style=swiglu_gptoss_style, + bias=bias, + activation=canonical_activation(activation_type), + routing=canonical_routing(routing), + ) + + +def infer_swiglu_gptoss_style( + *, + bias: bool = False, + swiglu_alpha: Optional[torch.Tensor] = None, + swiglu_beta: Optional[torch.Tensor] = None, + activation_type: Optional[Union[ActivationType, int]] = None, +) -> bool: + """True for the gpt-oss / MiniMax SwiGLU package (bias, alpha/beta, or SwigluBias). + + ``swiglu_limit`` alone is not enough — DeepSeek-V4 uses a plain clamp and + must not be treated as gpt-oss. + + ``activation_type`` is normalized because ``MoE`` stores the activation as a + plain ``int``, which no identity check against an enum member can match. + """ + if activation_type is not None and ActivationType(activation_type) is ActivationType.SwigluBias: + return True + return bool(bias or swiglu_alpha is not None or swiglu_beta is not None) + + +def build_moe_deployment( + model_config: ModelConfig, + *, + num_experts: Optional[int] = None, + environment: Optional[MoEEnvironment] = None, +) -> MoEDeployment: + """Assemble the deployment half, reading the same mapping ``MoE.__init__`` does.""" + mapping = model_config.mapping + # MoE._init_load_balancer only adopts the config's slot count once a + # balancer is registered; without one the layer keeps num_slots == + # num_experts, so the deployment must say the same. + eplb_enabled = get_moe_load_balancer() is not None + balancer_config = getattr(model_config, "moe_load_balancer", None) + num_slots = getattr(balancer_config, "num_slots", None) if eplb_enabled else None + if num_slots is None: + num_slots = num_experts if num_experts is not None else 0 + lora_config = getattr(model_config, "lora_config", None) + return MoEDeployment( + ep_size=mapping.moe_ep_size, + tp_size=mapping.moe_tp_size, + # Same meaning as MoE.parallel_size (mapping.tp_size). + parallel_size=mapping.tp_size, + cluster_size=mapping.moe_cluster_size, + use_dp=mapping.enable_attention_dp, + num_slots=num_slots, + env=environment if environment is not None else collect_moe_environment(), + # Registered balancer only; config alone does not enable EPLB. + eplb_enabled=eplb_enabled, + # Routed-expert LoRA only; attention-only LoRA stays False. + moe_lora_enabled=has_moe_lora_targets(lora_config), + ) + + +# --------------------------------------------------------------------------- +# Resolution +# --------------------------------------------------------------------------- + + +# Backends whose whole point is to be the one that runs: silently degrading +# them to Cutlass would hand back the very numbers the caller asked to compare +# against. They fail with the rejection trail instead. +NO_FALLBACK_BACKENDS: FrozenSet[str] = frozenset({"VANILLA"}) + + +def _candidates_for(backend: str) -> List[Type]: + """Requested family in priority order, then fallback.""" + normalized = backend.upper() + if normalized == "WIDEEP": + raise ValueError(WIDEEP_DEPRECATION_MESSAGE) + family = BACKEND_FAMILY.get(normalized) + if family is None: + raise ValueError(f"Unsupported moe backend: {backend}") + candidates = [impl_cls for impl_cls in IMPL_PRIORITY if impl_cls in family] + if FALLBACK_IMPL not in family and normalized not in NO_FALLBACK_BACKENDS: + candidates.append(FALLBACK_IMPL) + return candidates + + +def resolve_moe_impl( + model_config: ModelConfig, + *, + problem: Optional[MoEProblem] = None, + deployment: Optional[MoEDeployment] = None, + override_quant_config: Optional[QuantConfig] = None, + dtype: Optional[torch.dtype] = None, + num_experts: Optional[int] = None, + hidden_size: Optional[int] = None, + intermediate_size: Optional[int] = None, + swiglu_gptoss_style: Optional[bool] = None, + bias: Optional[bool] = None, + activation_type: Optional[ActivationType] = None, + routing: Optional["BaseMoeRoutingMethod | RoutingMethodType"] = None, + layer_idx: Optional[int] = None, +) -> MoEResolutionReport: + """Resolve a MoE backend and return the full eligibility report. + + Raises ValueError for unknown or deprecated backend literals. + """ + if problem is None: + problem = build_moe_problem( + model_config, + override_quant_config=override_quant_config, + dtype=dtype, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + swiglu_gptoss_style=swiglu_gptoss_style, + bias=bias, + activation_type=activation_type, + routing=routing, + ) + if deployment is None: + deployment = build_moe_deployment(model_config, num_experts=problem.num_experts) + + requested = model_config.moe_backend + candidates = _candidates_for(requested) + in_family = BACKEND_FAMILY[requested.upper()] + + # Ask all candidates so the report lists alternatives. + rejected = [] + eligible: List[Type] = [] + for candidate in candidates: + if deployment.moe_lora_enabled and not candidate.capabilities.supports_moe_lora: + rejected.append( + MoERejection( + _legacy_backend_name(candidate), + MoERejectReason.LORA_UNSUPPORTED, + f"{candidate.__name__} does not fuse routed-expert LoRA", + ) + ) + continue + eligibility = candidate.can_implement(problem, deployment) + if eligibility.eligible: + eligible.append(candidate) + continue + rejected.append( + MoERejection( + _legacy_backend_name(candidate), + eligibility.reject_reason, + eligibility.detail, + ) + ) + + # candidates is already priority-ordered. + winner_cls = eligible[0] if eligible else None + + if winner_cls is None: + selected_by = "failed" + elif winner_cls in in_family: + # Another family member still counts as pinned. + selected_by = "pinned" + else: + selected_by = "heuristic" + + report = MoEResolutionReport( + problem=problem, + deployment=deployment, + winner=None if winner_cls is None else _legacy_backend_name(winner_cls), + rejected=tuple(rejected), + eligible=tuple(_legacy_backend_name(impl_cls) for impl_cls in eligible), + selected_by=selected_by, + requested=requested, + env_fingerprint=deployment.env.fingerprint(), + ) + + if report.degraded and winner_cls is not None: + cause = report.degraded_from + location = "" if layer_idx is None else f" [layer_idx={layer_idx}]" + logger.warning( + f"MoE backend {requested} cannot serve this layer{location} " + f"({cause.reason.value}: {cause.detail}); running " + f"{winner_cls.__name__} instead. Full trail: {report.describe()}" + ) + else: + logger.debug(report.describe()) + + return report + + +def impl_class_for(report: MoEResolutionReport) -> Type: + """The class a report's winner names, or raise with the whole trail.""" + if report.winner is None: + raise ValueError(f"no MoE implementation can serve this layer. {report.describe()}") + for candidate in IMPL_PRIORITY: + if _legacy_backend_name(candidate) == report.winner: + return candidate + raise ValueError( + f"resolution report names legacy backend {report.winner!r}, which no candidate class claims" + ) + + +def resolve_moe_cls(model_config: ModelConfig, **kwargs) -> Type: + """Resolve and return only the implementation class.""" + return impl_class_for(resolve_moe_impl(model_config, **kwargs)) diff --git a/tensorrt_llm/_torch/peft/lora/validation.py b/tensorrt_llm/_torch/peft/lora/validation.py index bbb4a7c3a597..5dbdea9b8a6b 100644 --- a/tensorrt_llm/_torch/peft/lora/validation.py +++ b/tensorrt_llm/_torch/peft/lora/validation.py @@ -3,11 +3,12 @@ """Validation helpers for routed-expert (MoE) LoRA. MoE LoRA is supported only on the Cutlass backend with unquantized fp16/bf16 or -per-tensor FP8 (qdq) base weights. This module provides a single helper, -`check_moe_lora_supported`, that callers (typically the MoE factory in -`create_moe.py`) can invoke at construction time so that unsupported -combinations fail loudly instead of silently dropping the LoRA contribution at -runtime. +per-tensor FP8 (qdq) base weights. Resolution owns that contract at select time: +``supports_moe_lora`` filters backends, and ``CutlassFusedMoE.can_implement`` +rejects unsupported base-weight quants when ``moe_lora_enabled`` is set. + +`check_moe_lora_supported` remains as a standalone assertion for unit tests and +any caller that wants an explicit ValueError without going through resolution. Runtime-only rejections (min-latency mode, alltoall, CUDA-graph without slot pointers) are enforced in the C++ thop / runtime call paths and are NOT diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 18dc19960e89..464570af4f95 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1942,6 +1942,12 @@ def test_fp8_block_scales(self, mtp, fp8kv, attention_dp, cuda_graph, task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @pytest.mark.skip( + reason="CuteDslFusedMoE declines FP8 block scales: it has no FP8 " + "block-scale kernel, only a torch.einsum reference. See the ‡ footnote " + "in tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md. " + "Re-enable against DEEPGEMM / TRTLLM once this checkpoint has an " + "owner backend on SM100.") @skip_pre_blackwell @parametrize_with_ids("torch_compile", [False]) @parametrize_with_ids( @@ -2129,6 +2135,10 @@ def test_fp8_block_scales_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @pytest.mark.skip( + reason="Same as test_cute_dsl_fp8_block_scales: CuteDslFusedMoE " + "declines FP8 block scales, so a CUTEDSL request on this checkpoint no " + "longer resolves.") @pytest.mark.skip_less_device(4) @skip_pre_blackwell @parametrize_with_ids("torch_compile", [False]) @@ -7270,8 +7280,8 @@ def test_nvfp4_marlin_adp_4gpus(self, mtp_nextn): The NVFP4 checkpoint is MIXED_PRECISION with deliberately-unquantized MTP draft layers, so this also guards the per-layer - MARLIN -> Cutlass fallback in ``create_moe.get_moe_cls`` while the - main-model expert layers stay on MARLIN. Attention DP exercises the + MARLIN -> Cutlass degradation in ``moe_resolution.resolve_moe_impl`` + while the main-model expert layers stay on MARLIN. Attention DP exercises the external-comm dispatch path with scheduler-precomputed routing. """ model_path = f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4" diff --git a/tests/microbenchmarks/bench_moe/backend.py b/tests/microbenchmarks/bench_moe/backend.py index 866eabdd3774..cc7ce9f990c1 100644 --- a/tests/microbenchmarks/bench_moe/backend.py +++ b/tests/microbenchmarks/bench_moe/backend.py @@ -86,8 +86,15 @@ def ensure_cute_dsl_importable_for_benchmark() -> None: class CuteDslFusedMoE: @classmethod - def can_implement(cls, *_args, **_kwargs): - return False, f"CUTLASS DSL is unavailable: {import_error}" + def can_implement(cls, p, d): + from tensorrt_llm._torch.modules.fused_moe.impl_contract import ( + MoEEligibility, + MoERejectReason, + ) + + return MoEEligibility.no( + MoERejectReason.DEP_MISSING, f"CUTLASS DSL is unavailable: {import_error}" + ) def __init__(self, *_args, **_kwargs): raise RuntimeError(f"CUTLASS DSL is unavailable: {import_error}") diff --git a/tests/microbenchmarks/bench_moe/search.py b/tests/microbenchmarks/bench_moe/search.py index 7fb8c40518c4..491ecc2fb256 100644 --- a/tests/microbenchmarks/bench_moe/search.py +++ b/tests/microbenchmarks/bench_moe/search.py @@ -24,6 +24,12 @@ import torch +from tensorrt_llm._torch.modules.fused_moe.impl_contract import ( + MoEDeployment, + MoEProblem, + canonical_quant, +) +from tensorrt_llm._torch.modules.fused_moe.impl_environment import collect_moe_environment from tensorrt_llm._utils import local_mpi_size from tensorrt_llm.models.modeling_utils import QuantAlgo @@ -58,19 +64,36 @@ def _check_backend_can_implement( dtype_activation: torch.dtype, swiglu_gptoss_style: bool, ) -> Tuple[bool, Optional[str]]: - """Resolve backend_str to its MoE class and forward to can_implement.""" + """Resolve backend_str to its MoE class and ask whether it can serve this. + + The topology-dependent gates are deliberately not exercised here: this runs + before ``_resolve_mapping_layout``, and the EP / comm constraints get their + own explicit checks in :func:`is_candidate_valid` with better messages. + """ try: backend_cls = get_backend_class(MoeBackendType(backend_str.upper())) except (ImportError, KeyError, RuntimeError, ValueError) as exc: return False, f"unknown MoE backend {backend_str!r}: {exc}" + problem = MoEProblem( + quant=canonical_quant(quant_algo), + dtype_act=dtype_activation, + swiglu_gptoss_style=swiglu_gptoss_style, + ) + deployment = MoEDeployment( + ep_size=1, + tp_size=1, + parallel_size=1, + use_dp=False, + num_slots=0, + env=collect_moe_environment(), + ) try: - return backend_cls.can_implement( - quant_algo=quant_algo, - dtype_activation=dtype_activation, - swiglu_gptoss_style=swiglu_gptoss_style, - ) + verdict = backend_cls.can_implement(problem, deployment) except Exception as exc: return False, (f"{backend_cls.__name__}.can_implement raised {type(exc).__name__}: {exc}") + if verdict.eligible: + return True, None + return False, f"{verdict.reject_reason.value}: {verdict.detail}" def _expand_axis(values: Iterable[Any], default: Any) -> Tuple[Any, ...]: diff --git a/tests/unittest/_torch/models/checkpoints/laguna/test_laguna_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/laguna/test_laguna_weight_mapper.py index 7f85274ce912..316fafee6966 100644 --- a/tests/unittest/_torch/models/checkpoints/laguna/test_laguna_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/laguna/test_laguna_weight_mapper.py @@ -20,6 +20,7 @@ from torch import nn from tensorrt_llm._torch.models.modeling_laguna import LagunaHfWeightMapper +from tensorrt_llm._torch.modules.fused_moe.impl_contract import MoEEligibility from tensorrt_llm._torch.modules.fused_moe.interface import MoE pytestmark = pytest.mark.cpu_only @@ -39,8 +40,8 @@ def has_fp8_block_scales(self): class _FakeMoE(MoE): @classmethod - def can_implement(cls, *args, **kwargs): - return True, None + def can_implement(cls, p, d): + return MoEEligibility.ok() def __init__(self): nn.Module.__init__(self) diff --git a/tests/unittest/_torch/models/test_qwen3_next_moe_quant.py b/tests/unittest/_torch/models/test_qwen3_next_moe_quant.py index 7b2370e60503..fa522e5a71d0 100644 --- a/tests/unittest/_torch/models/test_qwen3_next_moe_quant.py +++ b/tests/unittest/_torch/models/test_qwen3_next_moe_quant.py @@ -168,13 +168,23 @@ def test_missing_layer_idx_is_a_noop(): class _StopBlockInit(Exception): - """Raised after the test captures the arguments passed to ``create_moe``.""" + """Raised once the resolved MoE class is captured, to abort the build.""" -def _build_moe_block(moe_backend, exclude_modules, layer_idx, quant_config_dict=None): - """Run sparse-MoE initialization through its ``create_moe`` call.""" +def _build_moe_block(moe_backend, exclude_modules, layer_idx, *, sm, quant_config_dict=None): + """Run sparse-MoE initialization through its ``create_moe`` call. + + ``sm`` is declared, not probed. Which implementation wins is a function of + the machine, so pinning an expected class while the environment comes from + whatever GPU CI happens to hand out asserts a property of the runner rather + than of the code. Declaring it keeps the assertions about selection + strategy and lets these cases run on any device. An empty dependency set + goes with it, so an installed flashinfer cannot change the outcome either. + """ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextSparseMoeBlock + from tensorrt_llm._torch.modules.fused_moe.impl_contract import MoEEnvironment + from tensorrt_llm._torch.modules.fused_moe.impl_environment import override_moe_environment model_config = ModelConfig( pretrained_config=SimpleNamespace( @@ -200,23 +210,31 @@ def _build_moe_block(moe_backend, exclude_modules, layer_idx, quant_config_dict= ) captured = {} - def _capture(*args, **kwargs): - from tensorrt_llm._torch.modules.fused_moe.create_moe import resolve_moe_cls - - captured["moe_backend"] = kwargs["model_config"].moe_backend - captured["override"] = kwargs["override_quant_config"] - captured["moe_cls"] = resolve_moe_cls( - kwargs["model_config"], - kwargs["routing_method"], - kwargs["dtype"], - kwargs["override_quant_config"], - kwargs["layer_idx"], - ).__name__ - raise _StopBlockInit + # ``fused_moe/__init__`` re-exports create_moe as a function, which shadows + # the submodule of the same name, so ``import ...create_moe as m`` binds the + # function. Both forms below resolve the submodule through sys.modules. + from tensorrt_llm._torch.modules.fused_moe.create_moe import ( + resolve_moe_cls as real_resolve_moe_cls, + ) - import tensorrt_llm._torch.models.modeling_qwen3_next as qwen3_next + # Observe the resolution create_moe actually performs instead of + # reproducing its argument list here: a copy drifts, and a copy that omits + # e.g. swiglu_gptoss_style resolves to a different class than the layer + # will be built with, which the assertions below would not catch. + def _capture(model_config, **kwargs): + captured["moe_backend"] = model_config.moe_backend + captured["override"] = kwargs.get("override_quant_config") + captured["moe_cls"] = real_resolve_moe_cls(model_config, **kwargs).__name__ + raise _StopBlockInit - with patch.object(qwen3_next, "create_moe", _capture), pytest.raises(_StopBlockInit): + with ( + override_moe_environment(MoEEnvironment(sm=sm)), + patch( + "tensorrt_llm._torch.modules.fused_moe.create_moe.resolve_moe_cls", + _capture, + ), + pytest.raises(_StopBlockInit), + ): Qwen3NextSparseMoeBlock(model_config, aux_stream=None, layer_idx=layer_idx) return captured @@ -225,10 +243,14 @@ def _capture(*args, **kwargs): @pytest.mark.parametrize("layer_idx", [5, MTP_LAYER_IDX]) def test_excluded_layer_builds_bf16_on_cutlass(backend, layer_idx): per_layer_quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) + # The exclusion rewrites the request to CUTLASS and strips the quantization, + # and unquantized Cutlass covers SM80+, so the declared SM only has to be a + # real one -- this case is about the rewrite, not about hardware. captured = _build_moe_block( backend, [f"model.layers.{layer_idx}*"], layer_idx, + sm=100, quant_config_dict={ f"model.layers.{layer_idx}.mlp.experts": per_layer_quant_config, }, @@ -241,26 +263,40 @@ def test_excluded_layer_builds_bf16_on_cutlass(backend, layer_idx): assert captured["override"].kv_cache_quant_algo == QuantAlgo.FP8 -_UNEXCLUDED_EXPECTED_MOE_CLS = { - "CUTLASS": "CutlassFusedMoE", - "TRTLLM": "TRTLLMGenFusedMoE", - "DEEPGEMM": "DeepGemmFusedMoE", - "CUTEDSL": "CuteDslFusedMoE", -} - - -@pytest.mark.parametrize("backend", sorted(_UNEXCLUDED_EXPECTED_MOE_CLS)) -def test_unexcluded_layer_keeps_configured_backend_and_layer_quant_config(backend): - per_layer_quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) +# "Requested backend wins" only holds where that backend can actually serve the +# layer, so each row pairs the request with an environment and an algorithm that +# make its family eligible. A single (algorithm, SM) shared by all four does not +# exist: FP8 block scales is Cutlass on SM90 and DeepGemm / TRTLLM-Gen on SM100, +# and CuteDSL only claims NVFP4. +@pytest.mark.parametrize( + "backend,sm,quant_algo,expected_moe_cls", + [ + ("CUTLASS", 90, QuantAlgo.FP8_BLOCK_SCALES, "CutlassFusedMoE"), + ("TRTLLM", 100, QuantAlgo.FP8_BLOCK_SCALES, "TRTLLMGenFusedMoE"), + ("DEEPGEMM", 100, QuantAlgo.FP8_BLOCK_SCALES, "DeepGemmFusedMoE"), + ("CUTEDSL", 100, QuantAlgo.NVFP4, "CuteDslFusedMoE"), + ], + ids=[ + "cutlass_sm90_fp8_block", + "trtllm_sm100_fp8_block", + "deepgemm_sm100_fp8_block", + "cutedsl_sm100_nvfp4", + ], +) +def test_unexcluded_layer_keeps_configured_backend_and_layer_quant_config( + backend, sm, quant_algo, expected_moe_cls +): + per_layer_quant_config = QuantConfig(quant_algo=quant_algo) captured = _build_moe_block( backend, ["model.layers.7*"], 5, + sm=sm, quant_config_dict={ "model.layers.5.mlp.experts": per_layer_quant_config, }, ) assert captured["moe_backend"] == backend - assert captured["moe_cls"] == _UNEXCLUDED_EXPECTED_MOE_CLS[backend] + assert captured["moe_cls"] == expected_moe_cls assert captured["override"] is per_layer_quant_config diff --git a/tests/unittest/_torch/modules/moe/moe_test_utils.py b/tests/unittest/_torch/modules/moe/moe_test_utils.py index 17c5fcd948e1..b6447fbdf90a 100644 --- a/tests/unittest/_torch/modules/moe/moe_test_utils.py +++ b/tests/unittest/_torch/modules/moe/moe_test_utils.py @@ -12,20 +12,7 @@ # 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. -""" -Shared utilities for MoE test files (test_moe_backend.py and test_moe_module.py). - -This module contains common code extracted from both test files: -- MoeBackendType enum and get_backend_class() -- MoeModelConfig dataclass -- Skip logic functions (should_skip_trtllm, should_skip_cutedsl, should_skip_routing_method, etc.) -- get_quick_skip_reason() - unified version supporting both backend and module tests -- supports_autotuner_capture() -- replay_tactics_and_check() -- module_timer fixture -- create_test_param() helper -- Common test parameter constants -""" +"""Shared MoE test utilities.""" import logging import os @@ -48,6 +35,12 @@ from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE 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.impl_contract import ( + MoEDeployment, + MoEProblem, + canonical_quant, +) +from tensorrt_llm._torch.modules.fused_moe.impl_environment import collect_moe_environment from tensorrt_llm._torch.modules.fused_moe.interface import MoE from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm from tensorrt_llm._torch.modules.fused_moe.mega_moe.mega_moe_cute_dsl import ( @@ -70,12 +63,7 @@ class MoeBackendType(str, Enum): CUTEDSL = "CUTEDSL" DEEPGEMM = "DEEPGEMM" DENSEGEMM = "DENSEGEMM" - # Two MegaMoE variants live side by side: the DeepGemm path and the - # CuteDSL path. Keep both keys explicit so ``value -> member`` lookup - # and grep are unambiguous (avoid an asymmetric pair where one variant - # has an alias and the other does not). The legacy - # ``MoeBackendType.MEGAMOE`` alias was removed; all call sites must - # spell out the variant explicitly. + # Keep the two MegaMoE variants explicit. MEGAMOE_DEEPGEMM = "MEGAMOE_DEEPGEMM" MEGAMOE_CUTEDSL = "MEGAMOE_CUTEDSL" CUTE_DSL_B12X = "CUTE_DSL_B12X" @@ -719,6 +707,7 @@ def should_skip_cutlass( model_config: "MoeModelConfig" = None, moe_tp_size: int = 1, dtype=None, + swiglu_gptoss_style: bool = False, ) -> Optional[str]: """ Check CUTLASS backend specific constraints for multi-GPU tests. @@ -729,6 +718,19 @@ def should_skip_cutlass( if backend_type != MoeBackendType.CUTLASS: return None + # W4A16_MXFP4 is the SM90 member of the MXFP4 family, and the family is the + # only one CutlassFusedMoE accepts for gpt-oss SwiGLU. Real gpt-oss + # checkpoints load fine on this path, but this harness's synthetic MXFP4 + # weights build a w3_w1 bias of 4 * intermediate_size while the loader + # supplies 2 * intermediate_size, so the copy raises a size mismatch. That + # is a pre-existing weight-generation gap, unrelated to backend selection. + if swiglu_gptoss_style and quant_algo == QuantAlgo.W4A16_MXFP4: + return ( + "CutlassFusedMoE W4A16_MXFP4 + gpt-oss SwiGLU: harness " + "synthetic MXFP4 bias shape mismatch (4x vs 2x " + "intermediate_size)" + ) + # TP per-shard alignment: W8A16, NVFP4, W4A8_AWQ, and MXFP8 require # 128-aligned per-shard intermediate_size. W8A16 fails in # preprocess_weights_for_mixed_gemm (num_rows % rows_per_tile != 0). NVFP4 @@ -935,13 +937,7 @@ def should_skip_cute_dsl_b12x( moe_tp_size: int = 1, parallel_mode: Optional[str] = None, ) -> Optional[str]: - """Check CuteDslB12xFusedMoE constraints not covered by can_implement(). - - can_implement() already gates SM version, quant_algo, dtype_activation, and - swiglu_gptoss_style. This helper covers the additional EP / alltoall hard - rejects enforced in __init__ (b12x has no expert-parallel dispatch/combine - kernel). - """ + """Check multi-rank constraints omitted from the capability query.""" if backend_type != MoeBackendType.CUTE_DSL_B12X: return None @@ -1170,17 +1166,7 @@ def get_quick_skip_reason( swiglu_gptoss_style: bool = False, seq_len: Optional[int] = None, ) -> Optional[str]: - """ - Fast skip check that calls backend's can_implement() method. - - Unified version supporting both backend-level and module-level tests: - - routing_method_cls: Used by test_moe_module.py for routing method compatibility checks - - swiglu_gptoss_style: Used by test_moe_backend.py for SwiGLU parameter checks - - seq_len: Optional sequence length for seq_len-sensitive skip checks - - Returns: - Skip reason string if test should be skipped, None otherwise - """ + """Return the first reason a test configuration is unsupported.""" import logging as _logging # Suppress logger warnings during parameter generation @@ -1189,24 +1175,29 @@ def get_quick_skip_reason( trtllm_logger.setLevel(_logging.ERROR) try: - # Call backend's can_implement for dtype/quant_algo checks backend_cls = get_backend_class(backend_type) - can_impl_kwargs = {"dtype_activation": dtype} - if swiglu_gptoss_style: - can_impl_kwargs["swiglu_gptoss_style"] = swiglu_gptoss_style - if ( - backend_type - in ( - MoeBackendType.MEGAMOE_DEEPGEMM, - MoeBackendType.MEGAMOE_CUTEDSL, - ) - 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 + problem = MoEProblem( + quant=canonical_quant(quant_algo), + dtype_act=dtype, + hidden_size=None if model_config is None else model_config.hidden_size, + intermediate_size=None if model_config is None else model_config.intermediate_size, + num_experts=None if model_config is None else model_config.num_experts, + top_k=None if model_config is None else model_config.top_k, + swiglu_gptoss_style=swiglu_gptoss_style, + bias=swiglu_gptoss_style, + ) + # Multi-rank constraints are checked by the helpers below. + deployment = MoEDeployment( + ep_size=1, + tp_size=1, + parallel_size=1, + use_dp=False, + num_slots=0 if model_config is None else model_config.num_experts, + env=collect_moe_environment(), + ) + verdict = backend_cls.can_implement(problem, deployment) + if not verdict.eligible: + return f"{verdict.reject_reason.value}: {verdict.detail}" # Chain skip checks: routing method, then per-backend constraints skip_checks = [ @@ -1220,7 +1211,11 @@ def get_quick_skip_reason( seq_len=seq_len, ), lambda: should_skip_cutlass( - backend_type, quant_algo=quant_algo, model_config=model_config, dtype=dtype + backend_type, + quant_algo=quant_algo, + model_config=model_config, + dtype=dtype, + swiglu_gptoss_style=swiglu_gptoss_style, ), lambda: should_skip_cutedsl( backend_type, quant_algo, model_config, routing_method_cls=routing_method_cls diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index 7e576b7d860a..1479d4546474 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -12,59 +12,92 @@ # 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. -"""Negative-path + dispatch tests for CuteDslB12xFusedMoE. - -These checks run without a GPU: they verify the can_implement() gating -matrix, the SM120/SM121 + NVFP4 selection in create_moe.get_moe_cls (the -backend is selected on the `moe_backend=CUTEDSL` path when flashinfer -is importable, never from `moe_backend=CUTLASS`), and the hybrid -CUTLASS-prefill / b12x-decode dispatch predicate. Functional -correctness of the b12x kernel is covered by end-to-end model tests on -SM120/SM121 hardware. -""" +"""CuteDslB12xFusedMoE gating and dispatch tests.""" import sys import types +from typing import Optional from unittest.mock import patch import pytest import torch -from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.fused_moe.create_moe import get_moe_cls -from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl import CuteDslFusedMoE from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE +from tensorrt_llm._torch.modules.fused_moe.impl_contract import ( + MoEDeployment, + MoEEnvironment, + MoEProblem, + MoERejectReason, + canonical_quant, +) +from tensorrt_llm._torch.modules.fused_moe.impl_environment import MoEDep from tensorrt_llm._torch.modules.fused_moe.quantization import ( NVFP4CuteDslB12xFusedMoEMethod, NVFP4CutlassFusedMoEMethod, ) from tensorrt_llm._torch.utils import ActivationType -from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig pytestmark = pytest.mark.cpu_only -_FUSED_MOE_MODULE = "tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl_b12x" +# Spelled out rather than read from _SUPPORTED_SM_VERSIONS: deriving the input +# from the value under test hides a narrowing of that set, which is the +# direction that silently drops hardware support. +SUPPORTED_SM = [120, 121] + + +def _deployment( + sm: int, + *, + flashinfer: bool = True, + ep_size: int = 1, + use_dp: bool = False, + parallel_size: Optional[int] = None, +) -> MoEDeployment: + """Declare the machine rather than patching the probes that read it.""" + return MoEDeployment( + ep_size=ep_size, + tp_size=1, + parallel_size=ep_size if parallel_size is None else parallel_size, + use_dp=use_dp, + num_slots=8, + env=MoEEnvironment( + sm=sm, + available_deps=(MoEDep.FLASHINFER.value,) if flashinfer else (), + ), + ) + + +def _problem( + quant_algo=QuantAlgo.NVFP4, dtype=torch.bfloat16, swiglu_gptoss_style=None +) -> MoEProblem: + return MoEProblem( + quant=canonical_quant(quant_algo), + dtype_act=dtype, + hidden_size=2048, + intermediate_size=2048, + num_experts=8, + top_k=2, + swiglu_gptoss_style=swiglu_gptoss_style, + ) @pytest.mark.parametrize("sm_version", [80, 89, 90, 100, 103]) def test_can_implement_rejects_unsupported_sm(sm_version): - """can_implement returns False on every SM outside the supported set.""" - with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version): - ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4) - assert not ok - assert reason is not None and f"SM{sm_version}" in reason + verdict = CuteDslB12xFusedMoE.can_implement(_problem(), _deployment(sm_version)) + assert not verdict.eligible + assert verdict.reject_reason is MoERejectReason.SM_UNSUPPORTED + assert f"SM{sm_version}" in verdict.detail -@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) +@pytest.mark.parametrize("sm_version", SUPPORTED_SM) @pytest.mark.parametrize("quant_algo", [QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4]) def test_can_implement_accepts_supported_sm(sm_version, quant_algo): - with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version): - ok, reason = CuteDslB12xFusedMoE.can_implement(quant_algo) - assert ok - assert reason is None + verdict = CuteDslB12xFusedMoE.can_implement(_problem(quant_algo), _deployment(sm_version)) + assert verdict.eligible + assert verdict.reject_reason is None @pytest.mark.parametrize( @@ -79,134 +112,46 @@ def test_can_implement_accepts_supported_sm(sm_version, quant_algo): ) def test_can_implement_rejects_non_nvfp4(quant_algo): """Only NVFP4 is supported; everything else must be turned away.""" - with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=120): - ok, reason = CuteDslB12xFusedMoE.can_implement(quant_algo) - assert not ok - assert reason is not None and "NVFP4" in reason + verdict = CuteDslB12xFusedMoE.can_implement(_problem(quant_algo), _deployment(120)) + assert not verdict.eligible + assert verdict.reject_reason is MoERejectReason.QUANT_UNSUPPORTED def test_can_implement_rejects_swiglu_gptoss_style(): - with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=120): - ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4, swiglu_gptoss_style=True) - assert not ok - assert reason is not None and "swiglu_gptoss_style" in reason + verdict = CuteDslB12xFusedMoE.can_implement( + _problem(swiglu_gptoss_style=True), _deployment(120) + ) + assert not verdict.eligible + assert verdict.reject_reason is MoERejectReason.ACTIVATION_UNSUPPORTED @pytest.mark.parametrize("dtype", [torch.float32, torch.float8_e4m3fn]) def test_can_implement_rejects_unsupported_activation_dtype(dtype): - with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=120): - ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4, dtype_activation=dtype) - assert not ok - assert reason is not None - - -def test_get_moe_cls_cutlass_path_never_auto_promotes(): - """Explicit ``moe_backend=CUTLASS`` always returns ``CutlassFusedMoE`` — - no silent override to the b12x backend even on eligible hardware. b12x - is opted into via ``moe_backend=CUTEDSL``.""" - cfg = ModelConfig() - cfg.moe_backend = "CUTLASS" - cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) - with patch("tensorrt_llm._utils.get_sm_version", return_value=120): - cls = get_moe_cls(cfg) - assert cls is CutlassFusedMoE - - -def test_get_moe_cls_cutedsl_falls_back_to_cutlass_on_unsupported_quant(): - """CUTEDSL + non-(fp8_block_scales|nvfp4) → warn + fall back to CutlassFusedMoE.""" - cfg = ModelConfig() - cfg.moe_backend = "CUTEDSL" - cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.FP8) - with patch("tensorrt_llm._utils.get_sm_version", return_value=120): - cls = get_moe_cls(cfg) - assert cls is CutlassFusedMoE - - -def test_get_moe_cls_cutedsl_falls_back_to_cutlass_on_missing_quant(): - cfg = ModelConfig() - cfg.moe_backend = "CUTEDSL" - cfg.quant_config = None - with patch("tensorrt_llm._utils.get_sm_version", return_value=120): - cls = get_moe_cls(cfg) - assert cls is CutlassFusedMoE - - -def test_get_moe_cls_cutedsl_returns_plain_cutedsl_on_unsupported_sm(): - """CUTEDSL + NVFP4 + non-SM120/121 → plain CuteDslFusedMoE (the SM100/103 - cuteDSL backend); the b12x branch is bypassed.""" - cfg = ModelConfig() - cfg.moe_backend = "CUTEDSL" - cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) - with patch("tensorrt_llm._utils.get_sm_version", return_value=100): - cls = get_moe_cls(cfg) - assert cls is CuteDslFusedMoE - - -def test_get_moe_cls_cutedsl_returns_cutlass_for_w4a16_nvfp4_on_unsupported_sm(): - """CUTEDSL + W4A16_NVFP4 + non-SM120/121 → CutlassFusedMoE.""" - cfg = ModelConfig() - cfg.moe_backend = "CUTEDSL" - cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) - with patch("tensorrt_llm._utils.get_sm_version", return_value=100): - cls = get_moe_cls(cfg) - assert cls is CutlassFusedMoE - - -@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) -@pytest.mark.parametrize("quant_algo", [QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4]) -def test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm(sm_version, quant_algo): - """CUTEDSL + NVFP4/W4A16_NVFP4 + SM120/121 + flashinfer importable → CuteDslB12xFusedMoE.""" - cfg = ModelConfig() - cfg.moe_backend = "CUTEDSL" - cfg.quant_config = QuantConfig(quant_algo=quant_algo) - with patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version): - cls = get_moe_cls(cfg) - assert cls is CuteDslB12xFusedMoE + verdict = CuteDslB12xFusedMoE.can_implement(_problem(dtype=dtype), _deployment(120)) + assert not verdict.eligible + assert verdict.reject_reason is MoERejectReason.DTYPE_UNSUPPORTED -@pytest.mark.parametrize( - "mapping", - [ - Mapping(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2), - Mapping( - world_size=2, - tp_size=2, - enable_attention_dp=True, - dwdp_size=2, - dwdp_rank=0, - ), - ], -) -def test_get_moe_cls_cutedsl_falls_back_to_cutlass_for_distributed_b12x(mapping): - cfg = ModelConfig(mapping=mapping) - cfg.moe_backend = "CUTEDSL" - cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) - - with patch("tensorrt_llm._utils.get_sm_version", return_value=120): - cls = get_moe_cls(cfg) +def test_can_implement_rejects_missing_flashinfer(): + verdict = CuteDslB12xFusedMoE.can_implement(_problem(), _deployment(120, flashinfer=False)) + assert not verdict.eligible + assert verdict.reject_reason is MoERejectReason.DEP_MISSING - assert cls is CutlassFusedMoE +def test_can_implement_rejects_expert_parallelism(): + verdict = CuteDslB12xFusedMoE.can_implement(_problem(), _deployment(120, ep_size=2)) + assert not verdict.eligible + assert verdict.reject_reason is MoERejectReason.TOPOLOGY_UNSUPPORTED -def test_get_moe_cls_cutedsl_falls_back_to_plain_cutedsl_when_flashinfer_missing(monkeypatch): - """CUTEDSL + NVFP4 + SM120/121 + flashinfer NOT importable → CuteDslFusedMoE.""" - import builtins - cfg = ModelConfig() - cfg.moe_backend = "CUTEDSL" - cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) - - real_import = builtins.__import__ - - def _raise_on_flashinfer(name, *args, **kwargs): - if name == "flashinfer": - raise ImportError("flashinfer not installed (simulated)") - return real_import(name, *args, **kwargs) - - monkeypatch.setattr(builtins, "__import__", _raise_on_flashinfer) - with patch("tensorrt_llm._utils.get_sm_version", return_value=120): - cls = get_moe_cls(cfg) - assert cls is CuteDslFusedMoE +def test_can_implement_rejects_attention_dp_without_expert_parallelism(): + """moe_tp == tp leaves ep_size at 1, so the EP gate alone would let this in.""" + verdict = CuteDslB12xFusedMoE.can_implement( + _problem(), _deployment(120, ep_size=1, use_dp=True, parallel_size=2) + ) + assert not verdict.eligible + assert verdict.reject_reason is MoERejectReason.TOPOLOGY_UNSUPPORTED + assert "attention-DP" in verdict.detail # -------------------------------------------------------------------------- @@ -220,9 +165,7 @@ def _raise_on_flashinfer(name, *args, **kwargs): class _RoutePredicateStub: - """Minimal carrier for ``_PREFILL_VIA_CUTLASS_THRESHOLD`` so we can call - the unbound ``_route_to_cutlass`` without instantiating the whole MoE - backend.""" + """Minimal carrier for the unbound dispatch predicate.""" _PREFILL_VIA_CUTLASS_THRESHOLD = CuteDslB12xFusedMoE._PREFILL_VIA_CUTLASS_THRESHOLD diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index a67703d75f82..2021a0f6b999 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -12,19 +12,7 @@ # 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 Backend Unit Tests - -This module provides a unified test framework for testing different MoE backends -through the backend-level interfaces (quantize_input + run_moe), rather than -the high-level forward() interface. - -Design Goals: -1. Test backend interfaces directly: routing_method.apply -> quantize_input -> run_moe -2. Cover all quantization + backend combinations -3. Use can_implement() interface to determine test skip logic -4. Support autotune and tactic capture testing -""" +"""MoE backend unit tests.""" import importlib import itertools @@ -59,16 +47,28 @@ DeepSeekV3MoeRoutingMethod, RenormalizeMoeRoutingMethod, ) -from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend, get_moe_cls +from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE from tensorrt_llm._torch.modules.fused_moe.fused_moe_marlin import MarlinFusedMoE -from tensorrt_llm._torch.modules.fused_moe.impl_contract import MoECommPlan, MoERunContext +from tensorrt_llm._torch.modules.fused_moe.impl_contract import ( + MoECommPlan, + MoEDeployment, + MoEEnvironment, + MoEProblem, + MoERejectReason, + MoERunContext, +) +from tensorrt_llm._torch.modules.fused_moe.impl_environment import ( + collect_moe_environment, + override_moe_environment, +) from tensorrt_llm._torch.modules.fused_moe.interface import ( MoE, MoESchedulerKind, MoEWeightLoadingMode, ) from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm +from tensorrt_llm._torch.modules.fused_moe.moe_resolution import impl_class_for, resolve_moe_impl from tensorrt_llm._torch.modules.fused_moe.quantization import ( FusedMoEMethodBase, NVFP4FusedMoEMethod, @@ -131,12 +131,7 @@ def test_fp8_block_scale_moe_fallback_tactic_is_explicit_and_deterministic(): def _ensure_single_proc_dist_for_megamoe(backend_type: MoeBackendType, rank: int) -> None: - """Every MegaMoE backend (DG + CuteDSL) resolves an EP ProcessGroup - at construction time via ``_resolve_ep_pg``. Single-process tests - must therefore initialise ``torch.distributed`` even when the test - only exercises ``ep_size == 1`` -- otherwise the constructor raises - ``MegaMoe*Unavailable``. Both MegaMoE backends need the same fixture - so the dist helper must accept the full set.""" + """Initialize the process group required by MegaMoE constructors.""" if backend_type not in _MEGAMOE_BACKEND_TYPES: return if not torch.cuda.is_available(): @@ -406,8 +401,17 @@ def _marlin_model_config(quant_algo=QuantAlgo.NVFP4): return cfg -def test_get_moe_cls_marlin_selects_marlin_for_nvfp4(): - assert get_moe_cls(_marlin_model_config()) is MarlinFusedMoE +def _marlin_environment(sm: int = 90) -> MoEEnvironment: + """Marlin's own SM window, so quantization stays the only variable.""" + return MoEEnvironment(sm=sm) + + +def test_marlin_is_selected_for_nvfp4(): + with override_moe_environment(_marlin_environment()): + report = resolve_moe_impl(_marlin_model_config()) + assert impl_class_for(report) is MarlinFusedMoE + assert report.selected_by == "pinned" + assert not report.degraded @pytest.mark.parametrize( @@ -417,21 +421,24 @@ def test_get_moe_cls_marlin_selects_marlin_for_nvfp4(): pytest.param(QuantAlgo.FP8, id="fp8"), ], ) -def test_get_moe_cls_marlin_falls_back_to_cutlass_on_non_nvfp4(quant_algo): - """MARLIN + non-NVFP4 layers (e.g. unquantized MTP draft layers in - MIXED_PRECISION checkpoints) fall back to CutlassFusedMoE instead of - raising, matching CUTEDSL/DENSEGEMM fallback behavior.""" - assert get_moe_cls(_marlin_model_config(quant_algo)) is CutlassFusedMoE +def test_marlin_degrades_to_cutlass_on_non_nvfp4(quant_algo): + with override_moe_environment(_marlin_environment()): + report = resolve_moe_impl(_marlin_model_config(quant_algo)) + assert impl_class_for(report) is CutlassFusedMoE + assert report.degraded + assert report.degraded_from.reason is MoERejectReason.QUANT_UNSUPPORTED -def test_get_moe_cls_marlin_override_quant_config_per_layer(): - """Per-layer override (the MTP draft-layer path): an unquantized per-layer - override falls back to Cutlass even though the global config is NVFP4.""" +def test_marlin_override_quant_config_degrades_per_layer(): cfg = _marlin_model_config() - assert ( - get_moe_cls(cfg, override_quant_config=QuantConfig(quant_algo=None), layer_idx=52) - is CutlassFusedMoE - ) + with override_moe_environment(_marlin_environment()): + report = resolve_moe_impl( + cfg, + override_quant_config=QuantConfig(quant_algo=None), + layer_idx=52, + ) + assert impl_class_for(report) is CutlassFusedMoE + assert report.degraded_from.reason is MoERejectReason.QUANT_UNSUPPORTED def test_megamoe_cutedsl_post_load_weights_uses_staged_hooks(): @@ -1057,7 +1064,7 @@ def generate_element_wise_test_params() -> List: # Skip Logic # ============================================================================= # Tests are automatically skipped for unsupported configurations using: -# - backend.can_implement(): Check dtype/quant_algo/swiglu_gptoss_style support +# - backend.can_implement(p, d): declared quant / dtype / SM / dependency support # - should_skip_trtllm(): TRTLLM-specific constraints (num_experts % 4, etc.) # - should_skip_cutedsl(): CuteDSL-specific accuracy issues # - 128-alignment requirements for quantization @@ -1339,15 +1346,33 @@ def test_trtllm_bf16_unquantized_moe( backend_type = MoeBackendType.TRTLLM dtype = torch.bfloat16 - can_impl, skip_reason = get_backend_class(backend_type).can_implement( - None, dtype_activation=dtype - ) - if not can_impl: - pytest.skip(skip_reason) - + num_experts = _BF16_UNQUANT_NUM_EXPERTS + top_k = _BF16_UNQUANT_TOP_K hidden_size = _BF16_UNQUANT_HIDDEN intermediate_size = _BF16_UNQUANT_INTERMEDIATE + # This test constructs the backend directly, so query it directly. + verdict = get_backend_class(backend_type).can_implement( + MoEProblem( + quant=None, + dtype_act=dtype, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + top_k=top_k, + ), + MoEDeployment( + ep_size=1, + tp_size=1, + parallel_size=1, + use_dp=False, + num_slots=num_experts, + env=collect_moe_environment(), + ), + ) + if not verdict.eligible: + pytest.skip(verdict.detail) + skip_if_insufficient_gpu_memory(num_experts, hidden_size, intermediate_size, dtype) mapping = Mapping() diff --git a/tests/unittest/_torch/modules/moe/test_moe_module.py b/tests/unittest/_torch/modules/moe/test_moe_module.py index d9ba6045e1a7..a7fae44529b5 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_module.py +++ b/tests/unittest/_torch/modules/moe/test_moe_module.py @@ -264,7 +264,7 @@ def _create_model_config( # CUTE_DSL_B12X is an internal-only MoeBackendType — it has no # corresponding user-facing MoeConfig.backend literal. Route through # "CUTEDSL" so the test exercises the cuteDSL-family selection path that - # users hit on SM120/121 + NVFP4 (where get_moe_cls returns the hybrid + # users hit on SM120/121 + NVFP4 (where resolve_moe_impl picks the hybrid # CuteDslB12xFusedMoE backend when flashinfer is importable). if moe_backend == MoeBackendType.CUTE_DSL_B12X.value: moe_backend = MoeBackendType.CUTEDSL.value @@ -1270,6 +1270,7 @@ def generate_multi_gpu_test_params( model_config=model_config, moe_tp_size=moe_tp_size, dtype=dtype, + swiglu_gptoss_style=swiglu_gptoss_style, ), should_skip_cutedsl( backend_type,