From b66da437e400a65f53a91f61416e2a12230505c0 Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Fri, 29 May 2026 10:15:57 +0800 Subject: [PATCH 1/3] [None][fix] Restore DSv4 NVFP4 routed swiglu_limit on TRTLLM-Gen (#14673) Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> (cherry picked from commit 0cebdcd98a88977fbf9c56015a187a5c89af2657) --- .../_torch/models/modeling_deepseekv4.py | 8 ----- .../_torch/moe/fused_moe/quantization.py | 35 +++++++++---------- 2 files changed, 17 insertions(+), 26 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 54a97a70c3d5..5b71c2064834 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -1571,13 +1571,6 @@ def __init__( 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 = ( @@ -1586,7 +1579,6 @@ def __init__( ) 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) diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index d3149adf14ec..52b0b6a092be 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -5246,6 +5246,23 @@ 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 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, 'swiglu_beta', None) is not None: + module.swiglu_beta.data.div_((module.fc31_alpha.data)) + if getattr(module, 'swiglu_limit', None) is not None: + module.swiglu_limit.data.div_((module.fc31_alpha.data)) + def _shuffle_shared_expert_tensors(self, module: torch.nn.Module, num_elts_per_sf: int = 16): @@ -5739,24 +5756,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 From c4076288e0fc9b99cced19fe2d055cadba155abf Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Wed, 26 Aug 2026 00:37:43 -0700 Subject: [PATCH 2/3] [None][chore] Apply pre-commit formatting Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> (cherry picked from commit 80a5ec4b99fc1b7789ffb8d58923a31e6f6603fb) --- tensorrt_llm/_torch/models/modeling_deepseekv4.py | 5 +---- tensorrt_llm/_torch/moe/fused_moe/quantization.py | 1 + 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 5b71c2064834..301ae7fe00ac 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -1577,10 +1577,7 @@ def __init__( moe_cls is TRTLLMGenFusedMoE and experts_quant_config.quant_mode.has_fp8_block_scales() ) - if ( - supports_swiglu_limit - and not requires_scalar_only_swiglu_limit - ): + if supports_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 diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index 52b0b6a092be..261264f65d69 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -5756,6 +5756,7 @@ def _shuffle_all_experts(self, module.w3_w1_bias.data[expert_idx]) self._shuffle_w2_weight(module.w2_bias.data[expert_idx]) + class W4A8NVFP4FP8TRTLLMGenFusedMoEMethod(NVFP4TRTLLMGenFusedMoEBaseMethod): eplb_support_status = EplbSupportStatus.NOT_VERIFIED From 8f0f18310183dffeb4896f511c0b2f4101809356 Mon Sep 17 00:00:00 2001 From: xxi Date: Thu, 27 Aug 2026 00:37:55 +0000 Subject: [PATCH 3/3] [TRTLLM-14959][refactor] declare MoE backend capabilities instead of checking exact classes The MoE factory decided what a backend could accept by naming it. Fifteen `moe_cls is ` / `moe_cls in [...]` / `isinstance` gates across `create_moe`, `moe_scheduler` and two model files spelled capability as class identity, so a subclass of a supported backend was silently ineligible, and the same eligibility rule lived in the factory, the constructor and a late runtime error, free to disagree. It already did: MEGAMOE_DEEPGEMM reached SiTU only because `activation="situ"` and `activation_type` were two parameters that contradicted each other -- the factory built its problem from the unset `activation_type` default (Swiglu), passed the gate, then handed the constructor `"situ"` -- while the identical MEGAMOE_CUTEDSL request raised. Each backend now declares what it accepts and the factory reads the declaration: - `MoEActivationSupport` (new `activation.py`) states which `ActivationType`s a backend implements and, per constant, the shape its kernel boundary requires (`PER_EXPERT_TENSOR` / `UNIFORM_SCALAR` / `UNSUPPORTED`). The bias and swiglu-parameter allow-lists are gone. - Selection reads that declaration in one place. `MoEProblem` carries `activation_constants` -- which registers the caller actually fills, a question the kind alone cannot answer, since clamped and unclamped SwiGLU share one `ActivationType` but not one set of eligible backends -- and `_reject_unsupported_activation` declines a candidate whose declaration cannot carry them. Previously such a mismatch raised from the adapter at construction, past the point where another candidate could still be chosen. - `MoEStaticCapability.supports_eplb` replaces `supported_load_balancer_backends`, the same shape of list one layer down. `CuteDslB12xFusedMoE` now rejects `eplb_enabled` in `can_implement` instead of failing at construction. - `MoEStaticCapability.supports_apply_router_weight_on_input` replaces the three `assert not apply_router_weight_on_input` gates keyed on `moe_cls`. The fold itself is MoEScheduler's (`x = x * token_final_scales`), so what a backend declares is whether it handles what the fold leaves behind: `None` scales, or all-ones under a DeepEP / NCCL comm strategy. `CuteDslB12xFusedMoE` declares False where its `CutlassFusedMoE` parent declares True, because only the NVFP4 prefill chunk reaches the parent's `run_moe` while the decode path hands `token_final_scales` to the flashinfer wrapper -- a combination the class-keyed gates never checked, since they listed only three classes by name. - `MoEImplBase.try_fused_route_quant` replaces the `isinstance(moe.backend, TRTLLMGenFusedMoE)` branch in `moe_scheduler`, so a backend opts into the fused route+quant path by overriding it. The hook drops the model name the method carried while it lived on one class, because the contract is generic -- return what `routing_method.apply` and `quantize_input` would have produced, or `None` -- and only the implementation is specialized: the op it calls hardcodes 896 experts, top-16, hidden 3584 and 64 tokens, so the `TRTLLMGenFusedMoE` override still says Kimi K3. The second `isinstance` in that same branch went the other way, to `input_requirement.routing_scales_dtype`. Both replaced a class name, but one asked which dtype the backend reads, which is data a declaration can carry, and this one asks whether a faster path exists, which is behavior only the backend can answer. - Kimi K3's explicit-backend requests pass `allow_backend_degradation=False` instead of asserting on the resolved class, so a silent substitution fails loudly rather than running a different kernel. The `moe_cls in (...)` branches that remain in `create_moe` select which kwargs each constructor takes, and with the last eligibility assert lifted out of them they carry nothing else. They are signature dispatch, not eligibility, and are the factory's job. A declaration is only worth reading if no backend inherits one it never made. `CuteDslFusedMoE`, `DeepGemmFusedMoE` and `MarlinFusedMoE` derived from `CutlassFusedMoE` for constructor reuse while overriding both entry points, so each silently inherited its parent's capability object -- including `supports_expert_bias=True`, which none of them implements and which this change would otherwise have turned from a dormant inconsistency into a selection input. They now derive from `MoEImplBase` and share the constructor through `apply_moe_impl_construction_state`, and every impl declares `capabilities`, `input_requirement` and `activation_support` itself, so no declared value travels by inheritance. `CuteDslB12xFusedMoE` keeps its `CutlassFusedMoE` parent, the one case where that is a real dependency rather than reuse: `_route_to_cutlass` sends every NVFP4 prefill chunk through `CutlassFusedMoE.quantize_input` / `run_moe`, which read the whole Cutlass execution state. This also closes TRTLLM-15649, because the exact-class gates were only needed while one activation was described by several parameters at once: - One carrier. `MoEActivation` (`SwigluActivation`, `SiTuActivation`, ...) holds type, alpha, beta and clamp together, and is what the layer hands to the factory and to selection. `trtllm_gen_activation_type/alpha/beta` and MegaMoE's `activation` string are deleted; `ActType_TrtllmGen` becomes a C++ ABI encoding rather than a second user-facing enum. `build_moe_problem` and `resolve_moe_impl` take only the carrier: their `activation_type` parameter had no caller left, since the kind is derivable and the constants are not. - One clamp, at every layer. `swiglu_limit` (per-expert tensor) and `swiglu_limit_scalar` merge into a single `clamp`, and the materialized view keeps one field too: which ABI a backend gets is already pinned by its declaration, so `materialize_activation_params` renders the shape and no backend chooses between two slots. - Kind-neutral names above the ABI. The installed slots are `act_alpha` / `act_beta` / `act_clamp`, because SiTU fills the same registers the C++ boundary spells `swiglu_*` and every non-SwiGLU backend was renaming it back on arrival. The op-call sites keep each kernel's own schema names. The MegaMoE CuteDSL op, whose signature we own, is renamed to match; DeepGEMM's `situ_beta` stays because `_import_deep_gemm` gates on the installed signature accepting exactly that. - SiTU travels as `ActivationType.SiTu` with alpha/beta, collapsing MegaMoE's two translation tables into the declaration. Gates that meant "FC1 is gated" and were written `== ActivationType.Swiglu` now say `is_gated_activation`, which stops them from rejecting SiTU. The slots are installed once per writer of an input the materialization reads, under the rule that whoever writes that input last re-installs. `apply_moe_impl_construction_state` establishes the invariant for an impl built directly, which unit tests and microbenchmarks do and which never sees the wrapper. `ConfigurableMoE.create_weights` re-installs after applying the layerwise quant config, since `TRTLLMGenFusedMoE` narrows its clamp to a uniform scalar on the FP8 block-scale path and would otherwise keep a tensor clamp that kernel silently ignores; being the last write before `create_weights`, it is also the value every reader in the wrapper path sees. The EPLB slot sync and DWDP's `fixup_moe_backends` re-install because they rewrite `expert_size_per_partition`, which sizes a per-expert constant. `MoEActivationSupport.limit_when_absent` exists because CuteDSL's epilogue always applies the clamp functor: at that boundary "no clamp" is `+inf`, not `None`. A constructor fixup cannot express it, because `ConfigurableMoE` re-installs the activation slots after syncing `expert_size_per_partition` and would overwrite it. `tests/unittest/_torch/moe/test_fused_moe.py` drops the 15 tests that were already permanently skipped as covered by `test_moe_backend.py` and `test_moe_module.py`, together with the six helpers only those tests called -- `run_fused_moe_nvfp4` alone was 240 lines -- and the module-level cloudpickle / `MPI.pickle` bootstrap, which existed so the deleted `MPIPoolExecutor` cases could pickle test functions across ranks. The file goes 2639 -> 466 lines. The three live tests and the two reference classes they use are byte-identical, and no test list or CI stage referenced any deleted name. Tested on GB300 (aga) and B300 at an earlier revision of this change: `test_moe_backend` 363 passed, `test_moe_module` single_gpu 267 / multi_gpu 438 / multi_gpu_eplb 18, `test_kimi_k3_situ_moe` 64 passed + 1 xfailed. multi_gpu gives byte-identical counts on both clusters. Accuracy on B300: `TestNemotronV3Nano::test_fp8`, `TestGPTOSS::test_w4_1gpu`, `TestDeepSeekV4Flash::test_nvfp4_4gpus_static_eplb` and `TestMiniMaxM3::test_nvfp4[use_msa=False]` all pass. Re-verified on GB300 (aga) after rebasing onto origin/main, which is the run that covers the selection-time activation gate, the two re-install points, the single-field clamp and the reparent: `test_moe_backend` 371 + 8 passed, `test_moe_module` single_gpu 278 / multi_gpu 470 / multi_gpu_eplb 18 / DEEPGEMM 14 / CUTEDSL-NVFP4 18 / DENSEGEMM 5, `test_kimi_k3_situ_moe` 69 passed + 1 xfailed, `test_configurable_moe` 2 passed, and the trimmed `test_fused_moe` 9 passed + 64 skipped. Every count is identical to the pre-rebase run of the same matrix, which is the signal that neither the reparent nor the upstream span the rebase pulled in -- the relocation of `_torch/modules/fused_moe` to `_torch/moe/fused_moe` among it -- changed anything observable. The 21 architectural gates in that job all pass, including `fc31_normalization_in_base_class` and `dsv4_swiglu_limit_war_removed`, which are what the two carried commits from #18239 exist to satisfy. `test_configurable_moe` goes 2 failed -> 2 passed relative to the pre-reparent run: its wrapper fixture builds a `ConfigurableMoE` through `__new__` and never set `ep_size`, which `_reject_non_divisible_ep_backend` reads. The accuracy cases run on B300 rather than aga and are covered by the run above. Signed-off-by: xxi --- .../_torch/models/modeling_deepseekv4.py | 61 +- tensorrt_llm/_torch/models/modeling_gemma4.py | 3 +- .../_torch/models/modeling_gpt_oss.py | 22 +- .../_torch/models/modeling_kimi_linear.py | 85 +- .../_torch/models/modeling_minimaxm3.py | 86 +- .../_torch/models/modeling_nemotron_h.py | 6 +- tensorrt_llm/_torch/modules/dwdp/setup.py | 11 + .../custom_ops/cute_dsl_megamoe_custom_op.py | 29 +- .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 88 +- tensorrt_llm/_torch/moe/fused_moe/__init__.py | 13 + .../_torch/moe/fused_moe/activation.py | 588 +++++ .../_torch/moe/fused_moe/configurable_moe.py | 55 +- .../_torch/moe/fused_moe/create_moe.py | 236 +- .../moe/fused_moe/fused_moe_cute_dsl.py | 88 +- .../moe/fused_moe/fused_moe_cute_dsl_b12x.py | 66 +- .../_torch/moe/fused_moe/fused_moe_cutlass.py | 60 +- .../moe/fused_moe/fused_moe_deepgemm.py | 54 +- .../moe/fused_moe/fused_moe_densegemm.py | 13 +- .../_torch/moe/fused_moe/fused_moe_marlin.py | 88 +- .../_torch/moe/fused_moe/fused_moe_triton.py | 73 +- .../moe/fused_moe/fused_moe_trtllm_gen.py | 254 +- .../_torch/moe/fused_moe/fused_moe_vanilla.py | 21 +- .../_torch/moe/fused_moe/impl_base.py | 41 +- .../_torch/moe/fused_moe/impl_contract.py | 23 + .../_torch/moe/fused_moe/interface.py | 28 +- .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 149 +- .../fused_moe/mega_moe/mega_moe_deepgemm.py | 120 +- .../_torch/moe/fused_moe/moe_resolution.py | 124 +- .../_torch/moe/fused_moe/moe_scheduler.py | 10 +- .../_torch/moe/fused_moe/quantization.py | 11 +- .../test_lists/test-db/l0_b200.yml | 8 +- tests/microbenchmarks/bench_moe/build.py | 124 +- .../moe/fused_moe/test_configurable_moe.py | 49 +- .../moe/test_cute_dsl_b12x_moe_backend.py | 9 + tests/unittest/_torch/moe/test_fused_moe.py | 2187 +---------------- .../_torch/moe/test_kimi_k3_situ_moe.py | 42 +- tests/unittest/_torch/moe/test_moe_backend.py | 181 +- tests/unittest/_torch/moe/test_moe_module.py | 14 +- 38 files changed, 1790 insertions(+), 3330 deletions(-) create mode 100644 tensorrt_llm/_torch/moe/fused_moe/activation.py diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 301ae7fe00ac..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,50 +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, - ) - # 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 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, @@ -1608,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 7e4c11060a59..e873b97e1e80 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 @@ -640,7 +641,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 261264f65d69..c947a316ab9d 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -5247,7 +5247,8 @@ def process_weights_after_loading(self, 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). + # 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 @@ -5258,10 +5259,10 @@ def process_weights_after_loading(self, 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, 'swiglu_beta', None) is not None: - module.swiglu_beta.data.div_((module.fc31_alpha.data)) - if getattr(module, 'swiglu_limit', None) is not None: - module.swiglu_limit.data.div_((module.fc31_alpha.data)) + 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, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 96bacf5215b7..342a20039964 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -145,12 +145,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, ):