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 1c1a553fe5c9..1d54d27fefac 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, @@ -85,6 +87,11 @@ class on the MRO does not change which kernels execute, only where the ``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}) @@ -216,29 +223,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 @@ -247,17 +246,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 70b87aa7d488..70daa39d561c 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: @@ -850,19 +859,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. @@ -870,22 +890,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( @@ -904,8 +924,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. @@ -918,7 +937,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 @@ -959,9 +978,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')) @@ -1034,16 +1050,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 5dc95306f55c..f9a7d1ad90dc 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))), @@ -147,20 +152,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 dc082eef24c6..85e2689e9905 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 f35d7020ad8d..8826bdaf3435 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, @@ -241,7 +246,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) @@ -317,7 +322,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) @@ -639,6 +644,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: @@ -655,12 +661,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) # ============================================================================ @@ -1237,6 +1256,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()