perf: optimize MXFP4xBF16 & INT4xFP8 CUTLASS MoE backend for SM90

csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu

@@ -29,6 +29,7 @@
 #include "moe_gemm_kernels.h"
 #include "tensorrt_llm/common/workspace.h"
 #include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h"
+#include "tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.h"
 
 namespace common = tensorrt_llm::common;
 namespace kernels = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE;
@@ -1274,4 +1275,48 @@ tvm::ffi::Module init(DLDataType activation_dtype, DLDataType weight_dtype, DLDa
   return tvm::ffi::Module(ptr);
 }
 
+// Interleave a 4-bit packed weight tensor into the layout required by the
+// SM90 mixed-input Hopper MoE GEMM. Expected input shape (num_experts, n,
+// k / 2) uint8 on CUDA. Writes into an output tensor of the same shape.
+// quant_type: 0 for INT4 (W4A8), 1 for FP4 (W4A16 / MXFP4).
+void interleave_moe_weights_for_Hopper_mixed_gemm(TensorView weight, TensorView weight_interleaved,
+                                                  int64_t quant_type) {
+  CHECK_INPUT_TYPE(weight, dl_uint8);
+  CHECK_INPUT_TYPE(weight_interleaved, dl_uint8);
+  CHECK_CONTIGUOUS(weight);
+  CHECK_CONTIGUOUS(weight_interleaved);
+  CHECK_DIM(3, weight);
+  CHECK_DIM(3, weight_interleaved);
+  TVM_FFI_ICHECK_EQ(weight.size(0), weight_interleaved.size(0))
+      << "weight and weight_interleaved must share num_experts dim";
+  TVM_FFI_ICHECK_EQ(weight.size(1), weight_interleaved.size(1))
+      << "weight and weight_interleaved must share n dim";
+  TVM_FFI_ICHECK_EQ(weight.size(2), weight_interleaved.size(2))
+      << "weight and weight_interleaved must share packed-k dim";
+  TVM_FFI_ICHECK(quant_type == 0 || quant_type == 1)
+      << "quant_type must be 0 (INT4) or 1 (FP4), got " << quant_type;
+
+  int64_t const num_experts = weight.size(0);
+  int64_t const n = weight.size(1);
+  int64_t const k = weight.size(2) * 2;
+  int64_t const per_expert_bytes = n * (k / 2);
+
+  auto stream = get_stream(weight.device());
+  auto* src = static_cast<uint8_t*>(weight.data_ptr());
+  auto* dst = static_cast<uint8_t*>(weight_interleaved.data_ptr());
+  for (int64_t e = 0; e < num_experts; ++e) {
+    uint8_t* src_e = src + e * per_expert_bytes;
+    uint8_t* dst_e = dst + e * per_expert_bytes;
+    if (quant_type == 1) {
+      tensorrt_llm::kernels::cutlass_kernels::interleave_fp4_weights_for_Hopper_mixed_gemm(
+          src_e, dst_e, static_cast<int>(n), static_cast<int>(k), stream);
+    } else {
+      tensorrt_llm::kernels::cutlass_kernels::interleave_int4_weights_for_Hopper_mixed_gemm(
+          src_e, dst_e, static_cast<int>(n), static_cast<int>(k), stream);
+    }
+  }
+}
+
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(init, init);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(interleave_moe_weights_for_Hopper_mixed_gemm,
+                              interleave_moe_weights_for_Hopper_mixed_gemm);

csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp

@@ -389,9 +389,16 @@ std::vector<CutlassGemmConfig> get_candidate_configs_sm90(
         (config & CutlassGemmConfig::WEIGHT_ONLY) && (config & CutlassGemmConfig::GROUPED_GEMM);
     if (has_w4afp8) {
       bool const has_coop_supported = sm90_supports_coop(tile_config);
-      std::set<MainloopScheduleType> mainloop_schedules{MainloopScheduleType::PINGPONG};
+      std::set<MainloopScheduleType> mainloop_schedules;
       if (has_coop_supported) {
+        // Due to the limitation on the number of registers on SM,
+        // cooperative scheduler does not support CtaShape128x128x128B
+        // for mixed-dtype (W4A16) grouped GEMM. Skip the tile entirely
+        // to avoid register overflow.
+        if (tile_config == CutlassTileConfigSM90::CtaShape128x128x128B) continue;
         mainloop_schedules.insert(MainloopScheduleType::COOPERATIVE);
+      } else {
+        mainloop_schedules.insert(MainloopScheduleType::PINGPONG);
       }
       auto const epilogue_schedule = EpilogueScheduleType::AUTO;
       for (auto const& mainloop_schedule : mainloop_schedules) {

csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl

@@ -26,6 +26,7 @@
 #include "cutlass/gemm/device/gemm_universal_adapter.h"
 #include "cutlass/gemm/dispatch_policy.hpp"
 #include "cutlass/gemm/kernel/gemm_universal.hpp"
+#include "cutlass/gemm/kernel/tile_scheduler_params.h"
 #include "cutlass/util/command_line.h"
 #include "cutlass/util/distribution.h"
 #include "cutlass/util/host_tensor.h"
@@ -202,6 +203,12 @@ void sm90_generic_mixed_moe_gemm_kernelLauncher(
        reinterpret_cast<StrideD*>(hopper_inputs.stride_d)},
       hw_info};
 
+  // Optimize tile scheduling for better L2 locality
+  using RasterOrderOptions =
+      typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90Params::RasterOrderOptions;
+  arguments.scheduler.max_swizzle_size = 2;
+  arguments.scheduler.raster_order = RasterOrderOptions::Heuristic;
+
   assert(group_size == int(inputs.groupwise_quant_group_size));
   if (workspace_size != nullptr) {
     *workspace_size = gemm.get_workspace_size(arguments);

csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.cu

@@ -0,0 +1,114 @@
+/*
+ * Copyright (c) 2020-2025, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "moe_gemm_mixed_utils.h"
+
+namespace tensorrt_llm {
+namespace kernels {
+namespace cutlass_kernels {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////////////
+
+__global__ void interleave_fp4_weights_for_Hopper_mixed_gemm_kernel(uint8_t* fp4_weight,
+                                                                    uint8_t* fp4_weight_interleaved,
+                                                                    int const rows,
+                                                                    int const cols) {
+  for (int block_id = blockIdx.x; block_id < rows / 2; block_id += gridDim.x) {
+    for (int partition_id = threadIdx.y; partition_id < cols / 64; partition_id += blockDim.y) {
+      int lane_id = threadIdx.x;
+      int row_id = block_id / 8 * 16 + block_id % 8;
+
+      int mma_id = lane_id / 8;
+      int dst_row_id = row_id + (mma_id % 2) * 8;
+
+      int interleaved_lane_id = lane_id / 16 * 16 + (lane_id % 4) * 4 + (lane_id % 8) / 4 * 2;
+
+      int col_id = partition_id * 32 + lane_id;
+      int dst_col_id = partition_id * 32 + interleaved_lane_id;
+
+      int index_a = row_id * cols / 2 + col_id;
+      int index_b = (row_id + 8) * cols / 2 + col_id;
+
+      uint8_t fp4x2_a = fp4_weight[index_a];
+      uint8_t fp4x2_b = fp4_weight[index_b];
+
+      uint8_t fp4_temp_a = (fp4x2_a & 0xF0U) >> 4;
+      uint8_t fp4_temp_b = (fp4x2_b & 0x0FU) << 4;
+
+      fp4x2_a = (fp4x2_a & 0x0FU) | fp4_temp_b;
+      fp4x2_b = (fp4x2_b & 0xF0U) | fp4_temp_a;
+
+      int dst_id = dst_row_id * cols / 2 + dst_col_id;
+
+      fp4_weight_interleaved[dst_id] = fp4x2_a;
+      fp4_weight_interleaved[dst_id + 1] = fp4x2_b;
+    }
+  }
+}
+
+__global__ void interleave_int4_weights_for_Hopper_mixed_gemm_kernel(
+    uint8_t* int4_weight, uint8_t* int4_weight_interleaved, int const rows, int const cols) {
+  uint16_t* uint16_ptr = reinterpret_cast<uint16_t*>(int4_weight);
+  uint16_t* uint16_interleaved_ptr = reinterpret_cast<uint16_t*>(int4_weight_interleaved);
+
+  for (int block_id = blockIdx.x; block_id < rows / 2; block_id += gridDim.x) {
+    for (int partition_id = threadIdx.y; partition_id < cols / 64; partition_id += blockDim.y) {
+      int lane_id = threadIdx.x;
+
+      int row_id = block_id / 8 * 16 + block_id % 8;
+      int dst_row_id = row_id + (lane_id % 8) / 4 * 8;
+
+      int mma_id = lane_id / 8;
+      int interleaved_lane_id = mma_id * 8 + lane_id % 4 * 2;
+
+      int col_id = partition_id * 16 + lane_id;
+      int dst_col_id = partition_id * 16 + interleaved_lane_id;
+
+      int src_id_a = row_id * cols / 4 + col_id;
+      int src_id_b = (row_id + 8) * cols / 4 + col_id;
+
+      uint16_t int4x2_a = uint16_ptr[src_id_a];
+      uint16_t int4x2_b = uint16_ptr[src_id_b];
+
+      int dst_id = dst_row_id * cols / 4 + dst_col_id;
+
+      uint16_interleaved_ptr[dst_id] = int4x2_a;
+      uint16_interleaved_ptr[dst_id + 1] = int4x2_b;
+    }
+  }
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void interleave_fp4_weights_for_Hopper_mixed_gemm(uint8_t* fp4_weight,
+                                                  uint8_t* fp4_weight_interleaved, int const rows,
+                                                  int const cols, cudaStream_t stream) {
+  dim3 block(32, 32);
+  interleave_fp4_weights_for_Hopper_mixed_gemm_kernel<<<1024, block, 0, stream>>>(
+      fp4_weight, fp4_weight_interleaved, rows, cols);
+}
+
+void interleave_int4_weights_for_Hopper_mixed_gemm(uint8_t* int4_weight,
+                                                   uint8_t* int4_weight_interleaved, int const rows,
+                                                   int const cols, cudaStream_t stream) {
+  dim3 block(16, 32);
+  interleave_int4_weights_for_Hopper_mixed_gemm_kernel<<<1024, block, 0, stream>>>(
+      int4_weight, int4_weight_interleaved, rows, cols);
+}
+
+}  // namespace cutlass_kernels
+}  // namespace kernels
+}  // namespace tensorrt_llm

csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.h

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+namespace tensorrt_llm {
+namespace kernels {
+namespace cutlass_kernels {
+
+void interleave_fp4_weights_for_Hopper_mixed_gemm(uint8_t* weight, uint8_t* weight_interleaved,
+                                                  int rows, int cols, cudaStream_t stream = 0);
+
+void interleave_int4_weights_for_Hopper_mixed_gemm(uint8_t* weight, uint8_t* weight_interleaved,
+                                                   int rows, int cols, cudaStream_t stream = 0);
+
+}  // namespace cutlass_kernels
+}  // namespace kernels
+}  // namespace tensorrt_llm

flashinfer/fused_moe/__init__.py

@@ -21,6 +21,8 @@
     WeightLayout,
     convert_to_block_layout,
     cutlass_fused_moe,
+    interleave_moe_scales_for_hopper_mixed_gemm,
+    interleave_moe_weights_for_hopper_mixed_gemm,
     gen_cutlass_fused_moe_sm120_module,
     gen_cutlass_fused_moe_sm103_module,
     gen_cutlass_fused_moe_sm100_module,
@@ -59,6 +61,8 @@
     "WeightLayout",
     "convert_to_block_layout",
     "cutlass_fused_moe",
+    "interleave_moe_scales_for_hopper_mixed_gemm",
+    "interleave_moe_weights_for_hopper_mixed_gemm",
     "gen_cutlass_fused_moe_sm120_module",
     "gen_cutlass_fused_moe_sm103_module",
     "gen_cutlass_fused_moe_sm100_module",

flashinfer/fused_moe/core.py

@@ -581,9 +581,122 @@ def _fake_cutlass_fused_moe(
     # Register the module
     return SimpleNamespace(
         cutlass_fused_moe=cutlass_fused_moe,
+        interleave_moe_weights_for_Hopper_mixed_gemm=(
+            module.interleave_moe_weights_for_Hopper_mixed_gemm
+        ),
     )
 
 
+@flashinfer_api
+def interleave_moe_scales_for_hopper_mixed_gemm(
+    scales: torch.Tensor,
+    group_size: int = 32,
+) -> torch.Tensor:
+    """Interleave MXFP4 block scales for the SM90 mixed-input Hopper GEMM.
+
+    The kernel expects scales in layout
+    ``(num_experts, K // (group_size * 4), rows * 4)`` rather than the natural
+    ``(num_experts, rows, K // group_size)`` produced by the MXFP4 quantizer.
+    This helper performs the reshape + permute equivalent to TensorRT-LLM's
+    ``WFP4A16FusedMoEMethod.load_quant_scales`` (PR #12451), with the fixed
+    interleave factor of ``128 // group_size`` used for MXFP4.
+
+    Parameters
+    ----------
+    scales:
+        ``[num_experts, rows, K // group_size]`` uint8 tensor of E8M0 block
+        scales.
+    group_size:
+        MXFP4 quantization group size (default 32).
+
+    Returns
+    -------
+    torch.Tensor
+        Contiguous uint8 tensor with shape
+        ``[num_experts, K // (group_size * factor), rows * factor]``
+        where ``factor = 128 // group_size``.
+    """
+    if scales.dim() != 3:
+        raise ValueError(
+            f"scales must be 3D (num_experts, rows, K/group_size); got {tuple(scales.shape)}"
+        )
+    if scales.dtype != torch.uint8:
+        raise ValueError(f"scales must be uint8 (E8M0); got {scales.dtype}")
+
+    factor = 128 // group_size
+    if factor < 1 or 128 % group_size != 0:
+        raise ValueError(
+            f"group_size={group_size} must divide 128 (interleave factor = 128 // group_size)"
+        )
+    e, rows, kgs = scales.shape
+    if kgs % factor != 0:
+        raise ValueError(
+            f"K/group_size={kgs} must be divisible by interleave factor {factor}"
+        )
+    tmp = (
+        scales.reshape(e, rows, kgs // factor, factor).permute(0, 2, 1, 3).contiguous()
+    )
+    return tmp.reshape(e, kgs // factor, rows * factor)
+
+
+@flashinfer_api
+def interleave_moe_weights_for_hopper_mixed_gemm(
+    weight: torch.Tensor,
+    quant_type: str = "fp4",
+) -> torch.Tensor:
+    """Interleave 4-bit packed MoE weights for the SM90 mixed-input GEMM.
+
+    The Hopper mixed-dtype MoE GEMM (used by ``cutlass_fused_moe`` with
+    ``use_w4_group_scaling=True``) expects weights in a specific interleaved
+    layout; without preprocessing, the LUT-based FP4→BF16 conversion reads
+    bytes from the wrong positions and the output diverges from a dequantized
+    reference for any K > 128. TensorRT-LLM's W4A16 MoE runs the equivalent
+    preprocessing at weight-load time (see ``trtllm::
+    interleave_4bit_weights_for_Hopper_mixed_gemm`` in PR #12451).
+
+    Parameters
+    ----------
+    weight:
+        ``[num_experts, n, k // 2]`` uint8 CUDA tensor (4-bit values packed
+        two-per-byte).
+    quant_type:
+        ``"fp4"`` for MXFP4 (the W4A16 path) or ``"int4"`` for INT4 (the
+        W4A8 path).
+
+    Returns
+    -------
+    torch.Tensor
+        A new uint8 tensor with the same shape as ``weight`` holding the
+        interleaved layout. Feed this directly as ``fc1_expert_weights`` /
+        ``fc2_expert_weights`` to :func:`cutlass_fused_moe`.
+    """
+    if weight.dim() != 3:
+        raise ValueError(
+            f"weight must be 3D (num_experts, n, k/2); got shape {tuple(weight.shape)}"
+        )
+    if weight.dtype != torch.uint8:
+        raise ValueError(f"weight must be uint8 (packed 4-bit); got {weight.dtype}")
+    if not weight.is_cuda:
+        raise ValueError("weight must live on CUDA")
+
+    qtype_map = {"fp4": 1, "int4": 0}
+    if quant_type not in qtype_map:
+        raise ValueError(
+            f"quant_type must be one of {list(qtype_map)}; got {quant_type!r}"
+        )
+
+    weight = weight.contiguous()
+    out = torch.empty_like(weight)
+
+    major, minor = get_compute_capability(weight.device)
+    device_arch = f"{major * 10 + minor}"
+    module = get_cutlass_fused_moe_module(device_arch)
+    module.interleave_moe_weights_for_Hopper_mixed_gemm(
+        weight, out, qtype_map[quant_type]
+    )
+    return out
+
+
 # ref: https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py#L121
 @flashinfer_api
 def cutlass_fused_moe(

flashinfer/jit/fused_moe.py

@@ -194,6 +194,8 @@ def gen_cutlass_fused_moe_module(
             jit_env.FLASHINFER_CSRC_DIR
             / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu",
             jit_env.FLASHINFER_CSRC_DIR
+            / "nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_mixed_utils.cu",
+            jit_env.FLASHINFER_CSRC_DIR
             / "fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu",
             jit_env.FLASHINFER_CSRC_DIR
             / "fused_moe/cutlass_backend/deepgemm_jit_setup.cu",