From 1a1b48a90118aaa18e8932e4c484db9eb35eab25 Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Wed, 22 Jul 2026 06:52:23 -0700 Subject: [PATCH 1/9] add SITU MOE --- csrc/trtllm_fused_moe_kernel_launcher.cu | 14 +- csrc/trtllm_fused_moe_runner.cu | 2 + flashinfer/fused_moe/core.py | 76 +++++ flashinfer/tllm_enums.py | 6 +- flashinfer/trace/templates/moe.py | 195 +++++++++++- .../trtllm/batched_gemm/KernelRunner.h | 16 +- include/flashinfer/trtllm/fused_moe/runner.h | 7 +- tests/moe/test_trtllm_gen_fused_moe.py | 286 ++++++++++++++++++ tests/moe/test_trtllm_gen_routed_fused_moe.py | 56 +++- tests/moe/trtllm_gen_fused_moe_utils.py | 141 ++++++--- tests/trace/test_trtllm_fp4_situ_moe_trace.py | 157 ++++++++++ tests/utils/test_tllm_enums.py | 8 + 12 files changed, 891 insertions(+), 73 deletions(-) create mode 100644 tests/trace/test_trtllm_fp4_situ_moe_trace.py diff --git a/csrc/trtllm_fused_moe_kernel_launcher.cu b/csrc/trtllm_fused_moe_kernel_launcher.cu index 02cc8a3ab79..19bee428188 100644 --- a/csrc/trtllm_fused_moe_kernel_launcher.cu +++ b/csrc/trtllm_fused_moe_kernel_launcher.cu @@ -1866,10 +1866,17 @@ 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. + bool const supports_tile_n_192 = + (dtype_weights == btg::Dtype::E2m1 && dtype_act == btg::Dtype::E2m1) || + (dtype_weights == btg::Dtype::MxE2m1 && dtype_act == btg::Dtype::MxE4m3); + if (supports_tile_n_192) { + tiles.push_back(192); + } tiles.push_back(256); } return tiles; @@ -2207,7 +2214,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); @@ -2761,7 +2768,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 diff --git a/csrc/trtllm_fused_moe_runner.cu b/csrc/trtllm_fused_moe_runner.cu index 83d985b7f3c..1ba5e50b6ab 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 ", diff --git a/flashinfer/fused_moe/core.py b/flashinfer/fused_moe/core.py index a7031946e77..09c08199b51 100644 --- a/flashinfer/fused_moe/core.py +++ b/flashinfer/fused_moe/core.py @@ -3178,6 +3178,46 @@ def _validate_bf16_gemm1_activation_params( ) +def _prepare_fp4_block_scale_gemm1_activation_params( + activation_type: int, + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + local_num_experts: int, + device: torch.device, +) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + if int(activation_type) != int(ActivationType.Situ): + return gemm1_alpha, gemm1_beta, gemm1_clamp_limit + + def validate_positive_finite(name: str, tensor: torch.Tensor) -> None: + check_shape_dtype_device( + tensor, (local_num_experts,), torch.float32, device, name + ) + # The warmup call validates values before CUDA graph capture. Avoid + # synchronizing a device tensor while the graph is being recorded. + if torch.cuda.is_current_stream_capturing(): + return + if not bool(torch.isfinite(tensor).all().item()): + raise ValueError(f"{name} must contain only finite values for SiTU.") + if not bool((tensor > 0).all().item()): + raise ValueError(f"{name} must contain only positive values for SiTU.") + + if gemm1_alpha is None: + gemm1_alpha = torch.ones(local_num_experts, dtype=torch.float32, device=device) + else: + validate_positive_finite("gemm1_alpha", gemm1_alpha) + + if gemm1_beta is None: + gemm1_beta = torch.ones(local_num_experts, dtype=torch.float32, device=device) + else: + validate_positive_finite("gemm1_beta", gemm1_beta) + + if gemm1_clamp_limit is not None: + validate_positive_finite("gemm1_clamp_limit", gemm1_clamp_limit) + + return gemm1_alpha, gemm1_beta, gemm1_clamp_limit + + def _validate_routing_replay_out( routing_replay_out: Optional[torch.Tensor], top_k: int ) -> None: @@ -4418,10 +4458,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``. @@ -4480,6 +4527,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] @@ -4507,6 +4555,16 @@ def trtllm_fp4_block_scale_moe( path and fp32 ``DeepSeekV3`` logits. """ _validate_routing_replay_out(routing_replay_out, top_k) + gemm1_alpha, gemm1_beta, gemm1_clamp_limit = ( + _prepare_fp4_block_scale_gemm1_activation_params( + activation_type, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, + local_num_experts, + hidden_states.device, + ) + ) return get_trtllm_moe_sm100_module().trtllm_fp4_block_scale_moe( RoutingInputMode.FromLogits, routing_logits, @@ -4618,10 +4676,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``. @@ -4679,6 +4744,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] @@ -4692,6 +4758,16 @@ def trtllm_fp4_block_scale_routed_moe( ``[output]`` when ``do_finalize`` is ``True``, otherwise ``[gemm2_output, expert_weights, expanded_idx_to_permuted_idx]``. """ + gemm1_alpha, gemm1_beta, gemm1_clamp_limit = ( + _prepare_fp4_block_scale_gemm1_activation_params( + activation_type, + gemm1_alpha, + gemm1_beta, + gemm1_clamp_limit, + local_num_experts, + hidden_states.device, + ) + ) # Determine routing mode based on input format if isinstance(topk_ids, tuple): # Unpacked format: (topk_ids, topk_weights) 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..3d0bbb794b8 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 @@ -1170,7 +1171,28 @@ def _unpack_fp4_e2m1(packed: torch.Tensor) -> torch.Tensor: Each byte stores two 4-bit values (low nibble = first element along the last axis). The returned tensor has twice the last-dim size of *packed*. """ - lut = torch.tensor(_E2M1_LUT_VALUES, dtype=torch.float32, device=packed.device) + lut = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=torch.float32, + device=packed.device, + ) p = packed.view(torch.uint8).to(torch.int64) lo = lut[p & 0x0F] hi = lut[(p >> 4) & 0x0F] @@ -1259,8 +1281,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 +1319,29 @@ 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_value = int(activation_type) + if activation_value not in (3, 10): + raise ValueError( + f"FP4 MoE trace reference supports SwiGLU (3) and SiTU (10), got {activation_value}" + ) + if activation_value == 10: + 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 +1355,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_value == 10: + 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 +1408,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 +1433,7 @@ def _trtllm_fp4_block_scale_moe_default_routing_reference( topk_idx, local_expert_offset, E_global, + **activation_kwargs, ) @@ -1378,6 +1452,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 +1477,7 @@ def _trtllm_fp4_block_scale_moe_renormalize_routing_reference( topk_idx, local_expert_offset, E_global, + **activation_kwargs, ) @@ -1422,6 +1498,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 +1551,7 @@ def _trtllm_fp4_block_scale_moe_ds_routing_reference( topk_idx, local_expert_offset, E_global, + **activation_kwargs, ) @@ -1492,6 +1570,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 +1594,7 @@ def _trtllm_fp4_block_scale_moe_llama4_routing_reference( topk_idx, local_expert_offset, E_global, + **activation_kwargs, ) @@ -1533,6 +1613,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 +1640,7 @@ def _trtllm_fp4_block_scale_moe_renormalize_naive_routing_reference( topk_idx, local_expert_offset, E_global, + **activation_kwargs, ) @@ -1577,6 +1659,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 +1686,7 @@ def _trtllm_fp4_block_scale_moe_topk_routing_reference( topk_idx, local_expert_offset, E_global, + **activation_kwargs, ) @@ -1634,6 +1718,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 +1758,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 +1808,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 +1842,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 +1860,29 @@ def _moe_fp4_block_scale_init( on by default for parity with the test). Requires SM100+ at runtime; CPU smoke tests skip. """ + situ_alpha = gemm1_alpha + situ_beta = gemm1_beta + situ_clamp = gemm1_clamp_limit + if int(activation_type) == 10: + 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 +1973,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 +1989,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 +2499,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,9 +2528,34 @@ 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, ) +_FP4_TRACE_REFERENCE_DEPENDENCIES = ( + _unpack_fp4_e2m1, + _ue8m0_to_float32, + _decode_block_scales, + _dequantize_fp4_tensor, + _dequantize_fp4_hidden_states, + _fp4_moe_run_experts, +) +for _fp4_reference in ( + _trtllm_fp4_block_scale_moe_default_routing_reference, + _trtllm_fp4_block_scale_moe_renormalize_routing_reference, + _trtllm_fp4_block_scale_moe_ds_routing_reference, + _trtllm_fp4_block_scale_moe_llama4_routing_reference, + _trtllm_fp4_block_scale_moe_renormalize_naive_routing_reference, + _trtllm_fp4_block_scale_moe_topk_routing_reference, + _trtllm_fp4_block_scale_routed_moe_reference, +): + _fp4_reference._trace_reference_dependencies = _FP4_TRACE_REFERENCE_DEPENDENCIES +del _fp4_reference + + @torch.no_grad() def _trtllm_mxint4_block_scale_moe_reference( routing_logits, @@ -3018,6 +3165,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 +3193,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 +3223,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..4246994b8d7 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, 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..83b04d20da9 100644 --- a/tests/moe/test_trtllm_gen_fused_moe.py +++ b/tests/moe/test_trtllm_gen_fused_moe.py @@ -20,6 +20,10 @@ from flashinfer.utils import get_compute_capability from flashinfer.fused_moe import fill_w_ptr +from flashinfer.fused_moe.core import ( + _prepare_fp4_block_scale_gemm1_activation_params, +) +from flashinfer.tllm_enums import DtypeTrtllmGen from tests.moe.trtllm_gen_fused_moe_utils import ( ActivationType, BF16Moe, @@ -36,6 +40,7 @@ pack_topk_for_routed_moe, routing_reference_renormalize, run_moe_test, + situ_activation_reference, trtllm_bf16_moe, trtllm_bf16_routed_moe, trtllm_fp8_block_scale_moe, @@ -2277,3 +2282,284 @@ def test_fp4_block_scale_deepseekv3_unfinalized_weight_dtype(cache_permute_indic routing_logits_dtype=torch.float32, verify_unfinalized_weight_dtype=True, ) + + +def test_situ_reference_default_and_branch_order(): + x0 = torch.tensor([[-3.0, 0.5, 2.0]], dtype=torch.float32) + x1 = torch.tensor([[0.25, -2.0, 4.0]], dtype=torch.float32) + + actual = situ_activation_reference(x0, x1) + expected = torch.tanh(x0) * torch.tanh(x1) * torch.sigmoid(x1) + branch_swapped = torch.tanh(x1) * torch.tanh(x0) * torch.sigmoid(x0) + + torch.testing.assert_close(actual, expected) + assert not torch.allclose(actual, branch_swapped) + + +@pytest.mark.parametrize( + "alpha_value,beta_value,clamp_value", + [ + pytest.param(4.0, 25.0, None, id="Alpha4Beta25"), + pytest.param(1.7, 1.0, 7.0, id="Alpha1p7Beta1Clamp7"), + ], +) +def test_situ_reference_asymmetric_alpha_beta_and_clamp( + alpha_value, beta_value, clamp_value +): + x0 = torch.tensor([[-9.0, 2.0, 8.0]], dtype=torch.float32) + x1 = torch.tensor([[6.0, -4.0, 1.0]], dtype=torch.float32) + alpha = torch.tensor(alpha_value) + beta = torch.tensor(beta_value) + limit = None if clamp_value is None else torch.tensor(clamp_value) + + actual = situ_activation_reference( + x0, x1, alpha=alpha, beta=beta, clamp_limit=limit + ) + if limit is None: + x0_bounded = x0 + x1_bounded = x1 + else: + x0_bounded = torch.clamp(x0, min=-limit, max=limit) + x1_bounded = torch.clamp(x1, max=limit) + expected = ( + beta + * torch.tanh(x0_bounded / beta) + * alpha + * torch.tanh(x1_bounded / alpha) + * torch.sigmoid(x1_bounded) + ) + swapped_parameters = ( + alpha + * torch.tanh(x0_bounded / alpha) + * beta + * torch.tanh(x1_bounded / beta) + * torch.sigmoid(x1_bounded) + ) + + torch.testing.assert_close(actual, expected) + assert not torch.allclose(actual, swapped_parameters) + + +def test_situ_activation_defaults_materialize_per_expert_ones() -> None: + alpha, beta, clamp = _prepare_fp4_block_scale_gemm1_activation_params( + ActivationType.Situ.value, None, None, None, 3, torch.device("cpu") + ) + + torch.testing.assert_close(alpha, torch.ones(3, dtype=torch.float32)) + torch.testing.assert_close(beta, torch.ones(3, dtype=torch.float32)) + assert clamp is None + + +@pytest.mark.parametrize( + ("alpha_value", "beta_value", "clamp_value"), + [ + pytest.param(4.0, 25.0, None, id="alpha4_beta25"), + pytest.param(1.7, 1.0, 7.0, id="alpha1p7_beta1_clamp7"), + ], +) +def test_situ_activation_accepts_requested_parameters( + alpha_value: float, beta_value: float, clamp_value: float | None +) -> None: + alpha = torch.full((2,), alpha_value, dtype=torch.float32) + beta = torch.full((2,), beta_value, dtype=torch.float32) + clamp = ( + None + if clamp_value is None + else torch.full((2,), clamp_value, dtype=torch.float32) + ) + + actual = _prepare_fp4_block_scale_gemm1_activation_params( + ActivationType.Situ.value, alpha, beta, clamp, 2, torch.device("cpu") + ) + + assert actual[0] is alpha + assert actual[1] is beta + assert actual[2] is clamp + + +@pytest.mark.parametrize( + ("parameter", "values"), + [ + ("gemm1_alpha", [0.0, 1.0]), + ("gemm1_alpha", [-1.0, 1.0]), + ("gemm1_alpha", [float("nan"), 1.0]), + ("gemm1_beta", [0.0, 1.0]), + ("gemm1_beta", [float("inf"), 1.0]), + ("gemm1_clamp_limit", [-1.0, 1.0]), + ("gemm1_clamp_limit", [float("nan"), 1.0]), + ], +) +def test_situ_activation_rejects_nonpositive_or_nonfinite_parameters( + parameter: str, values: list[float] +) -> None: + params = { + "gemm1_alpha": torch.ones(2, dtype=torch.float32), + "gemm1_beta": torch.ones(2, dtype=torch.float32), + "gemm1_clamp_limit": None, + } + params[parameter] = torch.tensor(values, dtype=torch.float32) + + with pytest.raises(ValueError, match=parameter): + _prepare_fp4_block_scale_gemm1_activation_params( + ActivationType.Situ.value, + params["gemm1_alpha"], + params["gemm1_beta"], + params["gemm1_clamp_limit"], + 2, + torch.device("cpu"), + ) + + +def test_situ_activation_rejects_wrong_shape_and_dtype() -> None: + with pytest.raises(ValueError, match="Invalid shape of gemm1_alpha"): + _prepare_fp4_block_scale_gemm1_activation_params( + ActivationType.Situ.value, + torch.ones(1), + None, + None, + 2, + torch.device("cpu"), + ) + with pytest.raises(ValueError, match="Invalid dtype of gemm1_beta"): + _prepare_fp4_block_scale_gemm1_activation_params( + ActivationType.Situ.value, + None, + torch.ones(2, dtype=torch.float64), + None, + 2, + torch.device("cpu"), + ) + + +def test_non_situ_activation_parameter_behavior_is_unchanged() -> None: + assert _prepare_fp4_block_scale_gemm1_activation_params( + ActivationType.Swiglu.value, None, None, None, 2, torch.device("cpu") + ) == (None, None, None) + + +SITU_TEST_CASES = ( + pytest.param(None, None, None, id="Situ_Defaults"), + pytest.param(4.0, 25.0, None, id="Situ_Alpha4Beta25"), + pytest.param(1.7, 1.0, 7.0, id="Situ_Alpha1p7Beta1Clamp7"), +) + + +def _run_fp4_situ_test( + quant_mode, alpha_value, beta_value, clamp_value, cache_permute_indices +): + num_experts = 8 + + def per_expert(value): + return ( + None + if value is None + else torch.full((num_experts,), value, device="cuda", dtype=torch.float32) + ) + + run_moe_test( + num_tokens=32, + hidden_size=1024, + intermediate_size=512, + moe_impl=FP4Moe(quant_mode=quant_mode), + routing_config={ + "num_experts": num_experts, + "top_k": 2, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Renormalize, + "compatible_moe_impls": [FP4Moe], + "compatible_intermediate_size": [512], + "compatible_activation_types": [ActivationType.Situ], + "enable_autotune": True, + }, + weight_processing={ + "use_shuffled_weight": True, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP4Moe], + }, + activation_type=ActivationType.Situ, + cache_permute_indices=cache_permute_indices, + routing_logits_dtype=torch.bfloat16, + gemm1_alpha=per_expert(alpha_value), + gemm1_beta=per_expert(beta_value), + gemm1_clamp_limit=per_expert(clamp_value), + ) + + +@pytest.mark.parametrize("alpha_value,beta_value,clamp_value", SITU_TEST_CASES) +def test_mxfp4_mxfp8_moe_situ( + alpha_value, beta_value, clamp_value, cache_permute_indices +): + _run_fp4_situ_test( + QuantMode.FP4_MXFP4_MXFP8, + alpha_value, + beta_value, + clamp_value, + cache_permute_indices, + ) + + +@pytest.mark.parametrize("alpha_value,beta_value,clamp_value", SITU_TEST_CASES) +def test_nvfp4_moe_situ(alpha_value, beta_value, clamp_value, cache_permute_indices): + _run_fp4_situ_test( + QuantMode.FP4_NVFP4_NVFP4, + alpha_value, + beta_value, + clamp_value, + cache_permute_indices, + ) + + +def _get_situ_tactic_tile_ns(dtype_act, dtype_weights): + from flashinfer.jit.fused_moe import gen_trtllm_gen_fused_moe_sm100_module + + raw_moe_op = gen_trtllm_gen_fused_moe_sm100_module().build_and_load() + tactics = raw_moe_op.trtllm_get_valid_moe_configs( + dtype_act, + dtype_weights, + Fp8QuantizationType.NoneFp8, + 2, + 1024, + 512, + 8, + ActivationType.Situ.value, + True, + WeightLayout.MajorK.value, + False, + 512, + False, + ) + return {int(tactic[0]) for tactic in tactics} + + +def test_mxfp4_mxfp8_autotuner_includes_tile_n_192(): + tile_ns = _get_situ_tactic_tile_ns(DtypeTrtllmGen.MxE4m3, DtypeTrtllmGen.MxE2m1) + assert 192 in tile_ns + + +def test_nvfp4_autotuner_includes_tile_n_192(): + tile_ns = _get_situ_tactic_tile_ns(DtypeTrtllmGen.E2m1, DtypeTrtllmGen.E2m1) + assert 192 in tile_ns + + +@pytest.mark.parametrize( + ("dtype_act", "dtype_weights"), + [ + pytest.param( + DtypeTrtllmGen.E2m1, + DtypeTrtllmGen.MxE2m1, + id="mxfp4_weights_nvfp4_activations", + ), + pytest.param( + DtypeTrtllmGen.MxE4m3, + DtypeTrtllmGen.E2m1, + id="nvfp4_weights_mxfp8_activations", + ), + ], +) +def test_unsupported_fp4_dtype_pairs_exclude_tile_n_192(dtype_act, dtype_weights): + tile_ns = _get_situ_tactic_tile_ns(dtype_act, dtype_weights) + assert 192 not in tile_ns 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"], diff --git a/tests/trace/test_trtllm_fp4_situ_moe_trace.py b/tests/trace/test_trtllm_fp4_situ_moe_trace.py new file mode 100644 index 00000000000..c6bfebc1fb4 --- /dev/null +++ b/tests/trace/test_trtllm_fp4_situ_moe_trace.py @@ -0,0 +1,157 @@ +"""Trace coverage for TRT-LLM Gen MXFP4 x MXFP8 SiTU MoE.""" + +import pytest +import torch + +from flashinfer import ActivationType + + +def _mxfp4_mxfp8_identity_inputs(): + seq_len = 4 + hidden_size = 32 + intermediate_size = 32 + + hidden_bf16 = torch.zeros(seq_len, hidden_size, dtype=torch.bfloat16) + hidden_bf16[:, 0] = torch.tensor([-3.0, -1.0, 3.0, -4.0]) + hidden_bf16[:, 1] = torch.tensor([0.5, -2.0, 4.0, 6.0]) + hidden_states = hidden_bf16.to(torch.float8_e4m3fn) + hidden_states_scale = torch.full((seq_len, 1), 127, dtype=torch.uint8) + + gemm1_weights = torch.zeros( + 1, 2 * intermediate_size, hidden_size // 2, dtype=torch.uint8 + ) + gemm1_weights[0, 0, 0] = 0x02 + gemm1_weights[0, intermediate_size, 0] = 0x20 + gemm1_weights_scale = torch.full( + (1, 2 * intermediate_size, 1), 127, dtype=torch.uint8 + ) + + gemm2_weights = torch.zeros( + 1, hidden_size, intermediate_size // 2, dtype=torch.uint8 + ) + gemm2_weights[0, 0, 0] = 0x02 + gemm2_weights_scale = torch.full((1, hidden_size, 1), 127, dtype=torch.uint8) + + return { + "routing_logits": torch.zeros(seq_len, 1, dtype=torch.float32), + "routing_bias": None, + "hidden_states": hidden_states, + "hidden_states_scale": hidden_states_scale, + "gemm1_weights": gemm1_weights, + "gemm1_weights_scale": gemm1_weights_scale, + "gemm1_bias": None, + "gemm2_weights": gemm2_weights, + "gemm2_weights_scale": gemm2_weights_scale, + "gemm2_bias": None, + "top_k": 1, + "local_expert_offset": 0, + "routed_scaling_factor": 1.0, + "hidden_size": hidden_size, + "intermediate_size": intermediate_size, + } + + +def _situ_params(): + return { + "activation_type": ActivationType.Situ.value, + "gemm1_alpha": torch.tensor([1.7], dtype=torch.float32), + "gemm1_beta": torch.tensor([1.0], dtype=torch.float32), + "gemm1_clamp_limit": torch.tensor([7.0], dtype=torch.float32), + } + + +def test_fp4_situ_trace_schema_records_activation(): + from flashinfer.trace.templates.moe import ( + trtllm_fp4_block_scale_moe_default_routing_trace, + trtllm_fp4_block_scale_routed_moe_trace, + ) + + inputs = _mxfp4_mxfp8_identity_inputs() + params = _situ_params() + common = { + **inputs, + **params, + "num_experts": 1, + "local_num_experts": 1, + } + default_defn = trtllm_fp4_block_scale_moe_default_routing_trace.build_fi_trace_fn( + "flashinfer.fused_moe.core.trtllm_fp4_block_scale_moe" + )(**common) + routed_inputs = dict(common) + routed_inputs.pop("routing_logits") + routed_inputs["topk_ids"] = torch.zeros(4, 1, dtype=torch.int32) + routed_defn = trtllm_fp4_block_scale_routed_moe_trace.build_fi_trace_fn( + "flashinfer.fused_moe.core.trtllm_fp4_block_scale_routed_moe" + )(**routed_inputs) + + for defn in (default_defn, routed_defn): + assert defn["axes"]["activation_type"]["value"] == ActivationType.Situ.value + assert defn["inputs"]["activation_type"]["dtype"] == "int32" + for name in ("gemm1_alpha", "gemm1_beta", "gemm1_clamp_limit"): + assert defn["inputs"][name]["shape"] == ["num_local_experts"] + assert defn["inputs"][name]["dtype"] == "float32" + assert defn["inputs"][name]["optional"] is True + + assert "activation_value == 10" in default_defn["reference"] + + +def test_fp4_situ_trace_reference_mxfp4_mxfp8_and_pre_routed_match(): + from flashinfer.trace.templates.moe import ( + trtllm_fp4_block_scale_moe_default_routing_trace, + trtllm_fp4_block_scale_routed_moe_trace, + ) + + inputs = _mxfp4_mxfp8_identity_inputs() + hidden_size = inputs.pop("hidden_size") + inputs.pop("intermediate_size") + params = _situ_params() + + default_out = trtllm_fp4_block_scale_moe_default_routing_trace.reference( + **inputs, **params + ).to(torch.float32) + routed_inputs = dict(inputs) + routed_inputs.pop("routing_logits") + routed_inputs.pop("routing_bias") + routed_out = trtllm_fp4_block_scale_routed_moe_trace.reference( + topk_ids=torch.zeros(4, 1, dtype=torch.int32), + num_experts=1, + **routed_inputs, + **params, + ).to(torch.float32) + + x0 = inputs["hidden_states"][:, :1].to(torch.float32).clamp(-7.0, 7.0) + x1 = inputs["hidden_states"][:, 1:2].to(torch.float32).clamp(max=7.0) + expected_first = torch.tanh(x0) * 1.7 * torch.tanh(x1 / 1.7) * torch.sigmoid(x1) + expected = torch.zeros(4, hidden_size, dtype=torch.float32) + expected[:, :1] = expected_first + expected = expected.to(torch.bfloat16).to(torch.float32) + + swiglu_out = trtllm_fp4_block_scale_moe_default_routing_trace.reference( + **inputs, activation_type=ActivationType.Swiglu.value + ).to(torch.float32) + + torch.testing.assert_close(default_out, expected, rtol=0, atol=0) + torch.testing.assert_close(routed_out, expected, rtol=0, atol=0) + assert not torch.allclose(default_out, swiglu_out) + + +@pytest.mark.parametrize( + ("parameter", "value"), + [ + pytest.param("gemm1_alpha", 0.0, id="zero_alpha"), + pytest.param("gemm1_beta", -1.0, id="negative_beta"), + pytest.param("gemm1_clamp_limit", float("nan"), id="nan_clamp"), + ], +) +def test_fp4_situ_trace_init_rejects_invalid_parameters(parameter, value): + from flashinfer.trace.templates.moe import ( + trtllm_fp4_block_scale_moe_default_routing_trace, + ) + + with pytest.raises(ValueError, match="finite and positive"): + trtllm_fp4_block_scale_moe_default_routing_trace.init( + seq_len=1, + num_local_experts=1, + activation_type=ActivationType.Situ.value, + **{parameter: value}, + ) diff --git a/tests/utils/test_tllm_enums.py b/tests/utils/test_tllm_enums.py index 83fc5051e79..6ff6ab6bb7e 100644 --- a/tests/utils/test_tllm_enums.py +++ b/tests/utils/test_tllm_enums.py @@ -5,3 +5,11 @@ def test_geglu_tanh_is_gated_activation(): assert is_gated_activation(ActivationType.GegluTanh) assert is_gated_activation(ActivationType.GegluTanh.value) assert ActivationType.GegluTanh.is_gated + + +def test_situ_activation_abi_and_gated_classification(): + assert ActivationType.Situ.value == 10 + assert ActivationType.InvalidType.value == 11 + assert is_gated_activation(ActivationType.Situ) + assert is_gated_activation(ActivationType.Situ.value) + assert ActivationType.Situ.is_gated From 6e5a783b0cc3a2358a63dfa68d555c4a7cbabc40 Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Mon, 27 Jul 2026 10:42:18 -0700 Subject: [PATCH 2/9] update tests --- csrc/trtllm_fused_moe_kernel_launcher.cu | 6 +- tests/moe/test_trtllm_gen_fused_moe.py | 157 ------------------ tests/trace/test_trtllm_fp4_situ_moe_trace.py | 157 ------------------ tests/utils/test_tllm_enums.py | 8 - 4 files changed, 2 insertions(+), 326 deletions(-) delete mode 100644 tests/trace/test_trtllm_fp4_situ_moe_trace.py diff --git a/csrc/trtllm_fused_moe_kernel_launcher.cu b/csrc/trtllm_fused_moe_kernel_launcher.cu index 19bee428188..3eaaedd5ab3 100644 --- a/csrc/trtllm_fused_moe_kernel_launcher.cu +++ b/csrc/trtllm_fused_moe_kernel_launcher.cu @@ -1871,10 +1871,8 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher { if (dtype_act != btg::Dtype::Bfloat16) { tiles.push_back(128); // Keep tactic enumeration aligned with the public BMM artifact. - bool const supports_tile_n_192 = - (dtype_weights == btg::Dtype::E2m1 && dtype_act == btg::Dtype::E2m1) || - (dtype_weights == btg::Dtype::MxE2m1 && dtype_act == btg::Dtype::MxE4m3); - if (supports_tile_n_192) { + 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); diff --git a/tests/moe/test_trtllm_gen_fused_moe.py b/tests/moe/test_trtllm_gen_fused_moe.py index 83b04d20da9..8a182adb530 100644 --- a/tests/moe/test_trtllm_gen_fused_moe.py +++ b/tests/moe/test_trtllm_gen_fused_moe.py @@ -20,9 +20,6 @@ from flashinfer.utils import get_compute_capability from flashinfer.fused_moe import fill_w_ptr -from flashinfer.fused_moe.core import ( - _prepare_fp4_block_scale_gemm1_activation_params, -) from flashinfer.tllm_enums import DtypeTrtllmGen from tests.moe.trtllm_gen_fused_moe_utils import ( ActivationType, @@ -40,7 +37,6 @@ pack_topk_for_routed_moe, routing_reference_renormalize, run_moe_test, - situ_activation_reference, trtllm_bf16_moe, trtllm_bf16_routed_moe, trtllm_fp8_block_scale_moe, @@ -2284,159 +2280,6 @@ def test_fp4_block_scale_deepseekv3_unfinalized_weight_dtype(cache_permute_indic ) -def test_situ_reference_default_and_branch_order(): - x0 = torch.tensor([[-3.0, 0.5, 2.0]], dtype=torch.float32) - x1 = torch.tensor([[0.25, -2.0, 4.0]], dtype=torch.float32) - - actual = situ_activation_reference(x0, x1) - expected = torch.tanh(x0) * torch.tanh(x1) * torch.sigmoid(x1) - branch_swapped = torch.tanh(x1) * torch.tanh(x0) * torch.sigmoid(x0) - - torch.testing.assert_close(actual, expected) - assert not torch.allclose(actual, branch_swapped) - - -@pytest.mark.parametrize( - "alpha_value,beta_value,clamp_value", - [ - pytest.param(4.0, 25.0, None, id="Alpha4Beta25"), - pytest.param(1.7, 1.0, 7.0, id="Alpha1p7Beta1Clamp7"), - ], -) -def test_situ_reference_asymmetric_alpha_beta_and_clamp( - alpha_value, beta_value, clamp_value -): - x0 = torch.tensor([[-9.0, 2.0, 8.0]], dtype=torch.float32) - x1 = torch.tensor([[6.0, -4.0, 1.0]], dtype=torch.float32) - alpha = torch.tensor(alpha_value) - beta = torch.tensor(beta_value) - limit = None if clamp_value is None else torch.tensor(clamp_value) - - actual = situ_activation_reference( - x0, x1, alpha=alpha, beta=beta, clamp_limit=limit - ) - if limit is None: - x0_bounded = x0 - x1_bounded = x1 - else: - x0_bounded = torch.clamp(x0, min=-limit, max=limit) - x1_bounded = torch.clamp(x1, max=limit) - expected = ( - beta - * torch.tanh(x0_bounded / beta) - * alpha - * torch.tanh(x1_bounded / alpha) - * torch.sigmoid(x1_bounded) - ) - swapped_parameters = ( - alpha - * torch.tanh(x0_bounded / alpha) - * beta - * torch.tanh(x1_bounded / beta) - * torch.sigmoid(x1_bounded) - ) - - torch.testing.assert_close(actual, expected) - assert not torch.allclose(actual, swapped_parameters) - - -def test_situ_activation_defaults_materialize_per_expert_ones() -> None: - alpha, beta, clamp = _prepare_fp4_block_scale_gemm1_activation_params( - ActivationType.Situ.value, None, None, None, 3, torch.device("cpu") - ) - - torch.testing.assert_close(alpha, torch.ones(3, dtype=torch.float32)) - torch.testing.assert_close(beta, torch.ones(3, dtype=torch.float32)) - assert clamp is None - - -@pytest.mark.parametrize( - ("alpha_value", "beta_value", "clamp_value"), - [ - pytest.param(4.0, 25.0, None, id="alpha4_beta25"), - pytest.param(1.7, 1.0, 7.0, id="alpha1p7_beta1_clamp7"), - ], -) -def test_situ_activation_accepts_requested_parameters( - alpha_value: float, beta_value: float, clamp_value: float | None -) -> None: - alpha = torch.full((2,), alpha_value, dtype=torch.float32) - beta = torch.full((2,), beta_value, dtype=torch.float32) - clamp = ( - None - if clamp_value is None - else torch.full((2,), clamp_value, dtype=torch.float32) - ) - - actual = _prepare_fp4_block_scale_gemm1_activation_params( - ActivationType.Situ.value, alpha, beta, clamp, 2, torch.device("cpu") - ) - - assert actual[0] is alpha - assert actual[1] is beta - assert actual[2] is clamp - - -@pytest.mark.parametrize( - ("parameter", "values"), - [ - ("gemm1_alpha", [0.0, 1.0]), - ("gemm1_alpha", [-1.0, 1.0]), - ("gemm1_alpha", [float("nan"), 1.0]), - ("gemm1_beta", [0.0, 1.0]), - ("gemm1_beta", [float("inf"), 1.0]), - ("gemm1_clamp_limit", [-1.0, 1.0]), - ("gemm1_clamp_limit", [float("nan"), 1.0]), - ], -) -def test_situ_activation_rejects_nonpositive_or_nonfinite_parameters( - parameter: str, values: list[float] -) -> None: - params = { - "gemm1_alpha": torch.ones(2, dtype=torch.float32), - "gemm1_beta": torch.ones(2, dtype=torch.float32), - "gemm1_clamp_limit": None, - } - params[parameter] = torch.tensor(values, dtype=torch.float32) - - with pytest.raises(ValueError, match=parameter): - _prepare_fp4_block_scale_gemm1_activation_params( - ActivationType.Situ.value, - params["gemm1_alpha"], - params["gemm1_beta"], - params["gemm1_clamp_limit"], - 2, - torch.device("cpu"), - ) - - -def test_situ_activation_rejects_wrong_shape_and_dtype() -> None: - with pytest.raises(ValueError, match="Invalid shape of gemm1_alpha"): - _prepare_fp4_block_scale_gemm1_activation_params( - ActivationType.Situ.value, - torch.ones(1), - None, - None, - 2, - torch.device("cpu"), - ) - with pytest.raises(ValueError, match="Invalid dtype of gemm1_beta"): - _prepare_fp4_block_scale_gemm1_activation_params( - ActivationType.Situ.value, - None, - torch.ones(2, dtype=torch.float64), - None, - 2, - torch.device("cpu"), - ) - - -def test_non_situ_activation_parameter_behavior_is_unchanged() -> None: - assert _prepare_fp4_block_scale_gemm1_activation_params( - ActivationType.Swiglu.value, None, None, None, 2, torch.device("cpu") - ) == (None, None, None) - - SITU_TEST_CASES = ( pytest.param(None, None, None, id="Situ_Defaults"), pytest.param(4.0, 25.0, None, id="Situ_Alpha4Beta25"), diff --git a/tests/trace/test_trtllm_fp4_situ_moe_trace.py b/tests/trace/test_trtllm_fp4_situ_moe_trace.py deleted file mode 100644 index c6bfebc1fb4..00000000000 --- a/tests/trace/test_trtllm_fp4_situ_moe_trace.py +++ /dev/null @@ -1,157 +0,0 @@ -"""Trace coverage for TRT-LLM Gen MXFP4 x MXFP8 SiTU MoE.""" - -import pytest -import torch - -from flashinfer import ActivationType - - -def _mxfp4_mxfp8_identity_inputs(): - seq_len = 4 - hidden_size = 32 - intermediate_size = 32 - - hidden_bf16 = torch.zeros(seq_len, hidden_size, dtype=torch.bfloat16) - hidden_bf16[:, 0] = torch.tensor([-3.0, -1.0, 3.0, -4.0]) - hidden_bf16[:, 1] = torch.tensor([0.5, -2.0, 4.0, 6.0]) - hidden_states = hidden_bf16.to(torch.float8_e4m3fn) - hidden_states_scale = torch.full((seq_len, 1), 127, dtype=torch.uint8) - - gemm1_weights = torch.zeros( - 1, 2 * intermediate_size, hidden_size // 2, dtype=torch.uint8 - ) - gemm1_weights[0, 0, 0] = 0x02 - gemm1_weights[0, intermediate_size, 0] = 0x20 - gemm1_weights_scale = torch.full( - (1, 2 * intermediate_size, 1), 127, dtype=torch.uint8 - ) - - gemm2_weights = torch.zeros( - 1, hidden_size, intermediate_size // 2, dtype=torch.uint8 - ) - gemm2_weights[0, 0, 0] = 0x02 - gemm2_weights_scale = torch.full((1, hidden_size, 1), 127, dtype=torch.uint8) - - return { - "routing_logits": torch.zeros(seq_len, 1, dtype=torch.float32), - "routing_bias": None, - "hidden_states": hidden_states, - "hidden_states_scale": hidden_states_scale, - "gemm1_weights": gemm1_weights, - "gemm1_weights_scale": gemm1_weights_scale, - "gemm1_bias": None, - "gemm2_weights": gemm2_weights, - "gemm2_weights_scale": gemm2_weights_scale, - "gemm2_bias": None, - "top_k": 1, - "local_expert_offset": 0, - "routed_scaling_factor": 1.0, - "hidden_size": hidden_size, - "intermediate_size": intermediate_size, - } - - -def _situ_params(): - return { - "activation_type": ActivationType.Situ.value, - "gemm1_alpha": torch.tensor([1.7], dtype=torch.float32), - "gemm1_beta": torch.tensor([1.0], dtype=torch.float32), - "gemm1_clamp_limit": torch.tensor([7.0], dtype=torch.float32), - } - - -def test_fp4_situ_trace_schema_records_activation(): - from flashinfer.trace.templates.moe import ( - trtllm_fp4_block_scale_moe_default_routing_trace, - trtllm_fp4_block_scale_routed_moe_trace, - ) - - inputs = _mxfp4_mxfp8_identity_inputs() - params = _situ_params() - common = { - **inputs, - **params, - "num_experts": 1, - "local_num_experts": 1, - } - default_defn = trtllm_fp4_block_scale_moe_default_routing_trace.build_fi_trace_fn( - "flashinfer.fused_moe.core.trtllm_fp4_block_scale_moe" - )(**common) - routed_inputs = dict(common) - routed_inputs.pop("routing_logits") - routed_inputs["topk_ids"] = torch.zeros(4, 1, dtype=torch.int32) - routed_defn = trtllm_fp4_block_scale_routed_moe_trace.build_fi_trace_fn( - "flashinfer.fused_moe.core.trtllm_fp4_block_scale_routed_moe" - )(**routed_inputs) - - for defn in (default_defn, routed_defn): - assert defn["axes"]["activation_type"]["value"] == ActivationType.Situ.value - assert defn["inputs"]["activation_type"]["dtype"] == "int32" - for name in ("gemm1_alpha", "gemm1_beta", "gemm1_clamp_limit"): - assert defn["inputs"][name]["shape"] == ["num_local_experts"] - assert defn["inputs"][name]["dtype"] == "float32" - assert defn["inputs"][name]["optional"] is True - - assert "activation_value == 10" in default_defn["reference"] - - -def test_fp4_situ_trace_reference_mxfp4_mxfp8_and_pre_routed_match(): - from flashinfer.trace.templates.moe import ( - trtllm_fp4_block_scale_moe_default_routing_trace, - trtllm_fp4_block_scale_routed_moe_trace, - ) - - inputs = _mxfp4_mxfp8_identity_inputs() - hidden_size = inputs.pop("hidden_size") - inputs.pop("intermediate_size") - params = _situ_params() - - default_out = trtllm_fp4_block_scale_moe_default_routing_trace.reference( - **inputs, **params - ).to(torch.float32) - routed_inputs = dict(inputs) - routed_inputs.pop("routing_logits") - routed_inputs.pop("routing_bias") - routed_out = trtllm_fp4_block_scale_routed_moe_trace.reference( - topk_ids=torch.zeros(4, 1, dtype=torch.int32), - num_experts=1, - **routed_inputs, - **params, - ).to(torch.float32) - - x0 = inputs["hidden_states"][:, :1].to(torch.float32).clamp(-7.0, 7.0) - x1 = inputs["hidden_states"][:, 1:2].to(torch.float32).clamp(max=7.0) - expected_first = torch.tanh(x0) * 1.7 * torch.tanh(x1 / 1.7) * torch.sigmoid(x1) - expected = torch.zeros(4, hidden_size, dtype=torch.float32) - expected[:, :1] = expected_first - expected = expected.to(torch.bfloat16).to(torch.float32) - - swiglu_out = trtllm_fp4_block_scale_moe_default_routing_trace.reference( - **inputs, activation_type=ActivationType.Swiglu.value - ).to(torch.float32) - - torch.testing.assert_close(default_out, expected, rtol=0, atol=0) - torch.testing.assert_close(routed_out, expected, rtol=0, atol=0) - assert not torch.allclose(default_out, swiglu_out) - - -@pytest.mark.parametrize( - ("parameter", "value"), - [ - pytest.param("gemm1_alpha", 0.0, id="zero_alpha"), - pytest.param("gemm1_beta", -1.0, id="negative_beta"), - pytest.param("gemm1_clamp_limit", float("nan"), id="nan_clamp"), - ], -) -def test_fp4_situ_trace_init_rejects_invalid_parameters(parameter, value): - from flashinfer.trace.templates.moe import ( - trtllm_fp4_block_scale_moe_default_routing_trace, - ) - - with pytest.raises(ValueError, match="finite and positive"): - trtllm_fp4_block_scale_moe_default_routing_trace.init( - seq_len=1, - num_local_experts=1, - activation_type=ActivationType.Situ.value, - **{parameter: value}, - ) diff --git a/tests/utils/test_tllm_enums.py b/tests/utils/test_tllm_enums.py index 6ff6ab6bb7e..83fc5051e79 100644 --- a/tests/utils/test_tllm_enums.py +++ b/tests/utils/test_tllm_enums.py @@ -5,11 +5,3 @@ def test_geglu_tanh_is_gated_activation(): assert is_gated_activation(ActivationType.GegluTanh) assert is_gated_activation(ActivationType.GegluTanh.value) assert ActivationType.GegluTanh.is_gated - - -def test_situ_activation_abi_and_gated_classification(): - assert ActivationType.Situ.value == 10 - assert ActivationType.InvalidType.value == 11 - assert is_gated_activation(ActivationType.Situ) - assert is_gated_activation(ActivationType.Situ.value) - assert ActivationType.Situ.is_gated From 5b2cef5813008d2dafe887b9c0268cee8f34746a Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Tue, 28 Jul 2026 00:39:14 -0700 Subject: [PATCH 3/9] address trace registration --- flashinfer/trace/templates/moe.py | 44 +------------------------------ 1 file changed, 1 insertion(+), 43 deletions(-) diff --git a/flashinfer/trace/templates/moe.py b/flashinfer/trace/templates/moe.py index 3d0bbb794b8..5e3ca8897a6 100644 --- a/flashinfer/trace/templates/moe.py +++ b/flashinfer/trace/templates/moe.py @@ -1171,28 +1171,7 @@ def _unpack_fp4_e2m1(packed: torch.Tensor) -> torch.Tensor: Each byte stores two 4-bit values (low nibble = first element along the last axis). The returned tensor has twice the last-dim size of *packed*. """ - lut = torch.tensor( - [ - 0.0, - 0.5, - 1.0, - 1.5, - 2.0, - 3.0, - 4.0, - 6.0, - -0.0, - -0.5, - -1.0, - -1.5, - -2.0, - -3.0, - -4.0, - -6.0, - ], - dtype=torch.float32, - device=packed.device, - ) + lut = torch.tensor(_E2M1_LUT_VALUES, dtype=torch.float32, device=packed.device) p = packed.view(torch.uint8).to(torch.int64) lo = lut[p & 0x0F] hi = lut[(p >> 4) & 0x0F] @@ -2535,27 +2514,6 @@ def _trtllm_fp4_block_scale_routed_moe_reference( ) -_FP4_TRACE_REFERENCE_DEPENDENCIES = ( - _unpack_fp4_e2m1, - _ue8m0_to_float32, - _decode_block_scales, - _dequantize_fp4_tensor, - _dequantize_fp4_hidden_states, - _fp4_moe_run_experts, -) -for _fp4_reference in ( - _trtllm_fp4_block_scale_moe_default_routing_reference, - _trtllm_fp4_block_scale_moe_renormalize_routing_reference, - _trtllm_fp4_block_scale_moe_ds_routing_reference, - _trtllm_fp4_block_scale_moe_llama4_routing_reference, - _trtllm_fp4_block_scale_moe_renormalize_naive_routing_reference, - _trtllm_fp4_block_scale_moe_topk_routing_reference, - _trtllm_fp4_block_scale_routed_moe_reference, -): - _fp4_reference._trace_reference_dependencies = _FP4_TRACE_REFERENCE_DEPENDENCIES -del _fp4_reference - - @torch.no_grad() def _trtllm_mxint4_block_scale_moe_reference( routing_logits, From 2f524007ff3c24ee3c68e29c40fcd80e7751cddf Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Tue, 28 Jul 2026 09:01:47 -0700 Subject: [PATCH 4/9] update cubin --- flashinfer/artifacts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 = ( From 45ce98be1750340835dd2104ee0854be59d03ba1 Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Tue, 28 Jul 2026 09:10:50 -0700 Subject: [PATCH 5/9] update test --- tests/moe/test_trtllm_gen_fused_moe.py | 159 +++++-------------------- 1 file changed, 27 insertions(+), 132 deletions(-) diff --git a/tests/moe/test_trtllm_gen_fused_moe.py b/tests/moe/test_trtllm_gen_fused_moe.py index 8a182adb530..487d36b7d6b 100644 --- a/tests/moe/test_trtllm_gen_fused_moe.py +++ b/tests/moe/test_trtllm_gen_fused_moe.py @@ -20,7 +20,6 @@ from flashinfer.utils import get_compute_capability from flashinfer.fused_moe import fill_w_ptr -from flashinfer.tllm_enums import DtypeTrtllmGen from tests.moe.trtllm_gen_fused_moe_utils import ( ActivationType, BF16Moe, @@ -448,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( @@ -469,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, @@ -483,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, ) @@ -2278,131 +2301,3 @@ def test_fp4_block_scale_deepseekv3_unfinalized_weight_dtype(cache_permute_indic routing_logits_dtype=torch.float32, verify_unfinalized_weight_dtype=True, ) - - -SITU_TEST_CASES = ( - pytest.param(None, None, None, id="Situ_Defaults"), - pytest.param(4.0, 25.0, None, id="Situ_Alpha4Beta25"), - pytest.param(1.7, 1.0, 7.0, id="Situ_Alpha1p7Beta1Clamp7"), -) - - -def _run_fp4_situ_test( - quant_mode, alpha_value, beta_value, clamp_value, cache_permute_indices -): - num_experts = 8 - - def per_expert(value): - return ( - None - if value is None - else torch.full((num_experts,), value, device="cuda", dtype=torch.float32) - ) - - run_moe_test( - num_tokens=32, - hidden_size=1024, - intermediate_size=512, - moe_impl=FP4Moe(quant_mode=quant_mode), - routing_config={ - "num_experts": num_experts, - "top_k": 2, - "padding": 8, - "n_groups": None, - "top_k_groups": None, - "routed_scaling": None, - "has_routing_bias": False, - "routing_method_type": RoutingMethodType.Renormalize, - "compatible_moe_impls": [FP4Moe], - "compatible_intermediate_size": [512], - "compatible_activation_types": [ActivationType.Situ], - "enable_autotune": True, - }, - weight_processing={ - "use_shuffled_weight": True, - "layout": WeightLayout.MajorK, - "compatible_moe_impls": [FP4Moe], - }, - activation_type=ActivationType.Situ, - cache_permute_indices=cache_permute_indices, - routing_logits_dtype=torch.bfloat16, - gemm1_alpha=per_expert(alpha_value), - gemm1_beta=per_expert(beta_value), - gemm1_clamp_limit=per_expert(clamp_value), - ) - - -@pytest.mark.parametrize("alpha_value,beta_value,clamp_value", SITU_TEST_CASES) -def test_mxfp4_mxfp8_moe_situ( - alpha_value, beta_value, clamp_value, cache_permute_indices -): - _run_fp4_situ_test( - QuantMode.FP4_MXFP4_MXFP8, - alpha_value, - beta_value, - clamp_value, - cache_permute_indices, - ) - - -@pytest.mark.parametrize("alpha_value,beta_value,clamp_value", SITU_TEST_CASES) -def test_nvfp4_moe_situ(alpha_value, beta_value, clamp_value, cache_permute_indices): - _run_fp4_situ_test( - QuantMode.FP4_NVFP4_NVFP4, - alpha_value, - beta_value, - clamp_value, - cache_permute_indices, - ) - - -def _get_situ_tactic_tile_ns(dtype_act, dtype_weights): - from flashinfer.jit.fused_moe import gen_trtllm_gen_fused_moe_sm100_module - - raw_moe_op = gen_trtllm_gen_fused_moe_sm100_module().build_and_load() - tactics = raw_moe_op.trtllm_get_valid_moe_configs( - dtype_act, - dtype_weights, - Fp8QuantizationType.NoneFp8, - 2, - 1024, - 512, - 8, - ActivationType.Situ.value, - True, - WeightLayout.MajorK.value, - False, - 512, - False, - ) - return {int(tactic[0]) for tactic in tactics} - - -def test_mxfp4_mxfp8_autotuner_includes_tile_n_192(): - tile_ns = _get_situ_tactic_tile_ns(DtypeTrtllmGen.MxE4m3, DtypeTrtllmGen.MxE2m1) - assert 192 in tile_ns - - -def test_nvfp4_autotuner_includes_tile_n_192(): - tile_ns = _get_situ_tactic_tile_ns(DtypeTrtllmGen.E2m1, DtypeTrtllmGen.E2m1) - assert 192 in tile_ns - - -@pytest.mark.parametrize( - ("dtype_act", "dtype_weights"), - [ - pytest.param( - DtypeTrtllmGen.E2m1, - DtypeTrtllmGen.MxE2m1, - id="mxfp4_weights_nvfp4_activations", - ), - pytest.param( - DtypeTrtllmGen.MxE4m3, - DtypeTrtllmGen.E2m1, - id="nvfp4_weights_mxfp8_activations", - ), - ], -) -def test_unsupported_fp4_dtype_pairs_exclude_tile_n_192(dtype_act, dtype_weights): - tile_ns = _get_situ_tactic_tile_ns(dtype_act, dtype_weights) - assert 192 not in tile_ns From 5fea7fd0040c349f05794edbe6d4ef997588aad0 Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Tue, 28 Jul 2026 09:53:54 -0700 Subject: [PATCH 6/9] fix per token --- csrc/trtllm_batched_gemm_runner.cu | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/csrc/trtllm_batched_gemm_runner.cu b/csrc/trtllm_batched_gemm_runner.cu index 277719feff8..263fb90f77b 100644 --- a/csrc/trtllm_batched_gemm_runner.cu +++ b/csrc/trtllm_batched_gemm_runner.cu @@ -117,14 +117,22 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner( if (options.mFusedBiasShuffleMode != mOptions.fusedBiasShuffleMode) continue; if (options.mBiasDtype != mOptions.biasDtype) continue; } - if (mOptions.usePerTokenScaling) { - if (options.mTransposeMmaOutput && !options.mUsePerTokenSfB) continue; - if (!options.mTransposeMmaOutput && !options.mUsePerTokenSfA) continue; - } + bool const isNonStandardSfC = (options.mDtypeSfC != Dtype::Void && + options.mDtypeSfC != dtypeGetBlockSfType(mOptions.dtypeC)); + // if use non-standard sfC, force sf layout to linear + if (isNonStandardSfC && options.mSfLayoutC != SfLayout::Linear) continue; if (mOptions.usePerChannelScaling) { if (options.mTransposeMmaOutput && !options.mUsePerTokenSfA) continue; if (!options.mTransposeMmaOutput && !options.mUsePerTokenSfB) continue; } + if (mOptions.usePerTokenScaling) { + if (options.mTransposeMmaOutput && !options.mUsePerTokenSfB) continue; + if (!options.mTransposeMmaOutput && !options.mUsePerTokenSfA) continue; + } else { + // Avoid using the cubins with non-standard dtype sfc. + // They should only be selected for the fast per-token scaling case. + if (isNonStandardSfC) continue; + } if (options.mFusedAct) { if (options.mActType != static_cast(mOptions.actType)) { continue; From d22e8c350de51a9420f5be6058becaa583953960 Mon Sep 17 00:00:00 2001 From: Siyuan Fu Date: Tue, 28 Jul 2026 19:11:03 +0000 Subject: [PATCH 7/9] cherry-pick 4168 Signed-off-by: Siyuan Fu --- csrc/trtllm_batched_gemm_runner.cu | 24 ++--- csrc/trtllm_fused_moe_kernel_launcher.cu | 13 ++- .../test_trtllm_gen_moe_autotune_tactics.py | 93 +++++++++++++++++++ 3 files changed, 114 insertions(+), 16 deletions(-) diff --git a/csrc/trtllm_batched_gemm_runner.cu b/csrc/trtllm_batched_gemm_runner.cu index 263fb90f77b..f7e916fa06a 100644 --- a/csrc/trtllm_batched_gemm_runner.cu +++ b/csrc/trtllm_batched_gemm_runner.cu @@ -117,22 +117,22 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner( if (options.mFusedBiasShuffleMode != mOptions.fusedBiasShuffleMode) continue; if (options.mBiasDtype != mOptions.biasDtype) continue; } - bool const isNonStandardSfC = (options.mDtypeSfC != Dtype::Void && - options.mDtypeSfC != dtypeGetBlockSfType(mOptions.dtypeC)); - // if use non-standard sfC, force sf layout to linear - if (isNonStandardSfC && options.mSfLayoutC != SfLayout::Linear) continue; + bool const usesPerTokenScaling = + options.mTransposeMmaOutput ? options.mUsePerTokenSfB : options.mUsePerTokenSfA; + if (mOptions.usePerTokenScaling && !usesPerTokenScaling) 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; if (!options.mTransposeMmaOutput && !options.mUsePerTokenSfB) continue; } - if (mOptions.usePerTokenScaling) { - if (options.mTransposeMmaOutput && !options.mUsePerTokenSfB) continue; - if (!options.mTransposeMmaOutput && !options.mUsePerTokenSfA) continue; - } else { - // Avoid using the cubins with non-standard dtype sfc. - // They should only be selected for the fast per-token scaling case. - if (isNonStandardSfC) continue; - } if (options.mFusedAct) { if (options.mActType != static_cast(mOptions.actType)) { continue; diff --git a/csrc/trtllm_fused_moe_kernel_launcher.cu b/csrc/trtllm_fused_moe_kernel_launcher.cu index 3eaaedd5ab3..49909c7d81e 100644 --- a/csrc/trtllm_fused_moe_kernel_launcher.cu +++ b/csrc/trtllm_fused_moe_kernel_launcher.cu @@ -1190,8 +1190,11 @@ 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); + // Keep tactic enumeration aligned with prepare_moe_common(). The + // one-way cubin filter still admits the exported FP8 FC1 kernels + // that optionally consume BF16 routing scales. + /*usePerTokenScalingGemm1*/ false, + /*usePerTokenScalingGemm2*/ false, false, false); auto cfgs = moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, num_local_experts, num_tokens); @@ -2225,9 +2228,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); 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]) From d0959b6edb6aa1c91f2bdd407ddff89cf94455bb Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Tue, 28 Jul 2026 17:59:09 -0700 Subject: [PATCH 8/9] address comment --- flashinfer/fused_moe/core.py | 60 ------------------------------- flashinfer/trace/templates/moe.py | 15 ++++---- 2 files changed, 9 insertions(+), 66 deletions(-) diff --git a/flashinfer/fused_moe/core.py b/flashinfer/fused_moe/core.py index 9c080b1e5d7..91447c54ac9 100644 --- a/flashinfer/fused_moe/core.py +++ b/flashinfer/fused_moe/core.py @@ -3188,46 +3188,6 @@ def _validate_bf16_gemm1_activation_params( ) -def _prepare_fp4_block_scale_gemm1_activation_params( - activation_type: int, - gemm1_alpha: Optional[torch.Tensor], - gemm1_beta: Optional[torch.Tensor], - gemm1_clamp_limit: Optional[torch.Tensor], - local_num_experts: int, - device: torch.device, -) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: - if int(activation_type) != int(ActivationType.Situ): - return gemm1_alpha, gemm1_beta, gemm1_clamp_limit - - def validate_positive_finite(name: str, tensor: torch.Tensor) -> None: - check_shape_dtype_device( - tensor, (local_num_experts,), torch.float32, device, name - ) - # The warmup call validates values before CUDA graph capture. Avoid - # synchronizing a device tensor while the graph is being recorded. - if torch.cuda.is_current_stream_capturing(): - return - if not bool(torch.isfinite(tensor).all().item()): - raise ValueError(f"{name} must contain only finite values for SiTU.") - if not bool((tensor > 0).all().item()): - raise ValueError(f"{name} must contain only positive values for SiTU.") - - if gemm1_alpha is None: - gemm1_alpha = torch.ones(local_num_experts, dtype=torch.float32, device=device) - else: - validate_positive_finite("gemm1_alpha", gemm1_alpha) - - if gemm1_beta is None: - gemm1_beta = torch.ones(local_num_experts, dtype=torch.float32, device=device) - else: - validate_positive_finite("gemm1_beta", gemm1_beta) - - if gemm1_clamp_limit is not None: - validate_positive_finite("gemm1_clamp_limit", gemm1_clamp_limit) - - return gemm1_alpha, gemm1_beta, gemm1_clamp_limit - - def _validate_routing_replay_out( routing_replay_out: Optional[torch.Tensor], top_k: int ) -> None: @@ -4565,16 +4525,6 @@ def trtllm_fp4_block_scale_moe( path and fp32 ``DeepSeekV3`` logits. """ _validate_routing_replay_out(routing_replay_out, top_k) - gemm1_alpha, gemm1_beta, gemm1_clamp_limit = ( - _prepare_fp4_block_scale_gemm1_activation_params( - activation_type, - gemm1_alpha, - gemm1_beta, - gemm1_clamp_limit, - local_num_experts, - hidden_states.device, - ) - ) return get_trtllm_moe_sm100_module().trtllm_fp4_block_scale_moe( RoutingInputMode.FromLogits, routing_logits, @@ -4768,16 +4718,6 @@ def trtllm_fp4_block_scale_routed_moe( ``[output]`` when ``do_finalize`` is ``True``, otherwise ``[gemm2_output, expert_weights, expanded_idx_to_permuted_idx]``. """ - gemm1_alpha, gemm1_beta, gemm1_clamp_limit = ( - _prepare_fp4_block_scale_gemm1_activation_params( - activation_type, - gemm1_alpha, - gemm1_beta, - gemm1_clamp_limit, - local_num_experts, - hidden_states.device, - ) - ) # Determine routing mode based on input format if isinstance(topk_ids, tuple): # Unpacked format: (topk_ids, topk_weights) diff --git a/flashinfer/trace/templates/moe.py b/flashinfer/trace/templates/moe.py index 5e3ca8897a6..aefc5718424 100644 --- a/flashinfer/trace/templates/moe.py +++ b/flashinfer/trace/templates/moe.py @@ -1298,12 +1298,14 @@ 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_value = int(activation_type) - if activation_value not in (3, 10): + activation_type = normalize_activation_type(activation_type) + if activation_type not in (ActivationType.Swiglu, ActivationType.Situ): raise ValueError( - f"FP4 MoE trace reference supports SwiGLU (3) and SiTU (10), got {activation_value}" + "FP4 MoE trace reference supports " + f"{ActivationType.Swiglu!r} and {ActivationType.Situ!r}, " + f"got {activation_type!r}" ) - if activation_value == 10: + if activation_type == ActivationType.Situ: for name, param in ( ("gemm1_alpha", gemm1_alpha), ("gemm1_beta", gemm1_beta), @@ -1336,7 +1338,7 @@ def _fp4_moe_run_experts( G1 = G1 + gemm1_bias[le].to(torch.float32) # TRTLLM-Gen convention: x0 is linear/first; x1 is gate/second. x0, x1 = G1[:, :I], G1[:, I:] - if activation_value == 10: + if activation_type == ActivationType.Situ: alpha = ( torch.tensor(1.0, dtype=torch.float32, device=device) if gemm1_alpha is None @@ -1839,10 +1841,11 @@ 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 int(activation_type) == 10: + if activation_type == ActivationType.Situ: materialized: dict[str, torch.Tensor | None] = {} for name, value, default in ( ("gemm1_alpha", gemm1_alpha, 1.0), From 61a2005a1456ded702de9a3d8d66763165bba72b Mon Sep 17 00:00:00 2001 From: jiahanc <173873397+jiahanc@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:42:46 -0700 Subject: [PATCH 9/9] fix per token --- csrc/trtllm_batched_gemm_runner.cu | 6 +++++- csrc/trtllm_fused_moe_kernel_launcher.cu | 18 ++++++------------ csrc/trtllm_fused_moe_runner.cu | 9 +++++++++ flashinfer/fused_moe/core.py | 1 + .../trtllm/batched_gemm/KernelRunner.h | 2 ++ 5 files changed, 23 insertions(+), 13 deletions(-) diff --git a/csrc/trtllm_batched_gemm_runner.cu b/csrc/trtllm_batched_gemm_runner.cu index f7e916fa06a..aa53bcd7ce4 100644 --- a/csrc/trtllm_batched_gemm_runner.cu +++ b/csrc/trtllm_batched_gemm_runner.cu @@ -119,7 +119,10 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner( } bool const usesPerTokenScaling = options.mTransposeMmaOutput ? options.mUsePerTokenSfB : options.mUsePerTokenSfA; - if (mOptions.usePerTokenScaling && !usesPerTokenScaling) continue; + if (mOptions.usePerTokenScaling) { + 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 @@ -172,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 49909c7d81e..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,10 +1187,7 @@ class Fp8PerTensorLauncher : public FusedMoeLauncher { static_cast(weight_layout), // FP8 per-tensor doesn't use Mn-bias (LoRA) cubins. /*gemm1BiasType*/ batchedGemm::gemm::BiasType::None, - // Keep tactic enumeration aligned with prepare_moe_common(). The - // one-way cubin filter still admits the exported FP8 FC1 kernels - // that optionally consume BF16 routing scales. - /*usePerTokenScalingGemm1*/ false, + /*usePerTokenScalingGemm1*/ use_routing_scales_on_input, /*usePerTokenScalingGemm2*/ false, false, false); auto cfgs = moe_runner->getValidConfigIndices(top_k, hidden_size, intermediate_size, @@ -2984,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 1ba5e50b6ab..0699d6c93e6 100644 --- a/csrc/trtllm_fused_moe_runner.cu +++ b/csrc/trtllm_fused_moe_runner.cu @@ -427,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; @@ -451,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; } @@ -562,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/fused_moe/core.py b/flashinfer/fused_moe/core.py index 91447c54ac9..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, ) diff --git a/include/flashinfer/trtllm/batched_gemm/KernelRunner.h b/include/flashinfer/trtllm/batched_gemm/KernelRunner.h index 4246994b8d7..d13bb842974 100644 --- a/include/flashinfer/trtllm/batched_gemm/KernelRunner.h +++ b/include/flashinfer/trtllm/batched_gemm/KernelRunner.h @@ -96,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}; };