perf: optimize MXFP4xBF16 & INT4xFP8 CUTLASS MoE backend for SM90

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

@@ -390,7 +390,10 @@ std::vector<CutlassGemmConfig> get_candidate_configs_sm90(
     if (has_w4afp8) {
       bool const has_coop_supported = sm90_supports_coop(tile_config);
       std::set<MainloopScheduleType> mainloop_schedules{MainloopScheduleType::PINGPONG};
-      if (has_coop_supported) {
+      if (has_coop_supported && tile_config != CutlassTileConfigSM90::CtaShape128x128x128B) {
+        // Due to register pressure on SM90, cooperative scheduler does not
+        // work with the 128x128x128B tile for mixed-dtype (W4A16) GEMM.
+        // Only PINGPONG is used for that tile to avoid register overflow.
         mainloop_schedules.insert(MainloopScheduleType::COOPERATIVE);
       }
       auto const epilogue_schedule = EpilogueScheduleType::AUTO;

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

@@ -43,6 +43,7 @@
 #include "cutlass_extensions/epilogue_helpers.h"
 #include "cutlass_extensions/gemm/collective/collective_builder_mixed_input.hpp"
 #include "cutlass_extensions/gemm_configs.h"
+#include "cutlass/gemm/kernel/tile_scheduler_params.h"
 
 #ifdef __GNUC__  // Check if the compiler is GCC or Clang
 #pragma GCC diagnostic pop
@@ -202,6 +203,11 @@ 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,113 @@
+/*
+ * 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,36 @@
+/*
+ * 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 <cstdint>
+#include <cuda_runtime.h>
+
+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