From 4e210c594fc55fc9e952984799087c758024d747 Mon Sep 17 00:00:00 2001 From: xxi Date: Thu, 6 Aug 2026 05:15:06 +0000 Subject: [PATCH] [TRTLLM-14955][refactor] declare MoE backend behaviour instead of comparing classes ExternalCommMoEScheduler and ConfigurableMoE decided what to prepare for a backend by naming its class: 13 sites spread over __class__ ==, isinstance, and a five-way chain in _get_backend_kwargs. Adding a backend therefore meant editing the scheduler, and the checks disagreed with each other about subclasses. Each site now reads a declaration off the backend class: - MoEStaticCapability for the multi-chunk LoRA guard and the DWDP gate. - MoEInputRequirement for routing-scale precision, run_moe workspace allocation, DeepEP expert-id sanitization, and the NVLink one-sided combine workspace dtype. - MoERunContext plus MoECommPlan for run_moe, which took 14 loosely-typed arguments assembled by the class chain. Backends deriving from another backend restate every field rather than inherit it, because the old predicates were inconsistent: the LoRA check was exact-class while the DWDP one was isinstance. Verified all 11 impl classes reproduce the old predicates for every declared field. That restatement rule is transitional and only has to hold while backends still derive from backends; TRTLLM-14960..14969 cut those inheritance edges. One collapse is intentional: the sanitize request drops its isinstance(moe.comm, DeepEP) half, because only DeepEP.dispatch reads the kwarg and every other strategy absorbs it through **kwargs. The routing-scale site keeps both of its halves instead. Its float32 assertion now fires whenever a backend asks for float32, rather than when it happens to be one of two named classes: several routing methods take an output_dtype, and one configured narrower has already dropped bits that the cast would silently widen back. MoEInputRequirement.requires_router_logits goes away unused. The design meant it to replace the scheduler's router-logits filter declaratively, but that filter also reads the routing method instance and an environment override, so a class-level bool cannot answer it. _routes_outside_the_kernel does. Removing the class chain also fixes latent defects, because it had no else branch: any impl outside it silently received empty kwargs and fell back to signature defaults. On CuteDslB12xFusedMoE the CUTLASS prefill path was told is_sf_swizzled=True while quantize_input had produced unswizzled scale factors under post-quant dispatch, and output_dtype was forced to bfloat16 for fp16 models. payload_in_workspace is now assigned on every path, so it can no longer be inherited from a previous forward. The five external-comm impls take the plan through require_comm_plan() rather than defaulting each field when it is absent. comm_plan stays optional on the context because a fused-comm impl owns the exchange and nothing outside its kernel decides anything, but an external-comm impl is only reachable through ExternalCommMoEScheduler, which builds a plan on every path. Defaulting there reintroduced the defect above in a new place: a wrong moe_output or enable_alltoall fails loudly, while a wrong input_sf_swizzled just makes the kernel read scale factors at the wrong stride. The three impls that had such a default disagreed on its value, which is what a guess rather than a decision looks like. Also drops the invalid_token_expert_id = -1 write, which every Communication.__init__ already performs. MoECommPlan is still produced by a scheduler helper rather than by the comm strategy. Relocating it needs combine() to read the flag off the plan instead of off the strategy, which changes the dispatch and combine signatures and stays with TRTLLM-14972. The router_logits filter that sat in the TRTLLMGen arm of the class chain moves into TRTLLMGenFusedMoE._routes_outside_the_kernel, beside the kernel that cares about it. Its three triggers are carried over unchanged. FORCE_SEPARATED_ROUTING moves to interface.py because both sides read it -- the scheduler to decide whether to precompute top-k at all, the backend to decide whether its kernel may route again -- and interface.py is the only module both already import, so neither gains a dependency. Reaching run_moe with router_logits set but no precomputed top-k now raises rather than dereferencing None. The scheduler cannot produce that pairing, since whatever forces the filter also forces the precompute; test_moe_backend calls run_moe directly and can, so it skips the combination. Signed-off-by: xxi --- .../modules/fused_moe/configurable_moe.py | 6 +- .../modules/fused_moe/fused_moe_cute_dsl.py | 38 +-- .../fused_moe/fused_moe_cute_dsl_b12x.py | 47 ++-- .../modules/fused_moe/fused_moe_cutlass.py | 86 +++--- .../modules/fused_moe/fused_moe_deepgemm.py | 43 ++- .../modules/fused_moe/fused_moe_densegemm.py | 28 +- .../modules/fused_moe/fused_moe_marlin.py | 31 ++- .../modules/fused_moe/fused_moe_trtllm_gen.py | 75 +++-- .../_torch/modules/fused_moe/impl_base.py | 5 + .../_torch/modules/fused_moe/impl_contract.py | 64 ++++- .../_torch/modules/fused_moe/interface.py | 78 ++++-- .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 16 +- .../fused_moe/mega_moe/mega_moe_deepgemm.py | 18 +- .../_torch/modules/fused_moe/moe_scheduler.py | 263 ++++++++---------- .../tools/layer_wise_benchmarks/runner.py | 32 +-- .../bench_moe/routing/native_logits.py | 27 +- .../_torch/lora/test_moe_lora_model_path.py | 50 ++-- .../_torch/modules/moe/test_moe_backend.py | 35 ++- 18 files changed, 546 insertions(+), 396 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index 11c648d59f49..ad714d3064e0 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -47,7 +47,6 @@ from tensorrt_llm.models.modeling_utils import QuantConfig from .communication import AllGatherReduceScatter, Communication, CommunicationFactory -from .fused_moe_cute_dsl import CuteDslFusedMoE from .moe_scheduler import MoEScheduler, create_moe_scheduler # Attributes that ConfigurableMoE owns (computed in MoE.__init__ from real @@ -410,8 +409,9 @@ def validate_config(self): ) def _should_enable_dwdp(self) -> bool: - # DWDP is currently supported only for CuteDslFusedMoE with NVFP4 quantization. - if not isinstance(self.backend, CuteDslFusedMoE): + # DWDP is currently supported only by CuteDSL backends, and only with + # NVFP4 quantization. + if not self.backend.capabilities.supports_dwdp: return False quant_config = getattr(self.backend, "quant_config", None) 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 35ac71b9f1b1..4e88d11848bd 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 @@ -37,6 +37,7 @@ 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 .quantization import MoEWeightLoadingMode, NVFP4CuteDslFusedMoEMethod from .routing import BaseMoeRoutingMethod @@ -353,6 +354,14 @@ class CuteDslFusedMoE(CutlassFusedMoE): model_config (ModelConfig): Configuration object for the model. """ + # ``supports_moe_lora`` is restated because CutlassFusedMoE declares True + # and the exact-class comparison it replaces answered False here. + # ``supports_dwdp`` is the capability this backend adds. CuteDslB12xFusedMoE + # derives from here and needs both, but must spell them out again: setting + # the attribute replaces the whole object rather than one field. + capabilities = MoEStaticCapability(supports_moe_lora=False, + supports_dwdp=True) + @classmethod def can_implement( cls, @@ -796,13 +805,9 @@ def run_moe_fp8_block_scales( def run_moe( self, - x: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: Optional[torch.Tensor], - x_sf: Optional[torch.Tensor] = None, - moe_output: Optional[torch.Tensor] = None, - enable_alltoall: bool = False, - **kwargs, + ctx: MoERunContext, + *, + workspace: Optional[dict] = None, ) -> torch.Tensor: """ Run MoE computation with CuteDSL backend. @@ -810,19 +815,18 @@ def run_moe( This method encapsulates the core MoE computation logic, handling different quantization schemes (fp8_block_scales and nvfp4). - Args: - # Standard MoE interface parameters: - x: Input hidden states (may be pre-quantized) - token_selected_experts: Expert IDs [num_tokens, top_k]. If EPLB is enabled, - this represents expert slots [num_tokens, top_k] instead. - token_final_scales: Final scaling factors for each token - x_sf: Input scale factors (optional, for certain quantization schemes) - moe_output: Pre-allocated MoE output buffer (optional, for NVLINK one-sided backend). - enable_alltoall: Whether alltoall communication is enabled. - Returns: final_hidden_states tensor. """ + del workspace # CuteDSL kernels allocate their own intermediates. + plan = require_comm_plan(self, ctx) + x = ctx.x + token_selected_experts = ctx.token_selected_experts + token_final_scales = ctx.token_final_scales + x_sf = ctx.x_sf + moe_output = plan.moe_output + enable_alltoall = plan.enable_alltoall + # Execute MoE computation if self.has_nvfp4: weight_view = self._build_local_weight_view() 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 f9cc59724ec1..962e5dc0fe76 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 @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import replace from typing import Optional, Tuple, Union import torch @@ -23,6 +24,7 @@ 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 # Shared MoE output buffer pool, keyed by (max_num_tokens, hidden_size, dtype, @@ -83,6 +85,11 @@ class CuteDslB12xFusedMoE(CuteDslFusedMoE): flashinfer-importable gates pass (see ``create_moe.get_moe_cls``). """ + # Restated rather than inherited: the LoRA gate this replaces compared the + # exact class and answered False here, while the DWDP gate used isinstance + # and answered True through CuteDslFusedMoE. + capabilities = MoEStaticCapability(supports_moe_lora=False, supports_dwdp=True) + # SM versions on which the FlashInfer b12x NVFP4 MoE kernel is available. # SM120 = desktop Blackwell (RTX 5090 / GB202); SM121 = GB10 / DGX Spark. _SUPPORTED_SM_VERSIONS = frozenset({120, 121}) @@ -211,29 +218,21 @@ def quantize_input( @nvtx_range("[b12x] run_moe") def run_moe( self, - x: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: torch.Tensor, - x_sf: Optional[torch.Tensor] = None, - is_sf_swizzled: bool = True, - output_dtype: Optional[torch.dtype] = None, - tuner_num_tokens: Optional[int] = None, - tuner_top_k: Optional[int] = None, - moe_output: Optional[torch.Tensor] = None, - enable_alltoall: Optional[bool] = None, + ctx: MoERunContext, + *, + workspace: Optional[dict] = None, ) -> torch.Tensor: + plan = require_comm_plan(self, ctx) + x = ctx.x if self._route_to_cutlass(x): # ``CutlassFusedMoE.run_moe`` forwards ``output_dtype`` straight # into the C++ ``trtllm::fused_moe`` op, which requires a concrete # high-precision ``ScalarType`` (uint8 / FP4-packed activations are # rejected at the kernel epilogue with "Invalid output type Byte"). - # Schedulers that drive ``run_moe`` directly (the KV-cache capacity - # probe, for one) leave ``output_dtype`` unset, so fall back to - # ``x.dtype`` if it is a real compute dtype, else bf16. Mirrors the - # ``forward_chunk`` convention while staying safe for the FP4 - # quant-input path (``x`` is uint8 after ``quantize_input``). + # ``ConfigurableMoE.forward`` always fills ``output_dtype``, so this + # only narrows the type for anything driving ``run_moe`` without it. _HIGH_PRECISION = {torch.float16, torch.bfloat16, torch.float32} - cutlass_output_dtype = output_dtype + cutlass_output_dtype = ctx.output_dtype if cutlass_output_dtype is None: cutlass_output_dtype = ( x.dtype @@ -242,17 +241,13 @@ def run_moe( ) return CutlassFusedMoE.run_moe( self, - x, - token_selected_experts=token_selected_experts, - token_final_scales=token_final_scales, - x_sf=x_sf, - is_sf_swizzled=is_sf_swizzled, - output_dtype=cutlass_output_dtype, - tuner_num_tokens=tuner_num_tokens, - tuner_top_k=tuner_top_k, - moe_output=moe_output, - enable_alltoall=enable_alltoall, + replace(ctx, output_dtype=cutlass_output_dtype), + workspace=workspace, ) + token_selected_experts = ctx.token_selected_experts + token_final_scales = ctx.token_final_scales + x_sf = ctx.x_sf + moe_output = plan.moe_output if self.b12x_wrapper is None or self._b12x_weights is None: raise RuntimeError( "CuteDslB12xFusedMoE.run_moe called before process_weights_after_loading completed." 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 cad3629f0d44..0c967d03a52f 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -28,6 +28,8 @@ from ...peft.lora.validation import has_moe_lora_targets from ...utils import (ActivationType, AuxStreamType, EventType, Fp4QuantizedTensor) +from .impl_contract import (MoEInputRequirement, MoERunContext, + MoEStaticCapability, require_comm_plan) from .interface import MoE from .quantization import UnquantizedFusedMoEMethod @@ -85,6 +87,13 @@ class CutlassFusedMoE(MoE): equals to: dynamic quant + routing(topK, etc.) [+ fp4_allgather] + scatter + gemm1 + swiglu + gemm2 + finalizeMoeRoute [no allreduce] + reducescatter """ + # Routed-expert MoE LoRA is fused into this backend's op only; the + # subclasses below each restate ``supports_moe_lora=False``. + capabilities = MoEStaticCapability(supports_moe_lora=True) + + # Inherited by every subclass, matching the isinstance check this replaces. + input_requirement = MoEInputRequirement(routing_scales_dtype=torch.float32) + # Quantization algorithm support table for can_implement() # Format: quant_algo -> {sm_constraint, dtypes} # sm_constraint types: @@ -848,19 +857,30 @@ def create_weights(self): def supports_moe_output_in_alltoall_workspace(self): return True + def _tuner_shapes( + self, + ctx: MoERunContext, + enable_alltoall: Optional[bool], + ) -> Tuple[Optional[int], Optional[int]]: + """Token/top-k shapes the profiling tuner should key on. + + Only meaningful under alltoall: the tuner must see pre-alltoall token + counts so tactics cached during the no-alltoall warmup still apply at + runtime. Without alltoall the kernel derives both from ``x`` itself. + """ + if not enable_alltoall: + return None, None + if ctx.all_rank_num_tokens is not None: + tuner_num_tokens = sum(ctx.all_rank_num_tokens) + else: + tuner_num_tokens = ctx.x.shape[0] * self.mapping.tp_size + return tuner_num_tokens, self.routing_method.top_k + def run_moe( self, - x: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: torch.Tensor, - x_sf: Optional[torch.Tensor] = None, - is_sf_swizzled: bool = True, - output_dtype: Optional[torch.dtype] = None, - tuner_num_tokens: Optional[int] = None, - tuner_top_k: Optional[int] = None, - moe_output: Optional[torch.Tensor] = None, - enable_alltoall: Optional[bool] = None, - lora_params: Optional[Dict] = None, + ctx: MoERunContext, + *, + workspace: Optional[dict] = None, ) -> torch.Tensor: """ Run MoE computation with Cutlass backend. @@ -868,22 +888,22 @@ def run_moe( This method encapsulates the core MoE computation logic, handling different quantization schemes. - Args: - x: Input hidden states (may be pre-quantized) - token_selected_experts: Expert IDs or expert slots [num_tokens, top_k] - If EPLB is enabled, represents expert slots; otherwise expert IDs - token_final_scales: Final scaling factors for each token - x_sf: Input scale factors (optional, for certain quantization schemes) - is_sf_swizzled: Whether scaling factors are swizzled - output_dtype: Output data type (optional) - tuner_num_tokens: Number of tokens for profiling tuner (optional) - tuner_top_k: Top-k value for profiling tuner (optional) - moe_output: Pre-allocated output buffer (optional) - enable_alltoall: Whether alltoall communication is enabled (optional). If None, defaults to self.enable_alltoall. - Returns: final_hidden_states: Output tensor from MoE computation """ + del workspace # Cutlass allocates its own intermediates. + plan = require_comm_plan(self, ctx) + x = ctx.x + token_selected_experts = ctx.token_selected_experts + token_final_scales = ctx.token_final_scales + x_sf = ctx.x_sf + output_dtype = ctx.output_dtype + lora_params = ctx.lora_params + is_sf_swizzled = plan.input_sf_swizzled + moe_output = plan.moe_output + enable_alltoall = plan.enable_alltoall + tuner_num_tokens, tuner_top_k = self._tuner_shapes(ctx, enable_alltoall) + # W4A16 NVFP4 fallback (SM<100). if isinstance(self.quant_method, W4A16NVFP4CutlassFusedMoEMethod): return self._run_moe_w4a16_nvfp4( @@ -902,8 +922,7 @@ def run_moe( if self.has_deepseek_fp8_block_scales and get_sm_version() == 120: from .fused_moe_triton_fp8_block_scale import \ run_triton_fp8_block_scale_moe - _use_alltoall = (enable_alltoall if enable_alltoall is not None else - self.enable_alltoall) + # forward_chunk sets token_final_scales=None when # apply_router_weight_on_input=True (weights already folded into x); # substitute ones so the Triton kernel's per-token scaling is a no-op. @@ -916,7 +935,7 @@ def run_moe( # (0 .. expert_size_per_partition-1), so remap and zero-scale any # non-local token-expert pairs to suppress their contribution. local_n = self.expert_size_per_partition - if _use_alltoall: + if enable_alltoall: # After alltoall dispatch, IDs are already local; padding = local_n local_ids = token_selected_experts.clamp(0, local_n - 1) is_local = token_selected_experts < local_n @@ -957,9 +976,6 @@ def run_moe( elif self.has_w4a16_mxfp4: weight_dtype = torch.uint8 - if enable_alltoall is None: - enable_alltoall = self.enable_alltoall - use_dynamic_fc2_scale = (self.has_nvfp4 and getattr( self, 'force_dynamic_quantization', False) and hasattr(self, 'fc2_weight_scale_2')) @@ -1032,16 +1048,20 @@ def _run_moe_w4a16_nvfp4( tuner_num_tokens: Optional[int] = None, tuner_top_k: Optional[int] = None, moe_output: Optional[torch.Tensor] = None, - enable_alltoall: Optional[bool] = None, + *, + enable_alltoall: bool, ) -> torch.Tensor: """W4A16 fallback for NVFP4 MoE on SM<100. Active-mask dequant into a static [E_total, N, K] bf16 workspace, then bf16 fused_moe with the original (global) token_selected_experts. CUDA-graph capturable. + + ``enable_alltoall`` has no default because it picks the expert-id remap + below, and either default is silently wrong for half the callers: the + ids are local after an alltoall dispatch and global otherwise, so a + wrong guess shifts every id by ``slot_start`` without failing. """ assert isinstance(self.quant_method, W4A16NVFP4CutlassFusedMoEMethod) - if enable_alltoall is None: - enable_alltoall = self.enable_alltoall if output_dtype is None: output_dtype = x.dtype 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 f030559b358e..7d3951762326 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py @@ -28,6 +28,8 @@ from ...model_config import ModelConfig from ...utils import AuxStreamType, Fp4QuantizedTensor from .fused_moe_cutlass import CutlassFusedMoE +from .impl_contract import (MoEInputRequirement, MoERunContext, + MoEStaticCapability) from .quantization import (DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm, MoEWeightLoadingMode, UnquantizedFusedMoEMethod) from .routing import BaseMoeRoutingMethod @@ -715,6 +717,25 @@ class DeepGemmFusedMoE(CutlassFusedMoE): model_config (ModelConfig): Configuration object for the model. """ + # Restated rather than inherited from CutlassFusedMoE: this backend does + # not fuse routed-expert LoRA, and the exact-class comparison this field + # replaces already answered False here. + capabilities = MoEStaticCapability(supports_moe_lora=False) + + # ``routing_scales_dtype`` is repeated from CutlassFusedMoE because setting + # any field here replaces the parent's object wholesale. + input_requirement = MoEInputRequirement( + routing_scales_dtype=torch.float32, + requires_run_moe_workspace=True, + ) + + def supports_moe_output_in_alltoall_workspace(self): + # Overrides the CutlassFusedMoE "True": run_moe emits into its own + # workspace buffers and never writes a caller-supplied output tensor, + # so a workspace-backed buffer would be left unfilled while combine() + # read from it. + return False + @classmethod def can_implement( cls, @@ -930,11 +951,9 @@ def quantize_input( def run_moe( self, - x: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: torch.Tensor, - x_sf: Optional[torch.Tensor] = None, - workspace: dict = None, + ctx: MoERunContext, + *, + workspace: Optional[dict] = None, ) -> torch.Tensor: """ Run MoE computation with DeepGemm backend. @@ -943,13 +962,9 @@ def run_moe( quantization with DeepGemm backend. Args: - # Standard MoE interface parameters: - x: Input hidden states (unquantized for DeepGemm) - token_selected_experts: Expert IDs [num_tokens, top_k]. If EPLB is enabled, - this represents expert slots [num_tokens, top_k] instead. - token_final_scales: Final scaling factors for each token - x_sf: Input scale factors (should be None for DeepGemm) - workspace: Workspace dictionary containing buffers for intermediate results + ctx: Run context; ``x`` is unquantized and ``x_sf`` must be None. + workspace: Buffers for intermediate results, allocated once per + chunk by the scheduler so the aux stream can reuse them. Required keys: 'workspace_0', 'workspace_1', 'workspace_sf' Returns: @@ -957,6 +972,10 @@ def run_moe( Note: Similar to CuteDslFusedMoE.run_moe_fp8_block_scales (fused_moe_cute_dsl.py:360-434) """ + x = ctx.x + token_selected_experts = ctx.token_selected_experts + token_final_scales = ctx.token_final_scales + x_sf = ctx.x_sf assert self.has_deepseek_fp8_block_scales assert x_sf is None assert workspace is not None, "workspace is required for DeepGemm backend" 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 6aac29f34ce6..1dad6dab2ec5 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py @@ -13,6 +13,7 @@ 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 .quantization import NVFP4CuteDslFusedMoEMethod from .routing import BaseMoeRoutingMethod @@ -101,6 +102,8 @@ class DenseGEMMFusedMoE(MoE): model_config (ModelConfig): Configuration object for the model. """ + input_requirement = MoEInputRequirement(routing_scales_dtype=torch.float32) + # Memory buffer pool for CUDA graph compatibility buffers = get_memory_buffers() @@ -500,28 +503,23 @@ def run_moe_nvfp4( def run_moe( self, - x: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: Optional[torch.Tensor], - x_sf: Optional[torch.Tensor] = None, - enable_alltoall: bool = False, - **kwargs, + ctx: MoERunContext, + *, + workspace: Optional[dict] = None, ) -> torch.Tensor: """ Run MoE computation with DenseGEMM backend (NVFP4 only). - Args: - x: Input hidden states (pre-quantized to NVFP4) - token_selected_experts: Expert IDs [num_tokens, top_k]. If EPLB is enabled, - this represents expert slots [num_tokens, top_k] instead. - token_final_scales: Final scaling factors for each token - x_sf: Input scale factors for NVFP4 - enable_alltoall: Whether alltoall communication is enabled. - **kwargs: Additional arguments for forward compatibility. - Returns: final_hidden_states tensor. """ + del workspace # DenseGEMM allocates its own intermediates. + plan = require_comm_plan(self, ctx) + x = ctx.x + token_selected_experts = ctx.token_selected_experts + token_final_scales = ctx.token_final_scales + x_sf = ctx.x_sf + enable_alltoall = plan.enable_alltoall assert self.has_nvfp4, ( f"{self.__class__.__name__} only supports nvfp4 quantization, " f"got {self.quant_config.quant_mode}." 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 c01a99da51d2..517bfe5346dd 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py @@ -31,6 +31,7 @@ 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 .quantization import NVFP4MarlinFusedMoEMethod @@ -52,6 +53,10 @@ class MarlinFusedMoE(CutlassFusedMoE): compatible. Requires the fused kernel to be built (no fallback path). """ + # Restated rather than inherited from CutlassFusedMoE, whose exact-class + # LoRA comparison answered False for this backend. + capabilities = MoEStaticCapability(supports_moe_lora=False) + _QUANT_SUPPORT_TABLE = { QuantAlgo.NVFP4: { "sm_constraint": ("in", set(range(89, 100))), @@ -143,20 +148,24 @@ def _ensure_workspace(self, device: torch.device): # Main entry point # ==================================================================== + def supports_moe_output_in_alltoall_workspace(self): + # Overrides the CutlassFusedMoE "True": this kernel always allocates + # and returns its own output tensor, so a workspace-backed buffer + # would be filled by nobody while combine() read from it. + return False + def run_moe( self, - x: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: torch.Tensor, - x_sf: Optional[torch.Tensor] = None, - is_sf_swizzled: bool = True, - output_dtype: Optional[torch.dtype] = None, - tuner_num_tokens: Optional[int] = None, - tuner_top_k: Optional[int] = None, - moe_output: Optional[torch.Tensor] = None, - enable_alltoall: Optional[bool] = None, - router_logits: Optional[torch.Tensor] = None, + ctx: MoERunContext, + *, + workspace: Optional[dict] = None, ) -> torch.Tensor: + del workspace # Marlin owns its own scratch (see _marlin_workspace). + x = ctx.x + token_selected_experts = ctx.token_selected_experts + token_final_scales = ctx.token_final_scales + router_logits = ctx.router_logits + output_dtype = ctx.output_dtype assert output_dtype is None or output_dtype == torch.bfloat16 assert _has_fused_moe_kernel(), ( "marlin_nvfp4_moe_gemm is not available. Rebuild TensorRT-LLM " 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 a1bf48be87a6..79c2f064e1ac 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 @@ -31,7 +31,8 @@ from ...utils import (ActivationType, ActType_TrtllmGen, AuxStreamType, Fp4QuantizedTensor) from ..gated_mlp import GatedMLP -from .interface import MoE, MoEWeightLoadingMode +from .impl_contract import MoEInputRequirement, MoERunContext, require_comm_plan +from .interface import FORCE_SEPARATED_ROUTING, MoE, MoEWeightLoadingMode from .moe_op_backend import MoEOpBackend, TRTLLMOpBackend, get_op_backend # isort: off @@ -87,6 +88,16 @@ class TRTLLMGenFusedMoE(MoE): There should be at lease `num_experts` slots in the model engine. More than that is OK, in that case, some experts may have multiple replicas. """ + # bfloat16 routing scales are what these kernels read, and the DeepEP + # dispatch has to mark unfilled rows before they reach them. + input_requirement = MoEInputRequirement( + routing_scales_dtype=torch.bfloat16, + requires_sanitized_expert_ids=True, + # The combine reduction runs in bf16 regardless of the model's output + # dtype, so the NVLink one-sided payload buffer must be bf16 too. + onesided_workspace_dtype=torch.bfloat16, + ) + # Supported quantization algorithms for TRTLLMGenFusedMoE _SUPPORTED_QUANT_ALGOS = { QuantAlgo.NVFP4, @@ -479,6 +490,18 @@ def _supports_load_balancer(self) -> bool: return True return self.use_dp and self.parallel_size > 1 + def _routes_outside_the_kernel(self) -> bool: + """Whether top-k is precomputed, so the kernel must not route again. + + Three independent triggers, none of which subsumes the others: a + kernel or parallel layout that forces it (both folded into + ``_supports_load_balancer``), a routing algorithm no C++ kernel + implements, and the host-routing override. + """ + return (self._supports_load_balancer() + or self.routing_method.requires_separated_routing + or FORCE_SEPARATED_ROUTING) + def _check_configs(self): assert not self.has_any_quant \ or self.has_deepseek_fp8_block_scales \ @@ -755,13 +778,9 @@ def fuse_shared_expert(self, shared_experts: GatedMLP): def run_moe( self, - x: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: Optional[torch.Tensor], - x_sf: Optional[torch.Tensor] = None, - router_logits: Optional[torch.Tensor] = None, - do_finalize: bool = True, - moe_output: Optional[torch.Tensor] = None, + ctx: MoERunContext, + *, + workspace: Optional[dict] = None, ) -> Union[torch.Tensor, tuple]: """ Run MoE computation with TRTLLMGen backend. @@ -770,26 +789,30 @@ def run_moe( quantization schemes (bf16, fp8_block_scales, nvfp4, w4a16_mxfp4, w4a8_nvfp4_fp8, w4a8_mxfp4_fp8, w4a8_mxfp4_mxfp8). - Args: - # Standard MoE interface parameters: - x: Input hidden states (may be pre-quantized) - token_selected_experts: Expert IDs [num_tokens, top_k]. If EPLB is enabled, - this represents expert slots [num_tokens, top_k] instead. - token_final_scales: Final scaling factors for each token - x_sf: Input scale factors (optional, for certain quantization schemes) - - # TRTLLMGen-specific additional parameters: - router_logits: Router logits for integrated routing in some kernels. - Should be None if routing has already been done (e.g., post_quant_comm). - do_finalize: Whether to finalize the output. If False, returns intermediate - results (tuple) for nvfp4 and w4a8_nvfp4_fp8 schemes. - moe_output: Pre-allocated output buffer from workspace (optional). - Used for mnnvlthroughput alltoall backend to avoid extra copies. - Returns: - If do_finalize=True: final_hidden_states tensor - If do_finalize=False: tuple of intermediate outputs (for nvfp4 and w4a8_nvfp4_fp8) + If ``ctx.do_finalize``: final_hidden_states tensor + Otherwise: tuple of intermediate outputs (for nvfp4 and w4a8_nvfp4_fp8) """ + del workspace # TRTLLMGen kernels allocate their own intermediates. + plan = require_comm_plan(self, ctx) + x = ctx.x + token_selected_experts = ctx.token_selected_experts + token_final_scales = ctx.token_final_scales + x_sf = ctx.x_sf + do_finalize = ctx.do_finalize + moe_output = plan.moe_output + # The caller used to apply this filter before handing over the kwargs. + if self._routes_outside_the_kernel(): + if ctx.router_logits is not None and token_selected_experts is None: + raise ValueError( + f"{type(self).__name__} requires separated routing for this " + "config, so ctx.router_logits is ignored, but " + "ctx.token_selected_experts is None -- there is nothing left " + "to route with. Supply precomputed top-k ids and scales.") + router_logits = None + else: + router_logits = ctx.router_logits + routing_params = self._extract_routing_params() top_k = routing_params.top_k routing_bias = routing_params.routing_bias if router_logits is not None else None diff --git a/tensorrt_llm/_torch/modules/fused_moe/impl_base.py b/tensorrt_llm/_torch/modules/fused_moe/impl_base.py index 1d3ac3eb9141..684a2452e770 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/impl_base.py +++ b/tensorrt_llm/_torch/modules/fused_moe/impl_base.py @@ -67,6 +67,11 @@ def quantize_input( self, x: "torch.Tensor | Fp4QuantizedTensor", **kwargs: object ) -> "tuple[torch.Tensor, torch.Tensor | None] | dict": ... + # Narrower than ``MoE.run_moe``, which also takes a keyword-only + # ``workspace``. Not a drift: the scheduler only hands out scratch because + # no impl allocates its own yet, and this signature is the state after + # ``get_workspaces`` below takes over. Impls moving onto this base + # (TRTLLM-14958, TRTLLM-14960..14969) drop the parameter as they arrive. @abc.abstractmethod def run_moe(self, ctx: MoERunContext) -> torch.Tensor: ... diff --git a/tensorrt_llm/_torch/modules/fused_moe/impl_contract.py b/tensorrt_llm/_torch/modules/fused_moe/impl_contract.py index 9a48c455c302..4c2d76d3b5b1 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/modules/fused_moe/impl_contract.py @@ -84,9 +84,8 @@ class MoEInputRequirement: # Legacy: the DeepEP-comm plus TRTLLMGenFusedMoE special case in # MoEScheduler. requires_sanitized_expert_ids: bool = False - requires_router_logits: bool = False - # Legacy: the bfloat16 combine workspace picked in - # ``MoEScheduler._get_nvlink_onesided_moe_output``. + # Overrides the NVLink one-sided combine payload dtype; None means the + # buffer follows the model output dtype. onesided_workspace_dtype: Optional[torch.dtype] = None # There is deliberately no sentinel field here. Every Communication sets @@ -94,6 +93,13 @@ class MoEInputRequirement: # kernels accept nothing else, so the value is a comm-side invariant rather # than something an impl needs the caller to supply. + # There is deliberately no ``requires_router_logits`` field either, though + # the design sketched one to replace the scheduler's router-logits filter. + # A class-level bool cannot express that condition: it also depends on the + # routing method instance and on an environment override, neither of which + # is known per class. ``TRTLLMGenFusedMoE._routes_outside_the_kernel`` + # answers it instead, next to the kernel whose contract it describes. + # --------------------------------------------------------------------------- # Selection inputs @@ -242,14 +248,23 @@ def eligible(self) -> Tuple["MoEImplId", ...]: class MoECommPlan: """What the comm layer decided for THIS forward. Facts, not capabilities. - Single producer: the comm strategy builds it at the end of dispatch, and - both ``run_moe`` and the following ``comm.combine()`` read the same object. + Single producer, so the value ``run_moe`` sees and the value the following + ``comm.combine()`` acts on cannot drift apart. That producer is currently + ``ExternalCommMoEScheduler._build_comm_plan`` rather than the comm strategy + itself; moving it onto the strategy needs ``combine()`` to read + ``payload_in_workspace`` off the plan instead of off the strategy, which + changes the dispatch and combine signatures and stays with TRTLLM-14972. """ input_sf_swizzled: bool # what quantize_input actually produced enable_alltoall: bool moe_output: Optional[torch.Tensor] # workspace-backed buffer, or None - payload_in_workspace: bool # combine() reads the same field + # No impl reads this one yet: ``combine()`` still takes the value off the + # strategy attribute, which the producer assigns from this same field so the + # two cannot disagree. It is here because the plan is where the decision is + # made; TRTLLM-14972 makes ``combine()`` read it from here and drops the + # attribute. + payload_in_workspace: bool @dataclass(frozen=True) @@ -261,12 +276,14 @@ class MoERunContext: than handed to it. """ - # produced by routing - token_selected_experts: torch.Tensor - token_final_scales: Optional[torch.Tensor] + # produced by routing. Expert IDs [num_tokens, top_k], or expert slots of + # the same shape when EPLB is enabled. None when the impl routes internally + # from ``router_logits`` instead. + token_selected_experts: Optional[torch.Tensor] + token_final_scales: Optional[torch.Tensor] # routing weights [num_tokens, top_k] # produced by quantize_input - x: torch.Tensor - x_sf: Optional[torch.Tensor] + x: torch.Tensor # activations [num_tokens, hidden_size] + x_sf: Optional[torch.Tensor] # scale factors, when the input is quantized # produced by the outer forward output_dtype: Optional[torch.dtype] = None do_finalize: bool = True @@ -278,6 +295,31 @@ class MoERunContext: comm_plan: Optional[MoECommPlan] = None +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 ctx.comm_plan + + @dataclass(frozen=True) class MoEEplbBinding: """Everything an impl needs to lay out and load its expert weights. diff --git a/tensorrt_llm/_torch/modules/fused_moe/interface.py b/tensorrt_llm/_torch/modules/fused_moe/interface.py index 1e7ecd06845f..73e28bb8beb2 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/interface.py +++ b/tensorrt_llm/_torch/modules/fused_moe/interface.py @@ -26,6 +26,20 @@ from tensorrt_llm.models.modeling_utils import QuantAlgo from ...distributed.ops import reducescatter +from .impl_contract import (MoEInputRequirement, MoERunContext, + MoEStaticCapability) + +# Route on the host (fused noaux_tc + post-topk pipeline) instead of inside +# the trtllm-gen cubin. The in-cubin top-k tier for large expert counts +# (896 experts / top-16) register-spills and costs ~33 us/layer at decode +# batch 5..64 vs ~10 us for the post-topk pipeline; the separated path is +# the same math the attention-DP deployments already run. +# +# Lives here rather than next to either reader because both need it: the +# scheduler decides whether to precompute top-k at all, and TRTLLMGenFusedMoE +# decides whether its kernel may route again. +FORCE_SEPARATED_ROUTING = os.environ.get( + "TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING", "0") == "1" def _warn_and_return(reason: str) -> Tuple[bool, Optional[str]]: @@ -227,6 +241,28 @@ 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. + 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. + input_requirement: MoEInputRequirement = MoEInputRequirement() + # Opt-in flag for non-divisible EP (num_experts % ep_size != 0). False by default # so backends whose dispatch/combine paths still assume uniform partitioning fail # fast with a clear error. Backends that fully exercise the ceil/floor partition @@ -927,32 +963,30 @@ def quantize_input( @abstractmethod def run_moe( self, - # ========== Common parameters (all backends use) ========== - x: torch.Tensor, - token_selected_experts: Optional[torch.Tensor], - token_final_scales: Optional[torch.Tensor], - x_sf: Optional[torch.Tensor] = None, - # ========== Backend-specific parameters (via kwargs) ========== - **kwargs + ctx: MoERunContext, + *, + workspace: Optional[dict] = None, ) -> torch.Tensor: """ - Unified MoE computation interface + Unified MoE computation interface. - NOTE: This is a TEMPORARY interface. In the future, this method should be moved - to the MoEBackend interface as part of the backend abstraction layer. + Every value the caller genuinely produces travels in ``ctx``; every + fact the comm layer decided for this forward travels in + ``ctx.comm_plan``. Backends read only the fields they need, so adding a + backend never requires touching the scheduler. - This method performs the core MoE computation. Different backends will implement - their specific computation logic while following this unified interface. - - Common parameters (all backends use): - x: Input activations [num_tokens, hidden_size] - token_selected_experts: Expert IDs [num_tokens, top_k] (used by DeepGemm/TRTLLMGen). - If EPLB is enabled, this represents expert slots [num_tokens, top_k]. - token_final_scales: Routing weights [num_tokens, top_k] - x_sf: Input scale factor (for quantization, if applicable) - - Backend-specific parameters (passed via kwargs, obtained from _get_backend_kwargs()): - TODO: This is not finalized, will be updated later. + Args: + ctx: Inputs for this forward. ``token_selected_experts`` holds + expert slots rather than expert IDs when EPLB is enabled. + workspace: Scratch buffers owned by the scheduler because they are + allocated once per chunk and reused across the aux stream. Only + backends declaring ``requires_run_moe_workspace`` receive one. + ``MoEImplBase.run_moe`` deliberately omits this parameter: it + describes the state after impls allocate their own scratch + through ``get_workspaces``, which happens as each impl moves + onto that base (TRTLLM-14958, TRTLLM-14960..14969). Keyword-only + here so that removing it is a mechanical change to named call + sites rather than a silent re-binding of a positional argument. Returns: torch.Tensor: MoE computation result [num_tokens, hidden_size] 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 00c67ecc4ea6..38f057f999eb 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 @@ -93,6 +93,7 @@ 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 ..quantization import NVFP4MegaMoECuteDslMethod from ..routing import BaseMoeRoutingMethod @@ -1111,13 +1112,9 @@ def quantize_input( def run_moe( self, - x: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: torch.Tensor, - x_sf: Optional[torch.Tensor] = None, + ctx: MoERunContext, *, - output_dtype: Optional[torch.dtype] = None, - **unused_kwargs, + workspace: Optional[dict] = None, ) -> torch.Tensor: """Run the fused MegaMoE CuteDSL kernel on pre-quantized inputs. @@ -1127,7 +1124,12 @@ def run_moe( to :meth:`_run_moe`, which returns the reduced ``(T, hidden)`` output. """ - del unused_kwargs + del workspace # The symmetric buffer is this backend's own workspace. + x = ctx.x + token_selected_experts = ctx.token_selected_experts + token_final_scales = ctx.token_final_scales + x_sf = ctx.x_sf + output_dtype = ctx.output_dtype if output_dtype is None: output_dtype = self.dtype or torch.bfloat16 if x_sf is None: 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 23b82228d196..530410cbdbd6 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 @@ -35,6 +35,7 @@ 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, @@ -651,13 +652,9 @@ def supports_fused_prepare(self) -> bool: def run_moe( self, - x: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: torch.Tensor, - x_sf: Optional[torch.Tensor] = None, + ctx: MoERunContext, *, - output_dtype: Optional[torch.dtype] = None, - **unused_kwargs, + workspace: Optional[dict] = None, ) -> torch.Tensor: """Run the fused kernel with either BF16 or pre-quantized activations. @@ -665,9 +662,12 @@ def run_moe( FP8+SF+topk SymmBuffer fields in one custom op. The fallback path keeps the original ``quantize_input`` + copy contract. """ - assert not unused_kwargs, ( - f"MegaMoEDeepGemm.run_moe got unexpected kwargs: {sorted(unused_kwargs)}" - ) + del workspace # The SymmBuffer is this backend's own workspace. + x = ctx.x + token_selected_experts = ctx.token_selected_experts + token_final_scales = ctx.token_final_scales + x_sf = ctx.x_sf + output_dtype = ctx.output_dtype if output_dtype is None: output_dtype = self.dtype or torch.bfloat16 dg = self._dg diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py index 72d355f97289..f86ecf887aee 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py @@ -42,7 +42,6 @@ from __future__ import annotations -import os from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union @@ -54,20 +53,9 @@ from .communication import DeepEP, DeepEPLowLatency, NcclEP, NVLinkOneSided, NVLinkTwoSided from .communication.nvlink_two_sided_flashinfer import NVLinkTwoSidedFlashinfer -from .fused_moe_cute_dsl import CuteDslFusedMoE -from .fused_moe_cutlass import CutlassFusedMoE, raise_moe_lora_multichunk_unsupported -from .fused_moe_deepgemm import DeepGemmFusedMoE -from .fused_moe_densegemm import DenseGEMMFusedMoE -from .fused_moe_marlin import MarlinFusedMoE -from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE -from .interface import MoESchedulerKind - -# Route on the host (fused noaux_tc + post-topk pipeline) instead of inside -# the trtllm-gen cubin. The in-cubin top-k tier for large expert counts -# (896 experts / top-16) register-spills and costs ~33 us/layer at decode -# batch 5..64 vs ~10 us for the post-topk pipeline; the separated path is -# the same math the attention-DP deployments already run. -FORCE_SEPARATED_ROUTING = os.environ.get("TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING", "0") == "1" +from .fused_moe_cutlass import raise_moe_lora_multichunk_unsupported +from .impl_contract import MoECommPlan, MoERunContext +from .interface import FORCE_SEPARATED_ROUTING, MoESchedulerKind __all__ = [ "MoEScheduler", @@ -163,7 +151,7 @@ def forward( if ( num_chunks > 1 - and moe.backend.__class__ == CutlassFusedMoE + and moe.backend.capabilities.supports_moe_lora and moe.backend._moe_lora_active(lora_params) ): raise_moe_lora_multichunk_unsupported(num_chunks) @@ -249,12 +237,12 @@ def _prepare_workspace_deepgemm( x: Union[torch.Tensor, Fp4QuantizedTensor], all_rank_num_tokens: List[int], ) -> Optional[torch.Tensor]: - """Single-chunk workspace for DeepGemmFusedMoE; otherwise ``None``. + """Single-chunk workspace for backends that ask for one; else ``None``. Multi-chunk execution uses ``_prepare_workspaces_for_chunk`` instead. """ moe = self.moe - if not isinstance(moe.backend, DeepGemmFusedMoE): + if not moe.backend.input_requirement.requires_run_moe_workspace: return None num_rows = x.shape[0] @@ -276,7 +264,7 @@ def _prepare_workspaces_for_chunk( chunk_size_list: List[int], use_multi_stream: bool, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - """Multi-chunk workspaces for DeepGemmFusedMoE; ``(None, None)`` otherwise. + """Multi-chunk workspaces for backends that ask for one; else ``(None, None)``. Single-chunk execution uses ``_prepare_workspace_deepgemm`` instead. """ @@ -284,7 +272,7 @@ def _prepare_workspaces_for_chunk( workspace_0 = None workspace_1 = None - if not isinstance(moe.backend, DeepGemmFusedMoE): + if not moe.backend.input_requirement.requires_run_moe_workspace: return workspace_0, workspace_1 # Always need at least workspace_0; reuse chunk_0 size for workspace_1 @@ -402,13 +390,27 @@ def _forward_chunk_impl( assert token_selected_experts.shape[1] == moe.routing_method.experts_per_token assert token_selected_experts.shape == token_final_scales.shape - # CutlassFusedMoE and DenseGEMMFusedMoE expect float32; TRTLLMGen expects bfloat16 - if isinstance(moe.backend, (CutlassFusedMoE, DenseGEMMFusedMoE)): - assert token_final_scales.dtype == torch.float32 assert token_selected_experts.dtype == torch.int32 - if token_final_scales is not None and isinstance(moe.backend, TRTLLMGenFusedMoE): - token_final_scales = token_final_scales.to(torch.bfloat16) + # Backends disagree on routing-scale precision, so the requirement + # names the dtype instead of the backend class. + scales_dtype = moe.backend.input_requirement.routing_scales_dtype + if scales_dtype is not None and token_final_scales is not None: + if scales_dtype == torch.float32: + # Asking for float32 is asking for routing's own + # full-precision output, so the cast below must be a no-op. + # Several routing methods take an ``output_dtype``, and one + # configured to a narrower type has already dropped mantissa + # bits that widening here cannot recover -- which is what + # this check catches. A narrower request (TRTLLM-Gen's + # bfloat16) is a deliberate conversion, not a loss. + assert token_final_scales.dtype == torch.float32, ( + f"{type(moe.backend).__name__} requires float32 routing " + f"scales, but {type(moe.routing_method).__name__} produced " + f"{token_final_scales.dtype}. Casting would widen a value " + "that already lost precision." + ) + token_final_scales = token_final_scales.to(scales_dtype) # apply_router_weight_on_input: fuse top-k weight onto x if moe.apply_router_weight_on_input: @@ -495,7 +497,9 @@ def _forward_chunk_impl( moe.dummy_allreduce() dispatch_kwargs = dict(eplb_dispatch_kwargs) - if isinstance(moe.comm, DeepEP) and isinstance(moe.backend, TRTLLMGenFusedMoE): + # Only DeepEP.dispatch reads this; every other strategy absorbs it + # through **kwargs, so the request does not need a comm-side test. + if moe.backend.input_requirement.requires_sanitized_expert_ids: dispatch_kwargs["enable_sanitize_expert_ids"] = True if supports_post_quant: @@ -537,19 +541,21 @@ def _forward_chunk_impl( # ========== Step 6: MoE computation ========== # If EPLB is enabled, token_selected_slots is slot ids; otherwise expert ids. - final_hidden_states = moe.backend.run_moe( + ctx = self._build_run_context( x=x, - token_selected_experts=token_selected_slots, - token_final_scales=token_final_scales, x_sf=x_sf, - **self._get_backend_kwargs( - router_logits, - do_finalize, - all_rank_num_tokens, - output_dtype, - x, - workspace, - lora_params=lora_params, + token_selected_slots=token_selected_slots, + token_final_scales=token_final_scales, + router_logits=router_logits, + do_finalize=do_finalize, + output_dtype=output_dtype, + all_rank_num_tokens=all_rank_num_tokens, + lora_params=lora_params, + ) + final_hidden_states = moe.backend.run_moe( + ctx, + workspace=( + workspace if moe.backend.input_requirement.requires_run_moe_workspace else None ), ) @@ -731,35 +737,30 @@ def _forward_multiple_chunks( return outputs # ------------------------------------------------------------------ - # Backend run_moe kwargs builder (external-comm only) + # Backend run_moe inputs (external-comm only) # ------------------------------------------------------------------ - def _get_nvlink_onesided_moe_output( + def _plan_onesided_workspace( self, all_rank_num_tokens: Optional[List[int]], output_dtype: Optional[torch.dtype], - ) -> Optional[torch.Tensor]: - """Workspace-backed output buffer for NVLinkOneSided combine, or None. + ) -> Tuple[Optional[torch.Tensor], bool]: + """Decide the NVLinkOneSided combine payload buffer for this forward. - Only meaningful when ``moe.comm`` is NVLinkOneSided AND the backend - supports payload-in-workspace combine. Returns None for all other - comm strategies; callers should always set the resulting kwarg - unconditionally and let backends ignore None. + Returns ``(moe_output, payload_in_workspace)``. Both are decided on + every path, including the ones that opt out, so the flag can never be + inherited from a previous forward. """ moe = self.moe if not isinstance(moe.comm, NVLinkOneSided): - return None + return None, False if not moe.backend.supports_moe_output_in_alltoall_workspace(): - # Backend opts out: keep payload off the workspace path. - moe.comm.payload_in_workspace = False - return None + # Backend emits its own output tensor; a workspace buffer would be + # left unfilled while combine() read from it. + return None, False - workspace_dtype = output_dtype - if isinstance(moe.backend, TRTLLMGenFusedMoE): - # TRTLLMGen sentinel for unfilled rows; bf16 workspace is the - # combine reduction precision used by the kernel. - moe.comm.invalid_token_expert_id = -1 - workspace_dtype = torch.bfloat16 + # None means "no override": the buffer matches the model output dtype. + workspace_dtype = moe.backend.input_requirement.onesided_workspace_dtype or output_dtype assert all_rank_num_tokens is not None, ( "all_rank_num_tokens must be provided for NVLinkOneSided backend" @@ -769,96 +770,68 @@ def _get_nvlink_onesided_moe_output( moe_output = moe.comm.get_combine_payload_tensor_in_workspace( runtime_max_tokens_per_rank, moe.hidden_size, workspace_dtype ) + return moe_output, True - # Toggle on for this forward; combine() reads this flag to decide - # whether to emit into the workspace tensor. - moe.comm.payload_in_workspace = True - return moe_output - - def _get_backend_kwargs( + def _build_run_context( self, - router_logits: Optional[torch.Tensor] = None, - do_finalize: bool = True, - all_rank_num_tokens: Optional[List[int]] = None, - output_dtype: Optional[torch.dtype] = None, - x: Optional[torch.Tensor] = None, - workspace: Optional[dict] = None, - lora_params: Optional[Dict] = None, - ) -> Dict: - """Backend-specific kwargs for ``backend.run_moe`` (external-comm only). - - ``FusedCommMoEScheduler`` constructs its own kwargs and never - calls this helper, so all branches here are EXTERNAL_COMM backends. - - Backend-specific kwargs: - - Cutlass: is_sf_swizzled, enable_alltoall, tuner_*, moe_output, lora_params - - CuteDSL: enable_alltoall, moe_output - - DeepGemm: workspace - - TRTLLMGen: router_logits, do_finalize, moe_output + *, + x: torch.Tensor, + x_sf: Optional[torch.Tensor], + token_selected_slots: Optional[torch.Tensor], + token_final_scales: Optional[torch.Tensor], + router_logits: Optional[torch.Tensor], + do_finalize: bool, + output_dtype: Optional[torch.dtype], + all_rank_num_tokens: Optional[List[int]], + lora_params: Optional[Dict], + ) -> MoERunContext: + """The single ``run_moe`` argument set, identical for every backend. - Only CutlassFusedMoE.run_moe accepts lora_params (routed-expert MoE LoRA - is fused there), so it is set on the Cutlass branch alone. + The only per-backend decision left here is dropping ``lora_params`` + for backends that do not fuse routed-expert LoRA: handing them one + would silently produce un-adapted output. """ moe = self.moe - kwargs: Dict = {} - - if moe.backend.__class__ == CutlassFusedMoE: - # Pre-quant dispatch: SFs arrive swizzled; post-quant dispatch: - # SFs arrive unswizzled. Backend uses this to skip a re-swizzle. - supports_post_quant = moe.comm is not None and moe.comm.supports_post_quant_dispatch() - kwargs["is_sf_swizzled"] = not supports_post_quant - kwargs["output_dtype"] = output_dtype - kwargs["lora_params"] = lora_params - - # Tuner sees pre-alltoall token shapes so cached tactics from the - # warmup (no-alltoall) phase still apply at runtime. - kwargs["enable_alltoall"] = moe.enable_alltoall - if moe.enable_alltoall: - if all_rank_num_tokens is not None: - kwargs["tuner_num_tokens"] = sum(all_rank_num_tokens) - else: - kwargs["tuner_num_tokens"] = ( - x.shape[0] * moe.mapping.tp_size if x is not None else None - ) - kwargs["tuner_top_k"] = moe.routing_method.top_k - - kwargs["moe_output"] = self._get_nvlink_onesided_moe_output( - all_rank_num_tokens=all_rank_num_tokens, output_dtype=output_dtype - ) - - elif moe.backend.__class__ == CuteDslFusedMoE: - kwargs["enable_alltoall"] = moe.enable_alltoall - kwargs["moe_output"] = self._get_nvlink_onesided_moe_output( - all_rank_num_tokens=all_rank_num_tokens, output_dtype=output_dtype - ) - - elif moe.backend.__class__ == DeepGemmFusedMoE: - if workspace is not None: - kwargs["workspace"] = workspace - - elif moe.backend.__class__ == TRTLLMGenFusedMoE: - # When the scheduler precomputes top-k for DP/load-balancer paths, - # the backend must not route again. Single-rank TRTLLMGen paths do - # not get precomputed top-k, so they still need router_logits. - router_logits_arg = ( - None - if ( - moe.backend._supports_load_balancer() - or moe.routing_method.requires_separated_routing - or FORCE_SEPARATED_ROUTING - ) - else router_logits - ) - kwargs["router_logits"] = router_logits_arg - kwargs["do_finalize"] = do_finalize - kwargs["moe_output"] = self._get_nvlink_onesided_moe_output( - all_rank_num_tokens=all_rank_num_tokens, output_dtype=output_dtype - ) + return MoERunContext( + token_selected_experts=token_selected_slots, + token_final_scales=token_final_scales, + x=x, + x_sf=x_sf, + output_dtype=output_dtype, + do_finalize=do_finalize, + lora_params=lora_params if moe.backend.capabilities.supports_moe_lora else None, + router_logits=router_logits, + all_rank_num_tokens=all_rank_num_tokens, + comm_plan=self._build_comm_plan(all_rank_num_tokens, output_dtype), + ) - elif moe.backend.__class__ == MarlinFusedMoE: - kwargs["router_logits"] = router_logits + def _build_comm_plan( + self, + all_rank_num_tokens: Optional[List[int]], + output_dtype: Optional[torch.dtype], + ) -> MoECommPlan: + """The comm-layer facts for this forward, derived once for every backend. - return kwargs + Backends read the fields they care about and ignore the rest, so the + set of facts no longer depends on which class is running. + """ + moe = self.moe + # Pre-quant dispatch: SFs arrive swizzled; post-quant dispatch: SFs + # arrive unswizzled. Backends use this to skip a re-swizzle. + supports_post_quant = moe.comm is not None and moe.comm.supports_post_quant_dispatch() + moe_output, payload_in_workspace = self._plan_onesided_workspace( + all_rank_num_tokens=all_rank_num_tokens, output_dtype=output_dtype + ) + if isinstance(moe.comm, NVLinkOneSided): + # combine() still reads the flag off the strategy; the plan stays + # the single place that decides its value. + moe.comm.payload_in_workspace = payload_in_workspace + return MoECommPlan( + input_sf_swizzled=not supports_post_quant, + enable_alltoall=moe.enable_alltoall, + moe_output=moe_output, + payload_in_workspace=payload_in_workspace, + ) # ============================================================================ @@ -1223,12 +1196,16 @@ def _forward_chunk( # ``token_selected_slots`` is in [0, num_slots), matching the kernel's # ``num_experts`` template parameter (SymmBuffer / weights sized to # num_slots in quantization.py). + # Fused-comm backends own the EP exchange, so there is no comm plan: + # nothing outside the fused kernel decided anything about this forward. out = moe.backend.run_moe( - x=moe_input, - token_selected_experts=token_selected_slots, - token_final_scales=token_final_scales, - x_sf=x_sf, - output_dtype=output_dtype, + MoERunContext( + token_selected_experts=token_selected_slots, + token_final_scales=token_final_scales, + x=moe_input, + x_sf=x_sf, + output_dtype=output_dtype, + ) ) # ----- EPLB: start/done CPU rebalance, AFTER run_moe ----- diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index 37e1bf239e72..2b3d61e4a802 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -4,6 +4,7 @@ import itertools import os import weakref +from dataclasses import replace from enum import IntEnum from typing import Optional @@ -297,19 +298,11 @@ def make_balanced_run_moe( dp_rank, ep_size, ): - def balanced_run_moe( - x, token_selected_experts, token_final_scales, x_sf, router_logits, do_finalize, moe_output - ): + def balanced_run_moe(ctx, *, workspace=None): if moe_module._routing_results_replaced_at is not None: - return run_moe_orig( - x, - token_selected_experts, - token_final_scales, - x_sf, - router_logits, - do_finalize, - moe_output, - ) + return run_moe_orig(ctx, workspace=workspace) + x = ctx.x + do_finalize = ctx.do_finalize logger.warning_once( 'Layer-wise benchmarks: Specifying routing results of "TRTLLM" MoE backend in TEP cases leads to different' " execution path around the topk kernel", @@ -355,15 +348,14 @@ def balanced_run_moe( token_final_scales = get_token_final_scales( token_selected_experts.shape, token_selected_experts.device ) - router_logits = None final_hidden_states = run_moe_orig( - x, - token_selected_experts, - token_final_scales, - x_sf, - router_logits, - do_finalize, - moe_output, + replace( + ctx, + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + router_logits=None, + ), + workspace=workspace, ) if not do_finalize: final_hidden_states = ( diff --git a/tests/microbenchmarks/bench_moe/routing/native_logits.py b/tests/microbenchmarks/bench_moe/routing/native_logits.py index b28c0a65bd5a..14e10554291d 100644 --- a/tests/microbenchmarks/bench_moe/routing/native_logits.py +++ b/tests/microbenchmarks/bench_moe/routing/native_logits.py @@ -26,6 +26,7 @@ from __future__ import annotations import contextlib +from dataclasses import replace from typing import Dict, Optional, Tuple import torch @@ -200,23 +201,23 @@ def _make_supplied_topk_run_moe( balanced/imbalanced selection helpers. """ - def supplied_run_moe( - x, token_selected_experts, token_final_scales, x_sf, router_logits, do_finalize, moe_output - ): + def supplied_run_moe(ctx, *, workspace=None): if getattr(moe_module, "_routing_results_replaced_at", None) is not None: - return run_moe_orig( - x, - token_selected_experts, - token_final_scales, - x_sf, - router_logits, - do_finalize, - moe_output, - ) + return run_moe_orig(ctx, workspace=workspace) + x = ctx.x + do_finalize = ctx.do_finalize local, scales = _align_topk_to_batch(materialized_ids, materialized_scales, x.shape[0]) local = local.to(device=x.device, dtype=torch.int32) scales = scales.to(device=x.device) - final_hidden_states = run_moe_orig(x, local, scales, x_sf, None, do_finalize, moe_output) + final_hidden_states = run_moe_orig( + replace( + ctx, + token_selected_experts=local, + token_final_scales=scales, + router_logits=None, + ), + workspace=workspace, + ) if not do_finalize: final_hidden_states = ( final_hidden_states[0], diff --git a/tests/unittest/_torch/lora/test_moe_lora_model_path.py b/tests/unittest/_torch/lora/test_moe_lora_model_path.py index 1900cc48118a..cc8bb0a1434a 100644 --- a/tests/unittest/_torch/lora/test_moe_lora_model_path.py +++ b/tests/unittest/_torch/lora/test_moe_lora_model_path.py @@ -12,8 +12,8 @@ 1. QwenMoE.forward to the routed self.experts call (legacy wrapper). 2. ConfigurableMoE.forward_impl to scheduler.forward. - 3. ExternalCommMoEScheduler._get_backend_kwargs to the CutlassFusedMoE - run_moe kwargs, and not to backends that cannot carry LoRA. + 3. ExternalCommMoEScheduler._build_run_context to the CutlassFusedMoE + run_moe context, and not to backends that cannot carry LoRA. """ from types import SimpleNamespace @@ -106,8 +106,8 @@ def test_configurable_moe_forward_impl_forwards_lora_params_to_scheduler(): def _make_external_comm_scheduler(backend_cls): """Build an ExternalCommMoEScheduler whose moe.backend is an uninitialized - instance of backend_cls, sufficient for _get_backend_kwargs class dispatch - without constructing weights.""" + instance of backend_cls, sufficient for building a run context without + constructing weights.""" backend = backend_cls.__new__(backend_cls) moe = SimpleNamespace( backend=backend, @@ -121,35 +121,39 @@ def _make_external_comm_scheduler(backend_cls): return scheduler -def test_scheduler_threads_lora_params_to_cutlass_run_moe_kwargs(): - """_get_backend_kwargs must thread lora_params into the - CutlassFusedMoE.run_moe kwargs.""" - scheduler = _make_external_comm_scheduler(CutlassFusedMoE) - - kwargs = scheduler._get_backend_kwargs( +def _build_run_context(scheduler): + return scheduler._build_run_context( + x=torch.randn(4, 8), + x_sf=None, + token_selected_slots=torch.zeros(4, 2, dtype=torch.int32), + token_final_scales=torch.ones(4, 2), + router_logits=None, + do_finalize=True, output_dtype=torch.bfloat16, + all_rank_num_tokens=None, lora_params=_LORA_PARAMS_SENTINEL, ) - assert kwargs.get("lora_params") is _LORA_PARAMS_SENTINEL, ( + +def test_scheduler_threads_lora_params_to_cutlass_run_context(): + """The run context handed to CutlassFusedMoE.run_moe must carry + lora_params.""" + ctx = _build_run_context(_make_external_comm_scheduler(CutlassFusedMoE)) + + assert ctx.lora_params is _LORA_PARAMS_SENTINEL, ( "Scheduler dropped lora_params before CutlassFusedMoE.run_moe; " "routed-expert MoE LoRA would be silently disabled." ) -def test_scheduler_does_not_thread_lora_params_to_non_cutlass_backend(): - """Only CutlassFusedMoE.run_moe accepts lora_params. Other backends must - not receive it, since it is not in their run_moe signature.""" - scheduler = _make_external_comm_scheduler(DeepGemmFusedMoE) - - kwargs = scheduler._get_backend_kwargs( - output_dtype=torch.bfloat16, - lora_params=_LORA_PARAMS_SENTINEL, - ) +def test_scheduler_does_not_thread_lora_params_to_non_lora_backend(): + """Backends that do not declare supports_moe_lora must get lora_params + cleared, so an adapter never silently no-ops inside their kernel.""" + ctx = _build_run_context(_make_external_comm_scheduler(DeepGemmFusedMoE)) - assert "lora_params" not in kwargs, ( - "lora_params must only be forwarded to CutlassFusedMoE.run_moe; " - f"DeepGemmFusedMoE.run_moe does not accept it. Got kwargs: {list(kwargs)}" + assert ctx.lora_params is None, ( + "lora_params must only reach backends declaring supports_moe_lora; " + f"DeepGemmFusedMoE does not fuse it. Got: {ctx.lora_params}" ) diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 00ec9070a70b..faa65f993e4e 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -61,7 +61,12 @@ from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend, get_moe_cls 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.interface import MoE, MoEWeightLoadingMode +from tensorrt_llm._torch.modules.fused_moe.impl_contract import MoECommPlan, MoERunContext +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.quantization import ( FusedMoEMethodBase, @@ -239,7 +244,7 @@ def load_weights(self, weights, allow_partial_loading=False): def quantize_input(self, x, **kwargs): return x, None - def run_moe(self, **kwargs): + def run_moe(self, ctx, *, workspace=None): raise NotImplementedError moe = HookTestMoE.__new__(HookTestMoE) @@ -315,7 +320,7 @@ class HookTestConfigurableMoE(ConfigurableMoE): def quantize_input(self, x, **kwargs): return x, None - def run_moe(self, **kwargs): + def run_moe(self, ctx, *, workspace=None): raise NotImplementedError configurable_moe = HookTestConfigurableMoE.__new__(HookTestConfigurableMoE) @@ -632,6 +637,7 @@ def run_backend_moe( token_final_scales=token_final_scales.to(torch.float32), x_sf=x_sf, ) + workspace = None # Backend-specific overrides if backend_type == MoeBackendType.CUTLASS: @@ -648,12 +654,25 @@ def run_backend_moe( import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils m_max = fp8_utils.align(x_quantized.shape[0], 128) - args["workspace"] = backend.get_workspace(m_max, 128) + workspace = backend.get_workspace(m_max, 128) elif backend_type in _MEGAMOE_BACKEND_TYPES: args["token_selected_experts"] = token_selected_experts.to(torch.int64) args["output_dtype"] = dtype - return backend.run_moe(**args) + # Mirror what each scheduler hands the backend: ExternalCommMoEScheduler + # builds a plan on every path, FusedCommMoEScheduler never does because the + # fused kernel owns the exchange. Single GPU with no comm strategy, so this + # is the no-comm plan: quantize_input ran locally and left the scale factors + # swizzled, no alltoall, and no workspace-backed output buffer. + if backend.scheduler_kind == MoESchedulerKind.EXTERNAL_COMM: + args["comm_plan"] = MoECommPlan( + input_sf_swizzled=True, + enable_alltoall=False, + moe_output=None, + payload_in_workspace=False, + ) + + return backend.run_moe(MoERunContext(**args), workspace=workspace) # ============================================================================ @@ -1227,6 +1246,12 @@ def test_trtllm_bf16_unquantized_moe( mapping=mapping, activation_type=activation_type, ) + if trtllm_use_router_logits and backend._routes_outside_the_kernel(): + # This config only routes outside the kernel, so run_moe drops + # router_logits and the scheduler never pairs the two. Asking for + # fused routing here tests a combination production cannot reach. + pytest.skip("routing happens outside the kernel; fused routing is unreachable") + backend.load_weights([weights]) backend.post_load_weights() backend.cuda()