diff --git a/csrc/trtllm_batched_gemm_runner.cu b/csrc/trtllm_batched_gemm_runner.cu index 277719feff8..aa53bcd7ce4 100644 --- a/csrc/trtllm_batched_gemm_runner.cu +++ b/csrc/trtllm_batched_gemm_runner.cu @@ -117,9 +117,20 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner( if (options.mFusedBiasShuffleMode != mOptions.fusedBiasShuffleMode) continue; if (options.mBiasDtype != mOptions.biasDtype) continue; } + bool const usesPerTokenScaling = + options.mTransposeMmaOutput ? options.mUsePerTokenSfB : options.mUsePerTokenSfA; if (mOptions.usePerTokenScaling) { - if (options.mTransposeMmaOutput && !options.mUsePerTokenSfB) continue; - if (!options.mTransposeMmaOutput && !options.mUsePerTokenSfA) continue; + if (!usesPerTokenScaling) continue; + if (options.mPerTokenSfDtype != mOptions.perTokenSfDtype) continue; + } + // The MoE pipeline allocates and consumes output scaling factors in the + // block format's default dtype. Reject cubins that override that + // contract, such as bmm_E2m1xFp32_* kernels that emit linear FP32 + // scaling factors into a buffer sized for E4M3 factors. + if (tg::dtypeIsBlockFmt(options.mDtypeC)) { + if (options.mDtypeSfC != tg::dtypeGetBlockSfType(options.mDtypeC)) continue; + } else if (options.mDtypeSfC != Dtype::Void) { + continue; } if (mOptions.usePerChannelScaling) { if (options.mTransposeMmaOutput && !options.mUsePerTokenSfA) continue; @@ -164,6 +175,7 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner( << ", mFusedBiasShuffleMode: " << (int64_t)mOptions.fusedBiasShuffleMode << ", mBiasDtype: " << tg::dtypeToString(mOptions.biasDtype) << ", mUsePerTokenScaling: " << mOptions.usePerTokenScaling + << ", mPerTokenSfDtype: " << tg::dtypeToString(mOptions.perTokenSfDtype) << ", mUsePerChannelScaling: " << mOptions.usePerChannelScaling; FLASHINFER_CHECK(!mPassingConfigIndices.empty(), error_msg.str()); } diff --git a/csrc/trtllm_fused_moe_kernel_launcher.cu b/csrc/trtllm_fused_moe_kernel_launcher.cu index 02cc8a3ab79..4fa81e69423 100644 --- a/csrc/trtllm_fused_moe_kernel_launcher.cu +++ b/csrc/trtllm_fused_moe_kernel_launcher.cu @@ -513,12 +513,7 @@ class FusedMoeLauncher { void prepare_moe_common(int64_t& moe_tactic) { using RunnerType = tensorrt_llm::kernels::trtllmgen_moe::MoE::Runner; - // FIXME(siyuan): check llama4 routing after the fp4 FC1 kernels with bf16 scale factors were - // generated - bool usePerTokenScalingGemm1 = - per_token_scales.has_value() /* || - static_cast(this->routing_method_type) == RoutingMethodType::Llama4*/ - ; + bool usePerTokenScalingGemm1 = per_token_scales.has_value() || args->mUseRoutingScalesOnInput; // FIXME(siyuan): currently only nvfp4 x nvfp4 uses per-token scaling in both FC1 and FC2 bool usePerTokenScalingGemm2 = per_token_scales.has_value() && mDtypeAct == btg::Dtype::E2m1; // For FP8 block-scale (E4m3 activations, E4m3 weights) with DeepSeek FP8 and no @@ -947,6 +942,7 @@ class Fp8PerTensorLauncher : public FusedMoeLauncher { int64_t weight_layout, bool use_routing_scales_on_input_param, ActivationType activation_type, bool norm_topk_prob = true) { this->use_routing_scales_on_input = use_routing_scales_on_input_param; + args->mUseRoutingScalesOnInput = use_routing_scales_on_input_param; auto dtype = hidden_states.dtype(); if (dtype == dl_float16) { @@ -1175,7 +1171,8 @@ class Fp8PerTensorLauncher : public FusedMoeLauncher { int64_t intermediate_size, int64_t num_local_experts, int64_t num_tokens, int64_t act_type, bool use_shuffled_weight, int64_t weight_layout, - btg::Dtype dtype_act, btg::Dtype dtype_weights) { + btg::Dtype dtype_act, btg::Dtype dtype_weights, + bool use_routing_scales_on_input) { Array> valid_configs; std::vector supported_tile_nums(mSupportedTileNums.begin(), mSupportedTileNums.end()); @@ -1190,8 +1187,8 @@ class Fp8PerTensorLauncher : public FusedMoeLauncher { static_cast(weight_layout), // FP8 per-tensor doesn't use Mn-bias (LoRA) cubins. /*gemm1BiasType*/ batchedGemm::gemm::BiasType::None, - true, // usePerTokenScalingGemm1. always true for per-tensor fp8 due to llama4 routing - false, false, false); + /*usePerTokenScalingGemm1*/ use_routing_scales_on_input, + /*usePerTokenScalingGemm2*/ false, false, false); auto cfgs = moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); @@ -1866,10 +1863,15 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher { public: static constexpr std::array mBaseSupportedTileNums = {8, 16, 32, 64}; - static std::vector getSupportedTileNums(btg::Dtype dtype_act) { + static std::vector getSupportedTileNums(btg::Dtype dtype_act, btg::Dtype dtype_weights) { std::vector tiles(mBaseSupportedTileNums.begin(), mBaseSupportedTileNums.end()); if (dtype_act != btg::Dtype::Bfloat16) { tiles.push_back(128); + // Keep tactic enumeration aligned with the public BMM artifact. + if ((dtype_weights == btg::Dtype::E2m1 && dtype_act == btg::Dtype::E2m1) || + (dtype_weights == btg::Dtype::MxE2m1 && dtype_act == btg::Dtype::MxE4m3)) { + tiles.push_back(192); + } tiles.push_back(256); } return tiles; @@ -2207,7 +2209,7 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher { bool use_per_token_scaling) { Array> valid_configs; - std::vector tile_sizes = getSupportedTileNums(dtype_act); + std::vector tile_sizes = getSupportedTileNums(dtype_act, dtype_weights); std::set selected_tile_nums = computeSelectedTileN(tile_sizes, num_tokens, top_k, num_local_experts); @@ -2220,9 +2222,11 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher { /*weight_layout*/ batchedGemm::gemm::MatrixLayout::MajorK, // FP4 MoE getValidConfigs doesn't exercise the Mn-bias (LoRA) cubins. /*gemm1BiasType*/ batchedGemm::gemm::BiasType::None, - // NOTE(siyuan): currently FP4 MoE always apply per-token scaling to both FC1 and FC2. /*usePerTokenScalingGemm1*/ use_per_token_scaling, - /*usePerTokenScalingGemm2*/ use_per_token_scaling, false, false); + // Match prepare_moe_common(): only NVFP4 uses the explicit + // per-token scale operand for FC2. + /*usePerTokenScalingGemm2*/ + use_per_token_scaling && dtype_act == btg::Dtype::E2m1, false, false); auto cfgs = moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); @@ -2761,7 +2765,8 @@ Array trtllm_fp4_block_scale_moe( } // Determine supported tile sizes - std::vector mSupportedTileN = FP4BlockScaleLauncher::getSupportedTileNums(mDtypeAct); + std::vector mSupportedTileN = + FP4BlockScaleLauncher::getSupportedTileNums(mDtypeAct, mDtypeWeights); // Build launchers for ALL supported tiles so autotuner-cached tactics always find their tile_N. // Create a map of launchers for each tile size @@ -2973,7 +2978,7 @@ Array> trtllm_get_valid_moe_configs( } return Fp8PerTensorLauncher::getValidConfigs( top_k, hidden_size, intermediate_size, num_local_experts, num_tokens, act_type, - use_shuffled_weight, weight_layout, dtype_act, dtype_weights); + use_shuffled_weight, weight_layout, dtype_act, dtype_weights, use_per_token_scaling); } else if (dtype_weights == btg::Dtype::E2m1 || dtype_weights == btg::Dtype::MxE2m1) { if (has_gemm1_lora_delta) { TVM_FFI_LOG_AND_THROW(NotImplementedError) diff --git a/csrc/trtllm_fused_moe_runner.cu b/csrc/trtllm_fused_moe_runner.cu index 83d985b7f3c..0699d6c93e6 100644 --- a/csrc/trtllm_fused_moe_runner.cu +++ b/csrc/trtllm_fused_moe_runner.cu @@ -357,6 +357,8 @@ static inline ActType activationTypeToGatedActType(ActivationType actType) { return ActType::SwiGlu; case ActivationType::Geglu: return ActType::GeGlu; + case ActivationType::Situ: + return ActType::SiTuGlu; default: FLASHINFER_CHECK(false, "Unsupported gated activation type ", serializeActivationType(actType), " of enum ", @@ -425,6 +427,9 @@ tensorrt_llm::kernels::TrtllmGenBatchedGemmRunnerOptions getOptions( .fusedBiasShuffleMode = fusedBiasShuffleMode, .biasDtype = biasDtype, .usePerTokenScaling = usePerTokenScaling, + .perTokenSfDtype = usePerTokenScaling ? (dtypeAct == btg::Dtype::E4m3 ? btg::Dtype::Bfloat16 + : btg::Dtype::Fp32) + : btg::Dtype::Void, .usePerChannelScaling = usePerChannelScaling, }; return options; @@ -449,6 +454,9 @@ tensorrt_llm::kernels::TrtllmGenBatchedGemmRunnerOptions getOptions( .fusedBiasShuffleMode = fusedBiasShuffleMode, .biasDtype = biasDtype, .usePerTokenScaling = usePerTokenScaling, + .perTokenSfDtype = usePerTokenScaling ? (dtypeAct == btg::Dtype::E4m3 ? btg::Dtype::Bfloat16 + : btg::Dtype::Fp32) + : btg::Dtype::Void, .usePerChannelScaling = usePerChannelScaling}; return options; } @@ -560,6 +568,9 @@ tensorrt_llm::kernels::TrtllmGenBatchedGemmRunnerOptions getOptions( .useShuffledMatrix = useShuffledMatrix, .weightLayout = weightLayout, .usePerTokenScaling = usePerTokenScaling, + .perTokenSfDtype = usePerTokenScaling ? (dtypeAct == btg::Dtype::E4m3 ? btg::Dtype::Bfloat16 + : btg::Dtype::Fp32) + : btg::Dtype::Void, .usePerChannelScaling = usePerChannelScaling}; return options; } diff --git a/flashinfer/artifacts.py b/flashinfer/artifacts.py index 0d0c9602bc2..7cabad5605b 100644 --- a/flashinfer/artifacts.py +++ b/flashinfer/artifacts.py @@ -137,7 +137,7 @@ class ArtifactPath: TRTLLM_GEN_FMHA: str = "158f6fa11ef139a098cfddcdddce73ca99d164ad/fmha/trtllm-gen/" TRTLLM_GEN_BMM: str = ( - "b368d003e8fdfe4b271bff7c788ac52ef789a81b/batched_gemm-da58956-b4ac80e/" + "5988e15c0e6d006c6a64c0f6c6748b4d3150c1af/batched_gemm-3d40263-3e19f0a/" ) TRTLLM_GEN_GEMM: str = ( "10f64528a1172dae8e29601a3b99ab9dc78d37be/gemm-91e0ba0-2710384/" @@ -160,7 +160,7 @@ class CheckSumHash: "c2d9399b2537be785882354a4f9902ed6c03136c0ea341e201eac40c3923e1dc" ) TRTLLM_GEN_BMM: str = ( - "d0178cd486be54e622386e88daba9c2aca654be7e6f3dcd1af7ecca3354492d2" + "b19ed6c8b1d3fc13ced823bd65ee764d35a19080aea97e742c82ee73ce4c19b0" ) DEEPGEMM: str = "1a2a166839042dbd2a57f48051c82cd1ad032815927c753db269a4ed10d0ffbf" TRTLLM_GEN_GEMM: str = ( diff --git a/flashinfer/fused_moe/core.py b/flashinfer/fused_moe/core.py index 661d643205b..71af28a04ec 100644 --- a/flashinfer/fused_moe/core.py +++ b/flashinfer/fused_moe/core.py @@ -2141,6 +2141,7 @@ def trtllm_fp8_per_tensor_scale_moe_op( weight_layout=WeightLayout.MajorK, use_shuffled_weight=True, activation_type=activation_type, + use_per_token_scaling=use_routing_scales_on_input, num_experts=num_experts, ) @@ -4428,10 +4429,17 @@ def trtllm_fp4_block_scale_moe( ``[num_experts, 2 * intermediate_size]`` FC1 bias, ``float32``. gemm1_alpha : Optional[torch.Tensor] ``[num_experts]`` swiglu alpha, ``float32``. + For SiTU this is ``[local_num_experts]``, finite and positive; + ``None`` materializes per-expert ``alpha=1``. + gemm1_beta : Optional[torch.Tensor] ``[num_experts]`` swiglu beta, ``float32``. + For SiTU this is ``[local_num_experts]``, finite and positive; + ``None`` materializes per-expert ``beta=1``. gemm1_clamp_limit : Optional[torch.Tensor] ``[num_experts]`` swiglu clamp limit, ``float32``. + For SiTU a provided limit is per-local-expert, finite, and positive; + it clamps ``x0`` to ``[-limit, limit]`` and ``x1`` from above. gemm2_weights : torch.Tensor ``[num_experts, hidden_size, intermediate_size]`` packed FP4 FC2 weights, dtype ``uint8``. @@ -4490,6 +4498,7 @@ def trtllm_fp4_block_scale_moe( activation_type : int Activation type (default ``3`` — Swiglu). ``3`` Swiglu; ``4`` Geglu; ``6`` Relu2; ``7`` Identity. + ``10`` SiTU uses ``beta*tanh(x0/beta) * alpha*tanh(x1/alpha)*sigmoid(x1)``. per_token_scale : Optional[torch.Tensor] ``[seq_len]`` per-token scaling factors, ``float32``. output : Optional[torch.Tensor] @@ -4628,10 +4637,17 @@ def trtllm_fp4_block_scale_routed_moe( ``[num_experts, 2 * intermediate_size]`` FC1 bias, float32. gemm1_alpha : Optional[torch.Tensor] ``[num_experts]`` swiglu alpha, float32. + For SiTU this is ``[local_num_experts]``, finite and positive; + ``None`` materializes per-expert ``alpha=1``. + gemm1_beta : Optional[torch.Tensor] ``[num_experts]`` swiglu beta, float32. + For SiTU this is ``[local_num_experts]``, finite and positive; + ``None`` materializes per-expert ``beta=1``. gemm1_clamp_limit : Optional[torch.Tensor] ``[num_experts]`` swiglu clamp limit, float32. + For SiTU a provided limit is per-local-expert, finite, and positive; + it clamps ``x0`` to ``[-limit, limit]`` and ``x1`` from above. gemm2_weights : torch.Tensor ``[num_experts, hidden_size, intermediate_size]`` packed FP4 FC2 weights, ``uint8``. @@ -4689,6 +4705,7 @@ def trtllm_fp4_block_scale_routed_moe( Whether to enable Programmatic Dependent Launch. activation_type : int Activation type (default ``3`` — Swiglu). + ``10`` SiTU uses ``beta*tanh(x0/beta) * alpha*tanh(x1/alpha)*sigmoid(x1)``. per_token_scale : Optional[torch.Tensor] ``[seq_len]`` per-token scaling factors, float32. output : Optional[torch.Tensor] diff --git a/flashinfer/tllm_enums.py b/flashinfer/tllm_enums.py index 1d4702e224d..7ac591dc9f6 100644 --- a/flashinfer/tllm_enums.py +++ b/flashinfer/tllm_enums.py @@ -74,7 +74,8 @@ class ActivationType(IntEnum): SwigluStep = 7 GegluTanh = 8 Identity = 9 - InvalidType = 10 + Situ = 10 + InvalidType = 11 # Eval-safe repr — see ``RoutingMethodType.__repr__``. def __repr__(self) -> str: @@ -92,6 +93,7 @@ def is_gated(self) -> bool: ActivationType.SwigluBias, ActivationType.SwigluStep, ActivationType.GegluTanh, + ActivationType.Situ, ) @@ -128,7 +130,7 @@ def is_gated_activation(activation_type: Union[int, ActivationType]) -> bool: ------- bool ``True`` if ``activation_type`` belongs to the gated activation family - (``Swiglu``, ``Geglu``, ``SwigluBias``, ``SwigluStep``, ``GegluTanh``); + (``Swiglu``, ``Geglu``, ``SwigluBias``, ``SwigluStep``, ``GegluTanh``, ``Situ``); ``False`` otherwise. Examples diff --git a/flashinfer/trace/templates/moe.py b/flashinfer/trace/templates/moe.py index 7587ace8ed0..aefc5718424 100644 --- a/flashinfer/trace/templates/moe.py +++ b/flashinfer/trace/templates/moe.py @@ -14,6 +14,7 @@ """TraceTemplates for Mixture-of-Experts operations.""" +import math import inspect import torch @@ -1259,8 +1260,13 @@ def _fp4_moe_run_experts( topk_idx, local_expert_offset, E_global, + activation_type=3, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + **_unused, ): - """FP4 dequantize + SwiGLU + GEMM for all routing types. + """FP4 dequantize + gated activation + GEMM for all routing types. ``weights`` : [T, TOP_K] float32 — per-token expert weights (normalised) ``topk_idx`` : [T, TOP_K] int64 — selected global expert indices @@ -1292,6 +1298,31 @@ def _fp4_moe_run_experts( T = A.shape[0] output = torch.zeros((T, H), dtype=torch.float32, device=device) local_start = int(local_expert_offset) + activation_type = normalize_activation_type(activation_type) + if activation_type not in (ActivationType.Swiglu, ActivationType.Situ): + raise ValueError( + "FP4 MoE trace reference supports " + f"{ActivationType.Swiglu!r} and {ActivationType.Situ!r}, " + f"got {activation_type!r}" + ) + if activation_type == ActivationType.Situ: + for name, param in ( + ("gemm1_alpha", gemm1_alpha), + ("gemm1_beta", gemm1_beta), + ("gemm1_clamp_limit", gemm1_clamp_limit), + ): + if param is None: + continue + if param.shape != (E_local,): + raise ValueError( + f"{name} must have shape [{E_local}], got {tuple(param.shape)}" + ) + if param.dtype != torch.float32: + raise ValueError(f"{name} must have dtype float32, got {param.dtype}") + if param.device != device: + raise ValueError(f"{name} must be on {device}, got {param.device}") + if not torch.isfinite(param).all() or not (param > 0).all(): + raise ValueError(f"{name} must contain only finite positive values") for le in range(E_local): ge = local_start + le @@ -1305,10 +1336,32 @@ def _fp4_moe_run_experts( G1 = A_e.matmul(W1[le].t()) # [N, 2*I] if gemm1_bias is not None: G1 = G1 + gemm1_bias[le].to(torch.float32) - # SwiGLU uses the trtllm-gen convention: silu(X2) * X1 with X1 first. - X1, X2 = G1[:, :I], G1[:, I:] - silu_X2 = X2 / (1.0 + torch.exp(-X2)) - activated = silu_X2 * X1 + # TRTLLM-Gen convention: x0 is linear/first; x1 is gate/second. + x0, x1 = G1[:, :I], G1[:, I:] + if activation_type == ActivationType.Situ: + alpha = ( + torch.tensor(1.0, dtype=torch.float32, device=device) + if gemm1_alpha is None + else gemm1_alpha[le] + ) + beta = ( + torch.tensor(1.0, dtype=torch.float32, device=device) + if gemm1_beta is None + else gemm1_beta[le] + ) + if gemm1_clamp_limit is not None: + limit = gemm1_clamp_limit[le] + x0 = torch.clamp(x0, min=-limit, max=limit) + x1 = torch.clamp(x1, max=limit) + activated = ( + beta + * torch.tanh(x0 / beta) + * alpha + * torch.tanh(x1 / alpha) + * torch.sigmoid(x1) + ) + else: + activated = x1 * torch.sigmoid(x1) * x0 O = activated.matmul(W2[le].t()) # [N, H] if gemm2_bias is not None: O = O + gemm2_bias[le].to(torch.float32) @@ -1336,6 +1389,7 @@ def _trtllm_fp4_block_scale_moe_default_routing_reference( top_k, local_expert_offset, routed_scaling_factor, + **activation_kwargs, ): """FP4 MoE with Default routing (Softmax → TopK).""" TOP_K = int(top_k) @@ -1360,6 +1414,7 @@ def _trtllm_fp4_block_scale_moe_default_routing_reference( topk_idx, local_expert_offset, E_global, + **activation_kwargs, ) @@ -1378,6 +1433,7 @@ def _trtllm_fp4_block_scale_moe_renormalize_routing_reference( top_k, local_expert_offset, routed_scaling_factor, + **activation_kwargs, ): """FP4 MoE with Renormalize routing (TopK on logits → Softmax).""" TOP_K = int(top_k) @@ -1402,6 +1458,7 @@ def _trtllm_fp4_block_scale_moe_renormalize_routing_reference( topk_idx, local_expert_offset, E_global, + **activation_kwargs, ) @@ -1422,6 +1479,7 @@ def _trtllm_fp4_block_scale_moe_ds_routing_reference( topk_group, local_expert_offset, routed_scaling_factor, + **activation_kwargs, ): """FP4 MoE with DeepSeek-V3 routing: sigmoid + groups + top_k.""" TOP_K = int(top_k) @@ -1474,6 +1532,7 @@ def _trtllm_fp4_block_scale_moe_ds_routing_reference( topk_idx, local_expert_offset, E_global, + **activation_kwargs, ) @@ -1492,6 +1551,7 @@ def _trtllm_fp4_block_scale_moe_llama4_routing_reference( top_k, local_expert_offset, routed_scaling_factor, + **activation_kwargs, ): """FP4 MoE with Llama4 routing (Top1 → Sigmoid). top_k is fixed at 1.""" E_global = routing_logits.shape[1] @@ -1515,6 +1575,7 @@ def _trtllm_fp4_block_scale_moe_llama4_routing_reference( topk_idx, local_expert_offset, E_global, + **activation_kwargs, ) @@ -1533,6 +1594,7 @@ def _trtllm_fp4_block_scale_moe_renormalize_naive_routing_reference( top_k, local_expert_offset, routed_scaling_factor, + **activation_kwargs, ): """FP4 MoE with RenormalizeNaive routing (Softmax → TopK → sum-to-1).""" TOP_K = int(top_k) @@ -1559,6 +1621,7 @@ def _trtllm_fp4_block_scale_moe_renormalize_naive_routing_reference( topk_idx, local_expert_offset, E_global, + **activation_kwargs, ) @@ -1577,6 +1640,7 @@ def _trtllm_fp4_block_scale_moe_topk_routing_reference( top_k, local_expert_offset, routed_scaling_factor, + **activation_kwargs, ): """FP4 MoE with TopK-only routing (uniform weights).""" TOP_K = int(top_k) @@ -1603,6 +1667,7 @@ def _trtllm_fp4_block_scale_moe_topk_routing_reference( topk_idx, local_expert_offset, E_global, + **activation_kwargs, ) @@ -1634,6 +1699,9 @@ def _trtllm_fp4_block_scale_moe_topk_routing_reference( description="Number of FP4 scale blocks along intermediate_size (intermediate_size // 16 for NvFP4).", abbrev="", ), + "activation_type": Const( + description="Fused activation type; 10 selects SiTU.", abbrev="act" + ), } _FP4_STANDARD_INPUTS: dict[str, Tensor | Scalar] = { @@ -1671,17 +1739,17 @@ def _trtllm_fp4_block_scale_moe_topk_routing_reference( ), "gemm1_alpha": Tensor( ["num_local_experts"], - description="Per-expert SwiGLU alpha (float32). Optional.", + description="Per-expert SiTU alpha or SwiGLU alpha (float32). Optional.", optional=True, ), "gemm1_beta": Tensor( ["num_local_experts"], - description="Per-expert SwiGLU beta (float32). Optional.", + description="Per-expert SiTU beta or SwiGLU beta (float32). Optional.", optional=True, ), "gemm1_clamp_limit": Tensor( ["num_local_experts"], - description="Per-expert SwiGLU clamp limit (float32). Optional.", + description="Per-expert gated-activation clamp limit (float32). Optional.", optional=True, ), "gemm2_weights": Tensor( @@ -1721,6 +1789,9 @@ def _trtllm_fp4_block_scale_moe_topk_routing_reference( optional=True, description="Scaling factor applied to routing weights. None for some routing methods.", ), + "activation_type": Scalar( + "int32", description="Fused activation type; 10 selects SiTU." + ), } _FP4_STANDARD_OUTPUTS = { @@ -1752,6 +1823,10 @@ def _moe_fp4_block_scale_init( num_fp4_intermediate_blocks: int = 0, # derived device: str = "cuda", seed: int = 0, + activation_type: int = 3, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, ): """Build inputs for ``trtllm_fp4_block_scale_moe`` (any routing variant). @@ -1766,6 +1841,30 @@ def _moe_fp4_block_scale_init( on by default for parity with the test). Requires SM100+ at runtime; CPU smoke tests skip. """ + activation_type = normalize_activation_type(activation_type) + situ_alpha = gemm1_alpha + situ_beta = gemm1_beta + situ_clamp = gemm1_clamp_limit + if activation_type == ActivationType.Situ: + materialized: dict[str, torch.Tensor | None] = {} + for name, value, default in ( + ("gemm1_alpha", gemm1_alpha, 1.0), + ("gemm1_beta", gemm1_beta, 1.0), + ("gemm1_clamp_limit", gemm1_clamp_limit, None), + ): + if value is None and default is None: + materialized[name] = None + continue + scalar = float(default if value is None else value) + if not math.isfinite(scalar) or scalar <= 0: + raise ValueError(f"{name} must be finite and positive for SiTU") + materialized[name] = torch.full( + (num_local_experts,), scalar, dtype=torch.float32, device=device + ) + situ_alpha = materialized["gemm1_alpha"] + situ_beta = materialized["gemm1_beta"] + situ_clamp = materialized["gemm1_clamp_limit"] + del gemm1_out_size, num_packed_hidden, num_fp4_hidden_blocks del num_packed_intermediate, num_fp4_intermediate_blocks from flashinfer import fp4_quantize # noqa: PLC0415 @@ -1856,9 +1955,9 @@ def _moe_fp4_block_scale_init( "gemm1_weights": gemm1_weights, "gemm1_weights_scale": gemm1_weights_scale, "gemm1_bias": None, - "gemm1_alpha": None, - "gemm1_beta": None, - "gemm1_clamp_limit": None, + "gemm1_alpha": situ_alpha, + "gemm1_beta": situ_beta, + "gemm1_clamp_limit": situ_clamp, "gemm2_weights": gemm2_weights, "gemm2_weights_scale": gemm2_weights_scale, "gemm2_bias": None, @@ -1872,6 +1971,7 @@ def _moe_fp4_block_scale_init( "local_num_experts": int(num_local_experts), "routed_scaling_factor": None, "routing_method_type": int(routing_method_type), + "activation_type": int(activation_type), } if routing_method_type == 2: result["n_group"] = int(n_group) if n_group else 8 @@ -2381,6 +2481,10 @@ def _trtllm_fp4_block_scale_routed_moe_reference( top_k, local_expert_offset, routed_scaling_factor=None, + activation_type=3, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, **_unused, ): """Reference for TRT-LLM FP4 block-scale routed MoE (precomputed topk_ids).""" @@ -2406,6 +2510,10 @@ def _trtllm_fp4_block_scale_routed_moe_reference( topk_ids.to(torch.int64), local_expert_offset, int(num_experts), + activation_type=activation_type, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, ) @@ -3018,6 +3126,9 @@ def _mono_moe_init( ), "num_packed_intermediate": Const(abbrev=""), "num_fp4_intermediate_blocks": Const(abbrev=""), + "activation_type": Const( + description="Fused activation type; 10 selects SiTU.", abbrev="act" + ), }, inputs={ "topk_ids": Tensor( @@ -3043,6 +3154,24 @@ def _mono_moe_init( ["num_local_experts", "gemm1_out_size", "num_fp4_hidden_blocks"], description="FC1 NvFP4 scale.", ), + "gemm1_alpha": Tensor( + ["num_local_experts"], + dtype="float32", + optional=True, + description="Per-expert SiTU alpha or SwiGLU alpha.", + ), + "gemm1_beta": Tensor( + ["num_local_experts"], + dtype="float32", + optional=True, + description="Per-expert SiTU beta or SwiGLU beta.", + ), + "gemm1_clamp_limit": Tensor( + ["num_local_experts"], + dtype="float32", + optional=True, + description="Per-expert gated-activation clamp limit.", + ), "gemm2_weights": Tensor( ["num_local_experts", "hidden_size", "num_packed_intermediate"], description="FC2 NvFP4 weights.", @@ -3055,6 +3184,9 @@ def _mono_moe_init( "top_k": Scalar("int32"), "local_expert_offset": Scalar("int32"), "routed_scaling_factor": Scalar("float32", optional=True), + "activation_type": Scalar( + "int32", description="Fused activation type; 10 selects SiTU." + ), }, outputs=dict(_TRTLLM_MOE_COMMON_OUTPUTS), tags=["status:experimental", "backend:trtllm", "quantization:nvfp4"], diff --git a/include/flashinfer/trtllm/batched_gemm/KernelRunner.h b/include/flashinfer/trtllm/batched_gemm/KernelRunner.h index b886d148c1f..d13bb842974 100644 --- a/include/flashinfer/trtllm/batched_gemm/KernelRunner.h +++ b/include/flashinfer/trtllm/batched_gemm/KernelRunner.h @@ -22,6 +22,7 @@ #include #include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/Enums.h" +#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h" #include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmOptions.h" #include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h" @@ -40,14 +41,25 @@ enum class ActType { // beta' = beta / scaleAb, scaleC' = scaleC * scaleAb. // // GatedSilu is a special case of SwiGlu where the alpha is 1.0 and the beta is 0.0. - SwiGlu, + SwiGlu = 0, // For ActType == GeGlu, we use the simplified version // gatedAct = scaleC' * (x0 + beta') * ((x1 * scaleGate) * phi(alpha * x1 * scaleGate)), // where x0 and x1 are the raw numbers from Gemm, while scaleC and scaleGate are input scales, // beta' = beta / scaleAb, scaleC' = scaleC * scaleAb. - GeGlu, + GeGlu = 1, + SiTuGlu = 2, + None = 3, }; +static_assert(static_cast(ActType::SwiGlu) == + static_cast(batchedGemm::gemmGatedAct::ActType::SwiGlu)); +static_assert(static_cast(ActType::GeGlu) == + static_cast(batchedGemm::gemmGatedAct::ActType::GeGlu)); +static_assert(static_cast(ActType::SiTuGlu) == + static_cast(batchedGemm::gemmGatedAct::ActType::SiTuGlu)); +static_assert(static_cast(ActType::None) == + static_cast(batchedGemm::gemmGatedAct::ActType::None)); + // Type of the element-wise activation to apply after the Gemm enum class EltwiseActType { None = 0, @@ -84,6 +96,8 @@ struct TrtllmGenBatchedGemmRunnerOptions { batchedGemm::trtllm::gen::Dtype biasDtype{batchedGemm::trtllm::gen::Dtype::Fp32}; // whether to apply row-wise scaling factors to the activations bool usePerTokenScaling{false}; + // dtype of the row-wise scaling factors when usePerTokenScaling is enabled + batchedGemm::trtllm::gen::Dtype perTokenSfDtype{batchedGemm::trtllm::gen::Dtype::Void}; // whether to apply row-wise scaling factors to the weights bool usePerChannelScaling{false}; }; diff --git a/include/flashinfer/trtllm/fused_moe/runner.h b/include/flashinfer/trtllm/fused_moe/runner.h index 1c43a0483d6..3e69d673ccf 100644 --- a/include/flashinfer/trtllm/fused_moe/runner.h +++ b/include/flashinfer/trtllm/fused_moe/runner.h @@ -168,7 +168,8 @@ enum class ActivationType : int64_t { SwigluStep = 7, GegluTanh = 8, Identity = 9, - InvalidType = 10, // Must be last + Situ = 10, + InvalidType = 11, // Must be last }; inline std::string serializeActivationType(ActivationType activationType) { @@ -193,6 +194,8 @@ inline std::string serializeActivationType(ActivationType activationType) { return "SwigluStep"; case ActivationType::GegluTanh: return "GegluTanh"; + case ActivationType::Situ: + return "Situ"; default: return "InvalidActivationType"; // TODO throw error }; @@ -202,7 +205,7 @@ inline bool isGatedActivation(ActivationType activationType) { return activationType == ActivationType::Swiglu || activationType == ActivationType::Geglu || activationType == ActivationType::SwigluBias || activationType == ActivationType::SwigluStep || - activationType == ActivationType::GegluTanh; + activationType == ActivationType::GegluTanh || activationType == ActivationType::Situ; } } // namespace MoE diff --git a/tests/moe/test_trtllm_gen_fused_moe.py b/tests/moe/test_trtllm_gen_fused_moe.py index 1a6aa1817c0..487d36b7d6b 100644 --- a/tests/moe/test_trtllm_gen_fused_moe.py +++ b/tests/moe/test_trtllm_gen_fused_moe.py @@ -447,10 +447,12 @@ def test_deepseekv3_routing( ], ) @pytest.mark.parametrize( - "activation_type", + ("activation_type", "alpha_value", "beta_value", "clamp_value"), [ - pytest.param(ActivationType.Swiglu, id="Swiglu"), - pytest.param(ActivationType.Geglu, id="Geglu"), + pytest.param(ActivationType.Swiglu, None, None, None, id="Swiglu"), + pytest.param(ActivationType.Geglu, None, None, None, id="Geglu"), + pytest.param(ActivationType.Situ, None, None, None, id="Situ_Defaults"), + pytest.param(ActivationType.Situ, 4.0, 25.0, None, id="Situ_kimi-k3"), ], ) @pytest.mark.parametrize( @@ -468,10 +470,29 @@ def test_topk_routing( routing_config, weight_processing, activation_type, + alpha_value, + beta_value, + clamp_value, routing_logits_dtype, cache_permute_indices, ): """Test TopK routing configuration.""" + num_experts = routing_config["num_experts"] + gemm1_alpha = ( + None + if alpha_value is None + else torch.full((num_experts,), alpha_value, device="cuda", dtype=torch.float32) + ) + gemm1_beta = ( + None + if beta_value is None + else torch.full((num_experts,), beta_value, device="cuda", dtype=torch.float32) + ) + gemm1_clamp_limit = ( + None + if clamp_value is None + else torch.full((num_experts,), clamp_value, device="cuda", dtype=torch.float32) + ) run_moe_test( num_tokens, hidden_size, @@ -482,6 +503,9 @@ def test_topk_routing( activation_type, cache_permute_indices, routing_logits_dtype, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, ) diff --git a/tests/moe/test_trtllm_gen_moe_autotune_tactics.py b/tests/moe/test_trtllm_gen_moe_autotune_tactics.py index 3e17249c615..d278cff64ce 100644 --- a/tests/moe/test_trtllm_gen_moe_autotune_tactics.py +++ b/tests/moe/test_trtllm_gen_moe_autotune_tactics.py @@ -333,6 +333,99 @@ def _enumerate_valid_tactics( ) +def test_nvfp4_per_tensor_small_shape_all_tactics_are_correct(): + """Every advertised small-shape tactic must honor the output-SF contract.""" + if get_compute_capability(torch.device(device="cuda"))[0] not in [10]: + pytest.skip("Only work on SM100 / SM103.") + + AutoTuner.get()._logged_file_hits.discard(_TEST_LOG_KEY_FP4) + + torch.manual_seed(42) + device = torch.device("cuda:0") + num_tokens = 32 + hidden_size = intermediate_size = 1024 + num_experts = 16 + top_k = 2 + inputs = _build_fp4_routed_moe_inputs( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + top_k=top_k, + num_experts=num_experts, + quant_mode="NvFP4xNvFP4", + routing_method_type=RoutingMethodType.Renormalize, + device=device, + ) + profile_shapes = _moe_profile_shapes(inputs, num_tokens, num_tokens) + + def _run(tactic: list[int] | None) -> torch.Tensor: + _force_tactic_in_autotuner_cache(profile_shapes, tactic, custom_op=_TEST_OP_FP4) + output = trtllm_fp4_block_scale_routed_moe( + topk_ids=inputs["packed_topk"], + routing_bias=None, + hidden_states=inputs["hidden_states"], + hidden_states_scale=inputs["hidden_states_scale"], + gemm1_weights=inputs["w13"], + gemm1_weights_scale=inputs["w13_scale"], + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=inputs["w2"], + gemm2_weights_scale=inputs["w2_scale"], + gemm2_bias=None, + output1_scale_scalar=inputs["output1_scale_scalar"], + output1_scale_gate_scalar=inputs["output1_scale_gate_scalar"], + output2_scale_scalar=inputs["output2_scale_scalar"], + num_experts=num_experts, + top_k=top_k, + n_group=None, + topk_group=None, + intermediate_size=intermediate_size, + local_expert_offset=0, + local_num_experts=num_experts, + routed_scaling_factor=None, + routing_method_type=RoutingMethodType.Renormalize.value, + do_finalize=True, + enable_pdl=device_support_pdl(device), + activation_type=ActivationType.Swiglu.value, + tune_max_num_tokens=num_tokens, + )[0] + torch.cuda.synchronize() + return output + + moe_op = gen_trtllm_gen_fused_moe_sm100_module().build_and_load() + valid_tactics = _enumerate_valid_tactics( + moe_op, + "NvFP4xNvFP4", + top_k, + hidden_size, + intermediate_size, + num_experts, + num_tokens, + ) + assert valid_tactics + assert any(tactic[0] == 8 for tactic in valid_tactics), ( + "the regression shape no longer exercises tile-N 8 tactics" + ) + reference = _run(None).float() + ref_max = reference.abs().max().item() + assert torch.isfinite(reference).all(), "heuristic reference output is not finite" + + failures = [] + for tactic in valid_tactics: + failure = _check_tactic(_run, tactic, reference, ref_max, n_iters=2) + if failure is not None: + failures.append(failure) + assert not failures, ( + f"{len(failures)} per-tensor NVFP4 tactics failed correctness; " + f"first failures: {failures[:10]}" + ) + assert _TEST_LOG_KEY_FP4 in AutoTuner.get()._logged_file_hits, ( + "the forced regression tactic was not dispatched through the autotuner cache" + ) + + @pytest.mark.parametrize("quant_mode", ["NvFP4xNvFP4", "MxFP4xMxFP8", "MxFP4xBf16"]) @pytest.mark.parametrize("num_tokens", [16, 23, 128]) @pytest.mark.parametrize("hidden_size", [4096, 7168]) diff --git a/tests/moe/test_trtllm_gen_routed_fused_moe.py b/tests/moe/test_trtllm_gen_routed_fused_moe.py index fd345e68742..73e159583e7 100644 --- a/tests/moe/test_trtllm_gen_routed_fused_moe.py +++ b/tests/moe/test_trtllm_gen_routed_fused_moe.py @@ -64,6 +64,10 @@ def _run_trtllm_gen_routed_fused_moe_case( routing_method_type: RoutingMethodType, quant_mode: Literal["NvFP4xNvFP4", "MxFP4xMxFP8", "MxFP4xBf16"], routing_format: Literal["packed", "unpacked", "unpacked_fp32"], + activation_type: ActivationType = ActivationType.Swiglu, + gemm1_alpha: torch.Tensor | None = None, + gemm1_beta: torch.Tensor | None = None, + gemm1_clamp_limit: torch.Tensor | None = None, ): compute_capability = get_compute_capability(torch.device(device="cuda")) if compute_capability[0] not in [10]: @@ -166,9 +170,9 @@ def _run_trtllm_gen_routed_fused_moe_case( w13, w13_scale, None, # w13_bias - None, # gemm1_alpha - None, # gemm1_beta - None, # gemm1_clamp_limit + gemm1_alpha, # gemm1_alpha + gemm1_beta, # gemm1_beta + gemm1_clamp_limit, # gemm1_clamp_limit w2, w2_scale, None, # w2_bias @@ -186,7 +190,7 @@ def _run_trtllm_gen_routed_fused_moe_case( routing_method_type.value, True, # do_finalize enable_pdl, - ActivationType.Swiglu.value, # act_type + activation_type.value, # act_type None, )[0].to(torch.float) @@ -229,9 +233,9 @@ def _run_trtllm_gen_routed_fused_moe_case( w13, w13_scale, None, # w13_bias - None, # gemm1_alpha - None, # gemm1_beta - None, # gemm1_clamp_limit + gemm1_alpha, # gemm1_alpha + gemm1_beta, # gemm1_beta + gemm1_clamp_limit, # gemm1_clamp_limit w2, w2_scale, None, # w2_bias @@ -249,7 +253,7 @@ def _run_trtllm_gen_routed_fused_moe_case( routing_method_type.value, True, # do_finalize enable_pdl, - ActivationType.Swiglu.value, # act_type + activation_type.value, # act_type None, )[0].to(torch.float) @@ -1415,3 +1419,39 @@ def test_fp8_block_scale_moe_routing_replay_custom_routing( f"Kernel wrote beyond active token rows " f"(kernel={kernel_tier}, routing={routing_method_type.name})" ) + + +@pytest.mark.parametrize( + "alpha_value,beta_value,clamp_value", + [ + pytest.param(4.0, 25.0, None, id="Situ_Alpha4Beta25"), + pytest.param(1.7, 1.0, 7.0, id="Situ_Alpha1p7Beta1Clamp7"), + ], +) +def test_situ_mxfp4_mxfp8_logits_match_pre_routed(alpha_value, beta_value, clamp_value): + num_experts = 8 + device = torch.device("cuda:0") + _run_trtllm_gen_routed_fused_moe_case( + num_tokens=32, + hidden_size=1024, + intermediate_size=512, + top_k=2, + num_experts=num_experts, + routing_method_type=RoutingMethodType.Renormalize, + quant_mode="MxFP4xMxFP8", + routing_format="unpacked", + activation_type=ActivationType.Situ, + gemm1_alpha=torch.full( + (num_experts,), alpha_value, device=device, dtype=torch.float32 + ), + gemm1_beta=torch.full( + (num_experts,), beta_value, device=device, dtype=torch.float32 + ), + gemm1_clamp_limit=( + None + if clamp_value is None + else torch.full( + (num_experts,), clamp_value, device=device, dtype=torch.float32 + ) + ), + ) diff --git a/tests/moe/trtllm_gen_fused_moe_utils.py b/tests/moe/trtllm_gen_fused_moe_utils.py index b4d1b8350d1..18618f96cbe 100644 --- a/tests/moe/trtllm_gen_fused_moe_utils.py +++ b/tests/moe/trtllm_gen_fused_moe_utils.py @@ -17,7 +17,7 @@ import math import pytest from abc import ABC, abstractmethod -from typing import Dict +from typing import Dict, Optional, Union import torch from cuda.bindings import runtime from torch.nn import functional as F @@ -212,9 +212,9 @@ def _run_moe_computation(self, runtime_args): gemm1_weights=self.static_data["gemm1_weights_fp4_shuffled"], gemm1_weights_scale=self.static_data["gemm1_scales_fp4_shuffled"], gemm1_bias=self.config["gemm1_bias"], - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, + gemm1_alpha=self.config.get("gemm1_alpha"), + gemm1_beta=self.config.get("gemm1_beta"), + gemm1_clamp_limit=self.config.get("gemm1_clamp_limit"), gemm2_weights=self.static_data["gemm2_weights_fp4_shuffled"], gemm2_weights_scale=self.static_data["gemm2_scales_fp4_shuffled"], gemm2_bias=self.config["gemm2_bias"], @@ -620,7 +620,15 @@ def prepare_static_weights_for_kernel( ) # Calculate scaling factors that depend on weights - if is_gated_activation(args.activation_type): + if args.activation_type == ActivationType.Situ: + # SiTU is nonlinear in both GEMM outputs, so applying the dequantization + # factor through scale_c_fc1 would move it inside tanh and change the + # activation. The kernel applies scale_gate_fc1 inside the activation; + # scale_c_fc1 must contain only the output quantization factor. + scale_c_fc1 = torch.full_like( + args.gemm1_scales_global, args_dequant.c_global_sf + ) + elif is_gated_activation(args.activation_type): scale_c_fc1 = ( args_dequant.c_global_sf * (1.0 / args.gemm1_scales_global) @@ -668,6 +676,18 @@ def call_moe( gemm1_bias = static_data["gemm1_bias_shuffled"] gemm2_bias = static_data["gemm2_bias_shuffled"] norm_topk_prob = kwargs.get("norm_topk_prob", True) + gemm1_alpha = kwargs.get("gemm1_alpha") + gemm1_beta = kwargs.get("gemm1_beta") + gemm1_clamp_limit = kwargs.get("gemm1_clamp_limit") + kernel_gemm1_clamp_limit = gemm1_clamp_limit + if ( + gemm1_clamp_limit is not None + and self.quant_mode == QuantMode.FP4_NVFP4_NVFP4 + and activation_type == ActivationType.Situ + ): + # trtllm-gen clamps the raw accumulator before applying the + # dequantization factor carried by scale_gate_fc1. + kernel_gemm1_clamp_limit = gemm1_clamp_limit / static_data["scale_gate_fc1"] # Create CUDA graph configuration config = { @@ -684,6 +704,9 @@ def call_moe( "gemm1_bias": gemm1_bias, "gemm2_bias": gemm2_bias, "norm_topk_prob": norm_topk_prob, + "gemm1_alpha": gemm1_alpha, + "gemm1_beta": gemm1_beta, + "gemm1_clamp_limit": kernel_gemm1_clamp_limit, } runtime_args = { @@ -2491,6 +2514,23 @@ def mxfp8_dequantize_batches(a, a_scales, is_swizzling=True): # ==================================================================================== +def situ_activation_reference( + x0: torch.Tensor, + x1: torch.Tensor, + *, + alpha: Union[float, torch.Tensor] = 1.0, + beta: Union[float, torch.Tensor] = 1.0, + clamp_limit: Optional[Union[float, torch.Tensor]] = None, +) -> torch.Tensor: + """Reference for TRTLLM-Gen SiTU v2 (linear x0, gate x1).""" + if clamp_limit is not None: + x0 = torch.clamp(x0, min=-clamp_limit, max=clamp_limit) + x1 = torch.clamp(x1, max=clamp_limit) + left = beta * torch.tanh(x0 / beta) + right = alpha * torch.tanh(x1 / alpha) * torch.sigmoid(x1) + return left * right + + def run_moe_dequant(args, quant_mode: QuantMode): """Common dequantized MoE reference implementation.""" # Permute @@ -2562,13 +2602,13 @@ def run_moe_dequant(args, quant_mode: QuantMode): (total_num_padded_tokens, args.intermediate_size), float("nan"), device="cuda" ).to(torch.float) - activation_type = args.activation_type + activation_type = ActivationType(args.activation_type) activation_type_to_func = { ActivationType.Swiglu: F.silu, ActivationType.Geglu: F.gelu, ActivationType.Relu2: lambda x: F.relu(x) ** 2, } - activation_func = activation_type_to_func[activation_type] + activation_func = activation_type_to_func.get(activation_type) i = 0 for expert_idx in range(args.num_experts): @@ -2576,45 +2616,54 @@ def run_moe_dequant(args, quant_mode: QuantMode): if my_num_tokens == 0: continue my_a = gemm1_output[i : i + my_num_tokens] - if is_gated_activation(args.activation_type): - my_x1 = my_a[:, : args.intermediate_size] - my_x2 = my_a[:, args.intermediate_size :] - if args.gemm1_clamp_limit is not None: - limit = args.gemm1_clamp_limit[expert_idx].to( - device=my_x1.device, dtype=torch.float + if is_gated_activation(activation_type): + x0 = my_a[:, : args.intermediate_size] + x1 = my_a[:, args.intermediate_size :] + alpha = ( + None + if args.gemm1_alpha is None + else args.gemm1_alpha[expert_idx].to( + device=x1.device, dtype=torch.float ) - my_x1 = torch.clamp(my_x1, min=-limit, max=limit) - my_x2 = torch.clamp(my_x2, max=limit) - if ( - args.gemm1_alpha is not None - or args.gemm1_beta is not None - or args.gemm1_clamp_limit is not None - ): - assert int(args.activation_type) == int(ActivationType.Swiglu) - alpha = ( - 1.0 - if args.gemm1_alpha is None - else args.gemm1_alpha[expert_idx].to( - device=my_x2.device, dtype=torch.float - ) - ) - beta = ( - 0.0 - if args.gemm1_beta is None - else args.gemm1_beta[expert_idx].to( - device=my_x1.device, dtype=torch.float - ) + ) + beta = ( + None + if args.gemm1_beta is None + else args.gemm1_beta[expert_idx].to(device=x0.device, dtype=torch.float) + ) + clamp_limit = ( + None + if args.gemm1_clamp_limit is None + else args.gemm1_clamp_limit[expert_idx].to( + device=x0.device, dtype=torch.float ) - activation_output[i : i + my_num_tokens] = ( - my_x2 * torch.sigmoid(alpha * my_x2) * (my_x1 + beta) + ) + if activation_type == ActivationType.Situ: + activation_output[i : i + my_num_tokens] = situ_activation_reference( + x0, + x1, + alpha=1.0 if alpha is None else alpha, + beta=1.0 if beta is None else beta, + clamp_limit=clamp_limit, ) else: - activation_output[i : i + my_num_tokens] = ( - activation_func(my_x2) * my_x1 - ) + if clamp_limit is not None: + x0 = torch.clamp(x0, min=-clamp_limit, max=clamp_limit) + x1 = torch.clamp(x1, max=clamp_limit) + if alpha is not None or beta is not None or clamp_limit is not None: + assert activation_type == ActivationType.Swiglu + alpha = 1.0 if alpha is None else alpha + beta = 0.0 if beta is None else beta + activation_output[i : i + my_num_tokens] = ( + x1 * torch.sigmoid(alpha * x1) * (x0 + beta) + ) + else: + assert activation_func is not None + activation_output[i : i + my_num_tokens] = activation_func(x1) * x0 else: - my_x1 = my_a[:, : args.intermediate_size] - activation_output[i : i + my_num_tokens] = activation_func(my_x1) + x0 = my_a[:, : args.intermediate_size] + assert activation_func is not None + activation_output[i : i + my_num_tokens] = activation_func(x0) i += my_num_tokens i = (i + args.padding - 1) // args.padding * args.padding @@ -2753,6 +2802,10 @@ def run_moe_reference_fp4(args, quant_mode: QuantMode): args.activation_type, gemm1_bias=args.gemm1_bias, gemm2_bias=args.gemm2_bias, + gemm1_lora_delta=args.gemm1_lora_delta, + gemm1_alpha=args.gemm1_alpha, + gemm1_beta=args.gemm1_beta, + gemm1_clamp_limit=args.gemm1_clamp_limit, ) return run_moe_dequant(args_dequant, quant_mode), args_dequant @@ -3591,9 +3644,9 @@ def run_moe_test( gemm1_weights=static_data["gemm1_weights_fp4_shuffled"], gemm1_weights_scale=static_data["gemm1_scales_fp4_shuffled"], gemm1_bias=static_data["gemm1_bias_shuffled"], - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=gemm1_clamp_limit, gemm2_weights=static_data["gemm2_weights_fp4_shuffled"], gemm2_weights_scale=static_data["gemm2_scales_fp4_shuffled"], gemm2_bias=static_data["gemm2_bias_shuffled"],