[TRTLLM-16404][fix] restore CUTLASS MoE compatibility and early feature checks

cpp/tensorrt_llm/CMakeLists.txt

@@ -222,6 +222,8 @@ endif()
 if(USING_OSS_CUTLASS_MOE_GEMM)
   set(TRTLLM_LINK_LIBS ${TRTLLM_LINK_LIBS} moe_gemm_src)
   message(STATUS "USING_OSS_CUTLASS_MOE_GEMM")
+else()
+  set(TRTLLM_LINK_LIBS ${TRTLLM_LINK_LIBS} moe_host_src)
 endif()
 
 if(USING_OSS_CUTLASS_ALLREDUCE_GEMM)

cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt

@@ -281,6 +281,14 @@ if(USING_OSS_CUTLASS_MOE_GEMM)
   target_include_directories(
     moe_gemm_src
     PUBLIC ${PROJECT_SOURCE_DIR}/tensorrt_llm/cutlass_extensions/include)
+else()
+  add_library(moe_host_src STATIC ${MOE_CUTLASS_DIR}/moe_kernels.cu)
+  target_compile_definitions(moe_host_src PRIVATE TLLM_MOE_NO_GEMM)
+  add_cuda_architectures(moe_host_src 89)
+  process_target(moe_host_src false false)
+  target_include_directories(
+    moe_host_src
+    PUBLIC ${PROJECT_SOURCE_DIR}/tensorrt_llm/cutlass_extensions/include)
 endif()
 
 if(USING_OSS_CUTLASS_ALLREDUCE_GEMM)
@@ -332,6 +340,10 @@ endforeach()
 
 list(APPEND ALL_KERNELS_LIB ${TARGET_LIB})
 
+if(NOT USING_OSS_CUTLASS_MOE_GEMM)
+  list(APPEND ALL_KERNELS_LIB moe_host_src)
+endif()
+
 add_custom_target(cutlass_kernels DEPENDS ${ALL_KERNELS_LIB})
 
 # foreach(target_name moe_gemm_src) process_target(${target_name} true true) #

cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu

@@ -3206,6 +3206,7 @@ void dequantFP8(OutputType* output, InputType const* input, int64_t const* num_v
         <<<blocks, threads, 0, stream>>>(output, input, num_valid_tokens_ptr, inter_size, scale, scale_is_dequant);
 }
 
+#if !defined(TLLM_MOE_NO_GEMM)
 // The DeepSeek FP8 block-scale MoE folds the pre-FC1 and pre-FC2 1x128 activation quant into the
 // row-expansion and activation-epilogue kernels, enabled automatically on supported hardware (Hopper / SM90).
 // Evaluated once at construction to pick the block-scale runner, so it is fixed for the runner's lifetime.
@@ -5860,6 +5861,7 @@ template class CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, __nv_bfloat16, _
 template class CutlassMoeFCRunner<__nv_bfloat16, __nv_fp4_e2m1>;
 #endif
 #endif
+#endif // !defined(TLLM_MOE_NO_GEMM)
 
 } // namespace kernels::cutlass_kernels
 

cpp/tests/CMakeLists.txt

@@ -99,7 +99,7 @@ endfunction()
 
 function(add_gtest test_name test_src)
   set(options NO_GTEST_MAIN NO_TLLM_LINKAGE)
-  set(oneValueArgs TIMEOUT)
+  set(oneValueArgs TIMEOUT TEST_PREFIX)
   cmake_parse_arguments(ARGS "${options}" "${oneValueArgs}" "${multiValueArgs}"
                         ${ARGN})
   add_executable(${test_name} ${test_src})
@@ -129,6 +129,7 @@ function(add_gtest test_name test_src)
 
   gtest_discover_tests(
     ${test_name}
+    TEST_PREFIX "${ARGS_TEST_PREFIX}"
     PROPERTIES ${test_properties}
                DISCOVERY_MODE
                PRE_TEST # WAR for DLL discovery on windows.

cpp/tests/unit_tests/kernels/CMakeLists.txt

@@ -33,14 +33,10 @@ if(USING_OSS_CUTLASS_MOE_GEMM)
   target_compile_definitions(mixtureOfExpertsTest
                              PUBLIC USING_OSS_CUTLASS_MOE_GEMM)
 
-  # The internal-path variant includes headers (quantization.h) that only ship
-  # with the internal cutlass kernels sources; the prebuilt tarball provides the
-  # library and a reduced header set only.
-  if(INTERNAL_CUTLASS_KERNELS_PATH)
-    add_gtest(mixtureOfExpertsInternalTest moe/mixtureOfExpertsTest.cu)
-    remove_compile_definition(mixtureOfExpertsInternalTest
-                              USING_OSS_CUTLASS_MOE_GEMM)
-  endif()
+  add_gtest(mixtureOfExpertsInternalTest moe/mixtureOfExpertsTest.cu
+            TEST_PREFIX "Internal.")
+  remove_compile_definition(mixtureOfExpertsInternalTest
+                            USING_OSS_CUTLASS_MOE_GEMM)
 endif()
 
 add_gtest(ropeTest ropeTest.cu)

cpp/tests/unit_tests/kernels/moe/mixtureOfExpertsTest.cu

@@ -17,6 +17,7 @@
 #include "tensorrt_llm/common/cudaUtils.h"
 #include "tensorrt_llm/common/memoryUtils.h"
 #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.h"
+#include "tensorrt_llm/kernels/quantization.h"
 #include "tensorrt_llm/runtime/cudaStream.h"
 
 #include <algorithm>
@@ -25,12 +26,11 @@
 
 #ifdef USING_OSS_CUTLASS_MOE_GEMM
 #include "tensorrt_llm/kernels/moe/cutlass/include/moe_kernels.h"
-#include <tensorrt_llm/kernels/quantization.h>
 #else
 #include "moe_kernels.h"
-#include "quantization.h"
 #endif
 #include "tensorrt_llm/kernels/cutlass_kernels/include/cutlass_kernel_selector.h"
+#include "tensorrt_llm/kernels/moe/cutlass/include/moe_gemm_kernels.h"
 
 #include "tensorrt_llm/common/tllmDataType.h"
 #include "tensorrt_llm/runtime/bufferManager.h"
@@ -42,7 +42,10 @@ using namespace tensorrt_llm::common;
 using namespace tensorrt_llm::runtime;
 
 using namespace CUTLASS_MOE_GEMM_KERNELS_NAMESPACE;
-using CUTLASS_MOE_GEMM_NAMESPACE::TmaWarpSpecializedGroupedGemmInput;
+using TmaWarpSpecializedGroupedGemmInput = tensorrt_llm::kernels::cutlass_kernels::TmaWarpSpecializedGroupedGemmInput;
+#if !defined(USING_OSS_CUTLASS_MOE_GEMM)
+using MoeGemmId = GemmProfilerBackend::GemmToProfile;
+#endif
 using CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::CutlassMoeFCRunner;
 using CUTLASS_MOE_GEMM_NAMESPACE::ActivationType;
 using CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::ActivationParams;
@@ -51,6 +54,15 @@ using CUTLASS_MOE_GEMM_NAMESPACE::isGatedActivation;
 constexpr static float FP8_MAX = 448.f;
 constexpr static float FP4_MAX = 6.f;
 
+constexpr bool supportsUnpaddedMoeOutput()
+{
+#if defined(USING_OSS_CUTLASS_MOE_GEMM)
+    return true;
+#else
+    return false;
+#endif
+}
+
 __host__ __device__ constexpr float applyExpertShift(float weight_value, int expert)
 {
     float lookup_table[] = {0.5f, 1.0f, 2.0f};
@@ -343,19 +355,12 @@ protected:
     using ElementSF = TmaWarpSpecializedGroupedGemmInput::ElementSF;
     constexpr static int FP4VecSize = MX_QUANT_WEIGHT ? TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaleVectorSize
                                                       : TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize;
-#ifdef USING_OSS_CUTLASS_MOE_GEMM
     constexpr static int MinNDimAlignmentFP4 = MX_QUANT_WEIGHT
         ? TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX
         : TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentNVFP4;
     constexpr static int MinKDimAlignmentFP4 = MX_QUANT_WEIGHT
         ? TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentMXFPX
         : TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentNVFP4;
-#else
-    constexpr static int MinNDimAlignmentFP4 = MX_QUANT_WEIGHT
-        ? TmaWarpSpecializedGroupedGemmInput::MinNumRowsAlignmentMXFPX
-        : TmaWarpSpecializedGroupedGemmInput::MiNumRowsAlignmentNVFP4;
-    constexpr static int MinKDimAlignmentFP4 = FP4VecSize * 4; // Hardcode the correct value
-#endif
     ElementSF* mFP4ScalingFactorsW1 = nullptr;
     ElementSF* mFP4ScalingFactorsW2 = nullptr;
 
@@ -800,15 +805,9 @@ protected:
         int64_t padded_in_dim = TmaWarpSpecializedGroupedGemmInput::alignToSfDim(in_shape, MinKDimAlignmentFP4);
         check_cuda_error(cudaMemsetAsync(scaling_factors, 0x00,
             num_experts * padded_out_dim * padded_in_dim / FP4VecSize * sizeof(ElementSF), mStream->get()));
-#ifdef USING_OSS_CUTLASS_MOE_GEMM
         invokeFP4Quantization<WeightRawType, FP4VecSize>(num_experts, out_shape, in_shape, raw_weights, global_scales,
             reinterpret_cast<int64_t*>(quant_weights), reinterpret_cast<int32_t*>(scaling_factors), MX_QUANT_WEIGHT,
             tensorrt_llm::QuantizationSFLayout::SWIZZLED, mMultiProcessorCount, mStream->get());
-#else
-        invokeBatchedFP4Quantization<WeightRawType, FP4VecSize>(num_experts, out_shape, in_shape, raw_weights,
-            global_scales, reinterpret_cast<int64_t*>(quant_weights), reinterpret_cast<int32_t*>(scaling_factors),
-            MX_QUANT_WEIGHT, mMultiProcessorCount, mStream->get());
-#endif
 
         // auto sf_data = getDataFromDevice<ElementSF>(scaling_factors, num_experts * padded_out_dim * padded_in_dim /
         // FP4VecSize); auto unquant_data = getDataFromDevice<WeightRawType>(raw_weights, num_experts * out_shape *
@@ -1237,7 +1236,11 @@ protected:
 
     auto getFilteredConfigs(int sm, MoeGemmId gemm_id)
     {
+#if defined(USING_OSS_CUTLASS_MOE_GEMM)
         auto tactics = mMoERunner.getTactics(gemm_id);
+#else
+        auto tactics = mMoERunner.getTactics();
+#endif
         if (sm == 89 || sm >= 120)
         {
             // Filter some unsupported configs for L40S
@@ -1307,6 +1310,18 @@ protected:
     using ConfigsToTestVec = std::vector<std::pair<tensorrt_llm::cutlass_extensions::CutlassGemmConfig,
         tensorrt_llm::cutlass_extensions::CutlassGemmConfig>>;
 
+    bool isFinalizeFused(tensorrt_llm::cutlass_extensions::CutlassGemmConfig const& gemm2) const
+    {
+#ifdef USING_OSS_CUTLASS_MOE_GEMM
+        return gemm2.epilogue_fusion_type
+            == tensorrt_llm::cutlass_extensions::CutlassGemmConfig::EpilogueFusionType::FINALIZE;
+#else
+        // The bundled runner controls finalize fusion with its runtime flag,
+        // rather than the tactic's epilogue_fusion_type metadata.
+        return gemm2.is_tma_warp_specialized && mUseFusedFinalizeInternal && !W4A8_AWQ;
+#endif
+    }
+
     auto getAllTileConfigsToTest()
     {
         if (mOverrideSelectedConfig1 && mOverrideSelectedConfig2)
@@ -1404,8 +1419,12 @@ protected:
             }
             else if constexpr (MXFP8_MXFP4)
             {
+#if defined(USING_OSS_CUTLASS_MOE_GEMM)
                 quant_params = QuantParams::MXFP8MXFP4(mFP4ScalingFactorsW1, static_cast<float const*>(scale1_ptr),
                     mFP4ScalingFactorsW2, static_cast<float const*>(scale3_ptr));
+#else
+                static_assert(!MXFP8_MXFP4, "MXFP8 x MXFP4 requires the open-source CUTLASS MoE kernels");
+#endif
             }
         }
 
@@ -1435,11 +1454,11 @@ protected:
             mNumExperts, mK, mWorkspace, mFinalOutput, mSourceToExpandedMap, parallelism_config, enable_alltoall,
             mUseLora, lora_params, useFp8BlockScales, minLatencyMode, min_latency_params, stream);
 #else
-        mMoERunner.runMoe(mInputTensor, nullptr, true, mSelectedExpert, mTokenFinalScales, weight1_ptr, bias1_ptr,
-            ActivationParams(mActType, mSwigluAlpha, mSwigluBeta, mSwigluLimit, mSwigluClampAfterSilu), weight2_ptr,
-            bias2_ptr, quant_params, mTotalTokens, mTotalTokens, mHiddenSize, mInterSize / parallelism_config.tp_size,
-            mNumExperts, mK, mWorkspace, mFinalOutput, mSourceToExpandedMap, parallelism_config, mUseLora, lora_params,
-            useFp8BlockScales, minLatencyMode, min_latency_params, stream);
+        mMoERunner.runMoe(mInputTensor, nullptr, mSelectedExpert, mTokenFinalScales, weight1_ptr, bias1_ptr,
+            ActivationParams(mActType, mSwigluAlpha, mSwigluBeta, mSwigluLimit), weight2_ptr, bias2_ptr, quant_params,
+            mTotalTokens, mHiddenSize, mInterSize / parallelism_config.tp_size, mNumExperts, mK, mWorkspace,
+            mFinalOutput, mSourceToExpandedMap, parallelism_config, mUseLora, lora_params, useFp8BlockScales,
+            minLatencyMode, min_latency_params, stream);
 #endif
 
         check_cuda_error(cudaStreamSynchronize(stream));
@@ -1731,8 +1750,10 @@ protected:
         // 1 expert per rank
         ParallelismTest(k, 1, num_experts, hidden_size, num_experts, num_tokens);
 
+#if defined(USING_OSS_CUTLASS_MOE_GEMM)
         // 2 expert per rank, enable alltoall optimised finalize
         ParallelismTest(k, 1, num_experts / 2, hidden_size, num_experts, num_tokens, true);
+#endif
     }
 
     // Tensor parallel tests default to inter_size_fraction = 1.0f so that all ranks have interesting values (i.e. a
@@ -1798,8 +1819,10 @@ using Types = ::testing::Types<
 
 #ifdef ENABLE_BF16
 #ifdef ENABLE_FP8
+#ifdef USING_OSS_CUTLASS_MOE_GEMM
     WeightParams<SafeFP8, cutlass::uint4b_t, __nv_bfloat16, void, __nv_bfloat16>,
 #endif
+#endif
 #endif
     WeightParams<half>, WeightParams<float>
 
@@ -1867,8 +1890,7 @@ void MixtureOfExpertsTest<TypeParam_>::BasicPermuteTest(
         auto [expected_experts, token_final_scales] = populateRouting(num_experts, num_tokens, k);
 
         runMoEPermute(hidden_input, expected_experts, token_final_scales, hidden_size, num_experts, k);
-        bool is_finalize_fusion = gemm2.epilogue_fusion_type
-            == tensorrt_llm::cutlass_extensions::CutlassGemmConfig::EpilogueFusionType::FINALIZE;
+        bool is_finalize_fusion = isFinalizeFused(gemm2);
         bool should_be_deterministic = !is_finalize_fusion || mK < 3;
         if (should_be_deterministic && !mIsLongTest)
         {
@@ -1988,6 +2010,9 @@ TYPED_TEST(MixtureOfExpertsTest, PermuteSwigluBias)
 
 TYPED_TEST(MixtureOfExpertsTest, PermuteSwigluPostSiluClamp)
 {
+#if !defined(USING_OSS_CUTLASS_MOE_GEMM)
+    GTEST_SKIP() << "Post-SiLU SwiGLU clamping requires the open-source CUTLASS MoE kernels";
+#endif
     if (this->W4A8_AWQ)
     {
         GTEST_SKIP() << "W4A8 does not support gated activations";
@@ -2196,8 +2221,7 @@ void MixtureOfExpertsTest<TypeParam_>::ParallelismTest(
                     // Only need to init the inputs on the first iteration
                     runMoEPermute(hidden_input, expected_experts, token_final_scales, hidden_size, num_experts, k,
                         MOEParallelismConfig{tp_size, i, ep_size, j}, enable_alltoall);
-                    bool is_finalize_fusion = gemm2.epilogue_fusion_type
-                        == tensorrt_llm::cutlass_extensions::CutlassGemmConfig::EpilogueFusionType::FINALIZE;
+                    bool is_finalize_fusion = isFinalizeFused(gemm2);
                     bool should_be_deterministic
                         = !is_finalize_fusion || mK < 3 || getSMVersion() < 90 || getSMVersion() >= 120;
                     if (should_be_deterministic && !mIsLongTest)
@@ -2214,8 +2238,7 @@ void MixtureOfExpertsTest<TypeParam_>::ParallelismTest(
                 else
                 {
                     runMoEPermute(MOEParallelismConfig{tp_size, i, ep_size, j}, enable_alltoall);
-                    bool is_finalize_fusion = gemm2.epilogue_fusion_type
-                        == tensorrt_llm::cutlass_extensions::CutlassGemmConfig::EpilogueFusionType::FINALIZE;
+                    bool is_finalize_fusion = isFinalizeFused(gemm2);
                     bool should_be_deterministic
                         = !is_finalize_fusion || mK < 3 || getSMVersion() < 90 || getSMVersion() >= 120;
                     if (should_be_deterministic && !mIsLongTest)
@@ -2353,6 +2376,10 @@ void MixtureOfExpertsTest<TypeParam_>::ParallelismTest(
     }                                                                                                                  \
     TYPED_TEST(MixtureOfExpertsTest, ParallelismType##GptOss120b)                                                      \
     {                                                                                                                  \
+        if (!supportsUnpaddedMoeOutput())                                                                              \
+        {                                                                                                              \
+            GTEST_SKIP() << "Unpadded MoE output requires the open-source CUTLASS MoE kernels";                        \
+        }                                                                                                              \
         this->mIsLongTest = true;                                                                                      \
         this->mUseBias = true;                                                                                         \
         this->mActType = ActivationType::Swiglu;                                                                       \
@@ -2659,12 +2686,16 @@ TEST_F(MixtureOfExpertsProfilerTest, TestGeneratedProfilerDistribution)
 #define GET_WS_PTR(type, name) auto* name = reinterpret_cast<type>(workspace + workspaces.at(#name).second)
 
             GET_WS_PTR(int64_t*, expert_first_token_offset);
+#ifdef USING_OSS_CUTLASS_MOE_GEMM
             GET_WS_PTR(int*, unpermuted_row_to_permuted_row);
             GET_WS_PTR(int*, permuted_row_to_unpermuted_row);
-#ifdef USING_OSS_CUTLASS_MOE_GEMM
             GET_WS_PTR(int*, token_selected_experts);
 #else
+            GET_WS_PTR(int*, source_to_dest);
+            GET_WS_PTR(int*, dest_to_source);
             GET_WS_PTR(int*, unpermuted_selected_experts);
+            auto* unpermuted_row_to_permuted_row = source_to_dest;
+            auto* permuted_row_to_unpermuted_row = dest_to_source;
 #endif
 #undef GET_WS_PTR
 
@@ -2742,11 +2773,7 @@ TEST_F(MixtureOfExpertsProfilerTest, TestGeneratedProfilerDistribution)
                         int64_t idx = token_idx * k + k_idx;
                         int64_t expert_idx = host_token_selected_experts[idx];
 
-#ifdef USING_OSS_CUTLASS_MOE_GEMM
                         if (expert_idx < num_experts_per_node)
-#else
-                        if (expert_idx < num_experts)
-#endif
                         {
                             int64_t unpermuted_row = k_idx * num_tokens + token_idx;
                             int64_t permuted_row = host_expert_first_token_offset_size[expert_idx]

tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md

@@ -657,6 +657,40 @@ and materialization enforces the same contract. CUTLASS opts in because its
 SwiGLU adaptor consumes the mode; its SM120 FP8-block-scale Triton fallback is
 rejected separately because that path does not.
 
+With `USING_OSS_CUTLASS_MOE_GEMM=OFF`, the Torch runner uses the bundled
+internal kernel archive's legacy interfaces. That build supports pre-SiLU
+SwiGLU clamping and rejects post-SiLU clamping, SiTu, Relu2, routed-expert LoRA,
+MXFP8 activations or weights, NVFP4 weight-only MoE, FP8-input W4A8,
+dynamic FC2 scaling, per-expert W4A8 activation scales,
+all-to-all, unpadded output, a separate valid-token count, and unswizzled
+input scales. These features require the open-source kernels. The shared
+MoE kernel tests use the public quantization utilities with either runner;
+post-SiLU, unpadded output, and MXFP8 x MXFP4 cases skip for the internal runner.
+The FP8-input W4A8 fixture and all-to-all variants require the OSS build. The Torch
+boundary translates public activation IDs to the internal enum values.
+When SwiGLU parameters are present, it supplies missing alpha/beta defaults
+as tensors so the legacy archive consumes the clamp limit.
+Public routing helpers remain available in either build.
+The build mode is collected once into `MoEEnvironment.oss_cutlass_moe` and
+included in its fingerprint. Libraries without the build-mode query preserve
+the OSS default. CUTLASS eligibility rejects unavailable activation and
+quantization formats, post-SiLU clamping, dynamic FC2 scaling, and routed-expert
+LoRA. W4A8 AWQ uses per-expert activation-scale buffers and is unavailable in
+the legacy build. The backend exposes `supports_alltoall` before communication
+selection. Unsupported auto-selection falls back to AllGather/ReduceScatter, and
+explicit all-to-all requests fail before allocation. Construction validation
+retains an all-to-all backstop. Input layout and token-count options supplied
+to native operators directly retain runtime guards.
+The neutral SwiGLU tensors are cached by device, expert count and CUDA stream;
+warm eager forwards reuse them without allocations or fill kernels. During CUDA
+graph capture, each forward creates its own defaults and records their fills in
+that graph instead of retaining graph-pool tensors in the runner cache.
+
+The CPU resolution tests run in `l0_cpu.yml`; routing and activation numerics
+run in the H100 and B200 single-GPU lists. The H100 C++ kernel stage discovers
+both OSS and bundled-internal MoE variants through CTest. With OSS kernels
+disabled, `mixtureOfExpertsTest` exercises the bundled runner directly.
+
 Declare `limit_when_absent` only when the ABI has no encoding for "no clamp":
 `CuteDslFusedMoE` passes `float("inf")` because its epilogue always applies the
 clamp functor. It is substituted *before* coercion, so a `PER_EXPERT_TENSOR`