csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/thread/fused_activations.h (1)

64-79: Drop the unused maximum instance.

maximum<T, PropagateNaN> mx; is never referenced; keeping it will trigger unused-variable warnings with stricter builds. You can remove it and rely on ReLu<T> directly.

flashinfer/fused_moe/core.py (1)

75-79: Review comment verified; refactoring to complete enum mirroring is recommended.

The verification confirms that Swiglu=3 and Relu2=6 correctly match the C++ header. However, the Python enum at flashinfer/fused_moe/core.py (lines 76-78) is incomplete—it defines only 2 of 9 members present in the C++ definition. The C++ header itself contains a comment on line 21 stating "Note update flashinfer/fused_moe/core.py to match", flagging this as a known sync issue.

The Python enum is missing: Gelu (0), Relu (1), Silu (2), Geglu (4), SwigluBias (5), Identity (7), and InvalidType (8). The suggestion to mirror all C++ values or add explicit comments is sound for reducing drift and improving maintainability.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

flashinfer/fused_moe/core.py (2)

76-87: Keep ActivationType in Python in lockstep with C++ source of truth.

Enum values currently match common.h, but this local copy can drift. Consider auto‑generating from the C++ header at build time or adding a cheap import‑time check that compares numeric values returned from a tiny C++ helper. At minimum, document the sync contract next to this enum.

---

572-601: Fake op signature missing activation_type (breaks meta/fake dispatch).

Add the new parameter with same name/default as the real op. Otherwise, fake/meta kernels will error on kwargs mismatch.

Apply this diff:

 @register_fake_op("flashinfer::cutlass_fused_moe")
 def _fake_cutlass_fused_moe(
     output: torch.Tensor,
     input: torch.Tensor,
     token_selected_experts: torch.Tensor,
     token_final_scales: torch.Tensor,
     fc1_expert_weights: torch.Tensor,
     fc1_expert_biases: Optional[torch.Tensor],
     fc2_expert_weights: torch.Tensor,
     fc2_expert_biases: Optional[torch.Tensor],
     output_dtype: torch.dtype,
     quant_scales: List[torch.Tensor],
     input_sf: Optional[torch.Tensor] = None,
     swiglu_alpha: Optional[torch.Tensor] = None,
     swiglu_beta: Optional[torch.Tensor] = None,
     swiglu_limit: Optional[torch.Tensor] = None,
     tp_size: int = 1,
     tp_rank: int = 0,
     ep_size: int = 1,
     ep_rank: int = 0,
     cluster_size: int = 1,
     cluster_rank: int = 0,
     enable_alltoall: bool = False,
     use_deepseek_fp8_block_scale: bool = False,
     use_w4_group_scaling: bool = False,
     use_mxfp8_act_scaling: bool = False,
     min_latency_mode: bool = False,
     tune_max_num_tokens: int = 8192,
     enable_pdl: Optional[bool] = None,
+    activation_type: ActivationType = ActivationType.Swiglu,
 ):

csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_sm100_binding.cu (1)

365-367: Quant-scale shape checks still assume gated (×2) for FC1; plumb activation to getQuantParams.

Shape validations for FC1 block scales hardcode inter_size * 2. That rejects valid Relu2 (non‑gated) layouts and mis-sizes workspaces. Pass ActivationType into getQuantParams and gate the multiplier.

Apply these diffs:

1. Thread activation into calls:

-    auto const quant_params =
-        getQuantParams(num_experts_on_rank, hidden_size, inter_size, quant_scales);
+    auto const quant_params =
+        getQuantParams(num_experts_on_rank, hidden_size, inter_size, quant_scales, base_activation_type);

-    auto const quant_params =
-        getQuantParams(num_experts_on_rank, hidden_size, inter_size, quant_scales);
+    auto const quant_params =
+        getQuantParams(num_experts_on_rank, hidden_size, inter_size, quant_scales, base_activation_type);

2. Update signature and gate multiplier:

-  kernels::QuantParams getQuantParams(int64_t num_experts_on_rank, int64_t hidden_size,
-                                      int64_t inter_size,
-                                      Optional<Array<Tensor>> quant_scales) const {
+  kernels::QuantParams getQuantParams(int64_t num_experts_on_rank, int64_t hidden_size,
+                                      int64_t inter_size, Optional<Array<Tensor>> quant_scales,
+                                      ActivationType base_activation_type) const {
+    bool const is_gated = isGatedActivation(base_activation_type);

3. Fix FC1 shape checks (examples shown for each block):

W4A8_MXFP4_FP8:

-      TVM_FFI_ICHECK(
+      TVM_FFI_ICHECK(
           fc1_weight_block->shape[0] == num_experts_on_rank &&
           fc1_weight_block->shape[1] ==
               TmaWarpSpecializedGroupedGemmInput::alignToSfDim(
-                  inter_size, TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX) *
-                  2 &&
+                  inter_size, TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX) *
+                  (is_gated ? 2 : 1) &&
           ...
-          << "fc1 weight block size must be (num_experts_on_rank, inter_size * 2, hidden_size // 4 "
+          << "fc1 weight block size must be (num_experts_on_rank, inter_size * (is_gated?2:1), hidden_size // 4 "
              "// block_scale_vector_size)";

W4A8_MXFP4_MXFP8:

-      TVM_FFI_ICHECK(
+      TVM_FFI_ICHECK(
           fc1_weight_block->shape[0] == num_experts_on_rank &&
           fc1_weight_block->shape[1] ==
               TmaWarpSpecializedGroupedGemmInput::alignToSfDim(
-                  inter_size, TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX) *
-                  2 &&
+                  inter_size, TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX) *
+                  (is_gated ? 2 : 1) &&
           ...
-          << "fc1 weight block size must be (num_experts_on_rank, inter_size * 2, hidden_size // 4 "
+          << "fc1 weight block size must be (num_experts_on_rank, inter_size * (is_gated?2:1), hidden_size // 4 "
              "// block_scale_vector_size)";

NVFP4:

-      TVM_FFI_ICHECK(
+      TVM_FFI_ICHECK(
           fc1_weight_block->shape[0] == num_experts_on_rank &&
           fc1_weight_block->shape[1] ==
               TmaWarpSpecializedGroupedGemmInput::alignToSfDim(
-                  inter_size, TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentNVFP4) *
-                  2 &&
+                  inter_size, TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentNVFP4) *
+                  (is_gated ? 2 : 1) &&
           ...
-          << "fc1 weight block size must be (num_experts_on_rank, inter_size * 2, hidden_size // 4 "
+          << "fc1 weight block size must be (num_experts_on_rank, inter_size * (is_gated?2:1), hidden_size // 4 "
              "// block_scale_vector_size)";

Also applies to: 542-544, 807-813, 871-882, 930-941, 1001-1011

flashinfer/fused_moe/core.py (1)

655-656: Docstring/type hint don’t match actual return shape.

Signature and docs promise torch.Tensor, but the op returns a list when min_latency_mode flips. After fixing the inversion, either:

• document Union[torch.Tensor, List[torch.Tensor]], or
• normalize to always return a torch.Tensor from this wrapper.

Pick one and update annotations/docs accordingly.

Also applies to: 759-761

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro