diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 54a97a70c3d5..55df32268e9e 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -74,16 +74,13 @@ from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm from ..moe.fused_moe import ( - CutlassFusedMoE, + DEFAULT_MOE_ACTIVATION, DeepSeekV4MoeRoutingMethod, MoEWeightLoadingMode, - TritonFusedMoE, - TRTLLMGenFusedMoE, + SwigluActivation, create_moe, is_moe_weight_owner, - resolve_moe_cls, ) -from ..moe.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE from ..peft.lora.layer import LoraLayer from ..speculative import SpecMetadata, get_num_extra_kv_tokens from ..utils import ( @@ -1545,61 +1542,15 @@ def __init__( if override_quant_config is not None and experts_quant_config is model_config.quant_config: experts_quant_config = override_quant_config + # One uniform clamp for the whole layer; the selected backend's + # ``activation_support`` decides whether its kernels take that as a + # per-expert ``float32`` buffer or a baked scalar. swiglu_limit = getattr(config, "swiglu_limit", None) - moe_swiglu_limit = None - if swiglu_limit is not None: - # `create_moe` only accepts swiglu_limit for these MoE classes; - # ask the resolver rather than the backend string so that a - # degradation (e.g. TRTLLM/CUTEDSL/DENSEGEMM dropping back to - # CutlassFusedMoE on unsupported quant) is accounted for here too. - moe_cls = resolve_moe_cls( - model_config, - override_quant_config=experts_quant_config, - dtype=dtype, - # Same routing object as create_moe below. - routing=self.gate.routing_method, - # create_moe below passes no bias and no swiglu alpha/beta, so - # it resolves with the plain SwiGLU package. Say so here too: - # leaving this unknown lets gates abstain that create_moe - # rejects, and the two calls would pick different backends. - swiglu_gptoss_style=False, - layer_idx=layer_idx, - ) - supports_swiglu_limit = moe_cls in ( - CutlassFusedMoE, - TritonFusedMoE, - TRTLLMGenFusedMoE, - DeepGemmFusedMoE, - ) - # NVFP4 routed-expert path: the TRTLLM-Gen fp4-block-scale fused-MoE - # cubin produces near-zero accuracy without bias even when - # swiglu_limit is supplied; drop the limit there until the cubin - # gains a no-bias clamp variant. MXFP4 variants are unaffected. - kernel_requires_bias_for_swiglu_limit = ( - moe_cls is TRTLLMGenFusedMoE and experts_quant_config.quant_mode.has_nvfp4() - ) - # DeepSeek-V4 supplies a uniform scalar limit. The TRTLLM-Gen FP8 - # path consumes it directly and rejects the redundant tensor. - requires_scalar_only_swiglu_limit = ( - moe_cls is TRTLLMGenFusedMoE - and experts_quant_config.quant_mode.has_fp8_block_scales() - ) - if ( - supports_swiglu_limit - and not kernel_requires_bias_for_swiglu_limit - and not requires_scalar_only_swiglu_limit - ): - moe_load_balancer_config = getattr(model_config, "moe_load_balancer", None) - num_slots = ( - moe_load_balancer_config.num_slots - if moe_load_balancer_config and moe_load_balancer_config.num_slots - else num_experts - ) - local_num_slots = num_slots // model_config.mapping.moe_ep_size - device = "cuda" if torch.cuda.is_available() else "cpu" - moe_swiglu_limit = torch.full( - (local_num_slots,), float(swiglu_limit), dtype=torch.float32, device=device - ) + moe_activation = ( + SwigluActivation(clamp=float(swiglu_limit)) + if swiglu_limit is not None + else DEFAULT_MOE_ACTIVATION + ) self.experts = create_moe( num_experts=num_experts, @@ -1619,8 +1570,7 @@ def __init__( if experts_quant_config.layer_quant_mode.is_int4_weight_only_per_group() else MoEWeightLoadingMode.VANILLA ), - swiglu_limit=moe_swiglu_limit, - swiglu_limit_scalar=(float(swiglu_limit) if swiglu_limit is not None else None), + activation=moe_activation, ) self.mapping = model_config.mapping diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index 900c3ef76f50..ba671795770a 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -27,6 +27,7 @@ from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper from tensorrt_llm._torch.modules.qk_norm_attention import QKNormRoPEAttention +from tensorrt_llm._torch.moe.fused_moe.activation import SimpleActivation from tensorrt_llm._torch.moe.fused_moe.create_moe import create_moe from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode from tensorrt_llm._torch.moe.fused_moe.routing import BaseMoeRoutingMethod @@ -669,7 +670,7 @@ def __init__( reduce_results=True, model_config=model_config, layer_idx=layer_idx, - activation_type=ActivationType.Geglu, + activation=SimpleActivation(kind=ActivationType.Geglu), # VANILLA mode: preprocess_weights splits 3D gate_up_proj into per-expert w1/w3 weight_loading_mode=MoEWeightLoadingMode.VANILLA, ) diff --git a/tensorrt_llm/_torch/models/modeling_gpt_oss.py b/tensorrt_llm/_torch/models/modeling_gpt_oss.py index 5808bcf386ba..4d48c3bbb641 100644 --- a/tensorrt_llm/_torch/models/modeling_gpt_oss.py +++ b/tensorrt_llm/_torch/models/modeling_gpt_oss.py @@ -25,7 +25,8 @@ # isort and yapf will fight against each other here, so we disable isort # isort: off from ..moe.fused_moe import (MoEWeightLoadingMode, RenormalizeMoeRoutingMethod, - TritonFusedMoE, create_moe, is_moe_weight_owner) + SwigluBiasActivation, TritonFusedMoE, create_moe, + is_moe_weight_owner) # isort: on from ..modules.linear import Linear, TensorParallelMode from ..modules.rms_norm import RMSNorm @@ -172,15 +173,12 @@ def __init__( output_dtype=torch.bfloat16 if config.moe_backend.upper() == "TRTLLM" else torch.float32) - self.swiglu_alpha = torch.tensor( - [1.702] * (self.num_slots // config.mapping.moe_ep_size), - dtype=torch.float32).cuda() - self.swiglu_beta = torch.tensor( - [1.0] * (self.num_slots // config.mapping.moe_ep_size), - dtype=torch.float32).cuda() - self.swiglu_limit = torch.tensor( - [7.0] * (self.num_slots // config.mapping.moe_ep_size), - dtype=torch.float32).cuda() + # gpt-oss constants are uniform across experts; the backend broadcasts + # them to whatever per-expert shape its kernels index, sized by the + # local slot count it resolved. + self.moe_activation = SwigluBiasActivation(gate_sigmoid_scale=1.702, + linear_offset=1.0, + clamp=7.0) # Prepare MoE creation parameters moe_params = { 'routing_method': self.routing_method, @@ -192,9 +190,7 @@ def __init__( 'model_config': config, 'weight_loading_mode': MoEWeightLoadingMode.FUSED_GATE_UP_PROJ, 'bias': True, - 'swiglu_alpha': self.swiglu_alpha, - 'swiglu_beta': self.swiglu_beta, - 'swiglu_limit': self.swiglu_limit, + 'activation': self.moe_activation, 'layer_idx': self.layer_idx, } diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index dce32f886909..91ddc5c5643b 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -117,10 +117,8 @@ from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm from ..modules.situ import SituAndMul -from ..moe.fused_moe import ConfigurableMoE, create_moe -from ..moe.fused_moe.interface import _compute_ep_partition +from ..moe.fused_moe import ConfigurableMoE, SiTuActivation, create_moe from ..moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod -from ..utils import ActivationType, ActType_TrtllmGen from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, register_auto_model, run_concurrently @@ -1120,6 +1118,16 @@ def __init__( layer_idx=layer_idx, # Let CommunicationFactory select the best available strategy. communication_method=None, + activation=SiTuActivation( + gate_softcap=situ_beta, + linear_softcap=situ_linear_beta, + ), + # A MegaMoE request that silently degraded to CUTLASS would be + # benchmarked as if it were MegaMoE, and the decline is easy to + # trigger (EP-only, own token / top-k limits). Fail in the resolver + # instead, which reports the rejection trail. + allow_backend_degradation=routed_moe_model_config.moe_backend + not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"), ) # trtllm-gen ships SiTu cubins for exactly one dtype combination # (``Bmm_MxE4m3_MxE2m1MxE4m3`` = MXFP8 act x MXFP4 weight) and has no @@ -1141,82 +1149,11 @@ def __init__( "CUTLASS or MEGAMOE_CUTEDSL." ) - if routed_moe_model_config.moe_backend == "CUTLASS": - # Size the per-expert SiTU constants with the same ceil/floor - # partition the MoE backend uses for ``expert_size_per_partition``. - # A plain ``num_experts // ep_size`` is one element short on the - # first ``num_experts % ep_size`` ranks, which trips the - # ``swiglu_alpha must have num_experts_on_rank elements`` check in - # moeOp.cpp. K3's 384 experts divide evenly at EP8/EP16, so the - # mismatch is latent there but real for any uneven split. - local_num_experts, _, _ = _compute_ep_partition( - self.num_experts, - routed_moe_model_config.mapping.moe_ep_size, - routed_moe_model_config.mapping.moe_ep_rank, - ) - device = torch.device("cuda", torch.cuda.current_device()) - self.routed_situ_alpha = torch.full( - (local_num_experts,), float(situ_beta), dtype=torch.float32, device=device - ) - self.routed_situ_beta = torch.full( - (local_num_experts,), - situ_linear_beta, - dtype=torch.float32, - device=device, - ) - routed_moe_kwargs.update( - activation_type=ActivationType.SiTu, - swiglu_alpha=self.routed_situ_alpha, - swiglu_beta=self.routed_situ_beta, - ) - elif routed_moe_model_config.moe_backend == "TRTLLM": - routed_moe_kwargs.update( - trtllm_gen_activation_type=ActType_TrtllmGen.SiTu, - # Cubin alpha is the gate-side SiTU beta; cubin beta is the - # linear-side SiTU beta. - trtllm_gen_activation_alpha=situ_beta, - trtllm_gen_activation_beta=situ_linear_beta, - ) - elif routed_moe_model_config.moe_backend == "MEGAMOE_DEEPGEMM": - routed_moe_kwargs.update( - activation="situ", - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, - ) - # MEGAMOE_CUTEDSL has no branch on purpose. ``MegaMoECuteDsl`` resolves - # SiTU from the pretrained config in ``_resolve_activation_config`` - # (``activation=None`` -> "situ" when ``activation_situ_beta`` is - # present), and ``create_moe`` currently rejects the explicit - # ``activation``/``situ_beta``/``situ_linear_beta`` trio for anything - # other than ``MegaMoEDeepGemm``, so passing them here would raise. - # Unifying that plumbing is tracked in TRTLLM-15649. self.routed_experts = create_moe(**routed_moe_kwargs) if not isinstance(self.routed_experts, ConfigurableMoE): raise RuntimeError( "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled." ) - if routed_moe_model_config.moe_backend == "MEGAMOE_DEEPGEMM": - from ..moe.fused_moe.mega_moe import MegaMoEDeepGemm - - if not isinstance(self.routed_experts.backend, MegaMoEDeepGemm): - raise RuntimeError( - "Kimi K3 explicitly requested MEGAMOE_DEEPGEMM, but the " - f"MoE factory selected {type(self.routed_experts.backend).__name__}." - ) - if routed_moe_model_config.moe_backend == "MEGAMOE_CUTEDSL": - from ..moe.fused_moe.mega_moe import MegaMoECuteDsl - - # Same guard as MEGAMOE_DEEPGEMM above, and for the same reason: - # create_moe silently falls back when a backend declines the - # config, and for MegaMoE the decline is easy to trigger (it is - # EP-only and has its own token/top-k limits), so an explicit - # request that quietly became CUTLASS would be measured as if it - # were MegaMoE. - if not isinstance(self.routed_experts.backend, MegaMoECuteDsl): - raise RuntimeError( - "Kimi K3 explicitly requested MEGAMOE_CUTEDSL, but the " - f"MoE factory selected {type(self.routed_experts.backend).__name__}." - ) if self.routed_experts.layer_load_balancer is not None: raise NotImplementedError( "Kimi K3 packed-checkpoint streaming does not yet support " diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 4e7162461694..0cab2ff3a95e 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -62,15 +62,9 @@ ) from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm -from ..moe.fused_moe import MiniMaxM3MoeRoutingMethod, create_moe +from ..moe.fused_moe import MiniMaxM3MoeRoutingMethod, SwigluBiasActivation, create_moe from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph -from ..utils import ( - ActivationType, - AuxStreamType, - EventType, - get_model_extra_attrs, - is_torch_compiling, -) +from ..utils import AuxStreamType, EventType, get_model_extra_attrs, is_torch_compiling from .checkpoints.base_weight_mapper import BaseWeightMapper from .checkpoints.hf.minimaxm3_weight_mapper import MINIMAX_M3_PARAMS_MAP, MiniMaxM3HfWeightMapper from .modeling_utils import ( @@ -300,33 +294,6 @@ def _build_swiglu_oai_dense_mlp( ) -def _resolve_minimax_m3_expert_size_per_partition( - num_experts: int, - mapping: Mapping, - moe_load_balancer_config, -) -> int: - """Compute the local expert/slot count the MoE module will resolve. - - Sizes the per-expert SwiGLU parameter tensors we hand to - ``create_moe`` to match what the backend sees. Some backends assert - ``swiglu_alpha.shape == (expert_size_per_partition,)`` at construct - time; a mismatch surfaces as an opaque CUDA-side shape failure. - - Priority: - 1. EPLB config → ``num_slots // moe_ep_size``. - 2. Plain EP → ``num_experts // moe_ep_size`` (with ``max(1, ...)`` - for tiny test configs). - """ - ep_size = mapping.moe_ep_size - if moe_load_balancer_config is not None and moe_load_balancer_config.num_slots: - # Mirror ``MoeLoadBalancerConfig.num_local_slots`` without - # requiring ``.setup(ep_rank, ep_size)`` to have been called - # (the ``num_local_slots`` property raises otherwise). - return moe_load_balancer_config.num_slots // ep_size - - return max(1, num_experts // ep_size) - - class MiniMaxM3Gate(nn.Module): """MiniMax-M3 router gate: float32 sigmoid scoring + per-expert bias correction. @@ -457,31 +424,14 @@ def __init__( self.swiglu_beta_value = 1.0 # SGLang's ``(up + 1)`` offset in swiglu_no_interleaved. self.swiglu_limit_value = float(getattr(config, "swiglu_limit", 7.0)) - # Size the per-expert SwiGLU parameter tensors from the same - # ``expert_size_per_partition`` the MoE module will resolve, not - # from a hand-rolled ``num_slots // ep_size`` guess. EPLB and - # DWDP can shift the local slot count (see - # :func:`_resolve_minimax_m3_expert_size_per_partition` for the - # priority order), and some backends assert - # ``swiglu_alpha.shape == (expert_size_per_partition,)``. - moe_load_balancer_config = model_config.moe_load_balancer - self.expert_size_per_partition = _resolve_minimax_m3_expert_size_per_partition( - num_experts=self.num_experts, - mapping=model_config.mapping, - moe_load_balancer_config=moe_load_balancer_config, + # One value per layer, not per expert: the MoE backend broadcasts these + # to whatever per-expert shape its kernels index (or bakes them as + # scalars), sized by the slot count it actually resolved. + self.moe_activation = SwigluBiasActivation( + gate_sigmoid_scale=self.swiglu_alpha_value, + linear_offset=self.swiglu_beta_value, + clamp=self.swiglu_limit_value, ) - self.swiglu_alpha = torch.tensor( - [self.swiglu_alpha_value] * self.expert_size_per_partition, - dtype=torch.float32, - ).cuda() - self.swiglu_beta = torch.tensor( - [self.swiglu_beta_value] * self.expert_size_per_partition, - dtype=torch.float32, - ).cuda() - self.swiglu_limit = torch.tensor( - [self.swiglu_limit_value] * self.expert_size_per_partition, - dtype=torch.float32, - ).cuda() # Router gate owns the float32 projection weight, the per-expert # ``e_score_correction_bias``, and the ``routing_method`` it @@ -508,24 +458,8 @@ def __init__( model_config=model_config, layer_idx=layer_idx, override_quant_config=experts_quant_config, - swiglu_alpha=self.swiglu_alpha, - swiglu_beta=self.swiglu_beta, - swiglu_limit=self.swiglu_limit, - activation_type=ActivationType.SwigluBias, + activation=self.moe_activation, ) - # Defensive: if a future MoE-resolution path (new load-balancer - # mode, new DWDP variant) shifts the local expert count in a - # way our resolver doesn't yet model, fail here with a - # diagnostic message instead of inside a CUDA-side shape - # assertion deep in a backend kernel dispatch. - resolved = getattr(self.experts, "expert_size_per_partition", None) - assert resolved is None or resolved == self.expert_size_per_partition, ( - f"MiniMax-M3 SwiGLU sizing mismatch: pre-create_moe estimate " - f"{self.expert_size_per_partition} != MoE-resolved {resolved}. " - f"Update _resolve_minimax_m3_expert_size_per_partition to " - f"match the MoE module's resolved layout." - ) - # Shared expert: dense MLP fused into MoE output. Constructed # with ``is_shared_expert=True`` so ``reduce_output=False`` # (the external AllReduce below performs the combined diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 13cf3a0905c8..c222240e5f21 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -47,7 +47,7 @@ from ..modules.mlp import MLP from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm -from ..moe.fused_moe import MoEWeightLoadingMode, create_moe +from ..moe.fused_moe import MoEWeightLoadingMode, SimpleActivation, create_moe from ..moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE from ..moe.fused_moe.quantization import (NVFP4CutlassFusedMoEMethod, W4A16NVFP4CutlassFusedMoEMethod) @@ -177,7 +177,7 @@ def __init__( # Import here to avoid circular dependency. from .modeling_deepseekv3 import DeepseekV3Gate - self.activation_type = ActivationType.Relu2 + self.moe_activation = SimpleActivation(kind=ActivationType.Relu2) self.reduce_results = False config = model_config.pretrained_config @@ -273,7 +273,7 @@ def _moe(name): layer_idx=self.layer_idx, weight_loading_mode=MoEWeightLoadingMode.VANILLA, bias=self.mlp_bias, - activation_type=self.activation_type, + activation=self.moe_activation, ) if reduce_output: diff --git a/tensorrt_llm/_torch/modules/dwdp/setup.py b/tensorrt_llm/_torch/modules/dwdp/setup.py index db857bbee818..4b239a802c3f 100644 --- a/tensorrt_llm/_torch/modules/dwdp/setup.py +++ b/tensorrt_llm/_torch/modules/dwdp/setup.py @@ -632,6 +632,17 @@ def fixup_moe_backends( target.initial_local_expert_ids = list(range(num_experts_total)) target.initial_global_assignments = list(range(num_experts_total)) + # ``expert_size_per_partition`` just changed, and it sizes every + # PER_EXPERT_TENSOR activation constant. This function is the layout's + # last writer, so re-materialize them here. Backend only: the wrapper + # declares no ``activation_support`` of its own. Imported locally because + # ``fused_moe`` reads the DWDP manager, so a module-scope import here + # would close that cycle. + from ..fused_moe.activation import install_activation_params + + if getattr(experts_module, "activation", None) is not None: + install_activation_params(experts_module) + logger.debug( f"[DWDP Setup] Layer {layer_idx}: patched " f"{'ConfigurableMoE + ' if len(targets) > 1 else ''}backend " diff --git a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py index 5b6805d9a25b..be75f148f165 100644 --- a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py +++ b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py @@ -1014,8 +1014,8 @@ def query_megamoe_shared_workspace_bytes( tactic: Optional[Tuple] = None, apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, - situ_beta: Optional[float] = None, - situ_linear_beta: Optional[float] = None, + act_alpha: Optional[float] = None, + act_beta: Optional[float] = None, in_kernel_fc2_reduce: bool = False, combine_format: str = "bf16", ) -> int: @@ -1073,8 +1073,8 @@ def query_megamoe_shared_workspace_bytes( epi_flag_batch=tuple(epi_flag_batch), apply_topk_in_fc1=bool(apply_topk_in_fc1), gate_up_clamp=(None if gate_up_clamp is None else float(gate_up_clamp)), - situ_beta=(None if situ_beta is None else float(situ_beta)), - situ_linear_beta=(None if situ_linear_beta is None else float(situ_linear_beta)), + situ_beta=(None if act_alpha is None else float(act_alpha)), + situ_linear_beta=(None if act_beta is None else float(act_beta)), **_LOCKED_KERNEL_KWARGS, ) # The probe MUST build the SAME kernel that runs (same combine_format): @@ -1702,8 +1702,8 @@ def cute_dsl_megamoe_nvfp4_blackwell( peer_offsets: List[int], apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, - situ_beta: Optional[float] = None, - situ_linear_beta: Optional[float] = None, + act_alpha: Optional[float] = None, + act_beta: Optional[float] = None, in_kernel_fc2_reduce: bool = False, combine_format: str = "bf16", tactic_autotune: bool = False, @@ -1730,6 +1730,12 @@ def cute_dsl_megamoe_nvfp4_blackwell( accumulation order). Both write ``combine_output`` with shape ``(T, 1, hidden)``. Perf knobs come from the tactic, not op arguments. + + ``act_alpha`` / ``act_beta`` also select the activation, because the + epilogue has no separate kind argument: both None runs SwiGLU, both set + runs SiTU with them as the gate-side and linear-side soft-caps + (``epilogue_refactor.py``'s ``situ_beta is None`` branch). They are + codegen-time constants, so a change recompiles the kernel. """ sm_version = get_sm_version() if sm_version not in (100, 103): @@ -1758,8 +1764,11 @@ def cute_dsl_megamoe_nvfp4_blackwell( output_dtype=combine_output.dtype, apply_topk_in_fc1=apply_topk_in_fc1, gate_up_clamp=gate_up_clamp, - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, + # The runner and the kernel package below it name these registers + # after SiTU; this op's boundary does not, so the spelling changes + # here and nowhere else. + situ_beta=act_alpha, + situ_linear_beta=act_beta, in_kernel_fc2_reduce=in_kernel_fc2_reduce, combine_format=combine_format, tactic_autotune=tactic_autotune, @@ -1872,8 +1881,8 @@ def _( peer_offsets: List[int], apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, - situ_beta: Optional[float] = None, - situ_linear_beta: Optional[float] = None, + act_alpha: Optional[float] = None, + act_beta: Optional[float] = None, in_kernel_fc2_reduce: bool = False, combine_format: str = "bf16", tactic_autotune: bool = False, diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index f7f19f81e62b..2b410f1b90c8 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -140,6 +140,7 @@ Still on old path (standalone, with embedded communication): | `impl_identity.py` | `MoEImplId` / `MoEImplDescriptor` / registry — the stable one-id-per-leaf-class mechanism used after an implementation migrates | | `interface.py` | Complete-layer base `MoE` and enums (`MoEWeightLoadingMode`, `MoESchedulerKind`) | | `impl_base.py` | Execution-unit base `MoEImplBase` — weights + `run_moe`, no `forward`; plus `apply_moe_impl_construction_state()`, which every execution unit must call | +| `activation.py` | Activation vocabulary — the `MoEActivation` carrier a model builds, the `MoEActivationSupport` a backend declares, and `install_activation_params` / `materialize_activation_params`, the only place a semantic constant becomes a kernel register | | `impl_blocks.py` | The blocks `MoE` and `MoEImplBase` share — `MoEExecutionContractMixin` (scheduler-facing declarations, `forward_fake`) and `MoEWeightOwnerMixin` (`create_weights` / `load_weights` / `_check_configs`) | | `quantization.py` | Quantization method implementations (`FusedMoEMethod` subclasses: weight creation, loading, quant/dequant ops per quant mode) | | `routing.py` | Routing methods (`TopKRouting`, etc.) | @@ -157,7 +158,7 @@ Still on old path (standalone, with embedded communication): | `fused_moe_cute_dsl.py` | `CuteDslFusedMoE` | SM100/SM103 | High throughput NVFP4, generally faster than Cutlass | `EXTERNAL_COMM` | | `fused_moe_cute_dsl_b12x.py` | `CuteDslB12xFusedMoE` | SM120/SM121 | NVFP4 hybrid CUTLASS-prefill / FlashInfer NVFP4 MoE decode — best perf on RTX PRO 6000 (SM120) and DGX Spark (SM121); select via the `CUTEDSL` backend path (it heads that family's candidate list, so it wins on SM120/121 when flashinfer is present and yields to `CuteDslFusedMoE` otherwise); single-GPU-shaped topology only — it rejects both `ep_size > 1` and attention-DP, because it has no dispatch/combine kernel and has never been exercised behind a DP allgather | `EXTERNAL_COMM` | | `mega_moe/mega_moe_deepgemm.py` | `MegaMoEDeepGemm` | SM100/SM103 | W4A8_MXFP4_MXFP8 via DeepGEMM `fp8_fp4_mega_moe` fused dispatch+GEMM+act+GEMM+combine kernel; requires `hidden_size % 512 == 0` | `FUSED_COMM` | -| `mega_moe/mega_moe_cute_dsl.py` | `MegaMoECuteDsl` | SM100/SM103 | NVFP4 via ported CuteDSL `Sm100MegaMoEKernel` fused dispatch+FC1+act+FC2+combine kernel; requires CUDA 13 Cutlass DSL runtime (PR #14354) and NVSHMEM provider (hard gate); threads per-expert `fc31_alpha`/`fc2_alpha`/`fc1_norm_const` through the kernel ABI and supports SwiGLU clamp via `swiglu_limit`; default deepgemm graph (topk score folded before fc1-out quant, host `combine_output.sum(dim=1)`) | `FUSED_COMM` | +| `mega_moe/mega_moe_cute_dsl.py` | `MegaMoECuteDsl` | SM100/SM103 | NVFP4 via ported CuteDSL `Sm100MegaMoEKernel` fused dispatch+FC1+act+FC2+combine kernel; requires CUDA 13 Cutlass DSL runtime (PR #14354) and NVSHMEM provider (hard gate); threads per-expert `fc31_alpha`/`fc2_alpha`/`fc1_norm_const` through the kernel ABI and takes the SwiGLU clamp as a `UNIFORM_SCALAR` (`gate_up_clamp=self.act_clamp`); default deepgemm graph (topk score folded before fc1-out quant, host `combine_output.sum(dim=1)`) | `FUSED_COMM` | | `fused_moe_marlin.py` | `MarlinFusedMoE` | SM89-SM99 | W4A16 NVFP4 on Ada/Hopper (BF16 activations + FP4 weights, fused single-launch `marlin_nvfp4_moe_gemm` kernel); supports attention-DP + EP via external comm (scheduler precomputes routing; dispatch payload is plain BF16, no activation scales); non-NVFP4 layers (e.g. unquantized MTP draft layers) degrade to Cutlass in `resolve_moe_impl`, recorded in the layer's `MoEResolutionReport`; no dynamic EPLB | `EXTERNAL_COMM` | | `fused_moe_triton.py` | `TritonFusedMoE` | SM90 only | GPT-OSS on Hopper (requires `swiglu_gptoss_style=True`) | (legacy path) | | `fused_moe_vanilla.py` | `VanillaMoE` | All devices | Reference / debugging only | (legacy path) | @@ -367,12 +368,79 @@ again. ### Activation Support -The matrix above is quantization only; activation style is a separate axis. The -gpt-oss SwiGLU package (per-expert bias plus `swiglu_alpha` / `swiglu_beta` / -`swiglu_limit`, surfaced as `MoEProblem.swiglu_gptoss_style`) is rejected by -every specialized backend — `CuteDslFusedMoE`, `CuteDslB12xFusedMoE`, -`DeepGemmFusedMoE`, `DenseGEMMFusedMoE`, `MarlinFusedMoE` — while -`TRTLLMGenFusedMoE` accepts only the algorithms in its `_GPTOSS_SUPPORTED_ALGOS`. +Activation is **declared, never inspected**: no factory or selection code asks +what class a backend is in order to decide whether it can run a checkpoint's +activation. Three types in `activation.py` carry that instead. + +| Layer | Type | Written by | +|-------|------|------------| +| Carrier | `MoEActivation` = `SwigluActivation` \| `SwigluBiasActivation` \| `SiTuActivation` \| `SimpleActivation` | The **model**, passed as `create_moe(activation=...)`; the default `DEFAULT_MOE_ACTIVATION` is plain SwiGLU. One frozen dataclass per kind, so a kind and its constants cannot disagree, and a constant the kind does not have cannot be written down at all | +| Declaration | `MoEActivationSupport` — `kinds`, plus an `ActivationParamShape` for `alpha_beta` and one for `limit` | The **backend**, as an `activation_support` class attribute. Every backend declares one; `resolve_activation_support` raises if a module does not | +| Adapter | `install_activation_params` → `materialize_activation_params` | Not the backend: the `apply_moe_impl_construction_state()` every execution unit already calls installs the slots, and `ConfigurableMoE` re-installs after the EPLB sync and inside `create_weights`. A complete layer that owns kernels directly (`TritonFusedMoE`, `VanillaMoE`) calls it itself — `MoE.__init__` deliberately does not, because a wrapper's declaration is the backend's | + +The carrier names constants **semantically**, per kind; only the declaration and +the adapter speak the ABI. `SwigluActivation` has just `clamp`; +`SwigluBiasActivation` has `gate_sigmoid_scale` / `linear_offset` / `clamp`; +`SiTuActivation` has `gate_softcap` / `linear_softcap` and no clamp at all; +`SimpleActivation(kind)` covers the kinds that take no constants. `constants()` +is the single seam where those become the functor's `alpha` / `beta` / `limit` +registers — and it has to be a seam, because the registers are borrowed for +unrelated jobs: `SwigluBias` reads `alpha` as a scale inside the sigmoid and +`beta` as an additive offset (neutral `0.0`), while `SiTu` reads both as tanh +soft-cap magnitudes the kernel divides by, so they must be positive (neutral +`1.0`). A single nullable `beta` has no coherent default without first finding +the kind, which is exactly the lookup this split removes. + +`kinds` is what the kernels **execute**, not what the gate tolerates — several +backends quietly narrow an accepted kind to SwiGLU or SiLU, and those kinds must +stay out of the declaration. + +The adapter assigns exactly three slots — `act_alpha`, `act_beta`, `act_clamp` — +and those are the only names a forward path or a quantization method reads. The +op schemas keep their historical spelling, so a backend passes +`torch.ops.trtllm.fused_moe(swiglu_alpha=self.act_alpha, ...)`; only the Python +plumbing was renamed. The slots' *types* come from the backend's declaration, +never from the checkpoint: + +- `PER_EXPERT_TENSOR` → `float32[expert_size_per_partition]`. A scalar is + broadcast; a caller tensor of the right length is cloned, because + `quantization.py` divides these buffers in place by the FC31 scale. +- `UNIFORM_SCALAR` → one `float`. A per-expert tensor whose values are not + *exactly* equal is rejected rather than averaged: the kernel bakes one value, + so a tolerance here would discard the experts that differ instead of + approximating them. +- `UNSUPPORTED` → the candidate is declined during resolution + (`MoERejectReason.ACTIVATION_UNSUPPORTED`) whenever the layer's activation + fills that register, so a backend that can serve it is still reachable. + +Declare `limit_when_absent` only when the ABI has no encoding for "no clamp": +`CuteDslFusedMoE` passes `float("inf")` because its epilogue always applies the +clamp functor. It is substituted *before* coercion, so a `PER_EXPERT_TENSOR` +backend gets it broadcast like any other scalar, and `MoEActivationSupport` +rejects it outright next to `limit=UNSUPPORTED`. + +`TRTLLMGenFusedMoE` is the one instance-level exception: its clamp is a +per-expert tensor for the FP4 fused-activation cubins but a by-value `double` +for the FP8 block-scale kernel, which is not a property of the class, so it +defines `resolve_activation_support` and narrows the shape per instance from +`quant_config`. That is also why `ConfigurableMoE.create_weights` re-installs +the slots: `apply_layerwise_quant_config` can move one layer onto that path +after `__init__` has already run. + +Selection keeps both halves of the picture in the tuning key — +`MoEProblem.activation` (the kind) and `MoEProblem.activation_constants` (which +registers the carrier actually fills) — because the kind alone does not separate +a clamped layer from an unclamped one, and those compile to different kernels. + +Quantization support (the matrix above) is a separate axis, and a backend can +execute a kind while still rejecting the quant algorithm it arrives with. The +gpt-oss SwiGLU package — per-expert bias plus a filled `alpha` / `beta` / +`limit` triple, summarized for selection as `MoEProblem.swiglu_gptoss_style` — +is declined by every specialized backend (`CuteDslFusedMoE`, +`CuteDslB12xFusedMoE`, `DeepGemmFusedMoE`, `DenseGEMMFusedMoE`, +`MarlinFusedMoE`) for the plain reason that none of them lists `SwigluBias` in +`kinds`, while `TRTLLMGenFusedMoE` executes it and then accepts only the +algorithms in its `_GPTOSS_SUPPORTED_ALGOS`. Cutlass gates gpt-oss / MiniMax SwiGLU on unquantized, MXFP8, NVFP4, and the MXFP4 family (`CutlassFusedMoE._GPTOSS_SUPPORTED_ALGOS` = `None`, `MXFP8`, @@ -382,7 +450,8 @@ is not the constraint — `torch.ops.trtllm.fused_moe` takes `swiglu_alpha` / NVFP4 (`CutlassMoeFCRunner<__nv_fp4_e2m1, __nv_fp4_e2m1>`), and TMA-WS GEMM1 applies `SwigluBiasAdaptor` in `doActivation`. NVFP4 is eligible only when there is no expert bias (`MoEProblem.bias is not True`): MiniMax-M3 NVFP4 -passes `ActivationType.SwigluBias` + alpha/beta/limit with `bias=False`. +builds a `SwigluBiasActivation` (`gate_sigmoid_scale` / `linear_offset` / +`clamp`) with `bias=False`. gpt-oss 1-D bias still goes through `NVFP4CutlassFusedMoEMethod`'s 2-D weight pad and is rejected at selection. Unquantized and the MXFP4 family can load that 1-D bias. W8A16 / W4A8_AWQ stay rejected because they inherit @@ -427,7 +496,7 @@ When adding new components, use these reference implementations: | Task | Reference | Key methods to implement | |------|-----------|--------------------------| -| New `EXTERNAL_COMM` Backend | `fused_moe_cutlass.py` (`CutlassFusedMoE`) | Declare `MoEImplBase`; implement `capabilities`, `can_implement`, `_get_quant_method`, `quantize_input`, `run_moe`; call `apply_moe_impl_construction_state()` in `__init__` (`create_weights` / `load_weights` come from `MoEWeightOwnerMixin` — override only if allocation needs more); then add the class to `moe_resolution.IMPL_PRIORITY` and `BACKEND_FAMILY`, and add a branch in `create_moe_backend`. Add a fixed `descriptor.identity` only for a one-implementation leaf class | +| New `EXTERNAL_COMM` Backend | `fused_moe_cutlass.py` (`CutlassFusedMoE`) | Declare `MoEImplBase`; implement `capabilities`, `activation_support`, `can_implement`, `_get_quant_method`, `quantize_input`, `run_moe`; call `apply_moe_impl_construction_state()` in `__init__` (`create_weights` / `load_weights` come from `MoEWeightOwnerMixin` — override only if allocation needs more); then add the class to `moe_resolution.IMPL_PRIORITY` and `BACKEND_FAMILY`, and add a branch in `create_moe_backend`. Add a fixed `descriptor.identity` only for a one-implementation leaf class | | New `FUSED_COMM` Backend | `mega_moe/mega_moe_deepgemm.py` (`MegaMoEDeepGemm`), `mega_moe/mega_moe_cute_dsl.py` (`MegaMoECuteDsl`) | Same as above + override `scheduler_kind = MoESchedulerKind.FUSED_COMM` and `validate_configurable_moe` for backend-specific constraints. For NVFP4 CuteDSL specifically, mirror the `MegaMoECuteDsl` pattern: capability probe for the CUDA 13 Cutlass DSL runtime, JSON-friendly tactic dict, lazy kernel import via `cute_dsl_kernels/mega_moe_nvfp4/import_kernel()`, and `quantize_input` that short-circuits zero-token input. | | New Quantization Method | `quantization.py` → `FP8QDQFusedMoEMethod` | Subclass `FusedMoEMethod`, implement quant/dequant ops | | New Communication Strategy | `communication/nvlink_one_sided.py` (`NVLinkOneSided`) | Subclass `Communication`, implement `prepare_dispatch`, `dispatch`, `combine` | @@ -449,6 +518,7 @@ Five backends declare `MoEImplBase` directly — `CutlassFusedMoE`, `TRTLLMGenFu - **Do NOT add new tests to `test_fused_moe.py` or `test_moe.py`** — Use `test_moe_backend.py` and `test_moe_module.py` - **Do NOT skip `can_implement()` checks** — Every backend must declare what it supports; an unsupported combination returns `MoEEligibility.no(MoERejectReason., detail)`, never a bare `False` and never a free-form string a test would have to pattern-match - **Do NOT probe the machine inside `can_implement()`** — No `get_sm_version()`, no `import` as a presence test, no `os.environ`. Read `d.env`; add the probe to `impl_environment.py` if it does not exist yet +- **Do NOT gate an activation on a class name** — No `isinstance`, no `moe_cls in [...]`, no ad-hoc `swiglu_gptoss_style` branch in a factory. A backend states what its kernels execute in `activation_support`; `_reject_unsupported_activation` reads that one declaration for every candidate, and `materialize_activation_params` enforces it again at construction. A per-class check re-creates the `assert moe_cls in [...]` this replaced, in a place the backend's own author will not find - **Do NOT add a second selection entry point** — `resolve_moe_impl` is the only one. A helper that picks a class on the side is how `get_moe_cls` and the old `resolve_moe_cls` drifted apart in the first place - **Do NOT substitute a backend without recording it** — A degradation must be visible in the `MoEResolutionReport`, not only in a log line - **Do NOT pick `scheduler_kind` opportunistically** — Use `EXTERNAL_COMM` (default) unless your backend's fused kernel genuinely owns cross-rank exchange via SymmBuffer / equivalent in-kernel collective; `FUSED_COMM` brings hard invariants (no host comm, lockstep launches, no multi-stream overlap) diff --git a/tensorrt_llm/_torch/moe/fused_moe/__init__.py b/tensorrt_llm/_torch/moe/fused_moe/__init__.py index 5b76cb7ad928..7461f879056a 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/__init__.py +++ b/tensorrt_llm/_torch/moe/fused_moe/__init__.py @@ -1,3 +1,7 @@ +from .activation import (ACTIVATION_PAYLOAD, DEFAULT_MOE_ACTIVATION, + ActivationParamShape, MoEActivation, + MoEActivationSupport, SimpleActivation, SiTuActivation, + SwigluActivation, SwigluBiasActivation) from .configurable_moe import ConfigurableMoE from .create_moe import (MoEImplClass, create_moe, resolve_moe_cls, resolve_moe_impl) @@ -28,10 +32,19 @@ # yapf: enable __all__ = [ + "ACTIVATION_PAYLOAD", + "ActivationParamShape", "BaseMoeRoutingMethod", "ConfigurableMoE", "create_load_balanced_logits", "create_moe", + "DEFAULT_MOE_ACTIVATION", + "MoEActivation", + "MoEActivationSupport", + "SimpleActivation", + "SiTuActivation", + "SwigluActivation", + "SwigluBiasActivation", "CuteDslB12xFusedMoE", "CuteDslFusedMoE", "CutlassFusedMoE", diff --git a/tensorrt_llm/_torch/moe/fused_moe/activation.py b/tensorrt_llm/_torch/moe/fused_moe/activation.py new file mode 100644 index 000000000000..9c80a11d5fd5 --- /dev/null +++ b/tensorrt_llm/_torch/moe/fused_moe/activation.py @@ -0,0 +1,588 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""A MoE layer's activation: one carrier, one backend declaration, one adapter. + +``MoEActivation`` + What a model builds. A union of per-kind dataclasses, so a kind and its + constants cannot disagree, and a constant a kind does not have cannot be + written down at all. + +``MoEActivationSupport`` + What a backend publishes, as a class attribute. Which kinds its kernels + execute, and what shape each constant must reach them in. + +``materialize_activation_params`` + The only place a semantic name becomes a kernel register, and the only + place a scalar is broadcast to per-expert or a per-expert tensor reduced + to a scalar. + +The split matters because the two floats every gated activation carries are +not one shared concept: they are two registers in the activation functor that +unrelated kinds borrow for unrelated jobs. ``SwigluBias`` reads ``alpha`` as a +scale inside the sigmoid and ``beta`` as an additive offset (neutral ``0.0``); +``SiTu`` reads both as tanh soft-cap magnitudes that must be positive (neutral +``1.0``). See ``cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/ +moe_kernels.cuh`` (``SwigluBiasAdaptor`` / ``SiTuAdaptor``) and +``GemmGatedActOptions.h``. A single nullable ``beta`` slot therefore has no +coherent default and no readable meaning until you know the kind. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum, auto +from typing import ClassVar + +import torch + +from ...utils import ActivationType + +__all__ = [ + "ACTIVATION_CONSTANT_NAMES", + "activation_constant_names", + "install_activation_params", + "ACTIVATION_PAYLOAD", + "ActivationConstant", + "ActivationConstants", + "ActivationParamShape", + "DEFAULT_MOE_ACTIVATION", + "MaterializedActivation", + "MoEActivation", + "MoEActivationSupport", + "resolve_activation_support", + "SimpleActivation", + "SiTuActivation", + "SwigluActivation", + "SwigluBiasActivation", + "materialize_activation_params", +] + +#: One value for the whole layer, or one per local expert. Which form reaches +#: the kernel is the backend's call (``MoEActivationSupport``). +ActivationConstant = torch.Tensor | float + + +# --------------------------------------------------------------------------- +# Validators +# --------------------------------------------------------------------------- +# +# Above the carriers: their ``__post_init__`` calls these, and +# ``DEFAULT_MOE_ACTIVATION`` constructs one at module scope. + + +def _reject_non_positive_clamp(value: ActivationConstant | None) -> None: + """Reject a clamp the kernel ABI cannot distinguish from "no clamp". + + ``swiglu_impl`` derives presence from the value itself + (``HAS_SWIGLU_LIMIT = swiglu_limit is not None and swiglu_limit > 0.0`` in + ``torch_custom_ops.py``), and the CuteDSL epilogue spells absence as + ``+inf``. So a clamp of zero silently means "no clamp". + """ + if value is None: + return + smallest = ( + float(value.detach().min().item()) if isinstance(value, torch.Tensor) else float(value) + ) + if not smallest > 0: + raise ValueError( + f"activation clamp must be positive because the kernel ABI encodes an absent " + f"clamp as a non-positive value; got {smallest}." + ) + + +def _reject_non_positive(value: ActivationConstant, *, name: str) -> None: + smallest = ( + float(value.detach().min().item()) if isinstance(value, torch.Tensor) else float(value) + ) + if not smallest > 0: + raise ValueError( + f"SiTu {name} must be positive because the kernel divides by it; got {smallest}." + ) + + +# --------------------------------------------------------------------------- +# Carrier +# --------------------------------------------------------------------------- +# +# ``eq=False`` throughout: these hold tensors, and the generated ``__eq__`` +# would compare elementwise and then call ``bool()`` on the result, which raises. + + +@dataclass(frozen=True, eq=False) +class ActivationConstants: + """The activation functor's register triple, in kernel-ABI order. + + The one representation that speaks ``alpha`` / ``beta``, because at this + boundary those *are* the names: ``ActFn::alpha`` / ``::beta`` / ``::limit`` + in ``moe_kernels.cuh``, ``gemm1_alpha`` / ``gemm1_beta`` in trtllm-gen. + Produced only by ``MoEActivation.constants()``; never authored by a caller. + """ + + alpha: ActivationConstant | None = None + beta: ActivationConstant | None = None + limit: ActivationConstant | None = None + + +@dataclass(frozen=True, eq=False) +class SwigluActivation: + """``silu(gate) * linear``, optionally clamped.""" + + clamp: ActivationConstant | None = None + + kind: ClassVar[ActivationType] = ActivationType.Swiglu + + def __post_init__(self) -> None: + _reject_non_positive_clamp(self.clamp) + + def constants(self) -> ActivationConstants: + return ActivationConstants(limit=self.clamp) + + +@dataclass(frozen=True, eq=False) +class SwigluBiasActivation: + """gpt-oss / MiniMax gated SwiGLU. + + ``g*sigmoid(g*gate_sigmoid_scale)*(l + linear_offset)`` where ``g`` / ``l`` + are the gate / linear halves of FC1 clamped by ``clamp``. + + ``linear_offset`` is not the MoE ``bias``: the kernel adds the per-expert + bias pointer to both halves *before* clamping, then adds this constant to + the clamped linear half (``moe_kernels.cuh`` ``SwigluBiasAdaptor``). Naming + it ``linear_bias`` would collide with that. + """ + + gate_sigmoid_scale: ActivationConstant + linear_offset: ActivationConstant + clamp: ActivationConstant | None = None + + kind: ClassVar[ActivationType] = ActivationType.SwigluBias + + def __post_init__(self) -> None: + _reject_non_positive_clamp(self.clamp) + + def constants(self) -> ActivationConstants: + return ActivationConstants( + alpha=self.gate_sigmoid_scale, beta=self.linear_offset, limit=self.clamp + ) + + +@dataclass(frozen=True, eq=False) +class SiTuActivation: + """Kimi K3 SiTU: two independently soft-capped branches. + + ``softcap(gate, gate_softcap) * sigmoid(gate) * softcap(linear, + linear_softcap)`` with ``softcap(x, c) = c*tanh(x / c)``, i.e. a smooth + saturation of ``x`` to ``+-c``. There is no separate clamp. + + Both caps must be positive because the kernel divides by them + (``GemmGatedActOptions.h``: ``SiTuGlu`` uses ``1/alpha`` and ``1/beta``), so + zero is not a neutral value here the way it is for ``SwigluBias``. + + Checkpoint provenance: ``gate_softcap`` is Kimi's ``activation_situ_beta`` + and ``linear_softcap`` its ``activation_situ_linear_beta``. + """ + + gate_softcap: ActivationConstant + linear_softcap: ActivationConstant + + kind: ClassVar[ActivationType] = ActivationType.SiTu + + def __post_init__(self) -> None: + for name in ("gate_softcap", "linear_softcap"): + _reject_non_positive(getattr(self, name), name=name) + + def constants(self) -> ActivationConstants: + return ActivationConstants(alpha=self.gate_softcap, beta=self.linear_softcap) + + +@dataclass(frozen=True, eq=False) +class SimpleActivation: + """A kind whose kernels take no activation constants (Silu, Geglu, Relu2 ...).""" + + kind: ActivationType + + def __post_init__(self) -> None: + payload = ACTIVATION_PAYLOAD.get(ActivationType(self.kind)) + if payload is None: + raise ValueError( + f"{ActivationType(self.kind).name} is not a MoE activation kind; " + f"known kinds: {', '.join(sorted(k.name for k in ACTIVATION_PAYLOAD))}" + ) + if payload is not SimpleActivation: + raise ValueError( + f"{ActivationType(self.kind).name} takes activation constants; " + f"build a {payload.__name__} instead of SimpleActivation." + ) + + def constants(self) -> ActivationConstants: + return ActivationConstants() + + +MoEActivation = SwigluActivation | SwigluBiasActivation | SiTuActivation | SimpleActivation + + +#: Entry point for a config-driven caller that knows only a kind: the member's +#: ``__init__`` signature *is* the constant list for that kind. +ACTIVATION_PAYLOAD: dict[ActivationType, type[MoEActivation]] = { + ActivationType.Gelu: SimpleActivation, + ActivationType.Relu: SimpleActivation, + ActivationType.Silu: SimpleActivation, + ActivationType.Relu2: SimpleActivation, + ActivationType.Geglu: SimpleActivation, + ActivationType.Swiglu: SwigluActivation, + ActivationType.SwigluBias: SwigluBiasActivation, + ActivationType.SiTu: SiTuActivation, +} + +#: Plain SwiGLU, no constants -- the historical default of every MoE signature. +#: Safe to share: the carrier is frozen. +DEFAULT_MOE_ACTIVATION: MoEActivation = SwigluActivation() + + +# --------------------------------------------------------------------------- +# Declaration +# --------------------------------------------------------------------------- + + +class ActivationParamShape(Enum): + """Form a backend's kernel boundary requires for a per-expert constant.""" + + #: Kernel dereferences a ``float*`` indexed by local expert. + PER_EXPERT_TENSOR = auto() + #: Kernel bakes one value for the whole layer (constexpr, or passed by value). + UNIFORM_SCALAR = auto() + #: Kernel has no such parameter. + UNSUPPORTED = auto() + + +@dataclass(frozen=True) +class MoEActivationSupport: + """What one backend's kernels can actually do with an activation. + + ``kinds`` declares what the backend **executes**, not what its gate + happens to accept. Several backends silently narrow an accepted kind to + SwiGLU or SiLU; those kinds do not belong here. + + ``alpha_beta`` keeps the ABI names on purpose: it describes the shape of the + functor's register pair, and it is only ever applied to an + ``ActivationConstants`` produced by ``MoEActivation.constants()``. + + ``limit_when_absent`` exists because some clamp ABIs have no "absent" + encoding -- the CuteDSL epilogue always applies the clamp functor, so "no + clamp" has to be spelled as a value. A backend whose ABI accepts None + leaves this unset. + """ + + kinds: frozenset[ActivationType] + alpha_beta: ActivationParamShape = ActivationParamShape.UNSUPPORTED + limit: ActivationParamShape = ActivationParamShape.UNSUPPORTED + limit_when_absent: float | None = None + + def __post_init__(self) -> None: + if self.limit_when_absent is not None and self.limit is ActivationParamShape.UNSUPPORTED: + raise ValueError( + "limit_when_absent names the value a clamp-less layer must still pass, so it " + "is meaningless with limit=UNSUPPORTED: declare a shape, or drop the value." + ) + + +# --------------------------------------------------------------------------- +# Adapter +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True, eq=False) +class MaterializedActivation: + """The activation in exactly the forms one backend's kernels take. + + Kind-neutral names on purpose. The C++ boundary calls these registers + ``swiglu_alpha`` / ``swiglu_beta`` / ``swiglu_limit`` (``moe_kernels.h``'s + ``ActivationParams``, and the ``moe_op`` schema), but SiTU fills the same + slots with tanh soft-caps. The op-call sites keep the schema names; nothing + above them does. + + One field per register, not one per ABI form: a clamp reaches its kernel + either as a ``float*`` indexed by expert or as a value, but which one is + already pinned by the backend's ``MoEActivationSupport``. + """ + + activation_type: ActivationType + alpha: ActivationConstant | None = None + beta: ActivationConstant | None = None + clamp: ActivationConstant | None = None + + +def materialize_activation_params( + activation: MoEActivation, + support: MoEActivationSupport, + *, + num_local_experts: int, + device: torch.device | str | None = None, + owner: str, +) -> MaterializedActivation: + """Convert an activation to the exact forms ``support`` declares. + + Takes ``activation.constants()`` -- the sole semantic-name-to-register seam + -- then broadcasts scalars to ``float32[num_local_experts]`` for + ``PER_EXPERT_TENSOR`` and reduces uniform tensors to a float for + ``UNIFORM_SCALAR``. A non-uniform tensor bound for ``UNIFORM_SCALAR`` is + rejected: the kernel bakes the value, so no single scalar is faithful. + + Rejections name ``owner`` and the offending parameter so a mis-declared + backend is identifiable without reading the factory. + """ + kind = ActivationType(activation.kind) + if kind not in support.kinds: + supported = ", ".join(sorted(k.name for k in support.kinds)) + raise ValueError( + f"{owner} does not implement activation {kind.name}; it executes: {supported}." + ) + + constants = activation.constants() + alpha = _materialize_to_declared_shape( + constants.alpha, + support.alpha_beta, + name="alpha", + kind=kind, + num_local_experts=num_local_experts, + device=device, + owner=owner, + ) + beta = _materialize_to_declared_shape( + constants.beta, + support.alpha_beta, + name="beta", + kind=kind, + num_local_experts=num_local_experts, + device=device, + owner=owner, + ) + # Substituted before materialization, not after: it is a plain float standing in for + # a missing constant, so a backend declaring PER_EXPERT_TENSOR needs it + # broadcast to the per-expert buffer like any caller-supplied scalar. + limit = _materialize_to_declared_shape( + constants.limit if constants.limit is not None else support.limit_when_absent, + support.limit, + name="clamp", + kind=kind, + num_local_experts=num_local_experts, + device=device, + owner=owner, + ) + + return MaterializedActivation( + activation_type=kind, + alpha=alpha, + beta=beta, + clamp=limit, + ) + + +def resolve_activation_support(module: torch.nn.Module) -> MoEActivationSupport: + """The declaration that applies to ``module``, class attribute or override. + + Static for ten of the eleven backends. TRTLLM-Gen is the documented + exception: its clamp ABI is a per-expert tensor for the FP4 fused-activation + cubins but a by-value ``double`` for the FP8 block-scale separate-activation + kernel, which is not a property of the class, so it defines + ``resolve_activation_support`` and narrows the shape per instance. + """ + override = getattr(module, "resolve_activation_support", None) + if callable(override): + return override() + support = getattr(type(module), "activation_support", None) + if support is None: + raise TypeError( + f"{type(module).__name__} must declare an ``activation_support`` class " + f"attribute stating which activations its kernels execute." + ) + return support + + +def install_activation_params( + module: torch.nn.Module, *, device: torch.device | str | None = None +) -> None: + """Assign the ``act_*`` slots ``module``'s kernels read, from ``module.activation``. + + The one place a layer or execution unit turns its declared activation into + the three attributes the forward paths and the quantization layer read. Runs + at construction, and again once ``ConfigurableMoE`` has synced + ``expert_size_per_partition`` -- the only thing that changes the per-expert + length -- before any weight is created. + + Safe to repeat only before weights exist: it re-materializes from + ``module.activation``, undoing any in-place transform a quant method applied. + """ + support = resolve_activation_support(module) + params = materialize_activation_params( + module.activation, + support, + num_local_experts=module.expert_size_per_partition, + device=device, + owner=type(module).__name__, + ) + module.activation_params = params + _write_activation_slot(module, "act_alpha", params.alpha) + _write_activation_slot(module, "act_beta", params.beta) + _write_activation_slot(module, "act_clamp", params.clamp) + + +def _write_activation_slot( + module: torch.nn.Module, name: str, value: ActivationConstant | None +) -> None: + """Write one slot so the constant travels with the weights it is used with. + + A tensor constant is registered rather than assigned, because a plain + attribute is the one thing ``nn.Module.to()`` does not move: the weights + reach the execution device that way, and a constant left behind arrives at + the kernel on the wrong device. ``persistent=False`` keeps it out of the + state dict, which is where a plain attribute already was -- these are + backend configuration, not checkpoint values. + + A slot that is already a parameter stays one. ``TRTLLMGenFusedMoE`` promotes + the SiTu slots so they do travel in the state dict, and the exclude-modules + pass clears ``_weights_created`` without unregistering them, so this runs + again over a live parameter. + """ + if isinstance(value, torch.Tensor) and name in getattr(module, "_parameters", {}): + setattr(module, name, torch.nn.Parameter(value, requires_grad=False)) + return + # Installing happens more than once -- EPLB slot sync and the layerwise + # quant config each redo it -- so clear the previous binding, of whichever + # kind, before making the new one. ``delattr`` rather than popping + # ``_buffers``, so a slot going from tensor back to scalar also leaves + # ``_non_persistent_buffers_set``. + if hasattr(module, name): + delattr(module, name) + if isinstance(value, torch.Tensor): + module.register_buffer(name, value, persistent=False) + else: + setattr(module, name, value) + + +#: The ABI register names, in the order ``ActivationConstants`` declares them. +#: ``alpha`` and ``beta`` share one declared shape (``alpha_beta``); ``clamp`` +#: has its own (``limit``). +ACTIVATION_CONSTANT_NAMES: tuple[str, ...] = ("alpha", "beta", "clamp") + + +def activation_constant_names(activation: MoEActivation | None) -> frozenset[str]: + """Which of the three ABI registers this activation actually fills. + + The selection layer needs this and cannot get it from the kind alone: SwiGLU + with a clamp and SwiGLU without one are the same ``ActivationType`` but not + the same request, and a backend declaring ``limit=UNSUPPORTED`` can serve + only the second. Returning names rather than the values keeps ``MoEProblem`` + hashable and JSON-serializable, which a tuning key has to be. + """ + if activation is None: + return frozenset() + constants = activation.constants() + return frozenset( + name + for name in ACTIVATION_CONSTANT_NAMES + if getattr(constants, "limit" if name == "clamp" else name) is not None + ) + + +def _materialize_to_declared_shape( + value: ActivationConstant | None, + shape: ActivationParamShape, + *, + name: str, + kind: ActivationType, + num_local_experts: int, + device: torch.device | str | None, + owner: str, +) -> ActivationConstant | None: + """Put one constant in the form ``shape`` declares, or refuse to. + + Not a lenient conversion: a value bound for ``UNSUPPORTED`` raises rather + than being dropped, and a per-expert tensor bound for ``UNIFORM_SCALAR`` + raises unless its values are exactly equal. + """ + if value is None: + return None + if shape is ActivationParamShape.UNSUPPORTED: + raise ValueError( + f"{owner} kernels take no activation {name}, but {kind.name} supplied one. " + f"Either drop it or select a backend that declares it." + ) + if shape is ActivationParamShape.UNIFORM_SCALAR: + return _reduce_to_uniform_scalar(value, name=name, owner=owner) + return _broadcast_to_per_expert( + value, name=name, num_local_experts=num_local_experts, device=device, owner=owner + ) + + +def _reduce_to_uniform_scalar(value: ActivationConstant, *, name: str, owner: str) -> float: + """Reduce a per-expert constant to the one float the kernel can bake.""" + if not isinstance(value, torch.Tensor): + return float(value) + flat = value.detach().reshape(-1) + if flat.numel() == 0: + raise ValueError(f"{owner} received an empty activation {name}.") + first = flat[0] + # Exact equality, not allclose: the kernel bakes ``first`` and every other + # expert inherits it, so a tolerance would discard the values that differ + # rather than approximate them. + if flat.numel() > 1 and bool((flat != first).any()): + raise ValueError( + f"{owner} only supports a uniform (per-layer) activation {name} because the " + f"kernel bakes it as a compile-time scalar; got per-expert values " + f"{flat.cpu().tolist()}." + ) + return float(first.item()) + + +def _broadcast_to_per_expert( + value: ActivationConstant, + *, + name: str, + num_local_experts: int, + device: torch.device | str | None, + owner: str, +) -> torch.Tensor: + """Produce the ``float32[num_local_experts]`` buffer the kernel indexes. + + ``device=None`` stays ``None`` so the constant lands where ``create_weights`` + puts the weights it is combined with: both create with the ambient device + (``nn.Parameter(torch.empty(shape, dtype=...))``), and + ``install_activation_params`` registers this as a buffer so the same + ``.to()`` moves both. Naming a device here instead makes the two disagree -- + a module built and loaded on CPU, then moved, divides by the FC31 scale + while still on CPU. + + Meta is the one exception. These are configuration, not checkpoint values, + so nothing reloads them after materialization; a meta constant would just be + empty. + """ + if device is None and torch.get_default_device().type == "meta": + device = "cuda" if torch.cuda.is_available() else "cpu" + if not isinstance(value, torch.Tensor): + return torch.full((num_local_experts,), float(value), dtype=torch.float32, device=device) + flat = value.detach().reshape(-1).to(dtype=torch.float32) + if flat.numel() == num_local_experts: + # ``.clone()`` because ``quantization.py`` divides these buffers in place + # by the FC31 scale, and ``detach``/``reshape``/``to`` are no-ops for an + # already-matching tensor -- so returning it would divide the caller's own + # constant. + return flat.to(device=device).clone() + if flat.numel() == 1: + return flat.to(device=device).expand(num_local_experts).contiguous() + raise ValueError( + f"{owner} indexes activation {name} by local expert, so it must hold " + f"{num_local_experts} values (or one to broadcast); got {flat.numel()}." + ) diff --git a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py index bd44358d74fa..47c3a07f221b 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py @@ -44,15 +44,11 @@ from tensorrt_llm._torch.moe.fused_moe.interface import MoE, MoESchedulerKind, _reject from tensorrt_llm._torch.moe.fused_moe.routing import BaseMoeRoutingMethod from tensorrt_llm._torch.pyexecutor.dwdp import get_global_dwdp_manager -from tensorrt_llm._torch.utils import ( - ActType_TrtllmGen, - AuxStreamType, - EventType, - Fp4QuantizedTensor, -) +from tensorrt_llm._torch.utils import AuxStreamType, EventType, Fp4QuantizedTensor from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig +from .activation import install_activation_params from .communication import AllGatherReduceScatter, Communication, CommunicationFactory from .moe_scheduler import MoEScheduler, create_moe_scheduler @@ -162,12 +158,6 @@ def __init__( layer_idx: Optional[int] = None, override_quant_config: Optional["QuantConfig"] = None, moe_cls: Optional[Type] = None, - activation: Optional[str] = None, - situ_beta: Optional[float] = None, - situ_linear_beta: Optional[float] = None, - trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None, - trtllm_gen_activation_alpha: Optional[float] = None, - trtllm_gen_activation_beta: Optional[float] = None, communication_method: Optional[str] = None, **kwargs, ): @@ -201,12 +191,6 @@ def __init__( routing_method=routing_method, override_quant_config=override_quant_config, moe_cls=moe_cls, - activation=activation, - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, - trtllm_gen_activation_type=trtllm_gen_activation_type, - trtllm_gen_activation_alpha=trtllm_gen_activation_alpha, - trtllm_gen_activation_beta=trtllm_gen_activation_beta, **kwargs, ) @@ -291,12 +275,6 @@ def _create_and_sync_backend( routing_method: BaseMoeRoutingMethod, override_quant_config: Optional["QuantConfig"], moe_cls: Optional[Type] = None, - activation: Optional[str] = None, - situ_beta: Optional[float] = None, - situ_linear_beta: Optional[float] = None, - trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None, - trtllm_gen_activation_alpha: Optional[float] = None, - trtllm_gen_activation_beta: Optional[float] = None, **kwargs, ) -> None: """Build the MoE backend, mirror EPLB attrs, then create weights. @@ -330,12 +308,10 @@ def _create_and_sync_backend( intermediate_size=self.intermediate_size, swiglu_gptoss_style=infer_swiglu_gptoss_style( bias=kwargs.get("bias", False), - swiglu_alpha=kwargs.get("swiglu_alpha"), - swiglu_beta=kwargs.get("swiglu_beta"), activation_type=self.activation_type, ), bias=kwargs.get("bias", False), - activation_type=self.activation_type, + activation=self.activation, routing=self.routing_method, layer_idx=self.layer_idx, ) @@ -360,18 +336,8 @@ def _create_and_sync_backend( bias=kwargs.get("bias", False), apply_router_weight_on_input=self.apply_router_weight_on_input, layer_idx=None, - swiglu_alpha=kwargs.get("swiglu_alpha"), - swiglu_beta=kwargs.get("swiglu_beta"), - swiglu_limit=kwargs.get("swiglu_limit"), - swiglu_limit_scalar=kwargs.get("swiglu_limit_scalar"), init_load_balancer=False, - activation_type=self.activation_type, - activation=activation, - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, - trtllm_gen_activation_type=trtllm_gen_activation_type, - trtllm_gen_activation_alpha=trtllm_gen_activation_alpha, - trtllm_gen_activation_beta=trtllm_gen_activation_beta, + activation=self.activation, ) # Backend acceptance is validated at the end of ``__init__`` instead @@ -389,6 +355,9 @@ def _create_and_sync_backend( if self.backend is not None: for attr in _BACKEND_SYNC_ATTRS: setattr(self.backend, attr, getattr(self, attr)) + # ``expert_size_per_partition`` may have just changed (EPLB slots), + # and it sizes every per-expert activation constant. + install_activation_params(self.backend) # Sync done -- now the backend has enough info to allocate weight # tensors with the right shard / slot count. @@ -755,6 +724,16 @@ def create_weights(self): if self._override_quant_config is not None else self.quant_config ) + # The quant config just changed and a backend may resolve its activation + # ABI from it: TRTLLMGenFusedMoE narrows the clamp to a uniform scalar on + # the FP8 block-scale path, whose kernel takes one ``double`` by value. + # The install in __init__ ran against the pre-layerwise config, so a layer + # that layerwise quantization moved onto that path would otherwise keep a + # tensor clamp the kernel ignores. Guarded because a re-install + # re-materializes from ``activation``, undoing the in-place division + # NVFP4TRTLLMGenFusedMoEBaseMethod applies to beta and clamp. + if not self.backend._weights_created: + install_activation_params(self.backend) return self.backend.create_weights() def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False): diff --git a/tensorrt_llm/_torch/moe/fused_moe/create_moe.py b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py index 66d9871792d0..85f65ef53aec 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py @@ -5,10 +5,10 @@ import torch from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.utils import (ActivationType, ActType_TrtllmGen, - AuxStreamType) +from tensorrt_llm._torch.utils import AuxStreamType from tensorrt_llm.models.modeling_utils import QuantConfig +from .activation import DEFAULT_MOE_ACTIVATION, MoEActivation from .configurable_moe import ConfigurableMoE from .fused_moe_cute_dsl import CuteDslFusedMoE from .fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE @@ -55,18 +55,8 @@ def create_moe_backend( bias: bool = False, apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, - swiglu_alpha: Optional[torch.Tensor] = None, - swiglu_beta: Optional[torch.Tensor] = None, - swiglu_limit: Optional[torch.Tensor] = None, - swiglu_limit_scalar: Optional[float] = None, init_load_balancer: bool = False, - activation_type: ActivationType = ActivationType.Swiglu, - activation: Optional[str] = None, - situ_beta: Optional[float] = None, - situ_linear_beta: Optional[float] = None, - trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None, - trtllm_gen_activation_alpha: Optional[float] = None, - trtllm_gen_activation_beta: Optional[float] = None, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, ) -> MoE | MoEImplBase | VanillaMoE: """ Create a MoE backend or a self-contained MoE layer. @@ -91,17 +81,9 @@ def create_moe_backend( bias: Whether to use bias apply_router_weight_on_input: Whether to apply router weight on input layer_idx: Layer index - swiglu_alpha: SwiGLU alpha parameter - swiglu_beta: SwiGLU beta parameter - swiglu_limit: SwiGLU limit parameter (per-expert tensor; for NVFP4) - swiglu_limit_scalar: SwiGLU limit scalar (uniform across experts; for FP8) - activation_type: Activation type - activation: Optional MegaMoE DeepGEMM activation name - situ_beta: Optional MegaMoE DeepGEMM SiTU beta - situ_linear_beta: Optional MegaMoE DeepGEMM SiTU linear beta - trtllm_gen_activation_type: Optional TRTLLM-Gen backend-local activation type - trtllm_gen_activation_alpha: Optional backend-local activation alpha - trtllm_gen_activation_beta: Optional backend-local activation beta + activation: The layer's activation kind and its constants. Whether a + constant reaches the kernel per expert or as one baked scalar is the + backend's declaration (``activation_support``), not a caller choice. Returns: A ``MoEImplBase`` execution unit, a self-contained ``MoE`` layer, or @@ -129,49 +111,24 @@ def create_moe_backend( "intermediate_size must be provided or model_config.pretrained_config " "must expose moe_intermediate_size / intermediate_size") - moe_load_balancer = get_moe_load_balancer() - if moe_load_balancer is not None: - supported_load_balancer_backends = ( - CutlassFusedMoE, - TRTLLMGenFusedMoE, - CuteDslFusedMoE, - DeepGemmFusedMoE, - DenseGEMMFusedMoE, - MegaMoEDeepGemm, - MegaMoECuteDsl, - ) - assert moe_cls in supported_load_balancer_backends, ( - "MoE Load Balance is only supported in " - f"{', '.join(cls.__name__ for cls in supported_load_balancer_backends)}." - ) - - if bias: - assert moe_cls in [CutlassFusedMoE, TritonFusedMoE, TRTLLMGenFusedMoE - ], f"bias not supported in {moe_cls.__name__}." - - if swiglu_alpha is not None or swiglu_beta is not None: - assert moe_cls in [CutlassFusedMoE, TritonFusedMoE, TRTLLMGenFusedMoE], \ - f"swiglu_alpha and swiglu_beta are only supported in CutlassFusedMoE, TritonFusedMoE and TRTLLMGenFusedMoE, not in {moe_cls.__name__}." - assert swiglu_alpha is not None and swiglu_beta is not None, \ - "Both swiglu_alpha and swiglu_beta must be provided." + # Backstop for the direct-``moe_cls`` callers that bypass resolution: a + # backend that cannot run under a load balancer already declines in + # ``can_implement``. + eplb_enabled = get_moe_load_balancer() is not None + if eplb_enabled and not moe_cls.capabilities.supports_eplb: + raise ValueError( + f"{moe_cls.__name__} does not support the MoE load balancer.") - if swiglu_limit is not None: - assert moe_cls in [ - CutlassFusedMoE, TritonFusedMoE, TRTLLMGenFusedMoE, - DeepGemmFusedMoE, MegaMoECuteDsl - ], f"swiglu_limit is not supported in {moe_cls.__name__}." + if bias and not moe_cls.capabilities.supports_expert_bias: + raise ValueError(f"bias not supported in {moe_cls.__name__}.") - if swiglu_limit_scalar is not None: - # MegaMoECuteDsl uses the scalar only as a fallback when no per-expert - # tensor limit is given (see the MegaMoE branch below). - assert moe_cls in [ - CutlassFusedMoE, TRTLLMGenFusedMoE, DeepGemmFusedMoE, - MegaMoEDeepGemm, CuteDslFusedMoE, MegaMoECuteDsl - ], f"swiglu_limit_scalar is not supported in {moe_cls.__name__}." + if (apply_router_weight_on_input + and not moe_cls.capabilities.supports_apply_router_weight_on_input): + raise ValueError( + f"apply_router_weight_on_input not supported in {moe_cls.__name__}." + ) if moe_cls == TRTLLMGenFusedMoE: - assert not apply_router_weight_on_input, "apply_router_weight_on_input is not supported in TRTLLMGenFusedMoE." - return moe_cls( routing_method=routing_method, num_experts=num_experts, @@ -184,33 +141,13 @@ def create_moe_backend( weight_loading_mode=weight_loading_mode, bias=bias, layer_idx=layer_idx, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - swiglu_limit=swiglu_limit, - swiglu_limit_scalar=swiglu_limit_scalar, init_load_balancer=init_load_balancer, - activation_type=activation_type, - trtllm_gen_activation_type=trtllm_gen_activation_type, - trtllm_gen_activation_alpha=trtllm_gen_activation_alpha, - trtllm_gen_activation_beta=trtllm_gen_activation_beta, + activation=activation, ) - if any(value is not None - for value in (activation, situ_beta, - situ_linear_beta)) and moe_cls is not MegaMoEDeepGemm: - raise ValueError("MegaMoE DeepGEMM activation options require " - f"MegaMoEDeepGemm, got {moe_cls.__name__}") - - if any(value is not None for value in (trtllm_gen_activation_type, - trtllm_gen_activation_alpha, - trtllm_gen_activation_beta)): - raise ValueError( - "TRTLLM-Gen backend-local activation options are only supported " - f"by TRTLLMGenFusedMoE, got {moe_cls.__name__}") - elif moe_cls in (CutlassFusedMoE, MarlinFusedMoE): - # CuteDslFusedMoE, DeepGemmFusedMoE, and CuteDslB12xFusedMoE - # also subclass CutlassFusedMoE but have narrower constructors, so - # they take their own branches below. + if moe_cls in (CutlassFusedMoE, MarlinFusedMoE): + # The two whose constructor takes an expert-bias flag. Marlin declines + # the flag itself, so the check above already rejected a True. return moe_cls( routing_method=routing_method, num_experts=num_experts, @@ -224,16 +161,10 @@ def create_moe_backend( bias=bias, apply_router_weight_on_input=apply_router_weight_on_input, layer_idx=layer_idx, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - swiglu_limit=swiglu_limit, - swiglu_limit_scalar=swiglu_limit_scalar, init_load_balancer=init_load_balancer, - activation_type=activation_type, + activation=activation, ) elif moe_cls == VanillaMoE: - assert not apply_router_weight_on_input, "apply_router_weight_on_input is not supported in VanillaMoE." - return moe_cls( routing_method=routing_method, num_experts=num_experts, @@ -245,13 +176,10 @@ def create_moe_backend( weight_loading_mode=weight_loading_mode, apply_router_weight_on_input=apply_router_weight_on_input, layer_idx=layer_idx, - activation_type=activation_type, + activation=activation, ) elif moe_cls in (CuteDslFusedMoE, CuteDslB12xFusedMoE): - # Both are constructed through the narrower CuteDsl argument set (no - # bias / swiglu_alpha-beta-limit). CuteDslB12xFusedMoE now delegates to - # CutlassFusedMoE.__init__, which does accept those four, so widening - # this branch would need the allow-lists above to admit b12x first. + # The narrower CuteDsl argument set: these kernels take no expert bias. return moe_cls( routing_method=routing_method, num_experts=num_experts, @@ -264,9 +192,8 @@ def create_moe_backend( weight_loading_mode=weight_loading_mode, apply_router_weight_on_input=apply_router_weight_on_input, layer_idx=layer_idx, - swiglu_limit_scalar=swiglu_limit_scalar, init_load_balancer=init_load_balancer, - activation_type=activation_type, + activation=activation, ) elif moe_cls == DeepGemmFusedMoE: return moe_cls( @@ -282,12 +209,9 @@ def create_moe_backend( apply_router_weight_on_input=apply_router_weight_on_input, layer_idx=layer_idx, init_load_balancer=init_load_balancer, - swiglu_limit=swiglu_limit, - swiglu_limit_scalar=swiglu_limit_scalar, + activation=activation, ) elif moe_cls == TritonFusedMoE: - assert not apply_router_weight_on_input, "apply_router_weight_on_input is not supported in TritonFusedMoE." - return moe_cls( routing_method=routing_method, num_experts=num_experts, @@ -299,9 +223,7 @@ def create_moe_backend( weight_loading_mode=weight_loading_mode, bias=bias, layer_idx=layer_idx, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - swiglu_limit=swiglu_limit, + activation=activation, ) elif moe_cls == DenseGEMMFusedMoE: return moe_cls( @@ -317,15 +239,10 @@ def create_moe_backend( apply_router_weight_on_input=apply_router_weight_on_input, layer_idx=layer_idx, init_load_balancer=init_load_balancer, - activation_type=activation_type, + activation=activation, ) elif moe_cls in (MegaMoEDeepGemm, MegaMoECuteDsl): - # MegaMoE fused-comm backends share the same construction surface. - # ``mega_moe_deepgemm`` lazily resolves DG via ``_import_deep_gemm`` - # at runtime and ``mega_moe_cute_dsl`` lazily imports the CuteDSL - # kernel package, so a top-level import here doesn't pull either - # heavyweight dependency on boxes that don't use these backends. - megamoe_kwargs = dict( + return moe_cls( routing_method=routing_method, num_experts=num_experts, hidden_size=hidden_size, @@ -338,22 +255,8 @@ def create_moe_backend( apply_router_weight_on_input=apply_router_weight_on_input, layer_idx=layer_idx, init_load_balancer=init_load_balancer, - activation_type=activation_type, + activation=activation, ) - if moe_cls is MegaMoECuteDsl: - # ``_resolve_gate_up_clamp`` accepts tensor or scalar; fall back - # to the scalar form when only that was wired. - megamoe_kwargs["swiglu_limit"] = (swiglu_limit - if swiglu_limit is not None else - swiglu_limit_scalar) - else: - megamoe_kwargs.update( - swiglu_limit_scalar=swiglu_limit_scalar, - activation=activation, - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, - ) - return moe_cls(**megamoe_kwargs) else: raise ValueError(f"Unsupported moe backend: {moe_cls}") @@ -372,18 +275,9 @@ def create_moe( bias: bool = False, apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, - swiglu_alpha: Optional[torch.Tensor] = None, - swiglu_beta: Optional[torch.Tensor] = None, - swiglu_limit: Optional[torch.Tensor] = None, - swiglu_limit_scalar: Optional[float] = None, - activation_type: ActivationType = ActivationType.Swiglu, - activation: Optional[str] = None, - situ_beta: Optional[float] = None, - situ_linear_beta: Optional[float] = None, - trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None, - trtllm_gen_activation_alpha: Optional[float] = None, - trtllm_gen_activation_beta: Optional[float] = None, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, communication_method: Optional[str] = None, + allow_backend_degradation: bool = True, ) -> MoE | VanillaMoE: """ Create MoE instance with automatic parameter inference from model_config. @@ -402,18 +296,14 @@ def create_moe( bias: Whether to use bias apply_router_weight_on_input: Whether to apply router weight on input layer_idx: Layer index - swiglu_alpha: SwiGLU alpha parameter - swiglu_beta: SwiGLU beta parameter - swiglu_limit: SwiGLU limit parameter (per-expert tensor; for NVFP4) - swiglu_limit_scalar: SwiGLU limit scalar (uniform across experts; for FP8) - activation_type: Activation type - activation: Optional MegaMoE DeepGEMM activation name - situ_beta: Optional MegaMoE DeepGEMM SiTU beta - situ_linear_beta: Optional MegaMoE DeepGEMM SiTU linear beta - trtllm_gen_activation_type: Optional TRTLLM-Gen backend-local activation type - trtllm_gen_activation_alpha: Optional backend-local activation alpha - trtllm_gen_activation_beta: Optional backend-local activation beta + activation: The layer's activation kind and its constants, e.g. + ``SwigluActivation(clamp=7.0)`` or + ``SiTuActivation(gate_softcap=..., linear_softcap=...)`` communication_method: Optional ConfigurableMoE communication method + allow_backend_degradation: When False, a requested backend that cannot + serve this layer raises with the rejection trail instead of falling + back. For callers that must know they got the backend they asked + for, e.g. because they are measuring it. Returns: A complete MoE layer: a ``MoE`` (``ConfigurableMoE`` around an @@ -452,30 +342,14 @@ def create_moe( intermediate_size=intermediate_size, swiglu_gptoss_style=infer_swiglu_gptoss_style( bias=bias, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - activation_type=activation_type, + activation_type=activation.kind, ), bias=bias, - activation_type=activation_type, + activation=activation, routing=routing_method, layer_idx=layer_idx, + allow_degradation=allow_backend_degradation, ) - if (any(value is not None - for value in (activation, situ_beta, situ_linear_beta)) - and moe_cls is not MegaMoEDeepGemm): - raise ValueError( - "MegaMoE DeepGEMM activation options require " - "MegaMoEDeepGemm without backend fallback, but resolved " - f"{moe_cls.__name__}.") - if (any(value is not None for value in (trtllm_gen_activation_type, - trtllm_gen_activation_alpha, - trtllm_gen_activation_beta)) - and moe_cls is not TRTLLMGenFusedMoE): - raise ValueError( - "A TRTLLM-Gen backend-local activation requires " - "TRTLLMGenFusedMoE without backend fallback, but resolved " - f"{moe_cls.__name__}.") # This dispatch needs no per-class entry: inheriting ``MoEImplBase`` is # enough to be wrapped. Becoming *selectable* still needs the constructor @@ -497,17 +371,7 @@ def create_moe( layer_idx=layer_idx, override_quant_config=override_quant_config, bias=bias, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - swiglu_limit=swiglu_limit, - swiglu_limit_scalar=swiglu_limit_scalar, - activation_type=activation_type, activation=activation, - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, - trtllm_gen_activation_type=trtllm_gen_activation_type, - trtllm_gen_activation_alpha=trtllm_gen_activation_alpha, - trtllm_gen_activation_beta=trtllm_gen_activation_beta, communication_method=communication_method, ) @@ -529,15 +393,5 @@ def create_moe( bias=bias, apply_router_weight_on_input=apply_router_weight_on_input, layer_idx=layer_idx, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - swiglu_limit=swiglu_limit, - swiglu_limit_scalar=swiglu_limit_scalar, - activation_type=activation_type, activation=activation, - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, - trtllm_gen_activation_type=trtllm_gen_activation_type, - trtllm_gen_activation_alpha=trtllm_gen_activation_alpha, - trtllm_gen_activation_beta=trtllm_gen_activation_beta, ) diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py index b1d411b1649e..6b356c33fd38 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py @@ -32,10 +32,12 @@ Fp4QuantizedTensor, get_last_power_of_2_num_tokens_buckets, last_positive_power_of_2) -from .fused_moe_cutlass import CutlassFusedMoE -from .impl_contract import (MoEDeployment, MoEEligibility, MoEProblem, - MoERejectReason, MoERunContext, MoEStaticCapability, - require_comm_plan) +from .activation import (DEFAULT_MOE_ACTIVATION, ActivationParamShape, + MoEActivation, MoEActivationSupport) +from .impl_base import MoEImplBase, apply_moe_impl_construction_state +from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement, + MoEProblem, MoERejectReason, MoERunContext, + MoEStaticCapability, require_comm_plan) from .interface import _reject from .quantization import MoEWeightLoadingMode, NVFP4CuteDslFusedMoEMethod from .routing import BaseMoeRoutingMethod @@ -358,7 +360,7 @@ def runner_tactic_comb_checker( return True -class CuteDslFusedMoE(CutlassFusedMoE): +class CuteDslFusedMoE(MoEImplBase): # CuteDSL dispatch/combine path exercises the ceil/floor partition # (NVLinkOneSided alltoall with kernel-level remainder handling), so this # backend is the only opt-in for non-divisible EP today. @@ -376,13 +378,22 @@ 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) + capabilities = MoEStaticCapability( + supports_dwdp=True, + supports_eplb=True, + supports_apply_router_weight_on_input=True) + + input_requirement = MoEInputRequirement(routing_scales_dtype=torch.float32) + + # Kinds mirror the kernel's own SUPPORTED_ACTIVATION_TYPES in + # cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py. + # The clamp is a kernel-cache-key scalar and the epilogue has no + # "clamp absent" branch, so an absent clamp is +inf, not None. + activation_support = MoEActivationSupport( + kinds=frozenset({ActivationType.Swiglu, ActivationType.Relu2}), + limit=ActivationParamShape.UNIFORM_SCALAR, + limit_when_absent=float("inf"), + ) @classmethod def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: @@ -450,11 +461,12 @@ def __init__( VANILLA, apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, - swiglu_limit_scalar: Optional[float] = None, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, init_load_balancer: bool = False, - activation_type: ActivationType = ActivationType.Swiglu, ): - super().__init__( + super().__init__(eplb=None) + apply_moe_impl_construction_state( + self, routing_method=routing_method, num_experts=num_experts, hidden_size=hidden_size, @@ -464,24 +476,32 @@ def __init__( model_config=model_config, aux_stream_dict=aux_stream_dict, weight_loading_mode=weight_loading_mode, - apply_router_weight_on_input=apply_router_weight_on_input, layer_idx=layer_idx, - swiglu_limit_scalar=swiglu_limit_scalar, + activation=activation, init_load_balancer=init_load_balancer, - activation_type=activation_type, ) - self.swiglu_limit_scalar = swiglu_limit_scalar or float("inf") + self.apply_router_weight_on_input = apply_router_weight_on_input + # Read by run_moe_nvfp4* to pick the fused-finalize epilogue, which + # leaves no seam for a LoRA GEMM. + self.use_fused_finalize = (not model_config.moe_disable_finalize_fusion + and model_config.lora_config is None) + + # ``run_moe_nvfp4*`` overlaps the output memset on its own stream. This + # backend never chunks, so it needs no chunking stream. if self.aux_stream_dict is None: - self.aux_stream_dict = aux_stream_dict if aux_stream_dict is not None else {} + self.aux_stream_dict = {} if AuxStreamType.MoeOutputMemset not in self.aux_stream_dict: self.aux_stream_dict[ AuxStreamType.MoeOutputMemset] = torch.cuda.Stream() - if self.event_dict is None: - self.event_dict = {} - for key in [EventType.Main, EventType.MoeOutputMemset]: - if key not in self.event_dict: - self.event_dict[key] = torch.cuda.Event() + self.event_dict = { + key: torch.cuda.Event() + for key in [EventType.Main, EventType.MoeOutputMemset] + } + + self._weights_created = False + if not model_config.skip_create_weights_in_init: + self.create_weights() def _build_local_weight_view(self) -> NvFp4WeightView: """Build the weight view from this backend's per-layer weights.""" @@ -501,7 +521,19 @@ def _get_quant_method(self): exclude_kv_cache=True): if self.quant_config.layer_quant_mode.has_nvfp4(): return NVFP4CuteDslFusedMoEMethod() - return super()._get_quant_method() + # ``can_implement`` admits NVFP4 only, so selection never lands here. + # Raise rather than fall back: any other method owns a weight layout + # these kernels cannot read. + raise ValueError( + f"CuteDslFusedMoE only supports NVFP4, got {self.quant_config}") + + def _supports_load_balancer(self) -> bool: + return True + + def _check_configs(self): + assert self._weights_created + if self.apply_router_weight_on_input: + assert self.routing_method.top_k == 1, "Current walkaround only supports top-1 routing" def supports_moe_output_in_alltoall_workspace(self): return self.has_nvfp4 @@ -665,7 +697,7 @@ def run_moe_nvfp4_impl( local_expert_offset=slot_start, tile_size=tile_size, activation_type=self.activation_type, - swiglu_limit_scalar=self.swiglu_limit_scalar, + swiglu_limit_scalar=self.act_clamp, ) if self.use_fused_finalize: @@ -778,7 +810,7 @@ def run_moe_fp8_block_scales( b_sf=self.quant_scales[0], offset_array=expert_first_token_offset, ) - x = swiglu_fused_moe(x, self.swiglu_limit_scalar) + x = swiglu_fused_moe(x, self.act_clamp) x, x_sf = torch.ops.trtllm.fp8_quantize_1x128(x) x = cute_dsl_fp8_group_blockwise_gemm_ref( a=x, diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py index b2f70dba5734..ffb7fe925a0d 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py @@ -22,10 +22,12 @@ from tensorrt_llm.models.modeling_utils import QuantAlgo from ...utils import ActivationType, Fp4QuantizedTensor +from .activation import MoEActivationSupport from .fused_moe_cutlass import CutlassFusedMoE from .impl_contract import ( MoEDeployment, MoEEligibility, + MoEInputRequirement, MoEProblem, MoERejectReason, MoERunContext, @@ -58,32 +60,49 @@ class CuteDslB12xFusedMoE(CutlassFusedMoE): Large prefill chunks use CUTLASS; decode uses FlashInfer's b12x kernel. - Inherits ``CutlassFusedMoE`` rather than only the shared blocks -- the same - shortcut its siblings (CuteDsl, DeepGemm, Marlin) take, but here it is a - real dependency: ``_route_to_cutlass`` sends every - NVFP4 prefill chunk through ``CutlassFusedMoE.quantize_input`` / + The only subclass of ``CutlassFusedMoE``, and the only backend for which + that is a real dependency rather than a shortcut: ``_route_to_cutlass`` + sends every NVFP4 prefill chunk through ``CutlassFusedMoE.quantize_input`` / ``CutlassFusedMoE.run_moe``, which read the whole Cutlass execution state (chunking stream and events, ``use_fused_finalize``, the tuner flags, the LoRA slot helpers, ``_tuner_shapes``, ``_run_moe_w4a16_nvfp4``). - Two ``CuteDslFusedMoE.__init__`` side effects are deliberately dropped, not - restated: the ``AuxStreamType.MoeOutputMemset`` / ``EventType`` entries, and - the ``swiglu_limit_scalar or inf`` fallback. Both are read only by - ``CuteDslFusedMoE.run_moe_nvfp4*``, which the ``run_moe`` override below - never reaches. Restate them before routing any CuteDSL path through this - class -- ``event_dict`` can now be None and ``swiglu_limit_scalar`` unset. + ``CuteDslFusedMoE.run_moe_nvfp4*`` is never reached from here, so the + ``AuxStreamType.MoeOutputMemset`` / ``EventType`` entries it needs are not + set up and ``event_dict`` can be None. Restate them, and + ``limit_when_absent`` below, before routing any CuteDSL path through this + class. """ - # 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) + # Inherited wholesale from CutlassFusedMoE, so every field is restated. + # No code path here reads ``w3_w1_bias`` / ``w2_bias`` or fuses LoRA, and + # ``supports_eplb`` stays False -- which is why ``can_implement`` has to + # decline ``d.eplb_enabled`` explicitly, since the inherited + # ``_supports_load_balancer()`` answers True. + # ``supports_apply_router_weight_on_input`` is False where the parent says + # True: only the NVFP4 prefill chunk reaches ``CutlassFusedMoE.run_moe``, + # while the decode path hands ``token_final_scales`` straight to the + # flashinfer b12x wrapper, which has no declared behaviour for the ``None`` + # the scheduler's fold leaves there. + capabilities = MoEStaticCapability( + supports_moe_lora=False, + supports_dwdp=True, + supports_expert_bias=False, + supports_apply_router_weight_on_input=False, + ) - # This and ``supports_moe_output_in_alltoall_workspace`` came through - # ``CuteDslFusedMoE`` before the reparent and Cutlass answers differently on - # both, so both are restated to keep the declared values unchanged. Read by - # ``ConfigurableMoE._reject_non_divisible_ep_backend()``; moot for this class - # in practice because ``can_implement`` rejects ``ep_size != 1`` outright. + # Same value the parent declares, pinned so a change there cannot silently + # retarget this backend. + input_requirement = MoEInputRequirement(routing_scales_dtype=torch.float32) + + # The kinds ``_ACTIVATION_MAP`` above gates on. The b12x decode kernel takes + # no activation constants, so none are declared. + activation_support = MoEActivationSupport( + kinds=frozenset({ActivationType.Swiglu, ActivationType.Relu2}) + ) + + # Read by ``ConfigurableMoE._reject_non_divisible_ep_backend()``; moot in + # practice because ``can_implement`` rejects ``ep_size != 1`` outright. _supports_non_divisible_ep: bool = True def supports_moe_output_in_alltoall_workspace(self) -> bool: @@ -139,6 +158,15 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: MoERejectReason.DEP_MISSING, "CuteDslB12xFusedMoE requires the flashinfer package", ) + # The only backend the construction-time allow-list turned down that + # never said so during selection, so an EPLB run resolved to b12x and + # then died in the factory. Declining here degrades with the usual + # warning instead, matching VanillaMoE / TritonFusedMoE / MarlinFusedMoE. + if d.eplb_enabled: + return _reject( + MoERejectReason.EPLB_UNSUPPORTED, + "CuteDslB12xFusedMoE does not support the MoE load balancer", + ) # No expert-parallel dispatch/combine kernel: EP must stay at 1. if d.ep_size != 1: return _reject( diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py index ce7a8b3a85b5..1f6dfaf93548 100755 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py @@ -27,6 +27,8 @@ from ...peft.lora.validation import has_moe_lora_targets from ...utils import (ActivationType, AuxStreamType, EventType, Fp4QuantizedTensor) +from .activation import (DEFAULT_MOE_ACTIVATION, ActivationParamShape, + MoEActivation, MoEActivationSupport) from .impl_base import MoEImplBase, apply_moe_impl_construction_state from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement, MoEProblem, MoERejectReason, MoERunContext, @@ -88,13 +90,31 @@ class CutlassFusedMoE(MoEImplBase): 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) + # Routed-expert MoE LoRA is fused into this backend's op only. Subclasses + # inherit this object wholesale, so they must restate every field. + capabilities = MoEStaticCapability( + supports_moe_lora=True, + supports_expert_bias=True, + supports_eplb=True, + supports_apply_router_weight_on_input=True) - # Inherited by every subclass, matching the isinstance check this replaces. input_requirement = MoEInputRequirement(routing_scales_dtype=torch.float32) + # The CUTLASS epilogue has an adaptor per kind (``moe_kernels.cuh``) and + # takes all three constants as ``float*`` indexed by expert. + activation_support = MoEActivationSupport( + kinds=frozenset({ + ActivationType.Swiglu, + ActivationType.SwigluBias, + ActivationType.Geglu, + ActivationType.SiTu, + ActivationType.Relu2, + ActivationType.Silu, + }), + alpha_beta=ActivationParamShape.PER_EXPERT_TENSOR, + limit=ActivationParamShape.PER_EXPERT_TENSOR, + ) + # Quantization algorithm support table for can_implement() # Format: quant_algo -> {sm_constraint, dtypes} # sm_constraint types: @@ -282,12 +302,8 @@ def __init__( bias: bool = False, apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, - swiglu_alpha: Optional[torch.Tensor] = None, - swiglu_beta: Optional[torch.Tensor] = None, - swiglu_limit: Optional[torch.Tensor] = None, - swiglu_limit_scalar: Optional[float] = None, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, init_load_balancer: bool = False, - activation_type: ActivationType = ActivationType.Swiglu, ): super().__init__(eplb=None) @@ -303,18 +319,13 @@ def __init__( aux_stream_dict=aux_stream_dict, weight_loading_mode=weight_loading_mode, bias=bias, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - swiglu_limit=swiglu_limit, - swiglu_limit_scalar=swiglu_limit_scalar, + activation=activation, layer_idx=layer_idx, init_load_balancer=init_load_balancer, - activation_type=activation_type, ) - # Store original hidden size before any potential padding - self.unpadded_hidden_size = self.hidden_size - + # ``unpadded_hidden_size`` is captured by + # apply_moe_impl_construction_state() above, before the padding below. if model_config.quant_config and model_config.quant_config.layer_quant_mode.has_w4a16_mxfp4( ): self.hidden_size = ((self.hidden_size + 127) // 128) * 128 @@ -1012,9 +1023,12 @@ def run_moe( ([self.fc2_weight_scale_2] if use_dynamic_fc2_scale else []), input_sf=x_sf, swizzled_input_sf=is_sf_swizzled, - swiglu_alpha=self.swiglu_alpha, - swiglu_beta=self.swiglu_beta, - swiglu_limit=self.swiglu_limit, + # ``swiglu_*`` are the moe_op schema's names for these registers + # (``ActivationParams`` in moe_kernels.h); SiTU fills the same three + # with tanh soft-caps. + swiglu_alpha=self.act_alpha, + swiglu_beta=self.act_beta, + swiglu_limit=self.act_clamp, tp_size=self.tp_size, tp_rank=self.tp_rank, ep_size=self.ep_size, @@ -1108,9 +1122,9 @@ def _run_moe_w4a16_nvfp4( quant_scales=[], input_sf=None, swizzled_input_sf=False, - swiglu_alpha=self.swiglu_alpha, - swiglu_beta=self.swiglu_beta, - swiglu_limit=self.swiglu_limit, + swiglu_alpha=self.act_alpha, + swiglu_beta=self.act_beta, + swiglu_limit=self.act_clamp, tp_size=self.tp_size, tp_rank=self.tp_rank, ep_size=self.ep_size, diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py index 6c32d00df4aa..e478dabf9ef8 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_deepgemm.py @@ -29,7 +29,9 @@ from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantAlgo -from .fused_moe_cutlass import CutlassFusedMoE +from .activation import (DEFAULT_MOE_ACTIVATION, ActivationParamShape, + MoEActivation, MoEActivationSupport) +from .impl_base import MoEImplBase, apply_moe_impl_construction_state from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement, MoEProblem, MoERejectReason, MoERunContext, MoEStaticCapability) @@ -749,7 +751,7 @@ def set_strides(workspace: torch.Tensor, g: int, m: int, k: int): return workspace -class DeepGemmFusedMoE(CutlassFusedMoE): +class DeepGemmFusedMoE(MoEImplBase): """DeepGEMM flow of fused mixture of experts (MoE) Layer. Args: @@ -763,23 +765,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) + capabilities = MoEStaticCapability( + supports_eplb=True, supports_apply_router_weight_on_input=True) - # ``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, ) + # The DeepGEMM Triton activation kernel implements SwiGLU only, and takes + # the clamp by value -- a per-expert tensor would never be read. + activation_support = MoEActivationSupport( + kinds=frozenset({ActivationType.Swiglu}), + limit=ActivationParamShape.UNIFORM_SCALAR, + ) + 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. + # ``run_moe`` emits into its own workspace buffers and never writes a + # caller-supplied output tensor, so nothing would fill a workspace-backed + # buffer that ``combine()`` then reads. return False @classmethod @@ -846,8 +850,7 @@ def __init__( VANILLA, apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, - swiglu_limit: Optional[torch.Tensor] = None, - swiglu_limit_scalar: Optional[float] = None, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, init_load_balancer: bool = False, ): # moe_max_num_tokens is set in ModelConfig.__post_init__ if not specified @@ -863,7 +866,9 @@ def __init__( # does not size the DeepGEMM workspace deliberately. _configure_deepgemm_moe_max_num_tokens(model_config) - super().__init__( + super().__init__(eplb=None) + apply_moe_impl_construction_state( + self, routing_method=routing_method, num_experts=num_experts, hidden_size=hidden_size, @@ -873,12 +878,23 @@ def __init__( model_config=model_config, aux_stream_dict=aux_stream_dict, weight_loading_mode=weight_loading_mode, - apply_router_weight_on_input=apply_router_weight_on_input, layer_idx=layer_idx, - swiglu_limit=swiglu_limit, - swiglu_limit_scalar=swiglu_limit_scalar, + activation=activation, init_load_balancer=init_load_balancer, ) + self.apply_router_weight_on_input = apply_router_weight_on_input + + self._weights_created = False + if not model_config.skip_create_weights_in_init: + self.create_weights() + + def _supports_load_balancer(self) -> bool: + return True + + def _check_configs(self): + assert self._weights_created + if self.apply_router_weight_on_input: + assert self.routing_method.top_k == 1, "Current walkaround only supports top-1 routing" def get_workspace(self, m_max: int, group_size: int): capture_graph = torch.cuda.is_current_stream_capturing() @@ -1113,7 +1129,7 @@ def run_moe( quant_group_size=128, masked_m=masked_m, scale_ue8m0=True, - swiglu_limit=self.swiglu_limit_scalar) + swiglu_limit=self.act_clamp) # Grouped gemm 2 h3 = set_strides(workspace["workspace_1"], diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py index 199ba8accfc5..a2104fc1cf10 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py @@ -20,6 +20,7 @@ swizzle_sf, unswizzle_sf, ) +from .activation import DEFAULT_MOE_ACTIVATION, MoEActivation, MoEActivationSupport from .impl_base import MoEImplBase, apply_moe_impl_construction_state from .impl_contract import ( MoEDeployment, @@ -28,6 +29,7 @@ MoEProblem, MoERejectReason, MoERunContext, + MoEStaticCapability, require_comm_plan, ) from .interface import MoEWeightLoadingMode, _reject @@ -123,8 +125,15 @@ class DenseGEMMFusedMoE(MoEImplBase): model_config (ModelConfig): Configuration object for the model. """ + # Declared because the default is conservative: this backend registers its + # weights with the load balancer, matching ``_supports_load_balancer``. + capabilities = MoEStaticCapability(supports_eplb=True) + input_requirement = MoEInputRequirement(routing_scales_dtype=torch.float32) + # The dense-GEMM epilogue fuses plain SwiGLU and takes no constants. + activation_support = MoEActivationSupport(kinds=frozenset({ActivationType.Swiglu})) + # Memory buffer pool for CUDA graph compatibility buffers = get_memory_buffers() @@ -194,8 +203,8 @@ def __init__( weight_loading_mode: MoEWeightLoadingMode = MoEWeightLoadingMode.VANILLA, apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, init_load_balancer: bool = False, - activation_type: ActivationType = ActivationType.Swiglu, ): # Eligibility (SM / quant / SwiGLU / EP / intermediate alignment) is # owned by ``can_implement``; do not re-assert it here. @@ -225,8 +234,8 @@ def __init__( aux_stream_dict=aux_stream_dict, weight_loading_mode=weight_loading_mode, layer_idx=layer_idx, + activation=activation, init_load_balancer=init_load_balancer, - activation_type=activation_type, ) # Environment variable to control fc2_alpha fusion into FC1's alpha_post. diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_marlin.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_marlin.py index ff3551875d07..74d1ddc07524 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_marlin.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_marlin.py @@ -20,13 +20,15 @@ No activation quantization overhead. In-kernel topk_weights multiplication. """ -from typing import Optional, Tuple +from typing import Dict, Optional, Tuple import torch import torch.nn.functional as F +from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.utils import ( ActivationType, + AuxStreamType, Fp4QuantizedTensor, is_gated_activation, is_nvfp4_marlin_supported_sm, @@ -34,17 +36,20 @@ ) from tensorrt_llm.models.modeling_utils import QuantAlgo -from .fused_moe_cutlass import CutlassFusedMoE +from .activation import DEFAULT_MOE_ACTIVATION, MoEActivation, MoEActivationSupport +from .impl_base import MoEImplBase, apply_moe_impl_construction_state from .impl_contract import ( MoEDeployment, MoEEligibility, + MoEInputRequirement, MoEProblem, MoERejectReason, MoERunContext, MoEStaticCapability, ) from .interface import _reject -from .quantization import NVFP4MarlinFusedMoEMethod +from .quantization import MoEWeightLoadingMode, NVFP4MarlinFusedMoEMethod +from .routing import BaseMoeRoutingMethod # Block size for moe_align_block_size — must match TILE_M in the kernel _MOE_BLOCK_SIZE = 16 @@ -55,7 +60,7 @@ def _has_fused_moe_kernel() -> bool: return hasattr(torch.ops.trtllm, "marlin_nvfp4_moe_gemm") -class MarlinFusedMoE(CutlassFusedMoE): +class MarlinFusedMoE(MoEImplBase): """MoE backend using Marlin W4A16 NVFP4 GEMM for SM89-SM99. Uses ``marlin_nvfp4_moe_gemm`` with BF16 activations to process all experts @@ -64,9 +69,18 @@ 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) + # Sorted-token dispatch has no EPLB slot layout, which is also why + # ``can_implement`` rejects ``d.eplb_enabled`` below. + capabilities = MoEStaticCapability(supports_apply_router_weight_on_input=True) + + input_requirement = MoEInputRequirement(routing_scales_dtype=torch.float32) + + # The Marlin epilogue takes no activation constants, and + # ``_apply_activation`` only distinguishes these three -- any other gated + # kind would silently run SiLU, any other non-gated kind ReLU. + activation_support = MoEActivationSupport( + kinds=frozenset({ActivationType.Swiglu, ActivationType.Geglu, ActivationType.Relu2}) + ) _QUANT_SUPPORT_TABLE = { QuantAlgo.NVFP4: { @@ -118,6 +132,58 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: return MoEEligibility.ok() + def __init__( + self, + *, + routing_method: BaseMoeRoutingMethod, + num_experts: int, + hidden_size: int, + intermediate_size: int, + dtype: Optional[torch.dtype] = None, + reduce_results: bool = False, + model_config: ModelConfig = ModelConfig(), + aux_stream_dict: Optional[Dict[AuxStreamType, torch.cuda.Stream]] = None, + weight_loading_mode: MoEWeightLoadingMode = MoEWeightLoadingMode.VANILLA, + bias: bool = False, + apply_router_weight_on_input: bool = False, + layer_idx: Optional[int] = None, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, + init_load_balancer: bool = False, + ): + """Construct the backend. + + ``bias`` is accepted because ``create_moe`` passes it on the branch this + class shares with CutlassFusedMoE, but is always False here: the factory + rejects it against ``capabilities.supports_expert_bias``. + """ + super().__init__(eplb=None) + apply_moe_impl_construction_state( + self, + routing_method=routing_method, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + reduce_results=reduce_results, + model_config=model_config, + aux_stream_dict=aux_stream_dict, + weight_loading_mode=weight_loading_mode, + bias=bias, + layer_idx=layer_idx, + activation=activation, + init_load_balancer=init_load_balancer, + ) + self.apply_router_weight_on_input = apply_router_weight_on_input + + self._weights_created = False + if not model_config.skip_create_weights_in_init: + self.create_weights() + + def _check_configs(self): + assert self._weights_created + if self.apply_router_weight_on_input: + assert self.routing_method.top_k == 1, "Current walkaround only supports top-1 routing" + def quantize_input( self, x: torch.Tensor | Fp4QuantizedTensor, post_quant_comm: bool = True, **kwargs ) -> Tuple[torch.Tensor, torch.Tensor | None]: @@ -131,9 +197,6 @@ def _get_quant_method(self): return NVFP4MarlinFusedMoEMethod() raise ValueError(f"MarlinFusedMoE only supports NVFP4, got {self.quant_config}") - def _supports_load_balancer(self) -> bool: - return False - def _apply_activation(self, gemm1_out: torch.Tensor) -> torch.Tensor: """Apply the activation function to the gemm1 output. @@ -174,9 +237,8 @@ def _ensure_workspace(self, device: torch.device): # ==================================================================== 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. + # This kernel always allocates and returns its own output tensor, so + # nothing would fill a workspace-backed buffer that ``combine()`` reads. return False def run_moe( diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py index ac7ff83a586d..bfc9738d69cf 100755 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_triton.py @@ -41,8 +41,11 @@ from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm.models.modeling_utils import QuantAlgo +from .activation import (DEFAULT_MOE_ACTIVATION, ActivationParamShape, + MoEActivation, MoEActivationSupport, + install_activation_params) from .impl_contract import (MoEDeployment, MoEEligibility, MoEProblem, - MoERejectReason) + MoERejectReason, MoEStaticCapability) from .interface import MoE, _reject from .quantization import (FusedMoEMethodBase, MoEWeightLoadingMode, load_activation_scales_fp8_qdq, @@ -51,6 +54,19 @@ RenormalizeMoeRoutingMethod) +def _swiglu_scalars(module: nn.Module) -> tuple[float, float]: + """SwiGLU ``alpha`` / ``beta``, with the kernel-neutral value for an absent one. + + ``is None`` rather than ``or``: 0.0 is a legal value for both registers, and + ``or`` cannot tell "the caller passed zero" from "the caller passed nothing". + ``beta`` in particular selects the fused-activation path at every call site, + so conflating the two silently reroutes the layer. + """ + alpha = 1.0 if module.act_alpha is None else float(module.act_alpha) + beta = 0.0 if module.act_beta is None else float(module.act_beta) + return alpha, beta + + # Triton kernels has hardcoded beta = 1, so we use this implementation when beta is not 1 def swiglu_torch(a: torch.Tensor, alpha: float, beta: float, limit: Optional[float]) -> torch.Tensor: @@ -469,13 +485,12 @@ def apply(self, out_dtype=module.dtype) # Call the Triton gemm kernel, which also does permutation and activation - alpha = module.swiglu_alpha or 1.0 - beta = module.swiglu_beta or 0.0 + alpha, beta = _swiglu_scalars(module) if beta == 1.0: act = FusedActivation( FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit"), - reduction_n=2), (alpha, module.swiglu_limit)) + reduction_n=2), (alpha, module.act_clamp)) act_out = matmul( hidden_states, gemm1_weights, @@ -492,7 +507,7 @@ def apply(self, a_ragged_metadata=rdata.ragged_metadata if rdata else None, gather_indx=gather_indx, precision_config=pc1) - act_out = swiglu_torch(act_out, alpha, beta, module.swiglu_limit) + act_out = swiglu_torch(act_out, alpha, beta, module.act_clamp) # Step 3: Gemm2 # Setup quantization context @@ -714,13 +729,12 @@ def apply(self, out_dtype=module.dtype) # Call the Triton gemm kernel, which also does permutation and activation - alpha = module.swiglu_alpha or 1.0 - beta = module.swiglu_beta or 0.0 + alpha, beta = _swiglu_scalars(module) if beta == 1.0: act = FusedActivation( FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit"), - reduction_n=2), (alpha, module.swiglu_limit)) + reduction_n=2), (alpha, module.act_clamp)) act_out = matmul( hidden_states, gemm1_weights, @@ -737,7 +751,7 @@ def apply(self, a_ragged_metadata=rdata.ragged_metadata if rdata else None, gather_indx=gather_indx, precision_config=pc1) - act_out = swiglu_torch(act_out, alpha, beta, module.swiglu_limit) + act_out = swiglu_torch(act_out, alpha, beta, module.act_clamp) # Quantize the activation output manually since the Triton activation kernel doesn't support bf16 in fp8 out act_out, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( @@ -1454,14 +1468,13 @@ def _maybe_pad_activation(hidden_states): out_dtype=module.dtype) # Call the Triton gemm kernel, which also does permutation and activation - alpha = module.swiglu_alpha or 1.0 - beta = module.swiglu_beta or 0.0 + alpha, beta = _swiglu_scalars(module) hidden_states = _maybe_pad_activation(hidden_states) if beta == 1.0: act = FusedActivation( FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit"), - reduction_n=2), (alpha, module.swiglu_limit)) + reduction_n=2), (alpha, module.act_clamp)) act_out = matmul( hidden_states, @@ -1479,7 +1492,7 @@ def _maybe_pad_activation(hidden_states): a_ragged_metadata=rdata.ragged_metadata if rdata else None, gather_indx=gather_indx, precision_config=pc1) - act_out = swiglu_torch(act_out, alpha, beta, module.swiglu_limit) + act_out = swiglu_torch(act_out, alpha, beta, module.act_clamp) if self.activation_dtype == torch.float8_e4m3fn: # Quantize the activation output manually since the Triton activation kernel doesn't support bf16 in fp8 out @@ -1550,6 +1563,17 @@ def transform_weights(self, module: torch.nn.Module) -> None: class TritonFusedMoE(MoE): + capabilities = MoEStaticCapability(supports_expert_bias=True) + + # The Triton epilogue always runs the OAI gated SwiGLU functor, with the + # constants baked into the kernel launch and defaulting to + # alpha=1.0 / beta=0.0 (plain SwiGLU) when the activation carries none. + activation_support = MoEActivationSupport( + kinds=frozenset({ActivationType.Swiglu, ActivationType.SwigluBias}), + alpha_beta=ActivationParamShape.UNIFORM_SCALAR, + limit=ActivationParamShape.UNIFORM_SCALAR, + ) + @classmethod def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: """Triton MoE: SM90 only, SwiGLU-family activations only. @@ -1645,9 +1669,7 @@ def __init__( VANILLA, bias: bool = False, layer_idx: Optional[int] = None, - swiglu_alpha: Optional[torch.Tensor] = None, - swiglu_beta: Optional[torch.Tensor] = None, - swiglu_limit: Optional[torch.Tensor] = None, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, ): super().__init__( routing_method=routing_method, @@ -1659,6 +1681,7 @@ def __init__( model_config=model_config, weight_loading_mode=weight_loading_mode, layer_idx=layer_idx, + activation=activation, ) # Eligibility (SM / routing / smart_router / quant) is owned by # ``can_implement``; do not re-assert it here. @@ -1679,21 +1702,9 @@ def __init__( self.bias = bias - def _maybe_squeeze_act_param(p): - if p is None or isinstance(p, (int, float)): - return p - assert isinstance(p, torch.Tensor) - assert p.dtype == torch.float32 - assert p.shape == (self.expert_size_per_partition, ), p.shape - assert torch.all( - p == p[0] - ), "All experts must have the same swiglu alpha/beta for Triton kernel" - p = p[0].item() - return p - - self.swiglu_alpha = _maybe_squeeze_act_param(swiglu_alpha) - self.swiglu_beta = _maybe_squeeze_act_param(swiglu_beta) - self.swiglu_limit = _maybe_squeeze_act_param(swiglu_limit) + # Reduces the constants to the per-launch scalars this kernel bakes, and + # rejects a per-expert one instead of silently using expert 0's value. + install_activation_params(self) self._weights_created = False if not model_config.skip_create_weights_in_init: diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py index 2888daae3916..b51894bce1fc 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py @@ -14,7 +14,7 @@ # limitations under the License. import os -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Dict, List, Optional, Union import torch @@ -30,10 +30,13 @@ from ...modules.gated_mlp import GatedMLP from ...utils import (ActivationType, ActType_TrtllmGen, AuxStreamType, Fp4QuantizedTensor, MxFp8QuantizedTensor) +from .activation import (DEFAULT_MOE_ACTIVATION, ActivationParamShape, + MoEActivation, MoEActivationSupport, + materialize_activation_params) from .impl_base import MoEImplBase, apply_moe_impl_construction_state from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement, MoEProblem, MoERejectReason, MoERunContext, - require_comm_plan) + MoEStaticCapability, require_comm_plan) from .impl_environment import MoEDep from .interface import FORCE_SEPARATED_ROUTING, MoEWeightLoadingMode, _reject from .moe_op_backend import MoEOpBackend, TRTLLMOpBackend, get_op_backend @@ -91,6 +94,9 @@ class TRTLLMGenFusedMoE(MoEImplBase): 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. """ + capabilities = MoEStaticCapability(supports_expert_bias=True, + supports_eplb=True) + # bfloat16 routing scales are what these kernels read, and the DeepEP # dispatch has to mark unfilled rows before they reach them. input_requirement = MoEInputRequirement( @@ -112,9 +118,8 @@ class TRTLLMGenFusedMoE(MoEImplBase): } # Quantization algorithms that support full swiglu_gptoss_style. - # FP8_BLOCK_SCALES supports the DSV4-style uniform swiglu_limit_scalar - # through the DeepSeek FP8 separate-activation path, but not bias or - # swiglu_alpha/beta. + # FP8_BLOCK_SCALES is absent on purpose: its separate-activation path takes + # the DSV4-style scalar clamp, but no bias and no alpha/beta. _GPTOSS_SUPPORTED_ALGOS = { QuantAlgo.NVFP4, QuantAlgo.W4A16_MXFP4, @@ -138,6 +143,51 @@ class TRTLLMGenFusedMoE(MoEImplBase): QuantAlgo.W4A8_MXFP4_MXFP8, } + # The fused-activation cubins index alpha/beta/clamp by expert + # (``gemm1_alpha`` / ``gemm1_beta`` / per-expert clamp tensor). The clamp is + # quant-dependent, so ``resolve_activation_support`` narrows it per instance. + activation_support = MoEActivationSupport( + kinds=frozenset({ + ActivationType.Swiglu, + ActivationType.SwigluBias, + ActivationType.Relu2, + ActivationType.Silu, + ActivationType.SiTu, + }), + alpha_beta=ActivationParamShape.PER_EXPERT_TENSOR, + limit=ActivationParamShape.PER_EXPERT_TENSOR, + ) + + # ActivationType -> the batched-GEMM ``ActType`` encoding the cubins are + # keyed by. SwigluBias shares the SwiGlu kernel (the gpt-oss constants + # travel as separate per-expert tensors, not as a distinct act type). + _TRTLLM_GEN_ACT_TYPE = { + ActivationType.Swiglu: ActType_TrtllmGen.SwiGlu, + ActivationType.SwigluBias: ActType_TrtllmGen.SwiGlu, + ActivationType.Relu2: ActType_TrtllmGen.Relu2, + ActivationType.Silu: ActType_TrtllmGen.Silu, + ActivationType.SiTu: ActType_TrtllmGen.SiTu, + } + + def resolve_activation_support(self) -> MoEActivationSupport: + """Narrow the clamp ABI to what this instance's quant path reads. + + The DeepSeek FP8 block-scale path runs the clamp in a separate + activation kernel (``DevKernel.cu::activationDeepSeekKernel``) that + takes one ``double`` by value, so a per-expert tensor would be silently + ignored there. Every other quant path consumes the per-expert tensor + the class attribute declares. + """ + support = type(self).activation_support + # Reads ``quant_config`` rather than ``has_deepseek_fp8_block_scales``: + # this runs while construction state is being installed, before + # create_weights sets the ``_weights_created`` those properties assert. + quant_config = self.quant_config + if (quant_config is not None + and quant_config.layer_quant_mode.has_fp8_block_scales()): + return replace(support, limit=ActivationParamShape.UNIFORM_SCALAR) + return support + @classmethod def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: """TRTLLM-Gen kernels: the SM100 (Blackwell) family, bfloat16 activations. @@ -246,15 +296,8 @@ def __init__( VANILLA, layer_idx: Optional[int] = None, bias: bool = False, - swiglu_alpha: Optional[torch.Tensor] = None, - swiglu_beta: Optional[torch.Tensor] = None, - swiglu_limit: Optional[torch.Tensor] = None, - swiglu_limit_scalar: Optional[float] = None, init_load_balancer: bool = False, - activation_type: ActivationType = ActivationType.Swiglu, - trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None, - trtllm_gen_activation_alpha: Optional[float] = None, - trtllm_gen_activation_beta: Optional[float] = None, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, ): super().__init__(eplb=None) apply_moe_impl_construction_state( @@ -269,21 +312,12 @@ def __init__( aux_stream_dict=aux_stream_dict, weight_loading_mode=weight_loading_mode, bias=bias, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - swiglu_limit=swiglu_limit, - swiglu_limit_scalar=swiglu_limit_scalar, layer_idx=layer_idx, init_load_balancer=init_load_balancer, - activation_type=activation_type, + activation=activation, ) - self.trtllm_gen_activation_type = ( - ActType_TrtllmGen(trtllm_gen_activation_type) - if trtllm_gen_activation_type is not None else None) - self.trtllm_gen_activation_alpha = trtllm_gen_activation_alpha - self.trtllm_gen_activation_beta = trtllm_gen_activation_beta - self._validate_backend_local_activation() + self._validate_situ_activation() # Cached for autotune profile sizing (forward path passes # tune_max_num_tokens to the MoE op). @@ -337,43 +371,27 @@ def __init__( def _to_trtllm_gen_activation_type(self, activation_type: ActivationType) -> int: - if self.trtllm_gen_activation_type is not None: - return int(self.trtllm_gen_activation_type) - if activation_type == ActivationType.Swiglu: - return 0 - elif activation_type == ActivationType.SwigluBias: - # SwigluBias uses the same SwiGlu kernel path (ActType::SwiGlu == 0); - # the per-expert alpha/beta/clamp_limit are passed as separate tensors. - return 0 - elif activation_type == ActivationType.Relu2: - return 1 - elif activation_type == ActivationType.Silu: - return 2 - else: + act_type = self._TRTLLM_GEN_ACT_TYPE.get( + ActivationType(activation_type)) + if act_type is None: raise ValueError(f"Unsupported activation type: {activation_type}") + return int(act_type) @property def is_situ_activation(self) -> bool: - return self.trtllm_gen_activation_type == ActType_TrtllmGen.SiTu - - def _validate_backend_local_activation(self) -> None: - # Runs from __init__, before create_weights, so the swiglu_* attributes - # checked below are still the constructor-provided values. For SiTu, - # create_weights later reuses the swiglu_alpha/swiglu_beta storage for - # the backend-local activation parameters (SiTu and SwiGLU are mutually - # exclusive and feed the same gemm1_alpha/gemm1_beta op slots). - if self.trtllm_gen_activation_type is None: - if (self.trtllm_gen_activation_alpha is not None - or self.trtllm_gen_activation_beta is not None): - raise ValueError( - "TRTLLM-Gen backend-local activation alpha/beta require " - "trtllm_gen_activation_type.") - return + return self.activation.kind is ActivationType.SiTu + def _validate_situ_activation(self) -> None: + """Hardware / quant preconditions the SiTu cubins carry. + + Kind and constant shape are already settled: the activation carrier only + admits two positive soft-caps for SiTu and no clamp, and + ``install_activation_params`` has materialized them as the per-expert + ``gemm1_alpha`` / ``gemm1_beta`` buffers the cubin indexes. What remains + is that trtllm-gen ships SiTu for exactly one dtype/quant combination. + """ if not self.is_situ_activation: - raise ValueError( - "Only the SiTu TRTLLM-Gen backend-local activation is " - f"supported, got {self.trtllm_gen_activation_type.name}.") + return if self.dtype != torch.bfloat16: raise ValueError( "TRTLLM-Gen SiTu requires bfloat16 activations, got " @@ -416,28 +434,10 @@ def _validate_backend_local_activation(self) -> None: f"({self.tp_size}) with the per-rank shard a multiple of " f"{alignment}, got " f"{self.intermediate_size_per_partition}.") - if self.activation_type != ActivationType.Swiglu: - raise ValueError( - "TRTLLM-Gen SiTu must use generic SwiGLU geometry so FC1 " - "contains gate and up projections.") - if self.bias or any( - value is not None - for value in (self.swiglu_alpha, self.swiglu_beta, - self.swiglu_limit, self.swiglu_limit_scalar)): + if self.bias: raise ValueError( - "TRTLLM-Gen SiTu does not support bias or SwiGLU-specific " - "alpha/beta/limit parameters.") - if (self.trtllm_gen_activation_alpha is None - or self.trtllm_gen_activation_beta is None): - raise ValueError( - "TRTLLM-Gen SiTu requires both backend-local activation " - "alpha and beta.") - if (self.trtllm_gen_activation_alpha <= 0.0 - or self.trtllm_gen_activation_beta <= 0.0): - raise ValueError( - "TRTLLM-Gen SiTu activation alpha/beta must be positive, got " - f"{self.trtllm_gen_activation_alpha} and " - f"{self.trtllm_gen_activation_beta}.") + "TRTLLM-Gen SiTu does not support expert bias; the cubin adds " + "no FC1 bias before the soft-caps.") @staticmethod def _is_flashinfer_fused_moe_available() -> bool: @@ -549,29 +549,19 @@ def _check_configs(self): ("TRTLLMGenFusedMoE BF16 path only supports " f"{[a.name for a in self._BF16_SUPPORTED_ACTIVATIONS]} activations, " f"got {self.activation_type.name}.") - assert not self.bias and self.swiglu_alpha is None and self.swiglu_beta is None and self.swiglu_limit is None, \ + assert not self.bias and self.act_alpha is None and self.act_beta is None and self.act_clamp is None, \ "TRTLLMGenFusedMoE BF16 path does not support bias/swiglu custom parameters." - if self.bias or self.swiglu_alpha is not None or self.swiglu_beta is not None: + if self.bias or self.act_alpha is not None or self.act_beta is not None: assert self.has_nvfp4 or self.has_w4a16_mxfp4 or self.has_w4a8_mxfp4_fp8 or self.has_w4a8_mxfp4_mxfp8, \ - "TRTLLMGenFusedMoE supports bias/swiglu_alpha/swiglu_beta only for nvfp4 and mxfp4 variants." - if self.swiglu_limit is not None or self.swiglu_limit_scalar is not None: - # swiglu_limit additionally goes through the DeepSeek FP8 - # separate-activation path - # (DevKernel.cu::activationDeepSeekKernel) when - # has_deepseek_fp8_block_scales. The FP8 path consumes the scalar - # variant (uniform across experts); NVFP4/MXFP4 fused-activation - # cubins consume the per-expert tensor. + "TRTLLMGenFusedMoE supports bias and the alpha/beta activation constants only for nvfp4 and mxfp4 variants." + if self.act_clamp is not None: + # Whether the clamp arrives as a scalar or a per-expert tensor is + # settled by ``resolve_activation_support``; which algorithms have a + # clamp at all is a quant fact, so it stays here. assert self.has_nvfp4 or self.has_w4a16_mxfp4 or self.has_w4a8_mxfp4_fp8 \ or self.has_w4a8_mxfp4_mxfp8 or self.has_deepseek_fp8_block_scales, \ - "TRTLLMGenFusedMoE supports swiglu_limit only for nvfp4, mxfp4, and fp8_block_scale variants." - # The FP8 block-scale separate-activation kernel only consumes the - # uniform scalar (swiglu_limit_scalar); a per-expert swiglu_limit - # tensor would be silently ignored, so reject it explicitly. - if self.has_deepseek_fp8_block_scales: - assert self.swiglu_limit is None, \ - "TRTLLMGenFusedMoE FP8 block-scale path only supports the uniform " \ - "swiglu_limit_scalar, not a per-expert swiglu_limit tensor." + "TRTLLMGenFusedMoE supports an activation clamp only for nvfp4, mxfp4, and fp8_block_scale variants." if self.is_situ_activation: if not isinstance(self.op_backend, TRTLLMOpBackend): @@ -588,7 +578,7 @@ def _check_configs(self): f"mode, got {self.scaling_vector_size}.") # For SiTu these hold the backend-local activation parameters # (populated by create_weights, which runs before this check). - for name in ("swiglu_alpha", "swiglu_beta"): + for name in ("act_alpha", "act_beta"): value = getattr(self, name) if (value.dtype != torch.float32 or value.shape != (self.expert_size_per_partition, ) @@ -603,15 +593,15 @@ def _get_quant_method(self): if self.quant_config.layer_quant_mode.has_fp8_block_scales(): return DeepSeekFP8BlockScalesFusedMoEMethod() elif self.quant_config.layer_quant_mode.has_nvfp4(): - # ``is_situ_activation`` (not ``swiglu_alpha is not None``): - # SiTu fills the swiglu_alpha/swiglu_beta slots from + # ``is_situ_activation`` (not ``act_alpha is not None``): + # SiTu fills the act_alpha/act_beta slots from # create_weights, i.e. after this runs, so keying off the # tensor would make the selected method depend on *when* # _get_quant_method is called. Like the SwiGLU-alpha and # element-wise cases, SiTu needs the padded method's # alignment handling. needs_padded_method = ( - self.swiglu_alpha is not None or self.is_situ_activation + self.act_alpha is not None or self.is_situ_activation or self.activation_type in [ActivationType.Relu2, ActivationType.Silu]) return (NVFP4TRTLLMGenFusedMoEMethod() if needs_padded_method @@ -642,27 +632,13 @@ def create_weights(self): else: self.quant_method.create_weights(self) - # SiTu reuses the swiglu_alpha/swiglu_beta storage: SiTu and SwiGLU are - # mutually exclusive (constructor-provided SwiGLU parameters are - # rejected by _validate_backend_local_activation) and feed the same - # gemm1_alpha/gemm1_beta op slots. Safe with respect to the - # `swiglu_alpha is not None` gates: create_moe.py checks the - # constructor kwargs (None for SiTu); _get_quant_method's nvfp4 branch - # keys off is_situ_activation rather than swiglu_alpha, so it returns - # the same method before and after this point; _check_configs runs - # after this point and its swiglu gate admits both nvfp4 and - # w4a8_mxfp4_mxfp8. + # SiTu's two soft-caps ride in the same gemm1_alpha/gemm1_beta op slots + # SwiGLU's alpha/beta use -- the kinds are mutually exclusive. They are + # backend configuration, not checkpoint weights, so ``cache_derived_state`` + # refills them if meta-device materialization wipes them. if self.is_situ_activation: - self.swiglu_alpha = nn.Parameter(torch.full( - (self.expert_size_per_partition, ), - float(self.trtllm_gen_activation_alpha), - dtype=torch.float32), - requires_grad=False) - self.swiglu_beta = nn.Parameter(torch.full( - (self.expert_size_per_partition, ), - float(self.trtllm_gen_activation_beta), - dtype=torch.float32), - requires_grad=False) + self.act_alpha = nn.Parameter(self.act_alpha, requires_grad=False) + self.act_beta = nn.Parameter(self.act_beta, requires_grad=False) self._weights_created = True self._check_configs() @@ -685,12 +661,25 @@ def cache_derived_state(self) -> None: super().cache_derived_state() if self.is_situ_activation: # Reinitialize constants after meta-device materialization. These - # are backend configuration, not checkpoint weights. - self.swiglu_alpha.data.fill_(float( - self.trtllm_gen_activation_alpha)) - self.swiglu_beta.data.fill_(float(self.trtllm_gen_activation_beta)) + # are backend configuration, not checkpoint weights, so nothing + # else refills them. + # + # Re-materialized from ``self.activation`` rather than from a + # snapshot of the slots taken in create_weights: under meta init + # those slots are themselves meta at that point, so the snapshot + # would carry no values. ``self.activation`` is the declaration and + # is always real. + params = materialize_activation_params( + self.activation, + self.resolve_activation_support(), + num_local_experts=self.expert_size_per_partition, + device=self.act_alpha.device, + owner=type(self).__name__, + ) + self.act_alpha.data.copy_(params.alpha) + self.act_beta.data.copy_(params.beta) - def try_fused_kimi_route_quant( + def try_fused_route_quant( self, x: Union[torch.Tensor, MxFp8QuantizedTensor], router_logits: torch.Tensor, @@ -699,8 +688,11 @@ def try_fused_kimi_route_quant( """Fuse Kimi K3 no-aux routing and MXFP8 input quantization. This launch-overhead optimization is deliberately specialized to the - K3 decode shape. Returning ``None`` keeps every other model, shape, - architecture, and op backend on the existing unfused path. + K3 decode shape: the op below hardcodes 896 experts, top-16, hidden + 3584 and at most 64 tokens. The checks here mirror its ``TORCH_CHECK``s + so a miss declines quietly instead of raising. Returning ``None`` keeps + every other model, shape, architecture, and op backend on the existing + unfused path. """ if (os.environ.get("TLLM_K3_DISABLE_FUSED_ROUTE_QUANT", "0") == "1" or isinstance(x, MxFp8QuantizedTensor)): @@ -971,7 +963,7 @@ def run_moe( self.routing_method.routing_method_type, topk_weights=token_final_scales, topk_ids=token_selected_experts, - gemm1_clamp_limit=self.swiglu_limit_scalar, + gemm1_clamp_limit=self.act_clamp, output=moe_output, tune_max_num_tokens=self.max_num_tokens, use_dp=self.use_dp, @@ -988,7 +980,7 @@ def run_moe( # Holds SwiGLU's per-expert alpha/beta, or SiTu's backend-local # activation parameters (which reuse this storage; see # create_weights). - gemm1_alpha, gemm1_beta = self.swiglu_alpha, self.swiglu_beta + gemm1_alpha, gemm1_beta = self.act_alpha, self.act_beta output1_scale_scalar = self._get_data_or_none("fc31_scale_c") output1_scale_gate_scalar = self._get_data_or_none("fc31_alpha") @@ -1004,7 +996,7 @@ def run_moe( self.w3_w1_bias if self.bias else None, gemm1_alpha, gemm1_beta, - self.swiglu_limit, + self.act_clamp, self.w2_weight, self.w2_weight_scale, self.w2_bias if self.bias else None, @@ -1096,9 +1088,9 @@ def run_moe( self.w3_w1_weight, self.w3_w1_weight_scale, self.w3_w1_bias, - self.swiglu_alpha, - self.swiglu_beta, - self.swiglu_limit, + self.act_alpha, + self.act_beta, + self.act_clamp, self.w2_weight, self.w2_weight_scale, self.w2_bias, diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_vanilla.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_vanilla.py index 5c3d4b5375c9..203672da9176 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_vanilla.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_vanilla.py @@ -13,6 +13,8 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.utils import fp4_utils +from .activation import (DEFAULT_MOE_ACTIVATION, MoEActivation, + MoEActivationSupport, install_activation_params) from .impl_contract import (MoEDeployment, MoEEligibility, MoEProblem, MoERejectReason, MoEStaticCapability) from .interface import MoEWeightLoadingMode, _reject @@ -24,6 +26,11 @@ class VanillaMoE(nn.ModuleList): #: Declared explicitly because the resolver may return this ModuleList. capabilities = MoEStaticCapability() + # What ``forward_chunk`` executes: the gated arm via ``GatedMLP`` (SiLU), + # and the non-gated Relu2 arm. + activation_support = MoEActivationSupport( + kinds=frozenset({ActivationType.Swiglu, ActivationType.Relu2})) + #: Quantization labels supported by the Linear dispatcher. _SUPPORTED_QUANT_LABELS = frozenset({ "FP8", @@ -98,7 +105,7 @@ def __init__( apply_router_weight_on_input: bool = False, pack_weights: bool = False, layer_idx: Optional[int] = None, - activation_type: ActivationType = ActivationType.Swiglu, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, ): from tensorrt_llm._torch.distributed import AllReduce @@ -111,8 +118,9 @@ def __init__( self.pack_weights = pack_weights self.layer_idx = layer_idx - self.activation_type = activation_type - self.is_gated_activation = is_gated_activation(activation_type) + self.activation = activation + self.activation_type = ActivationType(activation.kind) + self.is_gated_activation = is_gated_activation(activation.kind) # Activation eligibility (gated vs Relu2-only non-gated) is owned by # ``can_implement``. ``pack_weights`` is a construction option that is # not part of the problem/deployment question, so keep it here. @@ -159,6 +167,8 @@ def __init__( self.num_experts) self.expert_size_per_partition = self.expert_end - self.expert_start + install_activation_params(self) + # moe_max_num_tokens is set in ModelConfig.__post_init__ if not specified # The default value is max_num_tokens * dp_size self.moe_max_num_tokens = model_config.moe_max_num_tokens @@ -167,7 +177,10 @@ def __init__( if not model_config.skip_create_weights_in_init: self.create_weights() - # If True, the router weight will be multiplied on the input rather than at the end of FC2 + # Stored but never read: no path here folds the router weight into the + # input, which is why ``capabilities`` leaves + # ``supports_apply_router_weight_on_input`` at False and the factory + # rejects a True before construction. self.apply_router_weight_on_input = apply_router_weight_on_input def create_experts(self, module_list: nn.ModuleList = None): diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_base.py b/tensorrt_llm/_torch/moe/fused_moe/impl_base.py index aa8687486542..fbc7498e74e9 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_base.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_base.py @@ -25,6 +25,7 @@ from ...model_config import ModelConfig from ...utils import ActivationType, AuxStreamType, is_gated_activation +from .activation import DEFAULT_MOE_ACTIVATION, MoEActivation, install_activation_params from .impl_blocks import MoEEplbWeightLayoutMixin, MoEExecutionContractMixin, MoEWeightOwnerMixin from .impl_contract import MoEDeployment, MoEEligibility, MoEEplbBinding, MoEProblem, MoERunContext from .interface import MoESchedulerKind, MoEWeightLoadingMode, _compute_ep_partition @@ -55,12 +56,8 @@ def apply_moe_impl_construction_state( aux_stream_dict: Optional[dict[AuxStreamType, torch.cuda.Stream]] = None, weight_loading_mode: MoEWeightLoadingMode = MoEWeightLoadingMode.VANILLA, bias: bool = False, - swiglu_alpha: Optional[torch.Tensor] = None, - swiglu_beta: Optional[torch.Tensor] = None, - swiglu_limit: Optional[torch.Tensor] = None, - swiglu_limit_scalar: Optional[float] = None, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, layer_idx: Optional[int] = None, - activation_type: ActivationType = ActivationType.Swiglu, init_load_balancer: bool = False, ) -> None: """Install the construction state backends used to get from ``MoE.__init__``. @@ -96,19 +93,21 @@ def apply_moe_impl_construction_state( module.routing_method = routing_method module.num_experts = num_experts module.hidden_size = hidden_size + # The pre-padding value: a backend that rounds ``hidden_size`` up for its + # weight layout reads this back to slice the padding off again. + module.unpadded_hidden_size = hidden_size module.intermediate_size = intermediate_size module.weight_loading_mode = weight_loading_mode module.bias = bias module.dtype = dtype module.reduce_results = reduce_results - module.swiglu_alpha = swiglu_alpha - module.swiglu_beta = swiglu_beta - module.swiglu_limit = swiglu_limit - module.swiglu_limit_scalar = swiglu_limit_scalar module.layer_idx = layer_idx module.layer_idx_str = str(layer_idx) if layer_idx is not None else None - module.activation_type = int(activation_type) - module.is_gated_activation = is_gated_activation(activation_type) + module.activation = activation + # ``ActivationType``, not a bare int: as an ``IntEnum`` it still reads as an + # int for op schemas and tuning keys, and rejection messages get ``.name``. + module.activation_type = ActivationType(activation.kind) + module.is_gated_activation = is_gated_activation(activation.kind) module.intermediate_size_expand_ratio = 2 if module.is_gated_activation else 1 module.quant_config = model_config.quant_config @@ -153,6 +152,11 @@ def apply_moe_impl_construction_state( module.initial_global_assignments = list(range(module.num_experts)) module.allreduce = None + # Last, because a PER_EXPERT_TENSOR constant is sized by + # ``expert_size_per_partition`` set just above. An impl built directly -- + # as unit tests and microbenchmarks do -- has no other installer. + install_activation_params(module) + class MoEImplBase( MoEExecutionContractMixin, MoEWeightOwnerMixin, MoEEplbWeightLayoutMixin, nn.Module, abc.ABC @@ -247,6 +251,21 @@ def quantize_input( @abc.abstractmethod def run_moe(self, ctx: MoERunContext) -> torch.Tensor: ... + def try_fused_route_quant( + self, x: "torch.Tensor", router_logits: "torch.Tensor" + ) -> "tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None": + """Offer to fuse routing and input quantization into one launch. + + Accepting returns what the scheduler otherwise assembles from two calls, + in that order: ``routing_method.apply``'s selected experts and scales, + then ``quantize_input``'s quantized activations and scaling factors. + + Declining is the default and the common case: the one implementation + that accepts is shape-, SM-, and quant-specialized and checks all of + that itself. Overriding this is how a backend opts in. + """ + return None + # ---- impl-owned resources: produced here, never passed in ------------- def get_workspaces(self, *args: object, **kwargs: object) -> "list[dict] | None": return None diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py index efa45f35009d..cb2d21339683 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py @@ -41,6 +41,23 @@ class MoEStaticCapability: supports_moe_lora: bool = False # Legacy gate: CuteDslFusedMoE isinstance check in ConfigurableMoE DWDP. supports_dwdp: bool = False + # Legacy gate: ``assert moe_cls in supported_load_balancer_backends`` in + # ``create_moe_backend``. Not the same question as the instance-level + # ``_supports_load_balancer()``, which TRTLLMGenFusedMoE overrides to mean + # "separated routing is used". + supports_eplb: bool = False + # Legacy gate: the ``assert moe_cls in [...]`` bias allow-list in + # ``create_moe_backend``. Per-expert FC bias from the checkpoint, added + # before the activation functor runs -- not an activation constant. + supports_expert_bias: bool = False + # Legacy gate: the three ``assert not apply_router_weight_on_input`` checks + # keyed on ``moe_cls`` in ``create_moe_backend``. The fold itself belongs to + # MoEScheduler (``x = x * token_final_scales``), so what a backend declares + # here is whether it handles what the fold leaves behind: ``None`` scales, + # or all-ones under a DeepEP / NCCL comm strategy. Backends that reject the + # flag in their own constructor keep doing so; that check guards direct + # construction, which never reaches the factory. + supports_apply_router_weight_on_input: bool = False @dataclass(frozen=True) @@ -94,6 +111,11 @@ class MoEProblem: bias: Optional[bool] = None #: ``ActivationType`` member name; omitted values canonicalize to SwiGLU. activation: str = "Swiglu" + #: Which of the ``alpha`` / ``beta`` / ``clamp`` ABI registers the caller's + #: activation fills; the kind alone does not say, since clamped and + #: unclamped SwiGLU share one ``ActivationType``. Empty if the call site + #: supplied no activation carrier. + activation_constants: frozenset[str] = frozenset() #: ``RoutingMethodType`` member name; None means the call site did not say. routing: Optional[str] = None @@ -367,6 +389,7 @@ def to_dict(self) -> Dict[str, object]: "swiglu_gptoss_style": self.problem.swiglu_gptoss_style, "bias": self.problem.bias, "activation": self.problem.activation, + "activation_constants": sorted(self.problem.activation_constants), "routing": self.problem.routing, }, "deployment": { diff --git a/tensorrt_llm/_torch/moe/fused_moe/interface.py b/tensorrt_llm/_torch/moe/fused_moe/interface.py index 632737785648..3ef31b3b1f02 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/interface.py +++ b/tensorrt_llm/_torch/moe/fused_moe/interface.py @@ -23,6 +23,7 @@ from torch import nn from ...distributed.ops import reducescatter +from .activation import DEFAULT_MOE_ACTIVATION, MoEActivation from .impl_blocks import (MoEEplbWeightLayoutMixin, MoEExecutionContractMixin, MoEWeightOwnerMixin) from .impl_contract import (MoEDeployment, MoEEligibility, MoEProblem, @@ -248,14 +249,10 @@ def __init__( weight_loading_mode: MoEWeightLoadingMode = MoEWeightLoadingMode. VANILLA, bias: bool = False, - swiglu_alpha: Optional[torch.Tensor] = None, - swiglu_beta: Optional[torch.Tensor] = None, - swiglu_limit: Optional[torch.Tensor] = None, - swiglu_limit_scalar: Optional[float] = None, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, layer_idx: Optional[int] = None, - activation_type: ActivationType = ActivationType.Swiglu, init_load_balancer: bool = True, - ): + ) -> None: from ...distributed import AllReduce super().__init__() @@ -267,22 +264,19 @@ def __init__( self.bias = bias self.dtype = dtype self.reduce_results = reduce_results - self.swiglu_alpha = swiglu_alpha - self.swiglu_beta = swiglu_beta - self.swiglu_limit = swiglu_limit - # Uniform-across-experts scalar variant of swiglu_limit, consumed by - # FP8 paths (DeepGEMM Triton kernel, TRTLLMGen FP8 separate-activation - # kernel) that don't actually need a per-expert tensor. Distinct from - # `swiglu_limit` (kept for NVFP4 fused-activation cubins that *do* - # consume per-expert values via fc31_alpha rescaling). - self.swiglu_limit_scalar = swiglu_limit_scalar self.layer_idx = layer_idx self.layer_idx_str = str(layer_idx) if layer_idx is not None else None - self.activation_type = int(activation_type) + # Intent only: a layer that owns kernels turns this into the ``act_*`` + # slots via ``install_activation_params`` against its own declaration. + # ``ConfigurableMoE`` forwards it to the backend instead. + self.activation = activation + # ``ActivationType``, not a bare int: as an ``IntEnum`` it still reads as an + # int for op schemas and tuning keys, and rejection messages get ``.name``. + self.activation_type = ActivationType(activation.kind) # Note: # - for gated activations, there should be with gate and up projections, so the intermediate size should be expanded by 2. # - for non-gated activations, there is only one up projection (no gate projection), so the intermediate size should not be expanded. - self.is_gated_activation = is_gated_activation(activation_type) + self.is_gated_activation = is_gated_activation(activation.kind) self.intermediate_size_expand_ratio = 2 if self.is_gated_activation else 1 self._register_layer(model_config) diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index ec526e78173d..768f22410df6 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -94,6 +94,12 @@ from ....cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ....model_config import ModelConfig from ....utils import ActivationType, AuxStreamType, Fp4QuantizedTensor +from ..activation import ( + DEFAULT_MOE_ACTIVATION, + ActivationParamShape, + MoEActivation, + MoEActivationSupport, +) from ..impl_base import MoEImplBase, apply_moe_impl_construction_state from ..impl_contract import ( MoEDeployment, @@ -101,6 +107,7 @@ MoEProblem, MoERejectReason, MoERunContext, + MoEStaticCapability, ) from ..impl_environment import MoEDep from ..interface import MoESchedulerKind, MoEWeightLoadingMode, _reject @@ -345,6 +352,18 @@ class MegaMoECuteDsl(MoEImplBase): _SUPPORTED_ACTIVATION_DTYPES = frozenset({torch.bfloat16}) + # Static and dynamic EPLB both work: see ``_supports_load_balancer`` below + # for why the MegaMoE-format derived parameters migrate atomically. + capabilities = MoEStaticCapability(supports_eplb=True) + + # The elementwise function and its constants are codegen-time literals, so + # only a uniform per-layer value is representable. + activation_support = MoEActivationSupport( + kinds=frozenset({ActivationType.Swiglu, ActivationType.SiTu}), + alpha_beta=ActivationParamShape.UNIFORM_SCALAR, + limit=ActivationParamShape.UNIFORM_SCALAR, + ) + # Legal combine wire formats; must stay in sync with # ``CombineFormat.parse`` in the kernel package's token_comm.py. _SUPPORTED_COMBINE_FORMATS = frozenset({"bf16", "32e4m3xe8m0", "16e2m1xbf16"}) @@ -411,10 +430,12 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: # The fused path also requires the ``trtllm::cute_dsl_megamoe_nvfp4_*`` # custom ops to be registered (strict import of every kernel symbol in # cute_dsl_megamoe_custom_op). - if p.activation_type != ActivationType.Swiglu: + if p.activation_type not in cls.activation_support.kinds: + supported = ", ".join(sorted(a.name for a in cls.activation_support.kinds)) return _reject( MoERejectReason.ACTIVATION_UNSUPPORTED, - f"MegaMoECuteDsl only supports ActivationType.Swiglu (got {p.activation}).", + f"MegaMoECuteDsl does not support activation {p.activation}; " + f"supported: {supported}", ) if d.tp_size != 1: return _reject( @@ -467,13 +488,7 @@ def __init__( apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, init_load_balancer: bool = False, - activation_type: ActivationType = ActivationType.Swiglu, - swiglu_limit: Optional[torch.Tensor] = None, - # ``activation=None`` infers Kimi K3 SiTU from the pretrained config and - # otherwise defaults to SwiGLU. Mirrors MegaMoEDeepGemm. - activation: Optional[str] = None, - situ_beta: Optional[float] = None, - situ_linear_beta: Optional[float] = None, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, **kwargs, ) -> None: # ``aux_stream_dict`` is accepted for ``create_moe_backend`` signature @@ -493,8 +508,7 @@ def __init__( aux_stream_dict=None, weight_loading_mode=weight_loading_mode, layer_idx=layer_idx, - activation_type=activation_type, - swiglu_limit=swiglu_limit, + activation=activation, init_load_balancer=init_load_balancer, ) @@ -505,22 +519,6 @@ def __init__( "MegaMoECuteDsl does not support apply_router_weight_on_input; " "the fused kernel applies routing weights on the MoE output." ) - # ``ActivationType.Swiglu`` describes the gated FC1 tensor geometry shared - # by SwiGLU and SiTU; the elementwise function is selected by - # ``activation`` below. Same reasoning as MegaMoEDeepGemm. - if activation_type != ActivationType.Swiglu: - raise ValueError( - f"MegaMoECuteDsl only supports ActivationType.Swiglu (got {activation_type})." - ) - activation, situ_beta, situ_linear_beta = self._resolve_activation_config( - model_config, - activation=activation, - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, - ) - self.activation = activation - self.situ_beta = situ_beta - self.situ_linear_beta = situ_linear_beta self.apply_router_weight_on_input = apply_router_weight_on_input # topk-score application point. v2 default is the deepgemm graph @@ -548,18 +546,6 @@ def __init__( # gates collective tuning behavior). self.tactic_autotune = os.environ.get("MEGAMOE_TACTIC_AUTOTUNE", "0") == "1" - # SwiGLU clamp: map the model-provided per-layer ``swiglu_limit`` tensor - # to the kernel's codegen-time scalar ``gate_up_clamp``. The MegaMoE - # kernel clamps the post-fc1_alpha real gate/up, so the model value is - # used directly (NO trtllm-gen-style div_(fc31_alpha) normalization). - # Reject non-uniform / per-expert clamp: the kernel bakes one constant. - self.gate_up_clamp = self._resolve_gate_up_clamp(swiglu_limit) - if self.activation == "situ" and self.gate_up_clamp is not None: - raise ValueError( - "MegaMoECuteDsl SiTU does not support a gate/up clamp; " - "drop swiglu_limit for SiTU checkpoints." - ) - # Buffer sizing. MoE layers execute serially per forward; one pool # sized to the worst-case per-rank tokens covers every layer. The # kernel compile takes this as the static ``max_tokens_per_rank``. @@ -638,83 +624,6 @@ def __init__( # ------------------------------------------------------------------ # Topology # ------------------------------------------------------------------ - @staticmethod - def _resolve_activation_config( - model_config, - *, - activation: Optional[str], - situ_beta: Optional[float], - situ_linear_beta: Optional[float], - ): - """Resolve the elementwise activation and its SiTU constants. - - ``activation=None`` infers SiTU from ``activation_situ_beta`` in the - pretrained config (Kimi K3 sets it) and otherwise selects SwiGLU. Kept - byte-for-byte equivalent to ``MegaMoEDeepGemm._resolve_activation_config`` - so the two MegaMoE backends cannot disagree about the same checkpoint. - """ - pretrained_config = getattr(model_config, "pretrained_config", None) - text_config = getattr(pretrained_config, "text_config", None) - cfg_beta = getattr(pretrained_config, "activation_situ_beta", None) - cfg_lbeta = getattr(pretrained_config, "activation_situ_linear_beta", None) - if cfg_beta is None: - cfg_beta = getattr(text_config, "activation_situ_beta", None) - if cfg_lbeta is None: - cfg_lbeta = getattr(text_config, "activation_situ_linear_beta", None) - if activation is None: - activation = "situ" if cfg_beta is not None else "swiglu" - activation = activation.lower() - if activation not in ("swiglu", "situ"): - raise ValueError( - f"MegaMoECuteDsl activation must be 'swiglu' or 'situ'; got {activation!r}." - ) - if activation == "swiglu": - if situ_beta is not None or situ_linear_beta is not None: - raise ValueError("SiTU beta parameters require activation='situ'.") - return activation, None, None - situ_beta = cfg_beta if situ_beta is None else situ_beta - situ_linear_beta = cfg_lbeta if situ_linear_beta is None else situ_linear_beta - if situ_beta is None or situ_linear_beta is None: - raise ValueError( - "MegaMoECuteDsl SiTU requires activation_situ_beta and " - "activation_situ_linear_beta in the pretrained config, or explicit " - "situ_beta and situ_linear_beta arguments." - ) - if situ_beta <= 0 or situ_linear_beta <= 0: - raise ValueError("MegaMoECuteDsl SiTU beta parameters must be positive.") - return activation, float(situ_beta), float(situ_linear_beta) - - @staticmethod - def _resolve_gate_up_clamp( - swiglu_limit: Optional[torch.Tensor], - ) -> Optional[float]: - """Reduce a per-layer ``swiglu_limit`` tensor to a single codegen-time - ``gate_up_clamp`` float, or ``None`` when no clamp is configured. - - The MegaMoE kernel bakes ``gate_up_clamp`` into the compiled kernel as - one scalar, so only a uniform (per-layer) clamp is representable. - Non-uniform / per-expert clamp is rejected with a clear ``ValueError`` - rather than silently using one element. GPT-OSS-style clamp is rejected - earlier in ``can_implement`` via ``swiglu_gptoss_style``. - """ - if swiglu_limit is None: - return None - if not isinstance(swiglu_limit, torch.Tensor): - # Accept a plain python scalar for robustness. - return float(swiglu_limit) - flat = swiglu_limit.detach().reshape(-1) - if flat.numel() == 0: - return None - first = flat[0] - if flat.numel() > 1 and not torch.allclose(flat, first.expand_as(flat), rtol=1e-5, atol=0): - raise ValueError( - "MegaMoECuteDsl only supports a uniform (per-layer) " - "swiglu_limit because the kernel bakes gate_up_clamp as a " - "codegen-time scalar; got a non-uniform / per-expert " - f"swiglu_limit with values {flat.cpu().tolist()}." - ) - return float(first.item()) - @staticmethod def _resolve_maxt_buckets(max_num_tokens: int) -> List[int]: """Resolve the adaptive ``max_tokens_per_rank`` bucket ladder. @@ -1529,9 +1438,11 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): max_tokens_per_rank=int(launch_max_T), peer_offsets=peer_offsets, apply_topk_in_fc1=bool(self.apply_topk_in_fc1), - gate_up_clamp=self.gate_up_clamp, - situ_beta=self.situ_beta, - situ_linear_beta=self.situ_linear_beta, + # The kernel clamps the post-fc1_alpha gate/up, so the model + # value goes in as-is -- no trtllm-gen-style div_(fc31_alpha). + gate_up_clamp=self.act_clamp, + act_alpha=self.act_alpha, + act_beta=self.act_beta, # Keep the deterministic standalone TopkReduce until form-B # has dedicated GPU correctness and performance coverage. in_kernel_fc2_reduce=False, diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py index fa459bd0d5f4..0258879b2651 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py @@ -35,6 +35,12 @@ from ....model_config import ModelConfig from ....utils import ActivationType, AuxStreamType +from ..activation import ( + DEFAULT_MOE_ACTIVATION, + ActivationParamShape, + MoEActivation, + MoEActivationSupport, +) from ..impl_base import MoEImplBase, apply_moe_impl_construction_state from ..impl_contract import ( MoEDeployment, @@ -42,6 +48,7 @@ MoEProblem, MoERejectReason, MoERunContext, + MoEStaticCapability, ) from ..impl_environment import MoEDep from ..interface import MoESchedulerKind, MoEWeightLoadingMode, _reject @@ -140,6 +147,26 @@ class MegaMoEDeepGemm(MoEImplBase): _SUPPORTED_ACTIVATION_DTYPES = frozenset({torch.bfloat16}) + # Static and dynamic EPLB both work: see ``_supports_load_balancer`` below + # for why slot-id routing and DG-tensor migration are safe here. + capabilities = MoEStaticCapability(supports_eplb=True) + + # ActivationType -> the DeepGEMM library's own ``activation`` argument. The + # only place that external vocabulary is spoken; keys must cover + # ``activation_support``, which ``__init__`` checks. + _ACTIVATION_MAP = { + ActivationType.Swiglu: "swiglu", + ActivationType.SiTu: "situ", + } + + # Both SiTU constants and the SwiGLU clamp are baked into the compiled + # kernel, so a per-expert tensor is rejected rather than silently reduced. + activation_support = MoEActivationSupport( + kinds=frozenset({ActivationType.Swiglu, ActivationType.SiTu}), + alpha_beta=ActivationParamShape.UNIFORM_SCALAR, + limit=ActivationParamShape.UNIFORM_SCALAR, + ) + # Kernel owns dispatch + GEMM1 + gated activation + GEMM2 + combine via NVLink # SymmBuffer; ConfigurableMoE must NOT layer host-side comm on top. scheduler_kind = MoESchedulerKind.FUSED_COMM @@ -192,10 +219,12 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: f"(DeepGEMM TMA-aligned packed-UE8M0 SF row); " f"got intermediate_size={p.intermediate_size}", ) - if p.activation_type != ActivationType.Swiglu: + if p.activation_type not in cls.activation_support.kinds: + supported = ", ".join(sorted(a.name for a in cls.activation_support.kinds)) return _reject( MoERejectReason.ACTIVATION_UNSUPPORTED, - f"MegaMoEDeepGemm only supports ActivationType.Swiglu (got {p.activation})", + f"MegaMoEDeepGemm does not support activation {p.activation}; " + f"supported: {supported}", ) if d.tp_size != 1: return _reject( @@ -255,15 +284,9 @@ def __init__( weight_loading_mode: MoEWeightLoadingMode = MoEWeightLoadingMode.VANILLA, apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, - activation_type: ActivationType = ActivationType.Swiglu, + activation: MoEActivation = DEFAULT_MOE_ACTIVATION, init_load_balancer: bool = False, - # DG tunables. ``activation=None`` infers Kimi K3 SiTU from the - # pretrained config and otherwise defaults to SwiGLU. - activation: Optional[str] = None, - swiglu_limit_scalar: Optional[float] = None, fast_math: bool = True, - situ_beta: Optional[float] = None, - situ_linear_beta: Optional[float] = None, **kwargs, ) -> None: super().__init__(eplb=None) @@ -279,7 +302,7 @@ def __init__( aux_stream_dict=aux_stream_dict, weight_loading_mode=weight_loading_mode, layer_idx=layer_idx, - activation_type=activation_type, + activation=activation, init_load_balancer=init_load_balancer, ) @@ -298,20 +321,18 @@ def __init__( # Also gated in ``can_implement``, but that only covers the resolution # path; this catches direct construction. _assert_num_slots_divisible_by_ep(self.num_slots, self.ep_size) - activation, situ_beta, situ_linear_beta = self._resolve_activation_config( - model_config, - activation=activation, - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, - ) - if activation == "situ" and swiglu_limit_scalar is not None: - raise ValueError("MegaMoEDeepGemm SiTU does not support activation_clamp.") self.apply_router_weight_on_input = apply_router_weight_on_input - self.activation = activation - self.swiglu_limit_scalar = swiglu_limit_scalar + # Checked here rather than at module scope, where a table drift would + # surface as ``import tensorrt_llm`` failing for every model. + kind = ActivationType(self.activation.kind) + if kind not in self._ACTIVATION_MAP: + raise ValueError( + f"MegaMoEDeepGemm._ACTIVATION_MAP has no DeepGEMM spelling for " + f"{kind.name}, which activation_support declares as executable. Add " + f"the spelling, or drop the kind from activation_support." + ) + self.dg_activation = self._ACTIVATION_MAP[kind] self.fast_math = fast_math - self.situ_beta = situ_beta - self.situ_linear_beta = situ_linear_beta # Buffer sizing. MoE layers execute serially per forward; a single # process-level pool sized to worst-case per-rank tokens serves all. @@ -371,46 +392,6 @@ def __init__( if not model_config.skip_create_weights_in_init: self.create_weights() - @staticmethod - def _resolve_activation_config( - model_config: ModelConfig, - *, - activation: Optional[str], - situ_beta: Optional[float], - situ_linear_beta: Optional[float], - ) -> Tuple[str, Optional[float], Optional[float]]: - pretrained_config = model_config.pretrained_config - text_config = getattr(pretrained_config, "text_config", None) - config_situ_beta = getattr(pretrained_config, "activation_situ_beta", None) - config_situ_linear_beta = getattr(pretrained_config, "activation_situ_linear_beta", None) - if config_situ_beta is None: - config_situ_beta = getattr(text_config, "activation_situ_beta", None) - if config_situ_linear_beta is None: - config_situ_linear_beta = getattr(text_config, "activation_situ_linear_beta", None) - if activation is None: - activation = "situ" if config_situ_beta is not None else "swiglu" - activation = activation.lower() - if activation not in ("swiglu", "situ"): - raise ValueError( - f"MegaMoEDeepGemm activation must be 'swiglu' or 'situ'; got {activation!r}." - ) - if activation == "swiglu": - if situ_beta is not None or situ_linear_beta is not None: - raise ValueError("SiTU beta parameters require activation='situ'.") - return activation, None, None - - situ_beta = config_situ_beta if situ_beta is None else situ_beta - situ_linear_beta = config_situ_linear_beta if situ_linear_beta is None else situ_linear_beta - if situ_beta is None or situ_linear_beta is None: - raise ValueError( - "MegaMoEDeepGemm SiTU requires activation_situ_beta and " - "activation_situ_linear_beta in the pretrained config, or " - "explicit situ_beta and situ_linear_beta arguments." - ) - if situ_beta <= 0 or situ_linear_beta <= 0: - raise ValueError("MegaMoEDeepGemm SiTU beta parameters must be positive.") - return activation, float(situ_beta), float(situ_linear_beta) - def _supports_load_balancer(self) -> bool: # The DeepGEMM mega kernel routes by `topk_idx` interpreted as slot id # (range [0, num_slots)) once the SymmBuffer is sized to num_slots. @@ -610,7 +591,7 @@ def _alloc_symm_buffer(self) -> None: self.routing_method.experts_per_token, self.hidden_size, self.intermediate_size, - self.activation, + self.dg_activation, ) cached = _MEGA_MOE_SYMM_BUFFER_CACHE.get(key) if cached is None: @@ -623,7 +604,7 @@ def _alloc_symm_buffer(self) -> None: self.intermediate_size, num_shared_experts=0, mma_type="fp8xfp4", - activation=self.activation, + activation=self.dg_activation, ) _MEGA_MOE_SYMM_BUFFER_CACHE[key] = cached # Log only on the first layer; deeper layers reuse the cache @@ -758,10 +739,13 @@ def run_moe( self._t_l1, self._t_l2, buf, - activation=self.activation, - activation_clamp=self.swiglu_limit_scalar, + activation=self.dg_activation, + # ``situ_beta`` / ``situ_linear_beta`` are the bundled DeepGEMM's own + # parameter names: ``_import_deep_gemm`` gates on the installed + # signature accepting exactly these, so do not rename them to ours. + activation_clamp=self.act_clamp, fast_math=self.fast_math, - situ_beta=self.situ_beta, - situ_linear_beta=self.situ_linear_beta, + situ_beta=self.act_alpha, + situ_linear_beta=self.act_beta, ) return y.to(output_dtype) diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py index e374984cc6ba..13301a4617fe 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py @@ -28,6 +28,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig +from .activation import ActivationParamShape, MoEActivation, activation_constant_names from .fused_moe_cute_dsl import CuteDslFusedMoE from .fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE from .fused_moe_cutlass import CutlassFusedMoE @@ -203,14 +204,21 @@ def build_moe_problem( top_k: Optional[int] = None, swiglu_gptoss_style: Optional[bool] = None, bias: Optional[bool] = None, - activation_type: Optional[ActivationType] = None, + activation: Optional[MoEActivation] = None, routing: Optional["BaseMoeRoutingMethod | RoutingMethodType"] = None, ) -> MoEProblem: """Assemble the problem half of a selection question. Explicit args win over ``pretrained_config``. Missing fields stay ``None`` (unknown): shape gates abstain instead of rejecting on absent info. + + ``activation`` is the whole activation package, and both halves of it are + read: the kind, and *which constants* the caller supplies -- a question the + kind alone cannot answer, since clamped and unclamped SwiGLU share one + ``ActivationType``. Pass it wherever it is in hand; without it the problem + says "some activation" and the activation gate abstains. """ + activation_kind = None if activation is None else ActivationType(activation.kind) shapes = derive_moe_layer_shapes( model_config, num_experts=num_experts, @@ -232,7 +240,8 @@ def build_moe_problem( top_k=shapes.top_k, swiglu_gptoss_style=swiglu_gptoss_style, bias=bias, - activation=canonical_activation(activation_type), + activation=canonical_activation(activation_kind), + activation_constants=activation_constant_names(activation), routing=canonical_routing(routing), ) @@ -240,21 +249,22 @@ def build_moe_problem( def infer_swiglu_gptoss_style( *, bias: bool = False, - swiglu_alpha: Optional[torch.Tensor] = None, - swiglu_beta: Optional[torch.Tensor] = None, activation_type: Optional[Union[ActivationType, int]] = None, ) -> bool: - """True for the gpt-oss / MiniMax SwiGLU package (bias, alpha/beta, or SwigluBias). + """True for the gpt-oss / MiniMax SwiGLU package (expert bias, or SwigluBias). - ``swiglu_limit`` alone is not enough — DeepSeek-V4 uses a plain clamp and - must not be treated as gpt-oss. + Keyed off the kind alone. The older form also answered True whenever an + alpha/beta constant was merely *present*, which SiTU satisfies too, and that + silently downgraded a MegaMoE request to Cutlass. A clamp was never part of + it either -- DeepSeek-V4 uses a plain clamp. - ``activation_type`` is normalized because ``MoE`` stores the activation as a - plain ``int``, which no identity check against an enum member can match. + ``activation_type`` is still normalized even though ``MoE`` now stores an + ``ActivationType``: the parameter is also reached from call sites that pass a + bare int, and no identity check against an enum member would match one. """ if activation_type is not None and ActivationType(activation_type) is ActivationType.SwigluBias: return True - return bool(bias or swiglu_alpha is not None or swiglu_beta is not None) + return bool(bias) def build_moe_deployment( @@ -315,6 +325,59 @@ def _candidates_for(backend: str) -> List[MoEImplClass]: return candidates +#: ABI register name -> the ``MoEActivationSupport`` field declaring its shape. +_CONSTANT_SHAPE_FIELDS: Dict[str, str] = { + "alpha": "alpha_beta", + "beta": "alpha_beta", + "clamp": "limit", +} + + +def _reject_unsupported_activation( + candidate: MoEImplClass, problem: MoEProblem +) -> Optional[MoERejection]: + """Decline a candidate whose declaration cannot carry this activation. + + Central rather than repeated in eleven ``can_implement`` gates, because the + answer is already written down: ``activation_support`` is the same + declaration ``materialize_activation_params`` reads. A backend that forgot + to re-derive it would not run the layer anyway -- it would raise from the + adapter at construction, well past the point where another candidate could + still have been chosen. + + Reads the *class* declaration. TRTLLM-Gen overrides + ``resolve_activation_support`` per instance, but only to narrow + ``PER_EXPERT_TENSOR`` to ``UNIFORM_SCALAR``; no instance turns a shape into + ``UNSUPPORTED``, so no instance refuses a constant its class admits. + """ + support = getattr(candidate, "activation_support", None) + if support is None: + return None + + kind = problem.activation_type + if kind not in support.kinds: + executes = ", ".join(sorted(k.name for k in support.kinds)) + return MoERejection( + _legacy_backend_name(candidate), + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"{candidate.__name__} does not execute {kind.name} (executes: {executes})", + ) + + # Sorted so the rejection names the same register every run. + for constant in sorted(problem.activation_constants): + field = _CONSTANT_SHAPE_FIELDS.get(constant) + if field is None: + continue + if getattr(support, field) is ActivationParamShape.UNSUPPORTED: + return MoERejection( + _legacy_backend_name(candidate), + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"{candidate.__name__} kernels take no activation {constant}, " + f"which this layer's {kind.name} supplies", + ) + return None + + def resolve_moe_impl( model_config: ModelConfig, *, @@ -327,12 +390,19 @@ def resolve_moe_impl( intermediate_size: Optional[int] = None, swiglu_gptoss_style: Optional[bool] = None, bias: Optional[bool] = None, - activation_type: Optional[ActivationType] = None, + activation: Optional[MoEActivation] = None, routing: Optional["BaseMoeRoutingMethod | RoutingMethodType"] = None, layer_idx: Optional[int] = None, + allow_degradation: bool = True, ) -> MoEResolutionReport: """Resolve a MoE backend and return the full eligibility report. + ``allow_degradation=False`` turns the usual substitution warning into a + hard failure. A caller that is measuring one specific backend needs that: + silently running the fallback returns numbers attributed to the backend it + asked for. The rejection trail says which gate declined and why, so the + caller does not have to re-derive the winner and compare classes. + Raises ValueError for unknown or deprecated backend literals. """ if problem is None: @@ -345,7 +415,7 @@ def resolve_moe_impl( intermediate_size=intermediate_size, swiglu_gptoss_style=swiglu_gptoss_style, bias=bias, - activation_type=activation_type, + activation=activation, routing=routing, ) if deployment is None: @@ -369,16 +439,21 @@ def resolve_moe_impl( ) continue eligibility = candidate.can_implement(problem, deployment) - if eligibility.eligible: - eligible.append(candidate) - continue - rejected.append( - MoERejection( - _legacy_backend_name(candidate), - eligibility.reject_reason, - eligibility.detail, + if not eligibility.eligible: + rejected.append( + MoERejection( + _legacy_backend_name(candidate), + eligibility.reject_reason, + eligibility.detail, + ) ) - ) + continue + # After can_implement, so a backend's own more specific reason wins. + activation_rejection = _reject_unsupported_activation(candidate, problem) + if activation_rejection is not None: + rejected.append(activation_rejection) + continue + eligible.append(candidate) # candidates is already priority-ordered. winner_cls = eligible[0] if eligible else None @@ -402,6 +477,13 @@ def resolve_moe_impl( env_fingerprint=deployment.env.fingerprint(), ) + if report.degraded and not allow_degradation: + location = "" if layer_idx is None else f" [layer_idx={layer_idx}]" + raise ValueError( + f"MoE backend {requested} was requested with degradation disallowed " + f"but cannot serve this layer{location}. {report.describe()}" + ) + if report.degraded and winner_cls is not None: cause = report.degraded_from location = "" if layer_idx is None else f" [layer_idx={layer_idx}]" diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index c9b93a64945b..70b273700f15 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -54,7 +54,6 @@ from .communication import DeepEP, DeepEPLowLatency, NcclEP, NVLinkOneSided, NVLinkTwoSided from .communication.nvlink_two_sided_flashinfer import NVLinkTwoSidedFlashinfer from .fused_moe_cutlass import raise_moe_lora_multichunk_unsupported -from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE from .impl_contract import MoECommPlan, MoERunContext from .interface import FORCE_SEPARATED_ROUTING, MoESchedulerKind @@ -384,13 +383,16 @@ def _forward_chunk_impl( supports_post_quant = moe.comm is None or moe.comm.supports_post_quant_dispatch() used_fused_route_quant = False if requires_separated_routing: + # ``MoEImplBase`` declines a fused route+quant path by default. The + # conditions are the scheduler's: a fused result skips the separate + # dispatch and carries no per-token scale to fold a router weight or + # an EPLB layout into. if ( supports_post_quant - and isinstance(moe.backend, TRTLLMGenFusedMoE) and not moe._using_load_balancer() and not moe.apply_router_weight_on_input ): - fused_result = moe.backend.try_fused_kimi_route_quant(x, router_logits) + fused_result = moe.backend.try_fused_route_quant(x, router_logits) else: fused_result = None @@ -399,8 +401,6 @@ def _forward_chunk_impl( token_selected_experts, token_final_scales = moe.routing_method.apply( router_logits, input_ids ) - if token_final_scales is not None and isinstance(moe.backend, TRTLLMGenFusedMoE): - token_final_scales = token_final_scales.to(torch.bfloat16) else: token_selected_experts, token_final_scales, x, x_sf = fused_result used_fused_route_quant = True diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index d3149adf14ec..c947a316ab9d 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -5246,6 +5246,24 @@ def process_weights_after_loading(self, # Finalize shared expert alphas and fc31_scale_c for online EPLB self._finalize_shared_expert_alphas(module) + # Cubin clamp / GLU bias inputs are consumed in the pre-dequant GEMM + # output domain (i.e. divided by fc31_alpha / fc2_alpha). The ``act_*`` + # slots are per-expert tensors on every path this method class serves. + # + # The bias division is gated on module.bias to stay consistent with + # _shuffle_all_experts() below. Hoisting this block into the base method + # deliberately widens it to W4A8NVFP4FP8TRTLLMGenFusedMoEMethod, which + # auto-creates all-zero bias tensors that its kernel + # (fp8_fp4_block_scale_moe_runner) never consumes; those must not be + # divided, because a zero alpha would turn them into NaN. + if module.bias: + module.w3_w1_bias.data.div_((module.fc31_alpha.data).view(-1, 1)) + module.w2_bias.data.div_((module.fc2_alpha.data).view(-1, 1)) + if getattr(module, 'act_beta', None) is not None: + module.act_beta.data.div_((module.fc31_alpha.data)) + if getattr(module, 'act_clamp', None) is not None: + module.act_clamp.data.div_((module.fc31_alpha.data)) + def _shuffle_shared_expert_tensors(self, module: torch.nn.Module, num_elts_per_sf: int = 16): @@ -5739,23 +5757,6 @@ def _shuffle_all_experts(self, module.w3_w1_bias.data[expert_idx]) self._shuffle_w2_weight(module.w2_bias.data[expert_idx]) - def process_weights_after_loading(self, - module: torch.nn.Module, - num_elts_per_sf: int = 16): - super().process_weights_after_loading(module, - num_elts_per_sf=num_elts_per_sf) - - # Cubin clamp / GLU bias inputs are consumed in the pre-dequant GEMM - # output domain (i.e. divided by fc31_alpha / fc2_alpha). - if module.w3_w1_bias is not None: - module.w3_w1_bias.data.div_((module.fc31_alpha.data).view(-1, 1)) - if module.w2_bias is not None: - module.w2_bias.data.div_((module.fc2_alpha.data).view(-1, 1)) - if module.swiglu_beta is not None: - module.swiglu_beta.data.div_((module.fc31_alpha.data)) - if module.swiglu_limit is not None: - module.swiglu_limit.data.div_((module.fc31_alpha.data)) - class W4A8NVFP4FP8TRTLLMGenFusedMoEMethod(NVFP4TRTLLMGenFusedMoEBaseMethod): eplb_support_status = EplbSupportStatus.NOT_VERIFIED diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a3e8504ad21f..c54ac82cd567 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -157,12 +157,14 @@ l0_b200: - unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout - unittest/_torch/moe/test_moe_backend.py::test_import_deep_gemm_rejects_pre_situ_mega_moe_api - unittest/_torch/moe/test_moe_backend.py::test_megamoe_streaming_reload_resets_slot_claims - - unittest/_torch/moe/test_moe_backend.py::test_megamoe_deepgemm_infers_kimi_situ_from_pretrained_config - - unittest/_torch/moe/test_moe_backend.py::test_megamoe_deepgemm_defaults_to_swiglu_without_situ_config - - unittest/_torch/moe/test_moe_backend.py::test_create_moe_forwards_megamoe_activation_options - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_selects_padded_quant_method - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_fc31_scale_c_drops_dequant_scale - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_situ_rejects_quant_algos_without_fused_cubins + - unittest/_torch/moe/test_moe_backend.py::test_megamoe_bakes_situ_softcaps_as_uniform_scalars + - unittest/_torch/moe/test_moe_backend.py::test_megamoe_plain_swiglu_carries_no_constants + - unittest/_torch/moe/test_moe_backend.py::test_create_moe_forwards_situ_activation_as_one_carrier + - unittest/_torch/moe/test_moe_backend.py::test_create_moe_backend_rejects_apply_router_weight_on_input_by_declaration + - unittest/_torch/moe/test_moe_backend.py::test_apply_router_weight_on_input_support_is_not_inherited # ------------- MoE: test_single_gpu (by backend) --------------- - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and not None" - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "TRTLLM" diff --git a/tests/microbenchmarks/bench_moe/build.py b/tests/microbenchmarks/bench_moe/build.py index 1021349a25be..2c6dc5079df5 100644 --- a/tests/microbenchmarks/bench_moe/build.py +++ b/tests/microbenchmarks/bench_moe/build.py @@ -29,12 +29,15 @@ import torch -from tensorrt_llm._torch.moe.fused_moe.interface import ( - MoESchedulerKind, - MoEWeightLoadingMode, - _compute_ep_partition, +from tensorrt_llm._torch.moe.fused_moe.activation import ( + ACTIVATION_PAYLOAD, + MoEActivation, + SimpleActivation, + SiTuActivation, + SwigluBiasActivation, ) -from tensorrt_llm._torch.utils import ActivationType, ActType_TrtllmGen +from tensorrt_llm._torch.moe.fused_moe.interface import MoESchedulerKind, MoEWeightLoadingMode +from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo @@ -93,14 +96,12 @@ def _epilogue_activation_name(moe) -> str: The SiTU request can be dropped for reasons the spec cannot see (wrong backend, wrong quant, upstream fallback), so read it back rather than - reporting what was asked for. MegaMoE DeepGEMM / CuteDSL store the - resolved name; TRTLLM-Gen exposes a predicate; CUTLASS uses - ``ActivationType.SiTu``. + reporting what was asked for. Every impl records the resolved kind in + ``activation_type``; TRTLLM-Gen additionally exposes a predicate, which is + checked first because SiTU shares its constant slots with SwiGLU there. """ backend = getattr(moe, "backend", None) or moe - if getattr(backend, "activation", None) == "situ" or getattr( - backend, "is_situ_activation", False - ): + if getattr(backend, "is_situ_activation", False): return "situ" activation_type = getattr(backend, "activation_type", None) if activation_type is not None and ActivationType(activation_type) == ActivationType.SiTu: @@ -122,64 +123,56 @@ def _calculate_num_chunks_safe(moe, all_rank_num_tokens: List[int]) -> Optional[ return None +#: Backends whose kernels implement the SiTU epilogue, with the quant they +#: implement it on. Not a capability check -- ``create_moe`` does that, and +#: rejects rather than degrades; this is the bench choosing which cases are +#: worth asking for. +_SITU_PATHS = frozenset( + { + ("MEGAMOE_DEEPGEMM", QuantAlgo.W4A8_MXFP4_MXFP8), + ("MEGAMOE_CUTEDSL", QuantAlgo.NVFP4), + ("TRTLLM", QuantAlgo.W4A8_MXFP4_MXFP8), + ("CUTLASS", QuantAlgo.NVFP4), + } +) + + def _situ_kwargs( model: ModelSpec, moe_backend: str, quant_algo: Optional[QuantAlgo], - mapping: Optional[Mapping] = None, ) -> Dict: """``create_moe`` kwargs that switch the epilogue to SiTU, or ``{}``. - Each backend takes SiTU through its own parameters and ``create_moe`` - REJECTS them on any other backend, so this dispatches instead of passing - one set everywhere. Quant is paired with the backend that actually - implements the epilogue; a spec carrying SiTU constants falls back to - the SwiGLU proxy elsewhere rather than failing the case -- SiTU is - gated, so the GEMM shapes and comm volume are the same either way. + A spec carrying SiTU constants falls back to the SwiGLU proxy on any other + backend/quant pair rather than failing the case -- SiTU is gated, so the + GEMM shapes and comm volume are the same either way. + + The constants go out as scalars for every backend: each one's + ``activation_support`` decides whether its kernels read them as a baked + scalar or a per-expert ``float32`` buffer, so the bench no longer sizes + tensors against the EP partition to match a particular kernel ABI. """ if model.situ_beta is None: return {} - backend = moe_backend.upper() - mega_situ = { - "activation": "situ", - "situ_beta": model.situ_beta, - "situ_linear_beta": model.situ_linear_beta, + if (moe_backend.upper(), quant_algo) not in _SITU_PATHS: + return {} + return { + "activation": SiTuActivation( + gate_softcap=float(model.situ_beta), + linear_softcap=float(model.situ_linear_beta), + ) } - if backend == "MEGAMOE_DEEPGEMM" and quant_algo == QuantAlgo.W4A8_MXFP4_MXFP8: - return mega_situ - if backend == "MEGAMOE_CUTEDSL" and quant_algo == QuantAlgo.NVFP4: - return mega_situ - if backend == "TRTLLM" and quant_algo == QuantAlgo.W4A8_MXFP4_MXFP8: - # Cubin alpha is the gate-side beta, cubin beta the linear-side one. - return { - "trtllm_gen_activation_type": ActType_TrtllmGen.SiTu, - "trtllm_gen_activation_alpha": model.situ_beta, - "trtllm_gen_activation_beta": model.situ_linear_beta, - } - if backend == "CUTLASS" and quant_algo == QuantAlgo.NVFP4: - # CUTLASS takes SiTU as ActivationType plus per-rank alpha/beta - # (same packing as modeling_kimi_linear). Size the tensors with - # the ceil/floor EP partition so uneven splits stay valid. - ep_size = 1 if mapping is None else max(mapping.moe_ep_size, 1) - ep_rank = 0 if mapping is None else mapping.moe_ep_rank - local_num_experts, _, _ = _compute_ep_partition(model.num_experts, ep_size, ep_rank) - device = torch.device("cuda", torch.cuda.current_device()) - return { - "activation_type": ActivationType.SiTu, - "swiglu_alpha": torch.full( - (local_num_experts,), - float(model.situ_beta), - dtype=torch.float32, - device=device, - ), - "swiglu_beta": torch.full( - (local_num_experts,), - float(model.situ_linear_beta), - dtype=torch.float32, - device=device, - ), - } - return {} + + +def _plain_activation(kind: ActivationType) -> MoEActivation: + """The spec's ``activation_type`` as a carrier, for the no-constants case. + + Reads ``ACTIVATION_PAYLOAD`` rather than naming the two kinds ``specs.py`` + currently allows, so a third one added there needs no change here. + """ + payload = ACTIVATION_PAYLOAD[ActivationType(kind)] + return SimpleActivation(kind) if payload is SimpleActivation else payload() def _create_moe_for_benchmark(**kwargs): @@ -370,12 +363,17 @@ def _build_moe_module( model_config=model_config, weight_loading_mode=weight_loading_mode, bias=swiglu_gptoss_style, - swiglu_alpha=swiglu_tensors["swiglu_alpha"] if swiglu_tensors else None, - swiglu_beta=swiglu_tensors["swiglu_beta"] if swiglu_tensors else None, - swiglu_limit=swiglu_tensors["swiglu_limit"] if swiglu_tensors else None, - activation_type=activation_type, + activation=( + SwigluBiasActivation( + gate_sigmoid_scale=swiglu_tensors["swiglu_alpha"], + linear_offset=swiglu_tensors["swiglu_beta"], + clamp=swiglu_tensors["swiglu_limit"], + ) + if swiglu_tensors + else _plain_activation(activation_type) + ), ) - moe_kwargs.update(_situ_kwargs(model, moe_backend, quant_algo, mapping)) + moe_kwargs.update(_situ_kwargs(model, moe_backend, quant_algo)) moe = _create_moe_for_benchmark(**moe_kwargs) if quant_algo == QuantAlgo.W4A8_MXFP4_MXFP8: diff --git a/tests/unittest/_torch/moe/fused_moe/test_configurable_moe.py b/tests/unittest/_torch/moe/fused_moe/test_configurable_moe.py index 9a27486aec4d..515009621485 100644 --- a/tests/unittest/_torch/moe/fused_moe/test_configurable_moe.py +++ b/tests/unittest/_torch/moe/fused_moe/test_configurable_moe.py @@ -19,7 +19,13 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm._torch.moe.fused_moe.activation import ( + DEFAULT_MOE_ACTIVATION, + ActivationParamShape, + MoEActivationSupport, +) from tensorrt_llm._torch.moe.fused_moe.configurable_moe import _BACKEND_SYNC_ATTRS, ConfigurableMoE +from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -27,6 +33,13 @@ def _wrapper() -> ConfigurableMoE: wrapper = ConfigurableMoE.__new__(ConfigurableMoE) torch.nn.Module.__init__(wrapper) wrapper.num_experts = 8 + # Divides ``num_experts`` so ``_reject_non_divisible_ep_backend`` returns at + # the divisibility check. It must not fall past it: the branch below reads + # ``type(self.backend)._supports_non_divisible_ep``, and the backend here is + # a ``Mock`` *instance*, so that lookup lands on the ``Mock`` class and + # raises. Assigned because ``__new__`` skipped the ``MoE.__init__`` that + # normally derives it from the mapping. + wrapper.ep_size = 1 wrapper.hidden_size = 16 wrapper.intermediate_size = 32 wrapper.dtype = torch.bfloat16 @@ -34,19 +47,51 @@ def _wrapper() -> ConfigurableMoE: wrapper.aux_stream_dict = None wrapper.weight_loading_mode = None wrapper.apply_router_weight_on_input = False - wrapper.activation_type = None + # Both are read while the backend is being built: the carrier goes to + # ``resolve_moe_cls`` and to ``install_activation_params``, the kind to + # ``infer_swiglu_gptoss_style``. ``__new__`` skipped the ``__init__`` that + # normally assigns them, so a fixture that sets neither fails before + # reaching the quant-config behaviour these tests are about. + wrapper.activation = DEFAULT_MOE_ACTIVATION + wrapper.activation_type = ActivationType(DEFAULT_MOE_ACTIVATION.kind) + wrapper.routing_method = Mock() wrapper._override_quant_config = None for attr in _BACKEND_SYNC_ATTRS: setattr(wrapper, attr, None) return wrapper +def _backend_mock() -> Mock: + """A ``Mock`` backend that survives ``install_activation_params``. + + That call cannot be patched out here: ``ConfigurableMoE`` makes it on the + backend after the EPLB sync *and* inside ``create_weights``, which is the + code path under test. A bare ``Mock`` slips past + ``resolve_activation_support`` -- every attribute of a Mock is callable, so + the override branch is taken -- and then fails inside + ``materialize_activation_params``, which uses the returned Mock as a real + declaration. Declaring a real one keeps the failure surface at the quant + config. Still GPU-free: ``DEFAULT_MOE_ACTIVATION`` carries no constants, so + every register short-circuits to None without allocating. + """ + backend = Mock() + backend.activation = DEFAULT_MOE_ACTIVATION + backend.resolve_activation_support = Mock( + return_value=MoEActivationSupport( + kinds=frozenset({ActivationType.Swiglu}), + alpha_beta=ActivationParamShape.PER_EXPERT_TENSOR, + limit=ActivationParamShape.PER_EXPERT_TENSOR, + ) + ) + return backend + + def _create_backend( wrapper: ConfigurableMoE, model_config: ModelConfig, override_quant_config: QuantConfig | None = None, ) -> Mock: - backend = Mock() + backend = _backend_mock() with ( patch( "tensorrt_llm._torch.moe.fused_moe.create_moe.resolve_moe_cls", diff --git a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py index 672210ae2c2e..af73c0d0e302 100644 --- a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py @@ -55,6 +55,7 @@ def _deployment( ep_size: int = 1, use_dp: bool = False, parallel_size: Optional[int] = None, + eplb: bool = False, ) -> MoEDeployment: """Declare the machine rather than patching the probes that read it.""" return MoEDeployment( @@ -63,6 +64,7 @@ def _deployment( parallel_size=ep_size if parallel_size is None else parallel_size, use_dp=use_dp, num_slots=8, + eplb_enabled=eplb, env=MoEEnvironment( sm=sm, available_deps=(MoEDep.FLASHINFER.value,) if flashinfer else (), @@ -138,6 +140,13 @@ def test_can_implement_rejects_missing_flashinfer(): assert verdict.reject_reason is MoERejectReason.DEP_MISSING +def test_can_implement_rejects_eplb(): + """EPLB is its own reject class, not a topology one: the machine is fine.""" + verdict = CuteDslB12xFusedMoE.can_implement(_problem(), _deployment(120, eplb=True)) + assert not verdict.eligible + assert verdict.reject_reason is MoERejectReason.EPLB_UNSUPPORTED + + def test_can_implement_rejects_expert_parallelism(): verdict = CuteDslB12xFusedMoE.can_implement(_problem(), _deployment(120, ep_size=2)) assert not verdict.eligible diff --git a/tests/unittest/_torch/moe/test_fused_moe.py b/tests/unittest/_torch/moe/test_fused_moe.py index db3b348ea02f..3b169745b075 100644 --- a/tests/unittest/_torch/moe/test_fused_moe.py +++ b/tests/unittest/_torch/moe/test_fused_moe.py @@ -1,2081 +1,30 @@ -import os -import pickle -import sys -from contextlib import contextmanager -from itertools import product from typing import Dict, List, Optional -import _torch.helpers -import cloudpickle import pytest import torch import torch.nn as nn import torch.nn.functional as F -from _torch.helpers import (calc_woq_tolerence, per_block_cast_to_fp8, - per_block_cast_to_fp8_e8m0, - per_token_cast_to_fp8_e8m0) -from mpi4py import MPI -from mpi4py.futures import MPIPoolExecutor -from transformers.configuration_utils import PretrainedConfig -from utils.util import (check_accuracy, skip_blackwell, skip_blackwell_geforce, - skip_neither_ada_nor_hopper_unittest, skip_no_hopper, - skip_pre_blackwell, skip_pre_hopper) +from utils.util import check_accuracy, skip_no_hopper -from tensorrt_llm._torch.autotuner import AutoTuner, autotune from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl import CuteDslFusedMoE -from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode # isort and yapf will fight against each other here, so we disable isort # isort: off from tensorrt_llm._torch.moe.fused_moe import (BaseMoeRoutingMethod, - CutlassFusedMoE, - TRTLLMGenFusedMoE, - DefaultMoeRoutingMethod, RenormalizeMoeRoutingMethod, - TritonFusedMoE, create_moe) + TritonFusedMoE) from tensorrt_llm._torch.moe.fused_moe.quantization import \ NVFP4CutlassFusedMoEMethod # isort: on from tensorrt_llm._torch.modules.gated_mlp import GatedMLP -from tensorrt_llm._utils import get_sm_version, mpi_rank +from tensorrt_llm._utils import mpi_rank from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig -# NOTE: Most tests in this file are deprecated and skipped. They are now covered by the -# unified MoE test framework in tests/unittest/_torch/moe/test_moe_backend.py -# and test_moe_module.py. Add new MoE tests there instead of here. - -cloudpickle.register_pickle_by_value(sys.modules[__name__]) -cloudpickle.register_pickle_by_value(_torch.helpers) -MPI.pickle.__init__( - cloudpickle.dumps, - cloudpickle.loads, - pickle.HIGHEST_PROTOCOL, -) - - -@contextmanager -def moe_trtllm_debug_msg(enable=False): - TLLM_BATCHED_GEMM_PRINT_NAME = os.environ.get( - "TLLM_BATCHED_GEMM_PRINT_NAME", "0") - TLLM_BATCHED_GEMM_PRINT_CONFIGS = os.environ.get( - "TLLM_BATCHED_GEMM_PRINT_CONFIGS", "0") - if enable: - os.environ["TLLM_BATCHED_GEMM_PRINT_NAME"] = "1" - os.environ["TLLM_BATCHED_GEMM_PRINT_CONFIGS"] = "1" - try: - yield - finally: - os.environ[ - "TLLM_BATCHED_GEMM_PRINT_NAME"] = TLLM_BATCHED_GEMM_PRINT_NAME - os.environ[ - "TLLM_BATCHED_GEMM_PRINT_CONFIGS"] = TLLM_BATCHED_GEMM_PRINT_CONFIGS - - -def round_up(x, alignment): - return (x + alignment - 1) // alignment * alignment - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@pytest.mark.parametrize( - "moe_backend, dtype, experts, routing_cls, bias", - product(["CUTLASS", "VANILLA", "TRITON"], [torch.float16, torch.bfloat16], - [3, 8, 512], [DefaultMoeRoutingMethod, RenormalizeMoeRoutingMethod], - [True, False])) -def test_fused_moe(moe_backend, - dtype, - experts, - routing_cls, - bias, - mapping=None): - - if moe_backend == "TRITON": - if get_sm_version() != 90: - pytest.skip("TRITON moe backend is only supported on Hopper") - if dtype != torch.bfloat16: - pytest.skip("Unsupported for TritonFusedMoE") - if routing_cls != RenormalizeMoeRoutingMethod: - pytest.skip("Unsupported for TritonFusedMoE") - - if bias and moe_backend not in ["TRITON"]: - pytest.skip("Bias not supported.") - - mapping = mapping or Mapping() - mapping.rank = mpi_rank() - AutoTuner.get().setup_distributed_state(mapping) - - torch.cuda.set_device(mapping.rank) - - with torch.device(f'cuda:{mapping.rank}'): - SEQ_LEN = 8 - HIDDEN_SIZE = 64 - INTERMEDIATE_SIZE = 32 - NUM_EXPERTS = experts - TOP_K = 2 - routing_method = routing_cls(top_k=TOP_K) - - torch.manual_seed(0) - torch.cuda.manual_seed(0) - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype, device="cuda") - router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), - dtype=dtype, - device="cuda") - - weights = {} - for expert_id in range(NUM_EXPERTS): - if bias: - w1_bias = torch.randn((INTERMEDIATE_SIZE, ), dtype=dtype).cuda() - w2_bias = torch.randn((HIDDEN_SIZE, ), dtype=dtype).cuda() - w3_bias = torch.randn((INTERMEDIATE_SIZE, ), dtype=dtype).cuda() - weights[f"{expert_id}.w1.bias"] = w1_bias - weights[f"{expert_id}.w2.bias"] = w2_bias - weights[f"{expert_id}.w3.bias"] = w3_bias - w1_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=dtype, - device="cuda") - w2_weight = torch.randn((HIDDEN_SIZE, INTERMEDIATE_SIZE), - dtype=dtype, - device="cuda") - w3_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=dtype, - device="cuda") - weights[f"{expert_id}.w1.weight"] = w1_weight - weights[f"{expert_id}.w2.weight"] = w2_weight - weights[f"{expert_id}.w3.weight"] = w3_weight - - # Create pretrained_config with necessary parameters - pretrained_config = PretrainedConfig() - pretrained_config.num_experts = NUM_EXPERTS - pretrained_config.hidden_size = HIDDEN_SIZE - pretrained_config.intermediate_size = INTERMEDIATE_SIZE - pretrained_config.torch_dtype = dtype - - fused_moe = create_moe( - routing_method=routing_method, - reduce_results=True, - model_config=ModelConfig(pretrained_config=pretrained_config, - mapping=mapping, - moe_backend=moe_backend), - bias=bias, - ) - fused_moe.load_weights([weights]) - fused_moe.cuda() - - AutoTuner.get().clear_cache() - with torch.inference_mode(), autotune(): - fused_moe.forward(x, router_logits) - - ref_fused_moe = RefGatedMLPFusedMoE(num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - model_config=ModelConfig(), - bias=bias) - ref_fused_moe.load_weights([weights]) - ref_fused_moe.cuda() - - # Evaluate the outputs on a variant sequence length to cover all possible keys in Autotuner cache - m = SEQ_LEN - while m >= 2: - x = torch.randn((m, HIDDEN_SIZE), dtype=dtype, device="cuda") - router_logits = torch.randn((m, NUM_EXPERTS), - dtype=dtype, - device="cuda") - - with torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - ref_output = ref_fused_moe.forward(x, router_logits) - - # Evaluate outputs - torch.cuda.synchronize() - # There can be one off mismatch in the outputs due to different kernel implementations - # Here we check most of the outputs are within the tolerance - # The CutlassFusedMoE case fails as well without this change on H100 for bf16 - check_accuracy(output, ref_output, rtol=0.2, atol=0.2, percent=0.975) - m //= 2 - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@pytest.mark.skipif(torch.cuda.device_count() < 4, - reason="needs 4 GPUs to run this test") -@pytest.mark.parametrize("moe_cls", ["CUTLASS", "VANILLA"]) -@pytest.mark.parametrize("ep_size", [1, 2, 4]) -def test_fused_moe_multi_gpu(moe_cls, ep_size): - world_size = 4 - with MPIPoolExecutor(max_workers=world_size) as executor: - results = executor.map( - test_fused_moe, - *zip( - *[(moe_cls, torch.bfloat16, 512, DefaultMoeRoutingMethod, False, - Mapping(world_size=world_size, - tp_size=world_size, - moe_ep_size=ep_size, - moe_tp_size=world_size // ep_size))] * world_size), - ) - for r in results: - assert r is None - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@skip_pre_hopper -@pytest.mark.parametrize( - "moe_backend", - ["CUTLASS", pytest.param("TRITON", marks=skip_no_hopper)]) -@pytest.mark.parametrize("routing_cls", - [DefaultMoeRoutingMethod, RenormalizeMoeRoutingMethod]) -@pytest.mark.parametrize("bias", [True, False]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fused_moe_fp8(moe_backend, dtype, routing_cls, bias): - - if moe_backend == "TRITON": - if dtype != torch.bfloat16: - pytest.skip("Unsupported for TritonFusedMoE") - if routing_cls != RenormalizeMoeRoutingMethod: - pytest.skip("Unsupported for TritonFusedMoE") - - if bias and moe_backend not in ["TRITON"]: - pytest.skip("Bias not supported.") - - mapping = Mapping() - mapping.rank = mpi_rank() - - with torch.device(f'cuda:{mapping.rank}'): - SEQ_LEN = 4 - HIDDEN_SIZE = 64 - INTERMEDIATE_SIZE = 32 - NUM_EXPERTS = 3 - TOP_K = 2 - routing_method = routing_cls(top_k=TOP_K) - torch.manual_seed(0) - torch.cuda.manual_seed(0) - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype, device="cuda") - _, x_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(x) - x_scale = x_scale.float().squeeze() - router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), - dtype=dtype, - device="cuda") - - weights = {} - for expert_id in range(NUM_EXPERTS): - if bias: - w1_bias = torch.randn((INTERMEDIATE_SIZE, ), dtype=dtype).cuda() - w2_bias = torch.randn((HIDDEN_SIZE, ), dtype=dtype).cuda() - w3_bias = torch.randn((INTERMEDIATE_SIZE, ), dtype=dtype).cuda() - weights[f"{expert_id}.w1.bias"] = w1_bias - weights[f"{expert_id}.w2.bias"] = w2_bias - weights[f"{expert_id}.w3.bias"] = w3_bias - w1_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=dtype, - device="cuda") - w2_weight = torch.randn((HIDDEN_SIZE, INTERMEDIATE_SIZE), - dtype=dtype, - device="cuda") - w3_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=dtype, - device="cuda") - - w1_weight_fp8, w1_weight_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor( - w1_weight) - w1_weight_fp8 = w1_weight_fp8.view(torch.float8_e4m3fn).cuda() - - w2_weight_fp8, w2_weight_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor( - w2_weight) - w2_weight_fp8 = w2_weight_fp8.view(torch.float8_e4m3fn).cuda() - - w3_weight_fp8, w3_weight_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor( - w3_weight) - w3_weight_fp8 = w3_weight_fp8.view(torch.float8_e4m3fn).cuda() - - w1_input_scale = x_scale.cuda() - w2_input_scale = x_scale.cuda() - w3_input_scale = x_scale.cuda() - - weights[f"{expert_id}.w1.weight"] = w1_weight_fp8 - weights[f"{expert_id}.w2.weight"] = w2_weight_fp8 - weights[f"{expert_id}.w3.weight"] = w3_weight_fp8 - weights[f"{expert_id}.w1.weight_scale"] = w1_weight_scale.float() - weights[f"{expert_id}.w2.weight_scale"] = w2_weight_scale.float() - weights[f"{expert_id}.w3.weight_scale"] = w3_weight_scale.float() - weights[f"{expert_id}.w1.input_scale"] = w1_input_scale - weights[f"{expert_id}.w2.input_scale"] = w2_input_scale - weights[f"{expert_id}.w3.input_scale"] = w3_input_scale - - # Create pretrained_config with necessary parameters - pretrained_config = PretrainedConfig() - pretrained_config.num_experts = NUM_EXPERTS - pretrained_config.hidden_size = HIDDEN_SIZE - pretrained_config.intermediate_size = INTERMEDIATE_SIZE - pretrained_config.torch_dtype = dtype - - quant_config = QuantConfig(quant_algo=QuantAlgo.FP8) - fused_moe = create_moe(routing_method=routing_method, - reduce_results=False, - model_config=ModelConfig( - pretrained_config=pretrained_config, - quant_config=quant_config, - moe_backend=moe_backend), - bias=bias) - fused_moe.cuda() - fused_moe.load_weights([weights]) - - ref_fused_moe = RefGatedMLPFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - model_config=ModelConfig(quant_config=quant_config), - bias=bias) - ref_fused_moe.load_weights([weights]) - ref_fused_moe.cuda() - - AutoTuner.get().clear_cache() - with torch.inference_mode(): - ref_output = ref_fused_moe.forward(x, router_logits) - - with torch.inference_mode(), autotune(): - fused_moe.forward(x, router_logits) - - # TRITON backend uses Triton kernels which don't register with AutoTuner - if moe_backend == "TRITON": - with torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - check_accuracy(output, - ref_output, - rtol=0.04, - atol=0.1, - percent=0.99) - else: - # Explicitly capture context for kernel testing - with AutoTuner.get().capture() as all_tactics, torch.inference_mode( - ): - output = fused_moe.forward(x, router_logits) - - # Test all kernel tactics - for tactic in all_tactics: - with AutoTuner.get().replay(tactic), torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - check_accuracy(output, - ref_output, - rtol=0.04, - atol=0.1, - percent=0.99) - - -def set_tensor_value_2(x, num_row, num_cols): - # Create 2x2 base pattern matrix - pattern = torch.tensor([[0.2, -0.5], [-0.3, 0.1]], device=x.device) - - # Repeat pattern to cover entire matrix - repeated = pattern.repeat((num_row + 1) // 2, - (num_cols + 1) // 2)[:num_row, :num_cols] - - x.copy_(repeated) - - -def set_tensor_value_3(x, num_row, num_cols): - # Create 3x3 base pattern matrix - pattern = torch.tensor( - [[0.1, 0.21, 0.31], [0.3, 0.6, 0.1], [0.11, 0.51, 0.62]], - device=x.device) - - # Repeat pattern to cover entire matrix - repeated = pattern.repeat((num_row + 2) // 3, - (num_cols + 2) // 3)[:num_row, :num_cols] - - x.copy_(repeated) - - -def set_tensor_value_4(x, num_row, num_cols): - # Create 4x4 base pattern matrix - pattern = torch.tensor( - [ - [0.1, 0.21, 0.31, 0.41], - [0.3, 0.6, 0.1, 0.2], - [0.11, 0.51, 0.61, 0.71], - [0.11, 0.52, 0.62, 0.72], - ], - device=x.device, - ) - - # Repeat pattern to cover entire matrix - repeated = pattern.repeat((num_row + 3) // 4, - (num_cols + 3) // 4)[:num_row, :num_cols] - - x.copy_(repeated) - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@skip_pre_blackwell -@pytest.mark.parametrize( - "dtype, num_experts, seq_len, hidden_size, RoutingMethodCls", - product( - [torch.bfloat16], - [72], - [128, 256, 384, 512, 1024, 2048, 4096, 8192], - [2560], - [DefaultMoeRoutingMethod], - ), -) -def test_fused_moe_fp8_blockwise_deepgemm(dtype, - num_experts, - seq_len, - hidden_size, - RoutingMethodCls, - mapping=None): - - SEQ_LEN = seq_len - HIDDEN_SIZE = hidden_size - INTERMEDIATE_SIZE = 256 - NUM_EXPERTS = num_experts - TOP_K = 2 - - routing_method = RoutingMethodCls(top_k=TOP_K) - - mapping = mapping or Mapping() - mapping.rank = mpi_rank() - torch.cuda.set_device(mapping.rank) - torch.manual_seed(0) - torch.cuda.manual_seed(0) - - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype).cuda() - # Note: we use some special values init x and weight, otherwise the test will false positive failed. - set_tensor_value_2(x, SEQ_LEN, HIDDEN_SIZE) - - x = x.cuda() - router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), dtype=dtype).cuda() - - weights = {} - w3_w1_weight_scales = [] - w2_weight_scales = [] - for expert_id in range(NUM_EXPERTS): - w1_weight = torch.randn( - (INTERMEDIATE_SIZE, HIDDEN_SIZE), dtype=dtype).cuda() / HIDDEN_SIZE - w2_weight = torch.randn((HIDDEN_SIZE, INTERMEDIATE_SIZE), - dtype=dtype).cuda() - w3_weight = torch.randn( - (INTERMEDIATE_SIZE, HIDDEN_SIZE), dtype=dtype).cuda() / HIDDEN_SIZE - set_tensor_value_3(w1_weight, INTERMEDIATE_SIZE, HIDDEN_SIZE) - set_tensor_value_4(w2_weight, HIDDEN_SIZE, INTERMEDIATE_SIZE) - set_tensor_value_3(w3_weight, INTERMEDIATE_SIZE, HIDDEN_SIZE) - - w1_weight_fp8, w1_weight_scale = per_block_cast_to_fp8_e8m0(w1_weight) - w1_weight_fp8 = w1_weight_fp8.view(torch.float8_e4m3fn).cuda() - - w2_weight_fp8, w2_weight_scale = per_block_cast_to_fp8_e8m0(w2_weight) - w2_weight_fp8 = w2_weight_fp8.view(torch.float8_e4m3fn).cuda() - - w3_weight_fp8, w3_weight_scale = per_block_cast_to_fp8_e8m0(w3_weight) - w3_weight_fp8 = w3_weight_fp8.view(torch.float8_e4m3fn).cuda() - - weights[f"{expert_id}.w1.weight"] = w1_weight_fp8 - weights[f"{expert_id}.w2.weight"] = w2_weight_fp8 - weights[f"{expert_id}.w3.weight"] = w3_weight_fp8 - weights[f"{expert_id}.w1.weight_scale_inv"] = w1_weight_scale - weights[f"{expert_id}.w2.weight_scale_inv"] = w2_weight_scale - weights[f"{expert_id}.w3.weight_scale_inv"] = w3_weight_scale - weights[f"{expert_id}.w1.weight_scale"] = w1_weight_scale - weights[f"{expert_id}.w2.weight_scale"] = w2_weight_scale - weights[f"{expert_id}.w3.weight_scale"] = w3_weight_scale - - w3_w1_weight_scales.append( - torch.cat([w3_weight_scale, w1_weight_scale], dim=0)) - w2_weight_scales.append(w2_weight_scale) - - w3_w1_weight_scales = torch.stack(w3_w1_weight_scales, dim=0).cuda() - w2_weight_scales = torch.stack(w2_weight_scales, dim=0).cuda() - - quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) - - # Create pretrained_config with necessary parameters - pretrained_config = PretrainedConfig() - pretrained_config.num_experts = NUM_EXPERTS - pretrained_config.hidden_size = HIDDEN_SIZE - pretrained_config.intermediate_size = INTERMEDIATE_SIZE - pretrained_config.torch_dtype = dtype - - fused_moe = create_moe( - routing_method=routing_method, - reduce_results=True, - model_config=ModelConfig(pretrained_config=pretrained_config, - quant_config=quant_config, - mapping=mapping, - moe_backend="DEEPGEMM"), - ) - fused_moe.cuda() - fused_moe.load_weights([weights]) - fused_moe.post_load_weights() - - def swiglu_fused_moe(x): - x, gate = x.chunk(2, dim=-1) - return torch.nn.functional.silu(gate) * x - - def grouped_gemm(a: torch.Tensor, b: torch.Tensor, a_sf: torch.Tensor, - b_sf: torch.Tensor, - offset_array: torch.Tensor) -> torch.Tensor: - d = torch.empty((a.shape[0], b.shape[1]), - device=b.device, - dtype=torch.bfloat16) - m_indices = torch.empty(a.shape[0], device=b.device, dtype=torch.int32) - for idx in range(offset_array.numel() - 1): - m_indices[offset_array[idx]:offset_array[idx + 1]] = idx - - num_groups, n, k_ = b.shape - d = torch.empty((a.shape[0], b.shape[1]), - device=b.device, - dtype=torch.bfloat16) - m_indices = torch.empty(a.shape[0], device=b.device, dtype=torch.int32) - for idx in range(offset_array.numel() - 1): - m_indices[offset_array[idx]:offset_array[idx + 1]] = idx - - for g in range(num_groups): - aa = a[offset_array[g]:offset_array[g + 1], :].to(torch.bfloat16) - aa_sf = a_sf[offset_array[g]:offset_array[g + 1], :] - aa_dq = aa * aa_sf.repeat_interleave( - 128, dim=1)[:aa.shape[0], :aa.shape[1]] - bb = b[g, :, :].to(torch.bfloat16) - bb_sf = b_sf[g, :, :] - bb_dq = bb * bb_sf.repeat_interleave(128, dim=0).repeat_interleave( - 128, dim=1)[:bb.shape[0], :bb.shape[1]] - d[offset_array[g]:offset_array[g + 1], :] = (aa_dq @ bb_dq.t()) - return d - - token_selected_experts, token_final_scales = routing_method.apply( - router_logits) - t_idx = 0 - permuted_data_tensor = torch.empty((x.shape[0] * TOP_K, x.shape[1]), - device=x.device, - dtype=torch.bfloat16) - expert_first_token_offset_tensor = torch.zeros(NUM_EXPERTS + 1, - dtype=torch.int32) - unpermute_map = [] - scales = [] - for e_idx in range(NUM_EXPERTS): - for idx, token in enumerate(x): - for i, selected_expert in enumerate(token_selected_experts[idx]): - if e_idx == selected_expert: - permuted_data_tensor[t_idx, :] = token - unpermute_map.append(idx) - scales.append(token_final_scales[idx, i]) - t_idx += 1 - expert_first_token_offset_tensor[e_idx + 1] = t_idx - - act_input_fp8, act_input_sf = per_token_cast_to_fp8_e8m0( - permuted_data_tensor) - h1 = grouped_gemm( - a=act_input_fp8, - b=fused_moe.w3_w1_weight, - a_sf=act_input_sf, - b_sf=w3_w1_weight_scales, - offset_array=expert_first_token_offset_tensor, - ) - h2 = swiglu_fused_moe(h1) - act_input_fp8, act_input_sf = per_token_cast_to_fp8_e8m0(h2) - h3 = grouped_gemm( - a=act_input_fp8, - b=fused_moe.w2_weight, - a_sf=act_input_sf, - b_sf=w2_weight_scales, - offset_array=expert_first_token_offset_tensor, - ) - ref_output = torch.zeros_like(x) - for token_idx, h3_token in enumerate(h3): - original_idx = unpermute_map[token_idx] - ref_output[original_idx, :] += h3_token * scales[token_idx] - - with torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - - # compare - torch.cuda.synchronize() - torch.testing.assert_close(output, ref_output, rtol=1e-2, atol=0.1) - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@skip_pre_blackwell -@pytest.mark.parametrize( - "dtype, num_experts, seq_len, hidden_size, RoutingMethodCls, WeightLoadingMode", - product( - [torch.bfloat16], - [72], - [128, 256, 384, 512, 1024, 2048, 4096, 8192], - [2560], - [DefaultMoeRoutingMethod], - [MoEWeightLoadingMode.VANILLA, MoEWeightLoadingMode.FUSED_GATE_UP_PROJ], - ), -) -def test_fused_moe_fp8_blockwise_cute_dsl(dtype, - num_experts, - seq_len, - hidden_size, - RoutingMethodCls, - WeightLoadingMode, - mapping=None): - SEQ_LEN = seq_len - HIDDEN_SIZE = hidden_size - INTERMEDIATE_SIZE = 1536 - NUM_EXPERTS = num_experts - TOP_K = 6 - - routing_method = RoutingMethodCls(top_k=TOP_K) - - mapping = mapping or Mapping() - mapping.rank = mpi_rank() - torch.cuda.set_device(mapping.rank) - torch.manual_seed(0) - torch.cuda.manual_seed(0) - - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype, device="cuda") - # Note: we use some special values init x and weight, otherwise the test will false positive failed. - set_tensor_value_2(x, SEQ_LEN, HIDDEN_SIZE) - - x = x.cuda() - router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), - dtype=dtype, - device="cuda") - - weights = {} - - if WeightLoadingMode == MoEWeightLoadingMode.FUSED_GATE_UP_PROJ: - weights['gate_up_proj'] = {} - weights['down_proj'] = {} - weights['gate_up_proj_weight_scale'] = {} - weights['down_proj_weight_scale'] = {} - - for expert_id in range(NUM_EXPERTS): - w1_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=dtype, - device="cuda") - w2_weight = torch.randn((HIDDEN_SIZE, INTERMEDIATE_SIZE), - dtype=dtype, - device="cuda") - w3_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=dtype, - device="cuda") - set_tensor_value_3(w1_weight, INTERMEDIATE_SIZE, HIDDEN_SIZE) - set_tensor_value_4(w2_weight, HIDDEN_SIZE, INTERMEDIATE_SIZE) - set_tensor_value_3(w3_weight, INTERMEDIATE_SIZE, HIDDEN_SIZE) - - w1_weight_fp8, w1_weight_scale = per_block_cast_to_fp8(w1_weight) - w1_weight_fp8 = w1_weight_fp8.view(torch.float8_e4m3fn).cuda() - - w2_weight_fp8, w2_weight_scale = per_block_cast_to_fp8(w2_weight) - w2_weight_fp8 = w2_weight_fp8.view(torch.float8_e4m3fn).cuda() - - w3_weight_fp8, w3_weight_scale = per_block_cast_to_fp8(w3_weight) - w3_weight_fp8 = w3_weight_fp8.view(torch.float8_e4m3fn).cuda() - - weights[f"{expert_id}.w1.weight"] = w1_weight_fp8 - weights[f"{expert_id}.w2.weight"] = w2_weight_fp8 - weights[f"{expert_id}.w3.weight"] = w3_weight_fp8 - weights[f"{expert_id}.w1.weight_scale"] = w1_weight_scale - weights[f"{expert_id}.w2.weight_scale"] = w2_weight_scale - weights[f"{expert_id}.w3.weight_scale"] = w3_weight_scale - - if WeightLoadingMode == MoEWeightLoadingMode.FUSED_GATE_UP_PROJ: - weights['gate_up_proj'][expert_id] = torch.cat( - [w3_weight_fp8, w1_weight_fp8], - dim=-2).transpose(0, 1).contiguous() - weights['down_proj'][expert_id] = w2_weight_fp8.transpose( - 0, 1).contiguous() - weights['gate_up_proj_weight_scale'][expert_id] = torch.cat( - [w3_weight_scale, w1_weight_scale], - dim=-2).transpose(0, 1).contiguous() - weights['down_proj_weight_scale'][ - expert_id] = w2_weight_scale.transpose(0, 1).contiguous() - elif WeightLoadingMode == MoEWeightLoadingMode.VANILLA: - weights[f"{expert_id}.w1.weight_scale_inv"] = w1_weight_scale - weights[f"{expert_id}.w2.weight_scale_inv"] = w2_weight_scale - weights[f"{expert_id}.w3.weight_scale_inv"] = w3_weight_scale - - quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) - - fused_moe = CuteDslFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - reduce_results=True, - model_config=ModelConfig(quant_config=quant_config, mapping=mapping), - weight_loading_mode=WeightLoadingMode, - ) - fused_moe.cuda() - fused_moe.load_weights([weights]) - - ref_fused_moe = RefGatedMLPFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - model_config=ModelConfig(quant_config=quant_config), - # Note: use deepgemm mm will cause accuracy error, so we use trtllmgen mm here - use_cute_dsl_blockscaling_mm=True, - ) - ref_fused_moe.load_weights([weights]) - ref_fused_moe.cuda() - - with torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - ref_output = ref_fused_moe.forward(x, router_logits) - - # compare - torch.cuda.synchronize() - torch.testing.assert_close(output, ref_output, rtol=1e-2, atol=0.1) - return True - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@skip_no_hopper -@pytest.mark.parametrize( - "dtype, num_experts, seq_len, hidden_size, RoutingMethodCls, WeightLoadingMode", - product( - [torch.bfloat16], - [72], - [128, 256, 384, 512, 1024, 2048, 4096, 8192], - [2560], - [DefaultMoeRoutingMethod], - [MoEWeightLoadingMode.VANILLA, MoEWeightLoadingMode.FUSED_GATE_UP_PROJ], - ), -) -def test_fused_moe_fp8_blockwise_cutlass(dtype, - num_experts, - seq_len, - hidden_size, - RoutingMethodCls, - WeightLoadingMode, - mapping=None): - SEQ_LEN = seq_len - HIDDEN_SIZE = hidden_size - INTERMEDIATE_SIZE = 1536 - NUM_EXPERTS = num_experts - TOP_K = 6 - - routing_method = RoutingMethodCls(top_k=TOP_K) - - mapping = mapping or Mapping() - mapping.rank = mpi_rank() - torch.cuda.set_device(mapping.rank) - torch.manual_seed(0) - torch.cuda.manual_seed(0) - - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype, device="cuda") - # Note: we use some special values init x and weight, otherwise the test will false positive failed. - set_tensor_value_2(x, SEQ_LEN, HIDDEN_SIZE) - - x = x.cuda() - router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), - dtype=dtype, - device="cuda") - - weights = {} - - if WeightLoadingMode == MoEWeightLoadingMode.FUSED_GATE_UP_PROJ: - weights['gate_up_proj'] = {} - weights['down_proj'] = {} - weights['gate_up_proj_weight_scale'] = {} - weights['down_proj_weight_scale'] = {} - - for expert_id in range(NUM_EXPERTS): - w1_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=dtype, - device="cuda") - w2_weight = torch.randn((HIDDEN_SIZE, INTERMEDIATE_SIZE), - dtype=dtype, - device="cuda") - w3_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=dtype, - device="cuda") - set_tensor_value_3(w1_weight, INTERMEDIATE_SIZE, HIDDEN_SIZE) - set_tensor_value_4(w2_weight, HIDDEN_SIZE, INTERMEDIATE_SIZE) - set_tensor_value_3(w3_weight, INTERMEDIATE_SIZE, HIDDEN_SIZE) - - w1_weight_fp8, w1_weight_scale = per_block_cast_to_fp8(w1_weight) - w1_weight_fp8 = w1_weight_fp8.view(torch.float8_e4m3fn).cuda() - - w2_weight_fp8, w2_weight_scale = per_block_cast_to_fp8(w2_weight) - w2_weight_fp8 = w2_weight_fp8.view(torch.float8_e4m3fn).cuda() - - w3_weight_fp8, w3_weight_scale = per_block_cast_to_fp8(w3_weight) - w3_weight_fp8 = w3_weight_fp8.view(torch.float8_e4m3fn).cuda() - - weights[f"{expert_id}.w1.weight"] = w1_weight_fp8 - weights[f"{expert_id}.w2.weight"] = w2_weight_fp8 - weights[f"{expert_id}.w3.weight"] = w3_weight_fp8 - weights[f"{expert_id}.w1.weight_scale"] = w1_weight_scale - weights[f"{expert_id}.w2.weight_scale"] = w2_weight_scale - weights[f"{expert_id}.w3.weight_scale"] = w3_weight_scale - - if WeightLoadingMode == MoEWeightLoadingMode.FUSED_GATE_UP_PROJ: - weights['gate_up_proj'][expert_id] = torch.cat( - [w1_weight_fp8, w3_weight_fp8], - dim=-2).transpose(0, 1).contiguous() - weights['down_proj'][expert_id] = w2_weight_fp8.transpose( - 0, 1).contiguous() - weights['gate_up_proj_weight_scale'][expert_id] = torch.cat( - [w3_weight_scale, w1_weight_scale], - dim=-2).transpose(0, 1).contiguous() - weights['down_proj_weight_scale'][ - expert_id] = w2_weight_scale.transpose(0, 1).contiguous() - elif WeightLoadingMode == MoEWeightLoadingMode.VANILLA: - weights[f"{expert_id}.w1.weight_scale_inv"] = w1_weight_scale - weights[f"{expert_id}.w2.weight_scale_inv"] = w2_weight_scale - weights[f"{expert_id}.w3.weight_scale_inv"] = w3_weight_scale - - quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) - - fused_moe = CutlassFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - reduce_results=True, - model_config=ModelConfig(quant_config=quant_config, mapping=mapping), - weight_loading_mode=WeightLoadingMode, - ) - fused_moe.cuda() - fused_moe.load_weights([weights]) - fused_moe.post_load_weights() - - ref_fused_moe = RefGatedMLPFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - model_config=ModelConfig(quant_config=quant_config), - ) - ref_fused_moe.cuda() - ref_fused_moe.load_weights([weights]) - ref_fused_moe.post_load_weights() - - with torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - ref_output = ref_fused_moe.forward(x, router_logits) - - # compare - torch.cuda.synchronize() - torch.testing.assert_close(output, ref_output, rtol=1e-2, atol=0.1) - return True - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@skip_no_hopper -@pytest.mark.skipif(torch.cuda.device_count() < 4, - reason="needs 4 GPUs to run this test") -@pytest.mark.parametrize("ep_size", [1, 2, 4]) -@pytest.mark.parametrize("routing_method", [DefaultMoeRoutingMethod]) -@pytest.mark.parametrize( - "weight_loading_mode", - [MoEWeightLoadingMode.VANILLA, MoEWeightLoadingMode.FUSED_GATE_UP_PROJ]) -def test_fused_moe_fp8_blockwise_cutlass_multi_gpu(ep_size, routing_method, - weight_loading_mode): - world_size = 4 - with MPIPoolExecutor(max_workers=world_size) as executor: - results = executor.map( - test_fused_moe_fp8_blockwise_cutlass, - *zip(*[( - torch.bfloat16, - 72, - 384, - 384, - routing_method, - weight_loading_mode, - Mapping( - world_size=world_size, - tp_size=world_size, - moe_ep_size=ep_size, - moe_tp_size=world_size // ep_size, - ), - )] * world_size), - ) - for r in results: - assert r is True - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@skip_pre_blackwell -@pytest.mark.skipif(torch.cuda.device_count() < 4, - reason="needs 4 GPUs to run this test") -@pytest.mark.parametrize("ep_size", [1, 2, 4]) -@pytest.mark.parametrize("routing_method", [DefaultMoeRoutingMethod]) -@pytest.mark.parametrize( - "weight_loading_mode", - [MoEWeightLoadingMode.VANILLA, MoEWeightLoadingMode.FUSED_GATE_UP_PROJ]) -def test_fused_moe_fp8_blockwise_cute_dsl_multi_gpu(ep_size, routing_method, - weight_loading_mode): - world_size = 4 - with MPIPoolExecutor(max_workers=world_size) as executor: - results = executor.map( - test_fused_moe_fp8_blockwise_cute_dsl, - *zip(*[( - torch.bfloat16, - 72, - 384, - 384, - routing_method, - weight_loading_mode, - Mapping( - world_size=world_size, - tp_size=world_size, - moe_ep_size=ep_size, - moe_tp_size=world_size // ep_size, - ), - )] * world_size), - ) - for r in results: - assert r is True - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@skip_pre_blackwell -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("moe_backend", [ - pytest.param("TRTLLM", marks=skip_blackwell_geforce), "CUTLASS", "CUTEDSL" -]) -@pytest.mark.parametrize( - "finalize_fusion", [True, False], - ids=["enable_finalize_fusion", "disable_finalize_fusion"]) -def test_fused_moe_nvfp4(dtype, moe_backend, finalize_fusion): - - run_fused_moe_nvfp4(dtype, moe_backend, finalize_fusion) - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@skip_pre_blackwell -@pytest.mark.parametrize("hidden_size, intermediate_size", [(2880, 2880)]) -@pytest.mark.parametrize("swiglu_alpha", [1, 0.1], ids=lambda v: f"alpha{v}") -@pytest.mark.parametrize("swiglu_beta", [0, 1], ids=lambda v: f"beta{v}") -@pytest.mark.parametrize("swiglu_limit", [float("inf"), 1], - ids=lambda v: f"limit{v}") -def test_fused_moe_nvfp4_gptoss_style(hidden_size, intermediate_size, - swiglu_alpha, swiglu_beta, swiglu_limit): - run_fused_moe_nvfp4(dtype=torch.bfloat16, - moe_backend="TRTLLM", - finalize_fusion=False, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - num_experts=32, - top_k=4, - seq_len=256, - gptoss_style=True, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - swiglu_limit=swiglu_limit) - - -def run_fused_moe_nvfp4(dtype, - moe_backend, - finalize_fusion, - hidden_size=512, - intermediate_size=512, - num_experts=8, - top_k=2, - seq_len=4, - gptoss_style=False, - swiglu_alpha=None, - swiglu_beta=None, - swiglu_limit=None): - - if moe_backend == "TRTLLM": - if dtype == torch.float16: - pytest.skip("TRTLLM NVFP4 MoE backend does not support float16 yet") - if finalize_fusion: - pytest.skip( - "TRTLLM NVFP4 MoE backend does not support fused finalize yet") - if moe_backend == "CUTEDSL": - if dtype == torch.float16: - pytest.skip( - "CUTEDSL NVFP4 MoE backend does not support float16 yet") - if get_sm_version() not in (100, 103): - pytest.skip( - "CUTEDSL NVFP4 MoE backend supports SM 100 (B200) and SM 103 (B300) only" - ) - - test_all_kernels = True - if get_sm_version() == 120: - # Disable; got "Assertion failed: Failed to initialize cutlass TMA WS grouped gemm. Error: Error Internal" - test_all_kernels = False - - mapping = Mapping() - mapping.rank = mpi_rank() - - with torch.device(f"cuda:{mapping.rank}"): - SCALING_VECTOR_SIZE = 16 - - SEQ_LEN = seq_len - HIDDEN_SIZE = hidden_size - INTERMEDIATE_SIZE = intermediate_size - NUM_EXPERTS = num_experts - TOP_K = top_k - routing_method = RenormalizeMoeRoutingMethod(top_k=TOP_K) - torch.manual_seed(0) - torch.cuda.manual_seed(0) - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype, device="cuda") - x_sf_global = (448 * 6) / x.abs().max().float() - router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), - dtype=dtype, - device="cuda") - - weights = {} - for expert_id in range(NUM_EXPERTS): - w1_weight = torch.randn( - (INTERMEDIATE_SIZE, HIDDEN_SIZE), dtype=dtype, - device="cuda") * 0.05 - w1_sf_global = (448 * 6) / w1_weight.abs().max().float() - - w2_weight = torch.randn( - (HIDDEN_SIZE, INTERMEDIATE_SIZE), dtype=dtype, - device="cuda") * 0.05 - w2_sf_global = (448 * 6) / w2_weight.abs().max().float() - - w3_weight = torch.randn( - (INTERMEDIATE_SIZE, HIDDEN_SIZE), dtype=dtype, - device="cuda") * 0.05 - w3_sf_global = (448 * 6) / w3_weight.abs().max().float() - - if gptoss_style: - w1_bias = torch.randn(INTERMEDIATE_SIZE, - device='cuda', - dtype=torch.float) - w2_bias = torch.randn(HIDDEN_SIZE, - device='cuda', - dtype=torch.float) - w3_bias = torch.randn(INTERMEDIATE_SIZE, - device='cuda', - dtype=torch.float) - weights[f"{expert_id}.w1.bias"] = w1_bias - weights[f"{expert_id}.w2.bias"] = w2_bias - weights[f"{expert_id}.w3.bias"] = w3_bias - - w3_w1_global = min( - w1_sf_global, - w3_sf_global) # w3 global and w1 global must be the same - - w1_weight_nvfp4, w1_sf_block_unswizzled = torch.ops.trtllm.fp4_quantize( - w1_weight, w3_w1_global, SCALING_VECTOR_SIZE, False, False) - w1_sf_block_unswizzled = w1_sf_block_unswizzled.view( - INTERMEDIATE_SIZE, -1) - - w2_weight_nvfp4, w2_sf_block_unswizzled = torch.ops.trtllm.fp4_quantize( - w2_weight, w2_sf_global, SCALING_VECTOR_SIZE, False, False) - w2_sf_block_unswizzled = w2_sf_block_unswizzled.view( - HIDDEN_SIZE, -1) - - w3_weight_nvfp4, w3_sf_block_unswizzled = torch.ops.trtllm.fp4_quantize( - w3_weight, w3_w1_global, SCALING_VECTOR_SIZE, False, False) - w3_sf_block_unswizzled = w3_sf_block_unswizzled.view( - INTERMEDIATE_SIZE, -1) - - w1_input_scale = x_sf_global.cuda() - w2_input_scale = x_sf_global.cuda() - w3_input_scale = x_sf_global.cuda() - - weights[f"{expert_id}.w1.weight"] = w1_weight_nvfp4 - weights[f"{expert_id}.w2.weight"] = w2_weight_nvfp4 - weights[f"{expert_id}.w3.weight"] = w3_weight_nvfp4 - weights[ - f"{expert_id}.w1.weight_scale"] = w1_sf_block_unswizzled.view( - torch.float8_e4m3fn).cuda() - weights[ - f"{expert_id}.w2.weight_scale"] = w2_sf_block_unswizzled.view( - torch.float8_e4m3fn).cuda() - weights[ - f"{expert_id}.w3.weight_scale"] = w3_sf_block_unswizzled.view( - torch.float8_e4m3fn).cuda() - weights[f"{expert_id}.w1.input_scale"] = 1.0 / w1_input_scale - weights[f"{expert_id}.w2.input_scale"] = 1.0 / w2_input_scale - weights[f"{expert_id}.w3.input_scale"] = 1.0 / w3_input_scale - weights[f"{expert_id}.w1.weight_scale_2"] = 1.0 / w3_w1_global - weights[f"{expert_id}.w2.weight_scale_2"] = 1.0 / w2_sf_global - weights[f"{expert_id}.w3.weight_scale_2"] = 1.0 / w3_w1_global - - swiglu_alpha_tensor = None - swiglu_beta_tensor = None - swiglu_limit_tensor = None - if gptoss_style: - swiglu_alpha_tensor = torch.full((NUM_EXPERTS, ), - swiglu_alpha, - device='cuda', - dtype=torch.float) - swiglu_beta_tensor = torch.full((NUM_EXPERTS, ), - swiglu_beta, - device='cuda', - dtype=torch.float) - swiglu_limit_tensor = torch.full((NUM_EXPERTS, ), - swiglu_limit, - device='cuda', - dtype=torch.float) - - quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) - - # Create pretrained_config with necessary parameters - pretrained_config = PretrainedConfig() - pretrained_config.num_experts = NUM_EXPERTS - pretrained_config.hidden_size = HIDDEN_SIZE - pretrained_config.intermediate_size = INTERMEDIATE_SIZE - pretrained_config.torch_dtype = dtype - - fused_moe = create_moe( - routing_method=routing_method, - reduce_results=True, - model_config=ModelConfig( - pretrained_config=pretrained_config, - quant_config=quant_config, - moe_backend=moe_backend, - moe_disable_finalize_fusion=not finalize_fusion), - bias=gptoss_style, - swiglu_alpha=swiglu_alpha_tensor, - swiglu_beta=swiglu_beta_tensor, - swiglu_limit=swiglu_limit_tensor, - ) - fused_moe.cuda() - fused_moe.load_weights([weights]) - fused_moe.post_load_weights() - - # Evaluate the outputs on a variant sequence length to cover all possible keys in Autotuner cache - ref_fused_moe = RefGatedMLPFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - model_config=ModelConfig(quant_config=quant_config), - bias=gptoss_style, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - swiglu_limit=swiglu_limit) - ref_fused_moe.cuda() - ref_fused_moe.load_weights([weights]) - - AutoTuner.get().clear_cache() - with torch.inference_mode(): - ref_output = ref_fused_moe.forward(x, router_logits) - - if not gptoss_style: - with torch.inference_mode(), autotune(): - fused_moe.forward(x, router_logits) - else: - # We skip autotune for gptoss style to reduce memory usage since the input shape is already quite large. - with torch.inference_mode(): - fused_moe.forward(x, router_logits) - - output = fused_moe.forward(x, router_logits) - - if gptoss_style: - rtol = 0.1 - atol = 0.1 - percent = 0.95 - else: - rtol = 1e-2 - atol = 0.15 - percent = None - - if gptoss_style: - check_accuracy(output, - ref_output, - rtol=rtol, - atol=atol, - percent=percent) - else: - torch.testing.assert_close(output, ref_output, rtol=rtol, atol=atol) - - if not test_all_kernels: - return - - # flashinfer has no capture and replay mechanisms, so we skip test_all_kernels - if getattr(fused_moe, 'use_flashinfer', False): - return - # Explicitly capture context for kernel testing - with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - - # Test all kernel tactics - for tactic in all_tactics: - with AutoTuner.get().replay(tactic), torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - if gptoss_style: - check_accuracy(output, - ref_output, - rtol=rtol, - atol=atol, - percent=percent) - else: - torch.testing.assert_close(output, - ref_output, - rtol=rtol, - atol=atol) - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@skip_pre_blackwell -@pytest.mark.parametrize( - "moe_backend", - [pytest.param("TRTLLM", marks=skip_blackwell_geforce), "CUTLASS"]) -def test_fused_moe_w4a8_nvfp4_fp8(moe_backend): - dtype = torch.bfloat16 - mapping = Mapping() - mapping.rank = mpi_rank() - - with torch.device(f'cuda:{mapping.rank}'): - SCALING_VECTOR_SIZE = 32 - - SEQ_LEN = 4 - HIDDEN_SIZE = 512 - INTERMEDIATE_SIZE = 512 - NUM_EXPERTS = 4 - TOP_K = 2 - routing_method = RenormalizeMoeRoutingMethod(top_k=TOP_K) - torch.manual_seed(0) - torch.cuda.manual_seed(0) - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype, device="cuda") - x_sf_global = 448 / x.abs().max().float() - router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), - dtype=dtype, - device="cuda") - - weights = {} - for expert_id in range(NUM_EXPERTS): - w1_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=torch.float32, - device="cpu") - w1_sf_global = (448) / w1_weight.abs().max().float() - - w2_weight = torch.randn((HIDDEN_SIZE, INTERMEDIATE_SIZE), - dtype=torch.float32, - device="cpu") - w2_sf_global = (448) / w2_weight.abs().max().float() - - w3_weight = torch.randn((INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=torch.float32, - device="cpu") - w3_sf_global = (448) / w3_weight.abs().max().float() - - w3_w1_global = min( - w1_sf_global, - w3_sf_global) # w3 global and w1 global must be the same - - w1_weight_nvfp4, w1_sf_block, _ = torch.ops.tensorrt_llm.float_to_e2m1_and_ufp8sf_scale( - w1_weight * w3_w1_global, SCALING_VECTOR_SIZE, 1, False) - w1_sf_block_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - w1_sf_block.view(INTERMEDIATE_SIZE, -1)) - - w2_weight_nvfp4, w2_sf_block, _ = torch.ops.tensorrt_llm.float_to_e2m1_and_ufp8sf_scale( - w2_weight * w2_sf_global, SCALING_VECTOR_SIZE, 1, False) - w2_sf_block_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - w2_sf_block.view(HIDDEN_SIZE, -1)) - - w3_weight_nvfp4, w3_sf_block, _ = torch.ops.tensorrt_llm.float_to_e2m1_and_ufp8sf_scale( - w3_weight * w3_w1_global, SCALING_VECTOR_SIZE, 1, False) - w3_sf_block_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - w3_sf_block.view(INTERMEDIATE_SIZE, -1)) - - w1_weight_nvfp4 = w1_weight_nvfp4.cuda() - w1_sf_block_unswizzled = w1_sf_block_unswizzled.cuda() - w2_weight_nvfp4 = w2_weight_nvfp4.cuda() - w2_sf_block_unswizzled = w2_sf_block_unswizzled.cuda() - w3_weight_nvfp4 = w3_weight_nvfp4.cuda() - w3_sf_block_unswizzled = w3_sf_block_unswizzled.cuda() - - w1_input_scale = x_sf_global.cuda() - w2_input_scale = x_sf_global.cuda() - w3_input_scale = x_sf_global.cuda() - - weights[f"{expert_id}.w1.weight"] = w1_weight_nvfp4 - weights[f"{expert_id}.w2.weight"] = w2_weight_nvfp4 - weights[f"{expert_id}.w3.weight"] = w3_weight_nvfp4 - weights[ - f"{expert_id}.w1.weight_scale"] = w1_sf_block_unswizzled.view( - torch.float8_e4m3fn).cuda() - weights[ - f"{expert_id}.w2.weight_scale"] = w2_sf_block_unswizzled.view( - torch.float8_e4m3fn).cuda() - weights[ - f"{expert_id}.w3.weight_scale"] = w3_sf_block_unswizzled.view( - torch.float8_e4m3fn).cuda() - weights[f"{expert_id}.w1.input_scale"] = 1.0 / w1_input_scale - weights[f"{expert_id}.w2.input_scale"] = 1.0 / w2_input_scale - weights[f"{expert_id}.w3.input_scale"] = 1.0 / w3_input_scale - weights[f"{expert_id}.w1.weight_scale_2"] = 1.0 / w3_w1_global - weights[f"{expert_id}.w2.weight_scale_2"] = 1.0 / w2_sf_global - weights[f"{expert_id}.w3.weight_scale_2"] = 1.0 / w3_w1_global - - quant_config = QuantConfig(quant_algo=QuantAlgo.W4A8_NVFP4_FP8) - fused_moe = TRTLLMGenFusedMoE(num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - reduce_results=False, - model_config=ModelConfig( - quant_config=quant_config, - moe_backend=moe_backend)) - fused_moe.load_weights([weights]) - fused_moe.cuda() - - # Evaluate the outputs on a variant sequence length to cover all possible keys in Autotuner cache - ref_fused_moe = RefGatedMLPFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - model_config=ModelConfig(quant_config=quant_config)) - ref_fused_moe.load_weights([weights]) - ref_fused_moe.cuda() - - AutoTuner.get().clear_cache() - with torch.inference_mode(): - ref_output = ref_fused_moe.forward(x, router_logits) - - with torch.inference_mode(), autotune(): - fused_moe.forward(x, router_logits) - - # Explicitly capture context for kernel testing - with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - - # Test all kernel tactics - for tactic in all_tactics: - with AutoTuner.get().replay(tactic), torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - torch.testing.assert_close(output, - ref_output, - rtol=1e-1, - atol=0.5) - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@skip_neither_ada_nor_hopper_unittest -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize( - "weight_loading_mode", - [MoEWeightLoadingMode.VANILLA, MoEWeightLoadingMode.W4A8_CUSTOM]) -def test_fused_moe_w4afp8(dtype, weight_loading_mode): - mapping = Mapping() - mapping.rank = mpi_rank() - - with torch.device(f'cuda:{mapping.rank}'): - SEQ_LEN = 4 - HIDDEN_SIZE = 768 - INTERMEDIATE_SIZE = 640 - SCALING_GROUP_SIZE = 128 - NUM_EXPERTS = 3 - TOP_K = 2 - routing_method = RenormalizeMoeRoutingMethod(top_k=TOP_K) - torch.manual_seed(0) - torch.cuda.manual_seed(0) - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype, device="cuda") - router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), - dtype=dtype, - device="cuda") - - affine_coeff = 0.005 - - lut = { - "weight": - "weight", - "weight_scale": - ("weight_scale_inv" if weight_loading_mode - == MoEWeightLoadingMode.W4A8_CUSTOM else "weight_scale"), - "weight_scale_2": - "weight_scale_2", - "pre_quant_scale": - "pre_quant_scale", - "input_scale": - "input_scale", - } - - weights = {} - for expert_id in range(NUM_EXPERTS): - # ModelOpt W4A8 packs pairs of 4b weights in the output dimension into one 8b element. - if weight_loading_mode == MoEWeightLoadingMode.VANILLA: - w1_shape = (INTERMEDIATE_SIZE // 2, HIDDEN_SIZE) - w2_shape = (HIDDEN_SIZE // 2, INTERMEDIATE_SIZE) - w3_shape = (INTERMEDIATE_SIZE // 2, HIDDEN_SIZE) - # The custom W4A8 quantization script examples/quantization/quantize_mixed_precision_moe.py - # packs pairs of 4b weight in the input dimension into one 8b element. - if weight_loading_mode == MoEWeightLoadingMode.W4A8_CUSTOM: - w1_shape = (INTERMEDIATE_SIZE, HIDDEN_SIZE // 2) - w2_shape = (HIDDEN_SIZE, INTERMEDIATE_SIZE // 2) - w3_shape = (INTERMEDIATE_SIZE, HIDDEN_SIZE // 2) - - # The weights in int4 precision. - w1_weight = torch.randint(-128, 127, w1_shape, - dtype=torch.int8).cuda() - w2_weight = torch.randint(-128, 127, w2_shape, - dtype=torch.int8).cuda() - w3_weight = torch.randint(-128, 127, w3_shape, - dtype=torch.int8).cuda() - - # The pre-quant scale to be multiplied with the input activation. - # Use random pre-quant scales [0.95, 1.05] instead of fixed 1.0 to ensure the kernel handles - # non-uniform pre-quant scaling factors correctly - w1_pre_quant_scale = torch.rand( - HIDDEN_SIZE, dtype=dtype, device="cuda") * 0.1 + 0.95 - w2_pre_quant_scale = torch.rand( - INTERMEDIATE_SIZE, dtype=dtype, device="cuda") * 0.1 + 0.95 - w3_pre_quant_scale = torch.rand( - HIDDEN_SIZE, dtype=dtype, device="cuda") * 0.1 + 0.95 - - # The weight scale to dequantize int4 weights (by multiplication). - w1_scale = torch.randn( - (INTERMEDIATE_SIZE, HIDDEN_SIZE // SCALING_GROUP_SIZE), - dtype=dtype, - device="cuda") * affine_coeff - w2_scale = torch.randn( - (HIDDEN_SIZE, INTERMEDIATE_SIZE // SCALING_GROUP_SIZE), - dtype=dtype, - device="cuda") * affine_coeff - w3_scale = torch.randn( - (INTERMEDIATE_SIZE, HIDDEN_SIZE // SCALING_GROUP_SIZE), - dtype=dtype, - device="cuda") * affine_coeff - - # The input scale to quantize the input activation (by division). - w1_input_scale = torch.randn(1, dtype=torch.float32, - device="cuda") * 0.2 - w2_input_scale = w1_input_scale - w3_input_scale = w1_input_scale - - # The weight scale 2 to quantize the dequantized weights (by division). - w1_weight_scale_2 = torch.ones([1], - dtype=torch.float32, - device="cuda") - w2_weight_scale_2 = w1_weight_scale_2 - w3_weight_scale_2 = w1_weight_scale_2 - - # Prepare weights. - weights[f"{expert_id}.w1.{lut['weight']}"] = w1_weight - weights[f"{expert_id}.w2.{lut['weight']}"] = w2_weight - weights[f"{expert_id}.w3.{lut['weight']}"] = w3_weight - weights[f"{expert_id}.w1.{lut['input_scale']}"] = w1_input_scale - weights[f"{expert_id}.w2.{lut['input_scale']}"] = w2_input_scale - weights[f"{expert_id}.w3.{lut['input_scale']}"] = w3_input_scale - weights[f"{expert_id}.w1.{lut['weight_scale']}"] = w1_scale - weights[f"{expert_id}.w2.{lut['weight_scale']}"] = w2_scale - weights[f"{expert_id}.w3.{lut['weight_scale']}"] = w3_scale - weights[ - f"{expert_id}.w1.{lut['pre_quant_scale']}"] = w1_pre_quant_scale - weights[ - f"{expert_id}.w2.{lut['pre_quant_scale']}"] = w2_pre_quant_scale - weights[ - f"{expert_id}.w3.{lut['pre_quant_scale']}"] = w3_pre_quant_scale - weights[ - f"{expert_id}.w1.{lut['weight_scale_2']}"] = w1_weight_scale_2 - weights[ - f"{expert_id}.w2.{lut['weight_scale_2']}"] = w2_weight_scale_2 - weights[ - f"{expert_id}.w3.{lut['weight_scale_2']}"] = w3_weight_scale_2 - - quant_config = QuantConfig(quant_algo=QuantAlgo.W4A8_AWQ) - fused_moe = CutlassFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - reduce_results=False, - model_config=ModelConfig(quant_config=quant_config), - weight_loading_mode=weight_loading_mode) - fused_moe.load_weights([weights]) - fused_moe.cuda() - - def ref(): - results = torch.zeros_like(x) - selected_experts, final_scales = routing_method.apply(router_logits) - for e_idx in range(NUM_EXPERTS): - mask = selected_experts == e_idx - activated_tokens = mask.sum(1).bool() - act = x[activated_tokens, :] - if act.shape[0] == 0: - continue - final_scale = (final_scales * - mask).sum(1)[activated_tokens].unsqueeze(1) - - # weights - def unpack_weights(weight: torch.Tensor) -> torch.Tensor: - unpacker = torch.ops.trtllm.unpack_int4_packed_tensor_to_int8 - if weight_loading_mode == MoEWeightLoadingMode.VANILLA: - return unpacker(weight.cpu().T.contiguous()).cuda() - else: - return unpacker(weight.cpu()).T.contiguous().cuda() - - w1 = unpack_weights(weights[f"{e_idx}.w1.{lut['weight']}"]) - w2 = unpack_weights(weights[f"{e_idx}.w2.{lut['weight']}"]) - w3 = unpack_weights(weights[f"{e_idx}.w3.{lut['weight']}"]) - w3_w1 = torch.cat([w3, w1], dim=-1) - - # weight_scale - s1 = weights[f"{e_idx}.w1.{lut['weight_scale']}"].T.contiguous( - ).cuda() - s2 = weights[f"{e_idx}.w2.{lut['weight_scale']}"].T.contiguous( - ).cuda() - s3 = weights[f"{e_idx}.w3.{lut['weight_scale']}"].T.contiguous( - ).cuda() - s3_s1 = torch.cat([s3, s1], dim=-1) - - # input_scale - p1 = weights[f"{e_idx}.w1.{lut['input_scale']}"].cuda() - p2 = weights[f"{e_idx}.w2.{lut['input_scale']}"].cuda() - p3 = weights[f"{e_idx}.w3.{lut['input_scale']}"].cuda() - p3_p1 = torch.max(p1, p3) - - # pre_quant_scale - a1 = a2 = a3 = a1_a3 = None - if weight_loading_mode == MoEWeightLoadingMode.VANILLA: - a1 = weights[ - f"{e_idx}.w1.{lut['pre_quant_scale']}"].T.contiguous( - ).cuda() - a2 = weights[ - f"{e_idx}.w2.{lut['pre_quant_scale']}"].T.contiguous( - ).cuda() - a3 = weights[ - f"{e_idx}.w3.{lut['pre_quant_scale']}"].T.contiguous( - ).cuda() - a1_a3 = torch.max(a1, a3) - - # weight_scale_2 - q1 = q2 = q3 = q3_q1 = None - if weight_loading_mode == MoEWeightLoadingMode.VANILLA: - q1 = weights[f"{e_idx}.w1.{lut['weight_scale_2']}"].cuda() - q2 = weights[f"{e_idx}.w3.{lut['weight_scale_2']}"].cuda() - q3 = weights[f"{e_idx}.w2.{lut['weight_scale_2']}"].cuda() - q3_q1 = torch.max(q3, q1) - - # forward pass - def process_layer( - act, - weight, - weight_scale, - input_scale, - pre_quant_scale=None, - weight_scale_2=None, - ): - if pre_quant_scale is not None: - act = act * pre_quant_scale - act = (torch.clamp((act / input_scale), -448.0, - 448.0).to(torch.float8_e4m3fn).to(dtype)) - weight = (weight.float() * weight_scale.repeat_interleave( - 128, dim=0).float()).to(dtype) - if weight_scale_2 is not None: - weight /= weight_scale_2 - output = torch.matmul(act, weight) * input_scale - if weight_scale_2 is not None: - output *= weight_scale_2 - return output - - # fc13 - fc1 = process_layer( - act, - w3_w1, - s3_s1, - p3_p1, - pre_quant_scale=a1_a3, - weight_scale_2=q3_q1, - ) - fc1, gate = fc1.chunk(2, dim=-1) - fc1 = fc1 * torch.nn.functional.silu(gate) - - # fc2 - fc2 = process_layer(fc1, - w2, - s2, - p2, - pre_quant_scale=a2, - weight_scale_2=q2) - - results[activated_tokens, :] += (fc2 * final_scale).to( - results.dtype) - return results - - AutoTuner.get().clear_cache() - with torch.inference_mode(): - ref_output = ref() - - with torch.inference_mode(), autotune(): - fused_moe.forward(x, router_logits) - - # Explicitly capture context for kernel testing - with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - - # Test all kernel tactics - for tactic in all_tactics: - with AutoTuner.get().replay(tactic), torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - # assert that result does not contain NaN or is all 0s - assert not torch.isnan(output).any(), "output contains NaN" - assert torch.nonzero(output).numel() > 0, "output is empty" - torch.testing.assert_close(output, - ref_output, - rtol=1e-2, - atol=0.1) - - torch.cuda.synchronize() - assert not torch.isnan(ref_output).any(), "ref_output contains NaN" - assert not torch.isnan(output).any(), "output contains NaN" - assert torch.nonzero(output).numel() > 0, "output is empty" - assert torch.nonzero(ref_output).numel() > 0, "ref_output is empty" - # Final comparison - torch.testing.assert_close(output, ref_output, rtol=1e-2, atol=0.1) - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@skip_pre_blackwell -@pytest.mark.parametrize( - "moe_backend", - [pytest.param("TRTLLM", marks=skip_blackwell_geforce), "CUTLASS"]) -@pytest.mark.parametrize("hidden_unpadded", [64, 192, 256]) -@pytest.mark.parametrize("seq_len", [8, 128]) -@pytest.mark.parametrize("bias", [True, False]) -def test_fused_moe_mxfp4_mxfp8(moe_backend, hidden_unpadded, seq_len, bias): - SCALING_VECTOR_SIZE = 32 - dtype = torch.bfloat16 - SEQ_LEN = seq_len - HIDDEN_SIZE_UNPADDED = hidden_unpadded - INTERMEDIATE_SIZE_UNPADDED = hidden_unpadded - NUM_EXPERTS = 8 - TOP_K = 4 - routing_method = RenormalizeMoeRoutingMethod(top_k=TOP_K) - torch.manual_seed(0) - torch.cuda.manual_seed(0) - router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), dtype=dtype).cuda() - - quant_config = QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8) - - # Create pretrained_config with necessary parameters - pretrained_config = PretrainedConfig() - pretrained_config.num_experts = NUM_EXPERTS - pretrained_config.hidden_size = HIDDEN_SIZE_UNPADDED - pretrained_config.intermediate_size = INTERMEDIATE_SIZE_UNPADDED - pretrained_config.torch_dtype = dtype - - fused_moe = create_moe( - routing_method=routing_method, - reduce_results=True, - model_config=ModelConfig(pretrained_config=pretrained_config, - quant_config=quant_config, - moe_backend=moe_backend), - bias=bias, - ) - fused_moe.cuda() - fused_moe.create_weights() - - num_elts_per_dtype = torch.iinfo( - fused_moe.quant_method.weight_dtype).bits // 4 - - HIDDEN_SIZE_IN = fused_moe.w3_w1_weight.shape[ - -1] * num_elts_per_dtype # last dim packed type factor - HIDDEN_SIZE_OUT = fused_moe.w2_weight.shape[-2] - INTERMEDIATE_SIZE = fused_moe.w2_weight.shape[ - -1] * num_elts_per_dtype # last dim packed type factor - - def dist_to_alignment(size, alignment): - return round_up(size, alignment) - size - - x = torch.randn((SEQ_LEN, HIDDEN_SIZE_UNPADDED), dtype=dtype).cuda() * 0.1 - x = torch.nn.functional.pad( - x, (0, dist_to_alignment(HIDDEN_SIZE_UNPADDED, HIDDEN_SIZE_IN))) - - def prepare_weights(num_experts: int, - hidden_size_in: int, - hidden_size_out: int, - intermediate_size: int, - bias: bool, - hidden_size_unpadded: int, - intermediate_size_unpadded: int, - pad_zero_or_val: bool, - weight_alignment: int = 128, - input_hidden_alignment: int = 512): - # Ensures each call gives same outcome - torch.manual_seed(42) - # Contamination value - contam_val = 42 - intermediate_size_unpadded = intermediate_size_unpadded or intermediate_size - weights = {} - for expert_id in range(num_experts): - if bias: - w1_bias = torch.randn( - (intermediate_size_unpadded, ), dtype=dtype).cuda() * 0.1 - w2_bias = torch.randn( - (hidden_size_unpadded, ), dtype=dtype).cuda() * 0.1 - w3_bias = torch.randn( - (intermediate_size_unpadded, ), dtype=dtype).cuda() * 0.1 - # Pad to output dimension using contamination - w1_bias = torch.nn.functional.pad( - w1_bias, - (0, dist_to_alignment(w1_bias.shape[-1], - intermediate_size)), "constant", - 0 if pad_zero_or_val else contam_val) - w2_bias = torch.nn.functional.pad( - w2_bias, - (0, dist_to_alignment(hidden_size_unpadded, - hidden_size_out)), "constant", - 0 if pad_zero_or_val else contam_val) - w3_bias = torch.nn.functional.pad( - w3_bias, - (0, dist_to_alignment(w3_bias.shape[-1], - intermediate_size)), "constant", - 0 if pad_zero_or_val else contam_val) - weights[f"{expert_id}.w1.bias"] = w1_bias - weights[f"{expert_id}.w2.bias"] = w2_bias - weights[f"{expert_id}.w3.bias"] = w3_bias - - w1_weight = torch.randn( - (intermediate_size_unpadded, hidden_size_unpadded), - dtype=dtype).cuda() * 0.1 - w2_weight = torch.randn( - (hidden_size_unpadded, intermediate_size_unpadded), - dtype=dtype).cuda() * 0.1 - w3_weight = torch.randn( - (intermediate_size_unpadded, hidden_size_unpadded), - dtype=dtype).cuda() - # First padding step: pad weight tensors from unpadded dimensions to weight-aligned dimensions using 0s - w1_weight = torch.nn.functional.pad( - w1_weight, (0, - dist_to_alignment(hidden_size_unpadded, - input_hidden_alignment), 0, - dist_to_alignment(intermediate_size_unpadded, - weight_alignment))) - w2_weight = torch.nn.functional.pad( - w2_weight, (0, - dist_to_alignment(intermediate_size_unpadded, - weight_alignment))) - w3_weight = torch.nn.functional.pad( - w3_weight, (0, - dist_to_alignment(hidden_size_unpadded, - input_hidden_alignment), 0, - dist_to_alignment(intermediate_size_unpadded, - weight_alignment))) - # Second padding step: pad from aligned dimensions to final dimensions using contamination - w1_weight = torch.nn.functional.pad( - w1_weight, - (0, dist_to_alignment(w1_weight.shape[-1], hidden_size_in), 0, - dist_to_alignment(w1_weight.shape[-2], intermediate_size)), - "constant", 0 if pad_zero_or_val else contam_val) - w2_weight = torch.nn.functional.pad( - w2_weight, - (0, dist_to_alignment(w2_weight.shape[-1], intermediate_size), - 0, dist_to_alignment(w2_weight.shape[-2], hidden_size_out)), - "constant", 0 if pad_zero_or_val else contam_val) - w3_weight = torch.nn.functional.pad( - w3_weight, - (0, dist_to_alignment(w3_weight.shape[-1], hidden_size_in), 0, - dist_to_alignment(w3_weight.shape[-2], intermediate_size)), - "constant", 0 if pad_zero_or_val else contam_val) - - w1_weight_mxfp4, w1_sf_block = torch.ops.trtllm.fp4_quantize( - w1_weight, None, SCALING_VECTOR_SIZE, True) - w1_sf_block_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - w1_sf_block.cpu().view(intermediate_size, -1)) - - w2_weight_mxfp4, w2_sf_block = torch.ops.trtllm.fp4_quantize( - w2_weight, None, SCALING_VECTOR_SIZE, True) - w2_sf_block_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - w2_sf_block.cpu().view(hidden_size_out, -1)) - - w3_weight_mxfp4, w3_sf_block = torch.ops.trtllm.fp4_quantize( - w3_weight, None, SCALING_VECTOR_SIZE, True) - w3_sf_block_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - w3_sf_block.cpu().view(intermediate_size, -1)) - - weights[f"{expert_id}.w1.weight"] = w1_weight_mxfp4 - weights[f"{expert_id}.w2.weight"] = w2_weight_mxfp4 - weights[f"{expert_id}.w3.weight"] = w3_weight_mxfp4 - weights[ - f"{expert_id}.w1.weight_scale"] = w1_sf_block_unswizzled.view( - torch.uint8).cuda() - weights[ - f"{expert_id}.w2.weight_scale"] = w2_sf_block_unswizzled.view( - torch.uint8).cuda() - weights[ - f"{expert_id}.w3.weight_scale"] = w3_sf_block_unswizzled.view( - torch.uint8).cuda() - return weights - - weights_pad_unsanitized = prepare_weights( - NUM_EXPERTS, - HIDDEN_SIZE_IN, - HIDDEN_SIZE_OUT, - INTERMEDIATE_SIZE, - bias, - HIDDEN_SIZE_UNPADDED, - INTERMEDIATE_SIZE_UNPADDED, - pad_zero_or_val=False, - weight_alignment=fused_moe.quant_method.weight_alignment, - input_hidden_alignment=getattr(fused_moe.quant_method, - "input_hidden_alignment", - fused_moe.quant_method.weight_alignment), - ) - fused_moe.cuda() - fused_moe.load_weights([weights_pad_unsanitized]) - fused_moe.post_load_weights() - - if moe_backend == "TRTLLM": - # Check sizes match. Note shape[-1] is in number of uint8 elements each containing 2x mxfp4 elements. - assert (fused_moe.quant_method.intermediate_size_per_partition_lean == - INTERMEDIATE_SIZE_UNPADDED) - assert fused_moe.w2_weight.shape[-1] * num_elts_per_dtype == round_up( - INTERMEDIATE_SIZE, fused_moe.quant_method.weight_alignment) - assert fused_moe.w3_w1_weight.shape[-1] * num_elts_per_dtype == round_up( - HIDDEN_SIZE_IN, fused_moe.quant_method.input_hidden_alignment) - - weights_pad_sanitized = prepare_weights( - NUM_EXPERTS, - HIDDEN_SIZE_IN, - HIDDEN_SIZE_IN, - INTERMEDIATE_SIZE, - bias, - HIDDEN_SIZE_UNPADDED, - INTERMEDIATE_SIZE_UNPADDED, - pad_zero_or_val=True, - weight_alignment=fused_moe.quant_method.weight_alignment, - input_hidden_alignment=fused_moe.quant_method.weight_alignment, - ) - ref_fused_moe = RefGatedMLPFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE_IN, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - bias=bias, - model_config=ModelConfig(quant_config=quant_config), - ) - ref_fused_moe.cuda() - ref_fused_moe.load_weights([weights_pad_sanitized]) - - AutoTuner.get().clear_cache() - with torch.inference_mode(): - ref_output = ref_fused_moe.forward(x, - router_logits)[:, :HIDDEN_SIZE_OUT] - - with torch.inference_mode(), autotune(): - fused_moe.forward(x, router_logits) - - # Capture fused_moe underlying runners as autotuner context - with AutoTuner.get().capture() as all_tactics, torch.inference_mode( - ), moe_trtllm_debug_msg(enable=False): - output = fused_moe.forward(x, router_logits) - output = torch.nn.functional.pad( - output, (0, HIDDEN_SIZE_OUT - HIDDEN_SIZE_UNPADDED)) - assert not torch.isnan(output).any(), "output contains NaN" - torch.testing.assert_close(output, ref_output, rtol=1e-2, atol=0.15) - - for i, tactic in enumerate(all_tactics): - with AutoTuner.get().replay(tactic), torch.inference_mode( - ), moe_trtllm_debug_msg(enable=False): - output = fused_moe.forward(x, router_logits) - output = torch.nn.functional.pad( - output, (0, HIDDEN_SIZE_OUT - HIDDEN_SIZE_UNPADDED)) - assert not torch.isnan(output).any( - ), f"tactic {tactic} at index {i} output contains NaN" - torch.testing.assert_close(output, ref_output, rtol=1e-2, atol=0.15) - - -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("hidden_size", [768, 2880]) -@pytest.mark.parametrize( - "moe_backend", - [ - # smVersion - pytest.param("TRTLLM", - marks=[skip_blackwell_geforce, skip_pre_blackwell]), - pytest.param( - "CUTLASS", - marks=[skip_pre_hopper, skip_blackwell, skip_blackwell_geforce]), - ], -) -def test_fused_moe_wfp4a16(dtype, hidden_size, moe_backend): - mapping = Mapping() - mapping.rank = mpi_rank() - - with torch.device(f'cuda:{mapping.rank}'): - SEQ_LEN = 4 - HIDDEN_SIZE = hidden_size - INTERMEDIATE_SIZE = 640 - SCALING_GROUP_SIZE = 32 - NUM_EXPERTS = 4 - TOP_K = 2 - routing_method = RenormalizeMoeRoutingMethod(top_k=TOP_K) - torch.manual_seed(0) - torch.cuda.manual_seed(0) - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype).cuda() - router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), dtype=dtype).cuda() - - weights = {} - for expert_id in range(NUM_EXPERTS): - w1_weight = torch.randint(0, - 256, - (INTERMEDIATE_SIZE, HIDDEN_SIZE // 2), - dtype=torch.uint8, - device='cuda') - w2_weight = torch.randint(0, - 256, - (HIDDEN_SIZE, INTERMEDIATE_SIZE // 2), - dtype=torch.uint8, - device='cuda') - w3_weight = torch.randint(0, - 256, - (INTERMEDIATE_SIZE, HIDDEN_SIZE // 2), - dtype=torch.uint8, - device='cuda') - - w1_scale = torch.randint( - 118, - 123, (INTERMEDIATE_SIZE, HIDDEN_SIZE // SCALING_GROUP_SIZE), - dtype=torch.uint8, - device='cuda') - w2_scale = torch.randint( - 118, - 123, (HIDDEN_SIZE, INTERMEDIATE_SIZE // SCALING_GROUP_SIZE), - dtype=torch.uint8, - device='cuda') - w3_scale = torch.randint( - 118, - 123, (INTERMEDIATE_SIZE, HIDDEN_SIZE // SCALING_GROUP_SIZE), - dtype=torch.uint8, - device='cuda') - - weights[f"{expert_id}.w1.weight"] = w1_weight - weights[f"{expert_id}.w2.weight"] = w2_weight - weights[f"{expert_id}.w3.weight"] = w3_weight - # WFP4A16FusedMoEMethod - weights[f"{expert_id}.w1.weight_scale_inv"] = w1_scale - weights[f"{expert_id}.w2.weight_scale_inv"] = w2_scale - weights[f"{expert_id}.w3.weight_scale_inv"] = w3_scale - # MXFP4WeightFusedMoEMethod - weights[f"{expert_id}.w1.weight_scale"] = w1_scale - weights[f"{expert_id}.w2.weight_scale"] = w2_scale - weights[f"{expert_id}.w3.weight_scale"] = w3_scale - - quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_MXFP4) - - # Create pretrained_config with necessary parameters - pretrained_config = PretrainedConfig() - pretrained_config.num_experts = NUM_EXPERTS - pretrained_config.hidden_size = HIDDEN_SIZE - pretrained_config.intermediate_size = INTERMEDIATE_SIZE - pretrained_config.torch_dtype = dtype - - fused_moe = create_moe(routing_method=routing_method, - reduce_results=False, - model_config=ModelConfig( - pretrained_config=pretrained_config, - quant_config=quant_config, - moe_backend=moe_backend)) - fused_moe.load_weights([weights]) - fused_moe.cuda() - - def ref(): - results = torch.zeros_like(x) - selected_experts, final_scales = routing_method.apply(router_logits) - unpacker = torch.ops.trtllm.mxfp4_dequantize_unswizzled - for e_idx in range(NUM_EXPERTS): - mask = selected_experts == e_idx - activated_tokens = mask.sum(1).bool() - act = x[activated_tokens, :] - if act.shape[0] == 0: - continue - final_scale = (final_scales * - mask).sum(1)[activated_tokens].unsqueeze(1) - - # weights and scales - w1 = weights[f"{e_idx}.w1.weight"] - s1 = weights[f"{e_idx}.w1.weight_scale_inv"] - w2 = weights[f"{e_idx}.w2.weight"] - s2 = weights[f"{e_idx}.w2.weight_scale_inv"] - w3 = weights[f"{e_idx}.w3.weight"] - s3 = weights[f"{e_idx}.w3.weight_scale_inv"] - - # converted weights - w1 = unpacker(w1.cpu(), s1.cpu(), SCALING_GROUP_SIZE).to( - dtype=x.dtype, device=x.device).T.contiguous() - w2 = unpacker(w2.cpu(), s2.cpu(), SCALING_GROUP_SIZE).to( - dtype=x.dtype, device=x.device).T.contiguous() - w3 = unpacker(w3.cpu(), s3.cpu(), SCALING_GROUP_SIZE).to( - dtype=x.dtype, device=x.device).T.contiguous() - w3_w1 = torch.cat([w3, w1], dim=-1) - - fc1 = torch.matmul(act, w3_w1) - fc1, gate = fc1.chunk(2, dim=-1) - fc1 = fc1 * torch.nn.functional.silu(gate) - fc2 = torch.matmul(fc1, w2) - results[activated_tokens, :] += (fc2 * final_scale).to( - results.dtype) - return results - - AutoTuner.get().clear_cache() - with torch.inference_mode(): - ref_output = ref() - - with torch.inference_mode(), autotune(): - fused_moe.forward(x, router_logits) - - # Explicitly capture context for kernel testing - with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - - # Test all kernel tactics - for tactic in all_tactics: - with AutoTuner.get().replay(tactic), torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - check_accuracy(output, - ref_output, - rtol=1e-2, - atol=0.1, - percent=0.99) - - # compare - torch.cuda.synchronize() - check_accuracy(output, ref_output, rtol=1e-2, atol=0.1, percent=0.99) +# NOTE: This file is what is left after the deprecated, permanently-skipped MoE +# tests were removed; the unified MoE test framework in +# tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py +# covers them. Add new MoE tests there, not here. @skip_no_hopper @@ -2248,128 +197,6 @@ def mxfp4_to_fp32(tensor, scales): check_accuracy(output, ref_output, rtol=0.6, atol=0.6, percent=0.945) -@pytest.mark.skip( - reason= - "Deprecated: covered by tests/unittest/_torch/moe/test_moe_backend.py and test_moe_module.py. Add new tests there." -) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("weight_dtype", [torch.int8]) -def test_fused_moe_int8_woq_per_channel(dtype, weight_dtype): - - mapping = Mapping() - mapping.rank = mpi_rank() - - with torch.device(f'cuda:{mapping.rank}'): - SEQ_LEN = 4 - HIDDEN_SIZE = 768 - INTERMEDIATE_SIZE = 640 - NUM_EXPERTS = 3 - TOP_K = 2 - routing_method = RenormalizeMoeRoutingMethod(top_k=TOP_K) - torch.manual_seed(0) - torch.cuda.manual_seed(0) - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype, device="cuda") - - router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), - dtype=dtype, - device="cuda") - - weight_id = 1 # 1 for w8a16, 2 for w4a16 - quant_config = QuantConfig(quant_algo=QuantAlgo.W8A16) - weights = {} - for expert_id in range(NUM_EXPERTS): - w1_weight = torch.randint( - -128, - 127, (INTERMEDIATE_SIZE, HIDDEN_SIZE // weight_id), - dtype=torch.int8).cuda() - w2_weight = torch.randint( - -128, - 127, (HIDDEN_SIZE, INTERMEDIATE_SIZE // weight_id), - dtype=torch.int8).cuda() - w3_weight = torch.randint( - -128, - 127, (INTERMEDIATE_SIZE, HIDDEN_SIZE // weight_id), - dtype=torch.int8).cuda() - - w1_scale = torch.randn( - (INTERMEDIATE_SIZE), dtype=dtype, device="cuda") / HIDDEN_SIZE - w2_scale = torch.randn( - (HIDDEN_SIZE), dtype=dtype, device="cuda") / INTERMEDIATE_SIZE - w3_scale = torch.randn( - (INTERMEDIATE_SIZE), dtype=dtype, device="cuda") / HIDDEN_SIZE - - weights[f"{expert_id}.w1.weight"] = w1_weight - weights[f"{expert_id}.w2.weight"] = w2_weight - weights[f"{expert_id}.w3.weight"] = w3_weight - weights[f"{expert_id}.w1.weight_scale"] = w1_scale - weights[f"{expert_id}.w2.weight_scale"] = w2_scale - weights[f"{expert_id}.w3.weight_scale"] = w3_scale - - fused_moe = CutlassFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - reduce_results=False, - model_config=ModelConfig(quant_config=quant_config)) - fused_moe.load_weights([weights]) - fused_moe.cuda() - - def ref(): - results = torch.zeros_like(x) - selected_experts, final_scales = routing_method.apply(router_logits) - for e_idx in range(NUM_EXPERTS): - mask = selected_experts == e_idx - activated_tokens = mask.sum(1).bool() - act = x[activated_tokens, :] - if act.shape[0] == 0: - continue - final_scale = (final_scales * - mask).sum(1)[activated_tokens].unsqueeze(1) - # weights - w1 = weights[f"{e_idx}.w1.weight"].T.contiguous().cuda() - w2 = weights[f"{e_idx}.w2.weight"].T.contiguous().cuda() - w3 = weights[f"{e_idx}.w3.weight"].T.contiguous().cuda() - w3_w1 = torch.cat([w3, w1], dim=-1) - # scales - s1 = weights[f"{e_idx}.w1.weight_scale"].cuda() - s2 = weights[f"{e_idx}.w2.weight_scale"].cuda() - s3 = weights[f"{e_idx}.w3.weight_scale"].cuda() - s3_s1 = torch.cat([s3, s1], dim=-1) - # calculation - w3_w1 = (w3_w1.float() * s3_s1).to(dtype) - fc1 = torch.matmul(act, w3_w1) - fc1, gate = fc1.chunk(2, dim=-1) - act = fc1 * torch.nn.functional.silu(gate) - w2 = (w2.float() * s2).to(dtype) - fc2 = torch.matmul(act, w2) - results[activated_tokens, :] += (fc2 * final_scale).to( - results.dtype) - return results - - AutoTuner.get().clear_cache() - with torch.inference_mode(): - ref_output = ref() - - with torch.inference_mode(), autotune(): - fused_moe.forward(x, router_logits) - - # Explicitly capture context for kernel testing - with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - - # Test all kernel tactics - atol = calc_woq_tolerence(ref_output, weight_dtype) - for tactic in all_tactics: - with AutoTuner.get().replay(tactic), torch.inference_mode(): - output = fused_moe.forward(x, router_logits) - torch.testing.assert_close(output, - ref_output, - rtol=1e-7, - atol=atol) - - class RefGatedMLPFusedMoE(nn.Module): def __init__(self, diff --git a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py index 927fefbff87e..98d3756dd8e9 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py @@ -21,9 +21,9 @@ * the fused path fails loudly without loaded weights (no silent random-weight fallback). -The NVFP4 half of SiTU lives here too, on both FP4 backends: CUTLASS takes -SiTU as an ``ActivationType``, TRTLLM-Gen as an out-of-band -``trtllm_gen_activation_type`` served by the fused +The NVFP4 half of SiTU lives here too, on both FP4 backends. Both take it as +one ``SiTuActivation`` carrier; they differ only in what serves it -- CUTLASS +an ``ActivationType`` its kernels branch on, TRTLLM-Gen the fused ``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*`` FC1 cubins. Tests that apply to both are parametrized over ``moe_backend`` rather than duplicated. """ @@ -339,7 +339,9 @@ def test_fused_forward_launches_situ_kernel(fmt): env["TLLM_BATCHED_GEMM_PRINT_NAME"] = "1" env["TLLM_LOG_LEVEL"] = "INFO" this_dir = os.path.dirname(os.path.abspath(__file__)) - unittest_root = os.path.abspath(os.path.join(this_dir, "..", "..", "..")) + # The child process imports ``_torch.moe.kimi_k3_ref_moe``, so the root it + # needs on PYTHONPATH is tests/unittest, not the repo's tests/ directory. + unittest_root = os.path.abspath(os.path.join(this_dir, "..", "..")) env["PYTHONPATH"] = os.pathsep.join([this_dir, unittest_root, env.get("PYTHONPATH", "")]) result = subprocess.run( [sys.executable, "-c", _LAUNCH_EVIDENCE_SCRIPTS[fmt]], @@ -682,7 +684,7 @@ def _make_routed_moe( """Mirror KimiK3MoERuntime's create_moe call on a single-rank mapping.""" from transformers.configuration_utils import PretrainedConfig - from tensorrt_llm._torch.moe.fused_moe import ConfigurableMoE, create_moe + from tensorrt_llm._torch.moe.fused_moe import ConfigurableMoE, SiTuActivation, create_moe from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig pretrained_config = PretrainedConfig() @@ -712,30 +714,10 @@ def _make_routed_moe( ), layer_idx=0, communication_method=None, + # Mirror KimiK3MoERuntime exactly: one activation for every backend, + # naming the two soft-caps rather than the ABI registers they land in. + activation=SiTuActivation(gate_softcap=4.0, linear_softcap=25.0), ) - if moe_backend == "TRTLLM": - moe_kwargs.update( - trtllm_gen_activation_type=ActType_TrtllmGen.SiTu, - trtllm_gen_activation_alpha=4.0, - trtllm_gen_activation_beta=25.0, - ) - elif moe_backend == "CUTLASS": - # Mirror KimiK3MoERuntime exactly: CUTLASS is the one backend that - # takes SiTU as an ActivationType (the others carry it out of band), - # and that choice decides the FC1 weight geometry. - from tensorrt_llm._torch.utils import ActivationType - - moe_kwargs.update( - activation_type=ActivationType.SiTu, - swiglu_alpha=torch.full((num_experts,), 4.0, dtype=torch.float32, device="cuda"), - swiglu_beta=torch.full((num_experts,), 25.0, dtype=torch.float32, device="cuda"), - ) - else: - moe_kwargs.update( - activation="situ", - situ_beta=4.0, - situ_linear_beta=25.0, - ) moe = create_moe(**moe_kwargs).cuda() assert isinstance(moe, ConfigurableMoE) return moe @@ -996,8 +978,8 @@ def block_scales(*shape): def _make_nvfp4_moe(gate, num_experts=_TP_EXPERTS, moe_backend="CUTLASS"): """NVFP4 + SiTU routed MoE on either FP4 backend. - CUTLASS takes SiTU as an ``ActivationType``; TRTLLM-Gen carries it out of - band as ``trtllm_gen_activation_type`` and serves it with the fused + Both take the same ``SiTuActivation`` carrier; CUTLASS serves it as an + ``ActivationType`` its kernels branch on, TRTLLM-Gen with the fused ``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*`` FC1 cubins (group-16 block scales). ``_make_routed_moe`` already mirrors both of KimiK3MoERuntime's branches, so the backend is the only variable. diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index c03fe19da7b0..fe8d18d77c53 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -46,7 +46,16 @@ DeepSeekV3MoeRoutingMethod, RenormalizeMoeRoutingMethod, ) +from tensorrt_llm._torch.moe.fused_moe.activation import ( + DEFAULT_MOE_ACTIVATION, + SimpleActivation, + SiTuActivation, + SwigluActivation, + SwigluBiasActivation, + materialize_activation_params, +) from tensorrt_llm._torch.moe.fused_moe.create_moe import create_moe_backend +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE from tensorrt_llm._torch.moe.fused_moe.fused_moe_marlin import MarlinFusedMoE from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE @@ -57,6 +66,7 @@ MoEProblem, MoERejectReason, MoERunContext, + MoEStaticCapability, ) from tensorrt_llm._torch.moe.fused_moe.impl_environment import ( collect_moe_environment, @@ -75,12 +85,7 @@ W4A8MXFP4MXFP8MegaMoEDeepGemmMethod, W4A16NVFP4CutlassFusedMoEMethod, ) -from tensorrt_llm._torch.utils import ( - ActivationType, - ActType_TrtllmGen, - MxFp8QuantizedTensor, - is_gated_activation, -) +from tensorrt_llm._torch.utils import ActivationType, MxFp8QuantizedTensor, is_gated_activation from tensorrt_llm._utils import get_sm_version, is_sm_100f, mpi_rank from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -197,7 +202,7 @@ def test_kimi_fused_route_quant_skips_prequantized_input(monkeypatch) -> None: ) assert ( - backend.try_fused_kimi_route_quant(hidden_states, torch.empty(1, 896, dtype=torch.float32)) + backend.try_fused_route_quant(hidden_states, torch.empty(1, 896, dtype=torch.float32)) is None ) @@ -222,6 +227,34 @@ def test_kimi_mxfp8_quantized_tensor_handoff() -> None: assert torch.equal(scales, scaling_factor) +def build_test_activation( + activation_type: ActivationType, + swiglu_alpha: Optional[torch.Tensor] = None, + swiglu_beta: Optional[torch.Tensor] = None, + swiglu_limit: Optional[torch.Tensor] = None, +) -> "SimpleActivation | SwigluActivation | SwigluBiasActivation | SiTuActivation": + """Package the flat parameters these tests parametrize over as one activation. + + The tests still sweep alpha / beta / limit independently because that is + what ``quantize_util.get_swiglu_tensors`` produces for the reference + implementation. Presence of alpha or beta means the gpt-oss package, which + is the same rule the C++ op applied when it upgraded a bare ``Swiglu`` with + constants to ``SwigluBias``. + """ + kind = ActivationType(activation_type) + if kind is ActivationType.SiTu: + return SiTuActivation(gate_softcap=swiglu_alpha, linear_softcap=swiglu_beta) + if swiglu_alpha is not None or swiglu_beta is not None: + return SwigluBiasActivation( + gate_sigmoid_scale=swiglu_alpha, + linear_offset=swiglu_beta, + clamp=swiglu_limit, + ) + if kind in (ActivationType.Swiglu, ActivationType.SwigluBias): + return SwigluActivation(clamp=swiglu_limit) + return SimpleActivation(kind=kind) + + def create_test_backend( backend_type: MoeBackendType, routing_method: RenormalizeMoeRoutingMethod, @@ -282,11 +315,8 @@ def create_test_backend( model_config=model_config, init_load_balancer=False, bias=bias, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - swiglu_limit=swiglu_limit, weight_loading_mode=weight_loading_mode, - activation_type=activation_type, + activation=build_test_activation(activation_type, swiglu_alpha, swiglu_beta, swiglu_limit), ) if n_shared_experts > 0: backend.create_weights() @@ -486,10 +516,11 @@ def test_marlin_override_quant_config_degrades_per_layer(): # ============================================================================ # TRTLLM-Gen SiTu backend contract # ============================================================================ -# SiTu rides the generic SwiGLU geometry and is carried out of band -# (trtllm_gen_activation_type, not ActivationType), so the host-side wiring is -# easy to get wrong in ways no shape check catches. Kernel-level coverage -- -# tactic availability, launch evidence and SiTu-vs-SwiGLU numerics -- lives in +# SiTu rides the generic SwiGLU geometry, so the host-side wiring is easy to +# get wrong in ways no shape check catches: it reaches the cubin through the +# same ``gemm1_alpha`` / ``gemm1_beta`` slots SwiGLU's constants use, and only +# the activation kind separates them. Kernel-level coverage -- tactic +# availability, launch evidence and SiTu-vs-SwiGLU numerics -- lives in # test_kimi_k3_situ_moe.py (and thop/serial/test_moe.py for the runner). These # are the contracts that hold without running a cubin. @@ -525,15 +556,6 @@ def _make_trtllm_gen_moe( mapping=Mapping(world_size=1, tp_size=1, rank=0), moe_backend="TRTLLM", ) - situ_kwargs = ( - dict( - trtllm_gen_activation_type=ActType_TrtllmGen.SiTu, - trtllm_gen_activation_alpha=_SITU_GATE_ALPHA, - trtllm_gen_activation_beta=_SITU_LINEAR_BETA, - ) - if situ - else {} - ) return TRTLLMGenFusedMoE( routing_method=RenormalizeMoeRoutingMethod(top_k=top_k), num_experts=num_experts, @@ -544,10 +566,16 @@ def _make_trtllm_gen_moe( model_config=model_config, init_load_balancer=False, weight_loading_mode=MoEWeightLoadingMode.VANILLA, - # SiTu rides the generic SwiGLU geometry; the fused activation is - # selected by trtllm_gen_activation_type, not by activation_type. - activation_type=ActivationType.Swiglu, - **situ_kwargs, + # The soft-caps go out as scalars; the backend's declared + # PER_EXPERT_TENSOR shape is what broadcasts them per slot. + activation=( + SiTuActivation( + gate_softcap=_SITU_GATE_ALPHA, + linear_softcap=_SITU_LINEAR_BETA, + ) + if situ + else DEFAULT_MOE_ACTIVATION + ), ) @@ -587,9 +615,9 @@ def _make_loaded_nvfp4_trtllm_gen_moe(situ: bool) -> TRTLLMGenFusedMoE: def test_trtllm_gen_nvfp4_situ_selects_padded_quant_method() -> None: """``_get_quant_method`` must key off ``is_situ_activation``. - SiTu fills the swiglu_alpha/swiglu_beta slots from ``create_weights``, - i.e. after ``_get_quant_method`` has already run, so keying off - ``swiglu_alpha is not None`` would make the selected method depend on + SiTu fills the act_alpha/act_beta slots from ``create_weights``, i.e. + after ``_get_quant_method`` has already run, so keying off + ``act_alpha is not None`` would make the selected method depend on *when* it is resolved. Plain SwiGLU is the control: it still gets the unpadded base method, so this is not asserting a constant. """ @@ -844,44 +872,34 @@ def test_megamoe_deepgemm_cache_derived_state_allocates_symm_buffer(): quant_method.cache_derived_state.assert_called_once_with(moe) -def test_megamoe_deepgemm_infers_kimi_situ_from_pretrained_config(): - model_config = ModelConfig( - pretrained_config=SimpleNamespace( - text_config=SimpleNamespace( - activation_situ_beta=4.0, - activation_situ_linear_beta=25.0, - ) - ) - ) - - activation, situ_beta, situ_linear_beta = MegaMoEDeepGemm._resolve_activation_config( - model_config, - activation=None, - situ_beta=None, - situ_linear_beta=None, +def test_megamoe_bakes_situ_softcaps_as_uniform_scalars(): + # MegaMoE declares UNIFORM_SCALAR for alpha/beta because the kernels bake + # them at codegen time, so a per-expert tensor is reduced here. + params = materialize_activation_params( + SiTuActivation(gate_softcap=torch.full((8,), 4.0), linear_softcap=25.0), + MegaMoEDeepGemm.activation_support, + num_local_experts=8, + owner="MegaMoEDeepGemm", ) - assert activation == "situ" - assert situ_beta == 4.0 - assert situ_linear_beta == 25.0 - + assert params.activation_type is ActivationType.SiTu + assert params.alpha == 4.0 + assert params.beta == 25.0 -def test_megamoe_deepgemm_defaults_to_swiglu_without_situ_config(): - model_config = ModelConfig(pretrained_config=SimpleNamespace()) - activation, situ_beta, situ_linear_beta = MegaMoEDeepGemm._resolve_activation_config( - model_config, - activation=None, - situ_beta=None, - situ_linear_beta=None, +def test_megamoe_plain_swiglu_carries_no_constants(): + params = materialize_activation_params( + SwigluActivation(), + MegaMoEDeepGemm.activation_support, + num_local_experts=8, + owner="MegaMoEDeepGemm", ) - assert activation == "swiglu" - assert situ_beta is None - assert situ_linear_beta is None + assert (params.alpha, params.beta) == (None, None) + assert params.clamp is None -def test_create_moe_forwards_megamoe_activation_options(monkeypatch): +def test_create_moe_forwards_situ_activation_as_one_carrier(monkeypatch): create_moe_module = importlib.import_module("tensorrt_llm._torch.moe.fused_moe.create_moe") configurable_moe = MagicMock(return_value=object()) monkeypatch.setattr(create_moe_module, "ConfigurableMoE", configurable_moe) @@ -890,6 +908,7 @@ def test_create_moe_forwards_megamoe_activation_options(monkeypatch): "resolve_moe_cls", MagicMock(return_value=MegaMoEDeepGemm), ) + activation = SiTuActivation(gate_softcap=4.0, linear_softcap=25.0) result = create_moe_module.create_moe( routing_method=MagicMock(), @@ -898,15 +917,41 @@ def test_create_moe_forwards_megamoe_activation_options(monkeypatch): intermediate_size=512, dtype=torch.bfloat16, model_config=ModelConfig(), - activation="situ", - situ_beta=4.0, - situ_linear_beta=25.0, + activation=activation, ) assert result is configurable_moe.return_value - assert configurable_moe.call_args.kwargs["activation"] == "situ" - assert configurable_moe.call_args.kwargs["situ_beta"] == 4.0 - assert configurable_moe.call_args.kwargs["situ_linear_beta"] == 25.0 + assert configurable_moe.call_args.kwargs["activation"] is activation + + +def test_create_moe_backend_rejects_apply_router_weight_on_input_by_declaration(): + """The gate runs ahead of every ``moe_cls`` branch, so a class the factory + has never heard of still reaches it -- which is the point of reading a + declaration instead of matching a class.""" + + class _Undeclared: + capabilities = MoEStaticCapability() + + with pytest.raises(ValueError, match="apply_router_weight_on_input"): + create_moe_backend( + moe_cls=_Undeclared, + routing_method=MagicMock(), + num_experts=8, + hidden_size=512, + intermediate_size=512, + apply_router_weight_on_input=True, + ) + + +def test_apply_router_weight_on_input_support_is_not_inherited(): + """``CuteDslB12xFusedMoE`` is the one impl that keeps its ``CutlassFusedMoE`` + parent, and this is a field where the two disagree: only the NVFP4 prefill + chunk reaches the parent's ``run_moe``, while the decode path hands + ``token_final_scales`` to the flashinfer wrapper.""" + assert CutlassFusedMoE.capabilities.supports_apply_router_weight_on_input + assert MarlinFusedMoE.capabilities.supports_apply_router_weight_on_input + assert not CuteDslB12xFusedMoE.capabilities.supports_apply_router_weight_on_input + assert not TRTLLMGenFusedMoE.capabilities.supports_apply_router_weight_on_input def test_megamoe_init_rejects_uneven_num_slots_with_value_error(): @@ -997,7 +1042,7 @@ def test_megamoe_cutedsl_tactic_autotune_defaults_off( def test_enumerate_megamoe_candidate_tactics_curated_space() -> None: - from tensorrt_llm._torch.custom_ops import cute_dsl_megamoe_custom_op as megamoe_op + from tensorrt_llm._torch.moe.custom_ops import cute_dsl_megamoe_custom_op as megamoe_op decode = megamoe_op.enumerate_megamoe_candidate_tactics(1024) prefill = megamoe_op.enumerate_megamoe_candidate_tactics(16384) diff --git a/tests/unittest/_torch/moe/test_moe_module.py b/tests/unittest/_torch/moe/test_moe_module.py index 0ecd83776600..489c42b8ab01 100644 --- a/tests/unittest/_torch/moe/test_moe_module.py +++ b/tests/unittest/_torch/moe/test_moe_module.py @@ -73,6 +73,7 @@ from tensorrt_llm._torch.autotuner import AutoTuner, autotune from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.moe.fused_moe import ( + DEFAULT_MOE_ACTIVATION, DeepSeekV3MoeRoutingMethod, DefaultMoeRoutingMethod, Llama4RenormalizeMoeRoutingMethod, @@ -81,6 +82,7 @@ RenormalizeMoeRoutingMethod, RenormalizeNaiveMoeRoutingMethod, SigmoidRenormMoeRoutingMethod, + SwigluBiasActivation, create_moe, ) from tensorrt_llm._torch.moe.fused_moe.communication.deep_ep_low_latency import DeepEPLowLatency @@ -658,9 +660,15 @@ def _test_moe_worker_impl( reduce_results=True, model_config=model_cfg, bias=swiglu_gptoss_style, - swiglu_alpha=swiglu_tensors["swiglu_alpha"] if swiglu_tensors else None, - swiglu_beta=swiglu_tensors["swiglu_beta"] if swiglu_tensors else None, - swiglu_limit=swiglu_tensors["swiglu_limit"] if swiglu_tensors else None, + activation=( + SwigluBiasActivation( + gate_sigmoid_scale=swiglu_tensors["swiglu_alpha"], + linear_offset=swiglu_tensors["swiglu_beta"], + clamp=swiglu_tensors["swiglu_limit"], + ) + if swiglu_tensors + else DEFAULT_MOE_ACTIVATION + ), weight_loading_mode=weight_loading_mode, ) as fused_moe, ):