diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 25d1d5784c91..fddd8d80a327 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/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) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt index 0ff492623277..20aedcd1e321 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/CMakeLists.txt +++ b/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) # diff --git a/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu b/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu index 657516ed508f..1b4959485862 100644 --- a/cpp/tensorrt_llm/kernels/moe/cutlass/moe_kernels.cu +++ b/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 <<>>(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 diff --git a/cpp/tensorrt_llm/thop/moe/cutlass/moeOp.cpp b/cpp/tensorrt_llm/thop/moe/cutlass/moeOp.cpp index 5986ae85319f..275d12135ecb 100644 --- a/cpp/tensorrt_llm/thop/moe/cutlass/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moe/cutlass/moeOp.cpp @@ -63,9 +63,11 @@ namespace torch_ext namespace common = tensorrt_llm::common; namespace kernels = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE; -using ActivationParams = CUTLASS_MOE_GEMM_NAMESPACE::ActivationParams; +using ActivationParams = kernels::ActivationParams; using ActivationType = CUTLASS_MOE_GEMM_NAMESPACE::ActivationType; +#if defined(USING_OSS_CUTLASS_MOE_GEMM) using MoeGemmId = CUTLASS_MOE_GEMM_NAMESPACE::MoeGemmId; +#endif // Always use public header as it is just utility functions and types using TmaWarpSpecializedGroupedGemmInput = tensorrt_llm::kernels::cutlass_kernels::TmaWarpSpecializedGroupedGemmInput; using profiler_backend = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::GemmProfilerBackend; @@ -148,9 +150,67 @@ inline void moeLoraGroupedGemmRunImpl(::tensorrt_llm::kernels::cutlass_kernels:: sync_check_cuda_error(stream); } -static ActivationParams makeActivationParams(ActivationType activation_type, +namespace +{ +#if !defined(USING_OSS_CUTLASS_MOE_GEMM) +void validateInternalMoeInputs(int64_t num_rows, int64_t hidden_size, int64_t unpadded_hidden_size, + torch::optional const& num_valid_tokens, bool enable_alltoall, bool has_input_sf, bool swizzled_input_sf, + bool use_dynamic_fc2_scale) +{ + TORCH_CHECK( + unpadded_hidden_size == hidden_size, "Unpadded MoE output requires the open-source CUTLASS MoE kernels."); + TORCH_CHECK(!num_valid_tokens.has_value() || num_valid_tokens.value() == num_rows, + "A separate valid-token count requires the open-source CUTLASS MoE kernels."); + TORCH_CHECK(!enable_alltoall, "MoE all-to-all requires the open-source CUTLASS MoE kernels."); + TORCH_CHECK( + !has_input_sf || swizzled_input_sf, "Unswizzled MoE input scales require the open-source CUTLASS MoE kernels."); + TORCH_CHECK(!use_dynamic_fc2_scale, "Dynamic FC2 scaling requires the open-source CUTLASS MoE kernels."); +} +#endif + +ActivationType getActivationType(int64_t activation_type) +{ +#if defined(USING_OSS_CUTLASS_MOE_GEMM) + return static_cast(activation_type); +#else + using PublicActivationType = tensorrt_llm::kernels::cutlass_kernels::ActivationType; + switch (static_cast(activation_type)) + { + case PublicActivationType::InvalidType: + TORCH_CHECK(false, "Invalid MoE activation type."); + return ActivationType::InvalidType; + case PublicActivationType::Identity: return ActivationType::Identity; + case PublicActivationType::Gelu: return ActivationType::Gelu; + case PublicActivationType::Relu: return ActivationType::Relu; + case PublicActivationType::Silu: return ActivationType::Silu; + case PublicActivationType::Swiglu: return ActivationType::Swiglu; + case PublicActivationType::Geglu: return ActivationType::Geglu; + case PublicActivationType::SwigluBias: return ActivationType::SwigluBias; + default: TORCH_CHECK(false, "Unsupported activation type for internal CUTLASS MoE kernels."); + } + return ActivationType::InvalidType; +#endif +} + +bool isSiTuActivation(ActivationType activation_type) +{ +#if defined(USING_OSS_CUTLASS_MOE_GEMM) + return activation_type == ActivationType::SiTu; +#else + return false; +#endif +} + +struct ActivationParamBuffers +{ + std::map, std::pair> defaults; +}; + +ActivationParams makeActivationParams(ActivationType activation_type, torch::optional const& swiglu_alpha, torch::optional const& swiglu_beta, - torch::optional const& swiglu_limit, bool const swiglu_clamp_after_silu) + torch::optional const& swiglu_limit, bool const swiglu_clamp_after_silu, + [[maybe_unused]] ActivationParamBuffers& buffers, + [[maybe_unused]] std::pair& capture_buffers) { auto const* swiglu_alpha_ptr = reinterpret_cast(swiglu_alpha.has_value() ? swiglu_alpha.value().const_data_ptr() : nullptr); @@ -166,9 +226,43 @@ static ActivationParams makeActivationParams(ActivationType activation_type, // not implement this clamp order. Reject it instead of silently ignoring it. TORCH_CHECK( !swiglu_clamp_after_silu, "Post-SiLU SwiGLU clamping requires TensorRT-LLM's open-source CUTLASS MoE kernels."); + if ((swiglu_alpha_ptr || swiglu_beta_ptr || swiglu_limit_ptr) && (!swiglu_alpha_ptr || !swiglu_beta_ptr)) + { + auto const& reference = swiglu_alpha.has_value() ? swiglu_alpha.value() + : swiglu_beta.has_value() ? swiglu_beta.value() + : swiglu_limit.value(); + // The legacy adaptor ignores the limit unless alpha and beta are both + // present. Keep neutral tensors per device, expert count and stream so + // repeated forwards do not allocate or refill them, or cross streams. + auto const stream = at::cuda::getCurrentCUDAStream(reference.device().index()).stream(); + cudaStreamCaptureStatus capture_status; + TLLM_CUDA_CHECK(cudaStreamIsCapturing(stream, &capture_status)); + // Captured fills execute only on replay. Keep these tensors local to + // this forward so no later call can reuse an uninitialized graph buffer. + auto& defaults = capture_status != cudaStreamCaptureStatusNone + ? capture_buffers + : buffers.defaults[{reference.device().index(), reference.numel(), stream}]; + if (!swiglu_alpha_ptr) + { + if (!defaults.first.defined()) + { + defaults.first = torch::ones_like(reference); + } + swiglu_alpha_ptr = static_cast(defaults.first.const_data_ptr()); + } + if (!swiglu_beta_ptr) + { + if (!defaults.second.defined()) + { + defaults.second = torch::zeros_like(reference); + } + swiglu_beta_ptr = static_cast(defaults.second.const_data_ptr()); + } + } return ActivationParams(activation_type, swiglu_alpha_ptr, swiglu_beta_ptr, swiglu_limit_ptr); #endif } +} // namespace class FusedMoeRunner : public EagerMoeWorkspaceOwner { @@ -250,6 +344,10 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner mUseMxfp8WeightScaling = use_mxfp8_weight_scaling; mUseFusedFinalize = use_fused_finalize; mInnerDimMultiplier = 1; +#if !defined(USING_OSS_CUTLASS_MOE_GEMM) + TORCH_CHECK(!mUseMxfp8ActScaling, "MXFP8 activations require the open-source CUTLASS MoE kernels."); + TORCH_CHECK(!mUseMxfp8WeightScaling, "MXFP8 x MXFP8 requires the open-source CUTLASS MoE kernels."); +#endif // MXFP8xMXFP8 grouped MoE is only meaningful for the // template instantiation. Reject other (act, weight) dtype pairs at @@ -308,6 +406,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner } if (isWFP4A16Quant()) { +#if defined(USING_OSS_CUTLASS_MOE_GEMM) mInnerDimMultiplier = 2; if (mActivationDtype == c10::ScalarType::Half) { @@ -318,6 +417,9 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner { mKernelRunner = std::make_shared>(); } +#endif +#else + TORCH_CHECK(false, "NVFP4 weight-only MoE requires the open-source CUTLASS MoE kernels."); #endif } #endif @@ -334,8 +436,12 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner { if (isInt4Quant() and mUseW4GroupScaling) { +#if defined(USING_OSS_CUTLASS_MOE_GEMM) mKernelRunner = std::make_unique< kernels::CutlassMoeFCRunner<__nv_fp8_e4m3, cutlass::uint4b_t, __nv_bfloat16, __nv_fp8_e4m3>>(); +#else + TORCH_CHECK(false, "FP8-input W4A8 MoE requires the open-source CUTLASS MoE kernels."); +#endif } else { @@ -358,11 +464,18 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner } mKernelRunner->use_fused_finalize_ = mUseFusedFinalize; +#if defined(USING_OSS_CUTLASS_MOE_GEMM) mKernelRunner->use_mxfp8_weight_scaling_ = mUseMxfp8WeightScaling; +#endif mProfiler = std::make_shared(); +#if defined(USING_OSS_CUTLASS_MOE_GEMM) mGemm1Profiles = mKernelRunner->getTactics(MoeGemmId::GEMM_1); mGemm2Profiles = mKernelRunner->getTactics(MoeGemmId::GEMM_2); +#else + mGemm1Profiles = mKernelRunner->getTactics(); + mGemm2Profiles = mGemm1Profiles; +#endif replaceUnsupportedProfiles(mGemm1Profiles); replaceUnsupportedProfiles(mGemm2Profiles); cuInit(0); @@ -514,9 +627,8 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner TORCH_CHECK(fc1_expert_weights.sizes()[0] == fc2_expert_weights.sizes()[0], "fc1_expert_weights and fc2_expert_weights must have the same number of experts."); - ActivationType base_activation_type = activation_type.has_value() - ? static_cast(activation_type.value()) - : ActivationType::Swiglu; + ActivationType base_activation_type + = activation_type.has_value() ? getActivationType(activation_type.value()) : ActivationType::Swiglu; if (mUseINT8WoqPerChannel) { // Note: The weight shape for INT8 weight only quantization is different, e.g., fc2_expert_weights: @@ -552,8 +664,6 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner int experts_per_token = token_selected_experts.sizes()[1]; int64_t num_rows = input.sizes()[0]; int64_t hidden_size = fc2_expert_weights.sizes()[1]; - int64_t unpadded_hidden_size_val - = unpadded_hidden_size.has_value() ? unpadded_hidden_size.value() : hidden_size; int64_t inter_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; if (mUseINT8WoqPerChannel) { @@ -563,6 +673,9 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner inter_size = fc2_expert_weights.sizes()[1]; } + int64_t unpadded_hidden_size_val + = unpadded_hidden_size.has_value() ? unpadded_hidden_size.value() : hidden_size; + if (isWMxfp4AMxfp8Quant() || isWMxfp4AFp8Quant()) { // MXFP4 weights are required to bealigned to 128 bytes @@ -581,6 +694,10 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner } int const num_experts_on_rank = fc2_expert_weights.sizes()[0]; +#if !defined(USING_OSS_CUTLASS_MOE_GEMM) + validateInternalMoeInputs(num_rows, hidden_size, unpadded_hidden_size_val, num_valid_tokens, enable_alltoall, + input_sf.has_value(), swizzled_input_sf, use_dynamic_fc2_scale); +#endif auto const num_experts_total = static_cast(num_experts_on_rank * ep_size); auto parallelism_config = kernels::MOEParallelismConfig(tp_size, tp_rank, ep_size, ep_rank); @@ -589,7 +706,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner CHECK_INPUT(swiglu_alpha.value(), at::ScalarType::Float); TORCH_CHECK(swiglu_alpha.value().sizes()[0] == num_experts_on_rank, "swiglu_alpha must have num_experts_on_rank elements."); - if (base_activation_type != ActivationType::SiTu) + if (!isSiTuActivation(base_activation_type)) { base_activation_type = ActivationType::SwigluBias; } @@ -599,7 +716,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner CHECK_INPUT(swiglu_beta.value(), at::ScalarType::Float); TORCH_CHECK(swiglu_beta.value().sizes()[0] == num_experts_on_rank, "swiglu_beta must have num_experts_on_rank elements."); - if (base_activation_type != ActivationType::SiTu) + if (!isSiTuActivation(base_activation_type)) { base_activation_type = ActivationType::SwigluBias; } @@ -609,20 +726,20 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner CHECK_INPUT(swiglu_limit.value(), at::ScalarType::Float); TORCH_CHECK(swiglu_limit.value().sizes()[0] == num_experts_on_rank, "swiglu_limit must have num_experts_on_rank elements."); - if (base_activation_type != ActivationType::SiTu) + if (!isSiTuActivation(base_activation_type)) { base_activation_type = ActivationType::SwigluBias; } } - TORCH_CHECK( - base_activation_type != ActivationType::SiTu || (swiglu_alpha.has_value() && swiglu_beta.has_value()), + TORCH_CHECK(!isSiTuActivation(base_activation_type) || (swiglu_alpha.has_value() && swiglu_beta.has_value()), "SiTu requires both swiglu_alpha and swiglu_beta."); - TORCH_CHECK(base_activation_type != ActivationType::SiTu || !swiglu_limit.has_value(), + TORCH_CHECK(!isSiTuActivation(base_activation_type) || !swiglu_limit.has_value(), "SiTu does not support swiglu_limit."); // A SwiGLU clamp promotes base_activation_type to SwigluBias above, // which selects the only CUTLASS adaptor that consumes clampAfterSilu. - auto activation_params = makeActivationParams( - base_activation_type, swiglu_alpha, swiglu_beta, swiglu_limit, swiglu_clamp_after_silu); + std::pair capture_buffers; + auto activation_params = makeActivationParams(base_activation_type, swiglu_alpha, swiglu_beta, swiglu_limit, + swiglu_clamp_after_silu, mActivationParamBuffers, capture_buffers); // ===== Routed-expert LoRA activation flags ===== // LoRA is activated by the per-request (fc1_lora_ranks) or slot-indexed @@ -708,6 +825,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner = getQuantParams(num_experts_on_rank, hidden_size, inter_size, quant_scales, base_activation_type); // Dynamic fc2 scale: allocate workspace buffers +#if defined(USING_OSS_CUTLASS_MOE_GEMM) at::Tensor dynamic_fc2_amax_tensor, dynamic_fc2_alpha_tensor, dynamic_fc2_bf16_tensor; if (use_dynamic_fc2_scale && isNvfp4Quant() && quant_scales.has_value() && quant_scales.value().size() >= 7) { @@ -727,6 +845,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner = static_cast(quant_scales.value()[6].data_ptr()); } +#endif kernels::MoeMinLatencyParams min_latency_params{}; // Use the populated LoraParams when LoRA is active, otherwise a default-constructed empty one. @@ -749,7 +868,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner lora_params, mUseDeepSeekFP8BlockScaling, min_latency_mode, min_latency_params, stream); #else mKernelRunner->runMoe(input.const_data_ptr(), - input_sf.has_value() ? input_sf.value().const_data_ptr() : nullptr, swizzled_input_sf, + input_sf.has_value() ? input_sf.value().const_data_ptr() : nullptr, reinterpret_cast(token_selected_experts.const_data_ptr()), token_final_scales.has_value() ? reinterpret_cast(token_final_scales.value().const_data_ptr()) : nullptr, @@ -757,8 +876,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner fc1_expert_biases.has_value() ? fc1_expert_biases.value().const_data_ptr() : nullptr, activation_params, fc2_expert_weights.const_data_ptr(), fc2_expert_biases.has_value() ? fc2_expert_biases.value().const_data_ptr() : nullptr, quant_params, - num_rows, num_valid_tokens.has_value() ? num_valid_tokens.value() : num_rows, hidden_size, inter_size, - num_experts_total, static_cast(experts_per_token), + num_rows, hidden_size, inter_size, num_experts_total, static_cast(experts_per_token), static_cast(workspace_info.scratch.tensor().data_ptr()), output.data_ptr(), static_cast(workspace_info.src_to_dest_map), parallelism_config, lora_active, lora_params, mUseDeepSeekFP8BlockScaling, min_latency_mode, min_latency_params, stream); @@ -837,8 +955,6 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner int experts_per_token = token_selected_experts.sizes()[1]; int64_t num_rows = input.sizes()[0]; int64_t hidden_size = fc2_expert_weights.sizes()[1]; - int64_t unpadded_hidden_size_val - = unpadded_hidden_size.has_value() ? unpadded_hidden_size.value() : hidden_size; int64_t inter_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; if (mUseINT8WoqPerChannel) { @@ -847,19 +963,24 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner hidden_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; inter_size = fc2_expert_weights.sizes()[1]; } + int64_t unpadded_hidden_size_val + = unpadded_hidden_size.has_value() ? unpadded_hidden_size.value() : hidden_size; int const num_experts_on_rank = fc2_expert_weights.sizes()[0]; +#if !defined(USING_OSS_CUTLASS_MOE_GEMM) + validateInternalMoeInputs(num_rows, hidden_size, unpadded_hidden_size_val, num_valid_tokens, enable_alltoall, + input_sf.has_value(), swizzled_input_sf, /*use_dynamic_fc2_scale=*/false); +#endif auto const num_experts_total = static_cast(num_experts_on_rank * ep_size); auto parallelism_config = kernels::MOEParallelismConfig(tp_size, tp_rank, ep_size, ep_rank, cluster_size, cluster_rank); - ActivationType base_activation_type = activation_type.has_value() - ? static_cast(activation_type.value()) - : ActivationType::Swiglu; + ActivationType base_activation_type + = activation_type.has_value() ? getActivationType(activation_type.value()) : ActivationType::Swiglu; if (swiglu_alpha.has_value()) { CHECK_INPUT(swiglu_alpha.value(), at::ScalarType::Float); TORCH_CHECK(swiglu_alpha.value().sizes()[0] == num_experts_on_rank, "swiglu_alpha must have num_experts_on_rank elements."); - if (base_activation_type != ActivationType::SiTu) + if (!isSiTuActivation(base_activation_type)) { base_activation_type = ActivationType::SwigluBias; } @@ -869,7 +990,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner CHECK_INPUT(swiglu_beta.value(), at::ScalarType::Float); TORCH_CHECK(swiglu_beta.value().sizes()[0] == num_experts_on_rank, "swiglu_beta must have num_experts_on_rank elements."); - if (base_activation_type != ActivationType::SiTu) + if (!isSiTuActivation(base_activation_type)) { base_activation_type = ActivationType::SwigluBias; } @@ -879,18 +1000,18 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner CHECK_INPUT(swiglu_limit.value(), at::ScalarType::Float); TORCH_CHECK(swiglu_limit.value().sizes()[0] == num_experts_on_rank, "swiglu_limit must have num_experts_on_rank elements."); - if (base_activation_type != ActivationType::SiTu) + if (!isSiTuActivation(base_activation_type)) { base_activation_type = ActivationType::SwigluBias; } } - TORCH_CHECK( - base_activation_type != ActivationType::SiTu || (swiglu_alpha.has_value() && swiglu_beta.has_value()), + TORCH_CHECK(!isSiTuActivation(base_activation_type) || (swiglu_alpha.has_value() && swiglu_beta.has_value()), "SiTu requires both swiglu_alpha and swiglu_beta."); - TORCH_CHECK(base_activation_type != ActivationType::SiTu || !swiglu_limit.has_value(), + TORCH_CHECK(!isSiTuActivation(base_activation_type) || !swiglu_limit.has_value(), "SiTu does not support swiglu_limit."); - auto activation_params = makeActivationParams( - base_activation_type, swiglu_alpha, swiglu_beta, swiglu_limit, swiglu_clamp_after_silu); + std::pair capture_buffers; + auto activation_params = makeActivationParams(base_activation_type, swiglu_alpha, swiglu_beta, swiglu_limit, + swiglu_clamp_after_silu, mActivationParamBuffers, capture_buffers); // Validate the fc1/fc2 inter-size relationship now that the activation type (gated vs // non-gated) is finalized. INT8-woq uses a transposed weight layout, so its fc1/fc2 dim @@ -978,7 +1099,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner mUseDeepSeekFP8BlockScaling, min_latency_mode, min_latency_params, stream); #else mKernelRunner->runMoe(input.const_data_ptr(), - input_sf.has_value() ? input_sf.value().const_data_ptr() : nullptr, swizzled_input_sf, + input_sf.has_value() ? input_sf.value().const_data_ptr() : nullptr, reinterpret_cast(token_selected_experts.const_data_ptr()), token_final_scales.has_value() ? reinterpret_cast(token_final_scales.value().const_data_ptr()) : nullptr, @@ -986,8 +1107,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner fc1_expert_biases.has_value() ? fc1_expert_biases.value().const_data_ptr() : nullptr, activation_params, fc2_expert_weights.const_data_ptr(), fc2_expert_biases.has_value() ? fc2_expert_biases.value().const_data_ptr() : nullptr, quant_params, - num_rows, num_valid_tokens.has_value() ? num_valid_tokens.value() : num_rows, hidden_size, inter_size, - num_experts_total, static_cast(experts_per_token), + num_rows, hidden_size, inter_size, num_experts_total, static_cast(experts_per_token), static_cast(workspace_info.scratch.tensor().data_ptr()), output.data_ptr(), static_cast(workspace_info.src_to_dest_map), parallelism_config, false, lora_params, mUseDeepSeekFP8BlockScaling, min_latency_mode, min_latency_params, stream); @@ -1019,7 +1139,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner { return; } - ActivationType activation_type = static_cast(activation_type_int); + ActivationType activation_type = getActivationType(activation_type_int); int64_t const num_rows = input.sizes()[0]; int64_t hidden_size = fc2_expert_weights.sizes()[1]; @@ -1031,6 +1151,11 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner hidden_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; inter_size = fc2_expert_weights.sizes()[1]; } +#if !defined(USING_OSS_CUTLASS_MOE_GEMM) + validateInternalMoeInputs(num_rows, hidden_size, unpadded_hidden_size > 0 ? unpadded_hidden_size : hidden_size, + std::nullopt, enable_alltoall, + /*has_input_sf=*/false, /*swizzled_input_sf=*/true, /*use_dynamic_fc2_scale=*/false); +#endif int64_t const group_size_ = isInt4Quant() ? TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::int4_group_size : -1; int64_t const group_size = isWFP4A16Quant() @@ -1115,6 +1240,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner int64_t mInnerDimMultiplier; char* mProfileWorkspace = nullptr; std::map, WorkspaceInfo> mStreamWorkspaces; + ActivationParamBuffers mActivationParamBuffers; bool mUseDeepSeekFP8BlockScaling = false; bool mUseW4GroupScaling = false; @@ -1755,6 +1881,7 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner int64_t inter_size, c10::ScalarType act_dtype, int64_t lora_max_low_rank, bool is_gated_activation, cudaStream_t stream, int experts_per_token) { +#if defined(USING_OSS_CUTLASS_MOE_GEMM) bool const has_per_request = fc1_lora_ranks.has_value(); bool const has_slot_indexed = fc1_slot_lora_ranks.has_value(); if (!has_per_request && !has_slot_indexed) @@ -2131,6 +2258,11 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner } return lora_params; +#else + TORCH_CHECK(!fc1_lora_ranks.has_value() && !fc1_slot_lora_ranks.has_value(), + "Routed-expert LoRA requires the open-source CUTLASS MoE kernels."); + return std::nullopt; +#endif } kernels::QuantParams getQuantParams(int64_t const num_experts_on_rank, int64_t const hidden_size, @@ -2281,11 +2413,12 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner static_cast(fc2_weight_block.data_ptr()), static_cast(fc2_global.data_ptr())); #else - TORCH_CHECK(false, "MXFP8 x MXFP4 quantization is not supported in OSS Cutlass Moe Gemm"); + TORCH_CHECK(false, "MXFP8 x MXFP4 requires the open-source CUTLASS MoE kernels."); #endif } else if (isWMxfp8AMxfp8Quant()) { +#if defined(USING_OSS_CUTLASS_MOE_GEMM) // with MXFP8 1x32 UE8M0 block scales on both sides; // SF storage is int32-packed UE8M0 (same convention as MXFP4 MoE). TORCH_CHECK(quant_scales.has_value(), "Expecting quant scales for MXFP8 x MXFP8 quantization"); @@ -2303,6 +2436,9 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner return kernels::QuantParams::MXFP8MXFP8( static_cast(fc1_weight_block.data_ptr()), static_cast(fc2_weight_block.data_ptr())); +#else + TORCH_CHECK(false, "MXFP8 x MXFP8 requires the open-source CUTLASS MoE kernels."); +#endif } else if (isNvfp4Quant()) { @@ -2417,7 +2553,11 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner // Whether it is per-expert activation scale bool fc1_use_per_expert_act_scale = fc1_act_scales.numel() > hidden_size; bool fc2_use_per_expert_act_scale = fc2_act_scales.numel() > inter_size; - return kernels::QuantParams::GroupWise(group_size, +#if !defined(USING_OSS_CUTLASS_MOE_GEMM) + TORCH_CHECK(!fc1_use_per_expert_act_scale && !fc2_use_per_expert_act_scale, + "Per-expert W4A8 activation scales require the open-source CUTLASS MoE kernels."); +#endif + auto quant_params = kernels::QuantParams::GroupWise(group_size, static_cast(fc1_weight_scales.data_ptr()), static_cast(fc2_weight_scales.data_ptr()), static_cast(fc1_act_scales.numel() > 0 ? fc1_act_scales.data_ptr() : nullptr), @@ -2425,8 +2565,12 @@ class FusedMoeRunner : public EagerMoeWorkspaceOwner static_cast(fc1_weight_zeros.numel() > 0 ? fc1_weight_zeros.data_ptr() : nullptr), static_cast(fc2_weight_zeros.numel() > 0 ? fc2_weight_zeros.data_ptr() : nullptr), static_cast(fc1_alpha.numel() > 0 ? fc1_alpha.data_ptr() : nullptr), - static_cast(fc2_alpha.numel() > 0 ? fc2_alpha.data_ptr() : nullptr), - fc1_use_per_expert_act_scale, fc2_use_per_expert_act_scale); + static_cast(fc2_alpha.numel() > 0 ? fc2_alpha.data_ptr() : nullptr)); +#if defined(USING_OSS_CUTLASS_MOE_GEMM) + quant_params.groupwise.fc1.use_per_expert_act_scale = fc1_use_per_expert_act_scale; + quant_params.groupwise.fc2.use_per_expert_act_scale = fc2_use_per_expert_act_scale; +#endif + return quant_params; } else { @@ -2507,6 +2651,15 @@ TORCH_LIBRARY(trtllm, m) .def(torch::init()) .def("begin", &tensorrt_llm::torch_ext::EagerMoeWorkspaceScope::begin) .def("finish", &tensorrt_llm::torch_ext::EagerMoeWorkspaceScope::finish); + m.def("is_oss_cutlass_moe() -> bool", + []() + { +#if defined(USING_OSS_CUTLASS_MOE_GEMM) + return true; +#else + return false; +#endif + }); m.class_("FusedMoeRunner") .def(torch::init()) .def("run_gemm_profile", &tensorrt_llm::torch_ext::FusedMoeRunner::runGemmProfile) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 008c43846618..5ba58455a7a8 100644 --- a/cpp/tests/CMakeLists.txt +++ b/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. diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index 95b7bdc7d596..c0c393a2634a 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/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) diff --git a/cpp/tests/unit_tests/kernels/moe/mixtureOfExpertsTest.cu b/cpp/tests/unit_tests/kernels/moe/mixtureOfExpertsTest.cu index 6fb253cd365e..86b1f1cfa6b0 100644 --- a/cpp/tests/unit_tests/kernels/moe/mixtureOfExpertsTest.cu +++ b/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 @@ -25,12 +26,11 @@ #ifdef USING_OSS_CUTLASS_MOE_GEMM #include "tensorrt_llm/kernels/moe/cutlass/include/moe_kernels.h" -#include #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(num_experts, out_shape, in_shape, raw_weights, global_scales, reinterpret_cast(quant_weights), reinterpret_cast(scaling_factors), MX_QUANT_WEIGHT, tensorrt_llm::QuantizationSFLayout::SWIZZLED, mMultiProcessorCount, mStream->get()); -#else - invokeBatchedFP4Quantization(num_experts, out_shape, in_shape, raw_weights, - global_scales, reinterpret_cast(quant_weights), reinterpret_cast(scaling_factors), - MX_QUANT_WEIGHT, mMultiProcessorCount, mStream->get()); -#endif // auto sf_data = getDataFromDevice(scaling_factors, num_experts * padded_out_dim * padded_in_dim / // FP4VecSize); auto unquant_data = getDataFromDevice(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>; + 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(scale1_ptr), mFP4ScalingFactorsW2, static_cast(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, #endif +#endif #endif WeightParams, WeightParams @@ -1867,8 +1890,7 @@ void MixtureOfExpertsTest::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::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::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::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(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] diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index 68280e465725..40b139d83258 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/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` diff --git a/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py index bc757c5fb622..247e208f470b 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py @@ -95,6 +95,7 @@ def create_strategy( use_flashinfer: bool = False, hidden_size: Optional[int] = None, communication_method: Optional[str] = None, + allow_alltoall: bool = True, ) -> Optional[Communication]: """ Create the best communication method for the given configuration @@ -122,6 +123,7 @@ def create_strategy( Falls back to pretrained_config.hidden_size when not provided. communication_method: Optional model-selected communication method. ``TRTLLM_FORCE_COMM_METHOD`` takes precedence when set. + allow_alltoall: Whether the backend supports all-to-all dispatch. # TODO: Need a way to indicate whether EPLB is enabled. Returns: @@ -173,6 +175,17 @@ def create_strategy( else: force_method, force_source = communication_method, "communication_method" + if not allow_alltoall: + if force_method is not None and force_method.upper() != "ALLGATHER": + raise ValueError( + f"{force_source}={force_method} requires all-to-all, which the " + "selected MoE backend does not support; use ALLGATHER" + ) + logger.info( + "Selected communication strategy: AllGatherReduceScatter (backend capability)" + ) + return AllGatherReduceScatter(mapping) + if force_method is not None: strategy = CommunicationFactory._create_forced_method( force_method, diff --git a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py index 0909c8e7f528..d7b1dccffc98 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py @@ -250,8 +250,14 @@ def __init__( self.event_dict = None # Validate configuration - self.validate_config() - self.validate_backend(self.backend) + validated = False + try: + self.validate_config() + self.validate_backend(self.backend) + validated = True + finally: + if not validated and self.comm is not None: + self.comm.destroy() # Mark as _weights_removed to skip ConfigurableMoE's post_load_weights in model_loader # The backend's post_load_weights will be called directly by model_loader @@ -640,6 +646,7 @@ def _create_comm_strategy_auto(self) -> Optional[Communication]: use_flashinfer=self.use_flashinfer, hidden_size=self.hidden_size, communication_method=self.communication_method, + allow_alltoall=self.backend.supports_alltoall, ) def forward_impl( diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py index 09818bc08f62..95a7dd10f5a3 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py @@ -212,6 +212,11 @@ def supports_moe_output_in_alltoall_workspace(self) -> bool: # that might shrink prefill chunk size. _PREFILL_VIA_CUTLASS_THRESHOLD = 64 + @classmethod + def _prefill_uses_cutlass(cls, quant_algo: Optional[QuantAlgo]) -> bool: + """NVFP4 prefill uses CUTLASS; W4A16_NVFP4 stays on the B12x kernel.""" + return quant_algo == QuantAlgo.NVFP4 + @classmethod def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: sm_version = d.env.sm @@ -245,6 +250,18 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: f"CuteDslB12xFusedMoE does not support activation " f"{p.activation}; supported: {supported}", ) + if d.moe_lora_enabled: + return _reject( + MoERejectReason.LORA_UNSUPPORTED, + f"{cls.__name__} does not fuse routed-expert LoRA", + ) + if cls._prefill_uses_cutlass(p.quant_algo): + eligibility = cls._check_build_capabilities(p, d) + if not eligibility.eligible: + return replace( + eligibility, + detail="NVFP4 prefill routes through CUTLASS: " + eligibility.detail, + ) # The decode kernel ships in the FlashInfer wheel. if not d.env.has_dep(MoEDep.FLASHINFER): return _reject( @@ -322,7 +339,8 @@ def _route_to_cutlass(self, x) -> bool: always stay on the b12x path (which rejects them) so the existing error message is preserved.""" quant_config = getattr(self, "quant_config", None) - if quant_config is not None and quant_config.quant_algo == QuantAlgo.W4A16_NVFP4: + quant_algo = quant_config.quant_algo if quant_config is not None else None + if not self._prefill_uses_cutlass(quant_algo): return False return isinstance(x, torch.Tensor) and x.shape[0] >= self._PREFILL_VIA_CUTLASS_THRESHOLD diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py index 24d3ef6adaf7..43cc9bd1faf5 100755 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py @@ -35,6 +35,7 @@ MoEStaticCapability, nvfp4_fc1_row_alignment_rejection, require_comm_plan) +from .impl_environment import collect_moe_environment from .interface import _reject from .quantization import UnquantizedFusedMoEMethod @@ -204,6 +205,41 @@ class CutlassFusedMoE(MoEImplBase): NVFP4 weight pad only accepts 2-D tensors. """ + @classmethod + def _check_build_capabilities(cls, p: MoEProblem, + d: MoEDeployment) -> MoEEligibility: + """Check capabilities required by callers of the CUTLASS execution path.""" + quant_algo = p.quant_algo + if not d.env.oss_cutlass_moe: + if p.activation_type in (ActivationType.SiTu, ActivationType.Relu2): + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"{cls.__name__} {p.activation} requires a build with " + "USING_OSS_CUTLASS_MOE_GEMM=ON") + if quant_algo in (QuantAlgo.MXFP8, QuantAlgo.W4A8_MXFP4_MXFP8, + QuantAlgo.W4A16_MXFP4, QuantAlgo.W4A8_AWQ): + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, + f"{cls.__name__} {quant_algo.name} requires a build with " + "USING_OSS_CUTLASS_MOE_GEMM=ON") + if quant_algo == QuantAlgo.NVFP4 and d.force_dynamic_quantization: + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, + f"{cls.__name__} dynamic FC2 scaling requires a build with " + "USING_OSS_CUTLASS_MOE_GEMM=ON") + if p.clamp_after_silu: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"{cls.__name__} post-SiLU clamping requires a build with " + "USING_OSS_CUTLASS_MOE_GEMM=ON") + if d.moe_lora_enabled: + return _reject( + MoERejectReason.LORA_UNSUPPORTED, + f"{cls.__name__} routed-expert LoRA requires a build with " + "USING_OSS_CUTLASS_MOE_GEMM=ON") + + return MoEEligibility.ok() + @classmethod def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: """Cutlass grouped-GEMM MoE: the widest quant and SM coverage there is. @@ -214,6 +250,10 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: sm_version = d.env.sm quant_algo = p.quant_algo + eligibility = cls._check_build_capabilities(p, d) + if not eligibility.eligible: + return eligibility + # Check minimum SM version for Cutlass backend if sm_version < 80: return _reject( @@ -306,6 +346,15 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: return MoEEligibility.ok() + @property + def supports_alltoall(self) -> bool: + return collect_moe_environment().oss_cutlass_moe + + def validate_configurable_moe(self, moe: torch.nn.Module) -> None: + if moe.enable_alltoall and not self.supports_alltoall: + raise ValueError("CutlassFusedMoE all-to-all requires a build with " + "USING_OSS_CUTLASS_MOE_GEMM=ON") + def __init__( self, *, diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_blocks.py b/tensorrt_llm/_torch/moe/fused_moe/impl_blocks.py index d6d9bd243547..68c2127cd7b2 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_blocks.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_blocks.py @@ -84,6 +84,11 @@ def supports_unfinalized_output(self) -> bool: """ return False + @property + def supports_alltoall(self) -> bool: + """Whether the kernel can consume tokens dispatched with all-to-all.""" + return True + @property def uses_locality_domain(self) -> bool: """True when this impl runs partitioned across locality domains.""" diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py index 683302725064..09595d9aad07 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py @@ -241,6 +241,7 @@ class MoEEnvironment: sm: int # e.g. 100, 103 available_deps: Tuple[str, ...] = () env_flags: Tuple[Tuple[str, str], ...] = () # sorted (name, value) pairs + oss_cutlass_moe: bool = True def has_dep(self, name: str) -> bool: return name in self.available_deps @@ -259,7 +260,7 @@ def env_flag(self, name: str) -> Optional[str]: def fingerprint(self) -> str: """Return a stable fingerprint for the selection environment.""" - payload = repr((self.sm, sorted(self.available_deps), self.env_flags)) + payload = repr((self.sm, sorted(self.available_deps), self.env_flags, self.oss_cutlass_moe)) return hashlib.sha256(payload.encode()).hexdigest()[:16] @@ -286,6 +287,7 @@ class MoEDeployment: # ``model_config.locality_domain_policy.enabled``. Whether the machine can # actually serve it is ``env.has_dep(MoEDep.LOCALITY_DOMAIN)``. locality_domain_requested: bool = False + force_dynamic_quantization: bool = False @property def smart_router(self) -> bool: @@ -518,6 +520,8 @@ def to_dict(self) -> Dict[str, object]: "num_slots": self.deployment.num_slots, "eplb_enabled": self.deployment.eplb_enabled, "moe_lora_enabled": self.deployment.moe_lora_enabled, + "force_dynamic_quantization": self.deployment.force_dynamic_quantization, + "oss_cutlass_moe": self.deployment.env.oss_cutlass_moe, "sm": self.deployment.env.sm, "env_flags": dict(self.deployment.env.env_flags), }, diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py b/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py index d9640a60d996..c87c4107289a 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py @@ -17,6 +17,7 @@ import os from contextlib import contextmanager from enum import Enum +from functools import lru_cache from typing import Callable, Dict, Optional, Tuple from tensorrt_llm.logger import logger @@ -162,6 +163,19 @@ def _run_probe(dep: MoEDep, probe: DepProbe) -> bool: return available +@lru_cache(maxsize=1) +def _is_oss_cutlass_moe() -> bool: + """Query the native build mode, preserving the default for older libraries.""" + import torch + + try: + query = torch.ops.trtllm.is_oss_cutlass_moe + except AttributeError: + logger.debug("CUTLASS MoE build-mode query is unavailable; assuming OSS kernels") + return True + return bool(query()) + + def collect_moe_environment(force: bool = False) -> MoEEnvironment: """Collect and cache the frozen MoE selection environment.""" global _CACHED_ENVIRONMENT @@ -181,7 +195,14 @@ def collect_moe_environment(force: bool = False) -> MoEEnvironment: for flag, default in _ENV_FLAG_DEFAULTS.items() ) ) - environment = MoEEnvironment(sm=get_sm_version(), available_deps=available, env_flags=env_flags) + if force: + _is_oss_cutlass_moe.cache_clear() + environment = MoEEnvironment( + sm=get_sm_version(), + available_deps=available, + env_flags=env_flags, + oss_cutlass_moe=_is_oss_cutlass_moe(), + ) logger.debug( f"collected MoE environment: sm={environment.sm} deps={available} " f"flags={env_flags} ({environment.fingerprint()})" @@ -194,6 +215,7 @@ def reset_moe_environment_cache() -> None: """Drop the cached probe result. For tests that change probe outcomes.""" global _CACHED_ENVIRONMENT _CACHED_ENVIRONMENT = None + _is_oss_cutlass_moe.cache_clear() @contextmanager diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py index a132c9f28ff9..61044cfba365 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py @@ -379,6 +379,7 @@ def build_moe_deployment( eplb_enabled=eplb_enabled, # Routed-expert LoRA only; attention-only LoRA stays False. moe_lora_enabled=has_moe_lora_targets(lora_config), + force_dynamic_quantization=model_config.force_dynamic_quantization, # Same expression the impls use to set ``use_fused_finalize``; any # LoRA counts, not just routed-expert LoRA. fused_finalize_enabled=( diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 0061bf2389ad..f2a6bbaae7cc 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -194,6 +194,8 @@ l0_b200: - unittest/_torch/moe/fused_moe/test_deepgemm_fused_gather_finalize.py - unittest/_torch/moe/fused_moe/test_deepgemm_fused_expand_quant.py # ------------- MoE: test_moe_backend (by backend) --------------- + - unittest/_torch/moe/test_moe_backend.py::test_cutlass_moe_routing_round_trip + - unittest/_torch/moe/test_moe_backend.py::test_cutlass_moe_public_activation_ids - unittest/_torch/moe/test_kimi_k3_moe_gate.py::test_fused_route_quant_matches_unfused_chain - unittest/_torch/moe/test_moe_backend.py::test_kimi_fused_route_quant_skips_prequantized_input - unittest/_torch/moe/test_moe_backend.py::test_kimi_mxfp8_quantized_tensor_handoff diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index a42393e3d02a..074115ba4b27 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -12,6 +12,15 @@ l0_cpu: backend: generic orchestrator: mpi tests: + - unittest/_torch/moe/test_moe_impl.py::test_cutlass_build_mode_resolution + - unittest/_torch/moe/test_moe_impl.py::test_b12x_lora_rejection_is_build_independent + - unittest/_torch/moe/test_moe_impl.py::test_b12x_cutlass_prefill_routing + - unittest/_torch/moe/test_moe_impl.py::test_b12x_cutlass_build_mode_resolution + - unittest/_torch/moe/test_moe_impl.py::test_cutlass_dynamic_scale_resolution + - unittest/_torch/moe/test_moe_impl.py::test_cutlass_communication_validation + - unittest/_torch/moe/test_moe_impl.py::test_cutlass_communication_selection + - unittest/_torch/moe/test_moe_impl.py::test_cutlass_build_mode_probe + - unittest/_torch/moe/test_moe_impl.py::test_cutlass_build_mode_fingerprint - disaggregated/test_aiperf_gate.py::test_fires_on_error_storm - disaggregated/test_aiperf_gate.py::test_passes_healthy_run_with_cancellations - disaggregated/test_aiperf_gate.py::test_missing_export_raises diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 3305a7c62845..1684ffceee2b 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -16,6 +16,8 @@ l0_h100: backend: pytorch orchestrator: mpi tests: + - unittest/_torch/moe/test_moe_backend.py::test_cutlass_moe_routing_round_trip + - unittest/_torch/moe/test_moe_backend.py::test_cutlass_moe_public_activation_ids # Covers tests/unittest/_torch/attention/. Two sub-trees moved in here from elsewhere # under tests/unittest/_torch/ and have no entry of their own on any list, so this entry # is what picks them up: diff --git a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py index b5c3faca4edc..d49fb7174995 100644 --- a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py @@ -169,7 +169,7 @@ def test_can_implement_rejects_attention_dp_without_expert_parallelism(): # -------------------------------------------------------------------------- # Hybrid CUTLASS-prefill / b12x-decode dispatch predicate tests # -# ``_route_to_cutlass`` is a pure shape predicate on its input ``x``; we test +# ``_route_to_cutlass`` checks quantization and input shape; we test # it via a stub that holds the class constant, sidestepping the full # CutlassFusedMoE constructor (which needs a routing method, real model # config, etc.). @@ -182,6 +182,10 @@ class _RoutePredicateStub: _PREFILL_VIA_CUTLASS_THRESHOLD = CuteDslB12xFusedMoE._PREFILL_VIA_CUTLASS_THRESHOLD _route_to_cutlass = CuteDslB12xFusedMoE._route_to_cutlass + _prefill_uses_cutlass = CuteDslB12xFusedMoE._prefill_uses_cutlass + + def __init__(self): + self.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) def test_dispatch_routes_prefill_shape_via_cutlass(): diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index ed26f89c4268..cd331a733d91 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -141,6 +141,172 @@ logger = logging.getLogger(__name__) + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_cutlass_moe_routing_round_trip(dtype: torch.dtype) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUTLASS routing requires CUDA") + torch.manual_seed(17) + num_experts, hidden_size = 4, 128 + x = torch.randn(8, hidden_size, device="cuda", dtype=dtype) + selected = torch.arange(8, device="cuda", dtype=torch.int32).remainder(num_experts).unsqueeze(1) + scores = torch.ones(8, 1, device="cuda", dtype=torch.float32) + w1 = torch.empty(num_experts, hidden_size * 2, hidden_size, device="cuda", dtype=dtype) + w2 = torch.empty(num_experts, hidden_size, hidden_size, device="cuda", dtype=dtype) + + p2u, _, permuted, offsets, _, u2p = torch.ops.trtllm.moe_permute_op( + x, selected, scores, w1, w2, None, None, num_experts, 1, 0, 1, 0, 1, 0, False, False, False + ) + torch.testing.assert_close(permuted, x[p2u.long()], atol=0, rtol=0) + restored = torch.ops.trtllm.moe_finalize_scale_op( + permuted, + None, + scores, + u2p, + p2u, + selected, + offsets, + False, + x.shape[0], + hidden_size, + hidden_size, + 1, + num_experts, + 1, + 0, + 1, + 0, + ) + torch.testing.assert_close(restored, x, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "activation_type,clamp_limit,alpha_only,capture_first", + [ + (ActivationType.Identity, None, False, False), + (ActivationType.Gelu, None, False, False), + (ActivationType.Relu, None, False, False), + (ActivationType.Silu, None, False, False), + (ActivationType.Swiglu, None, False, False), + (ActivationType.Geglu, None, False, False), + (ActivationType.SwigluBias, None, False, False), + (ActivationType.Swiglu, 0.2, False, False), + (ActivationType.SwigluBias, 0.2, False, False), + pytest.param(ActivationType.Swiglu, 0.2, True, False, id="alpha-only-clamp"), + pytest.param(ActivationType.Swiglu, 0.2, False, True, id="capture-before-eager"), + ], +) +def test_cutlass_moe_public_activation_ids( + activation_type: ActivationType, + clamp_limit: Optional[float], + alpha_only: bool, + capture_first: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + if not torch.cuda.is_available() or get_sm_version() < 80: + pytest.skip("BF16 CUTLASS MoE requires CUDA and SM80 or newer") + + torch.manual_seed(17) + dtype = torch.bfloat16 + num_experts, hidden_size, inter_size = 4, 128, 128 + gated = is_gated_activation(activation_type) + x = torch.randn(8, hidden_size, device="cuda", dtype=dtype) * 0.5 + w1 = ( + torch.randn( + num_experts, inter_size * (2 if gated else 1), hidden_size, device="cuda", dtype=dtype + ) + * 0.1 + ) + w2 = torch.randn(num_experts, hidden_size, inter_size, device="cuda", dtype=dtype) * 0.1 + selected = torch.arange(8, device="cuda", dtype=torch.int32).remainder(num_experts).unsqueeze(1) + scores = torch.ones(8, 1, device="cuda", dtype=torch.float32) + alpha = ( + torch.ones(num_experts, device="cuda") + if activation_type == ActivationType.SwigluBias or alpha_only + else None + ) + beta = ( + torch.zeros(num_experts, device="cuda") + if activation_type == ActivationType.SwigluBias + else None + ) + limit = ( + torch.full((num_experts,), clamp_limit, device="cuda") if clamp_limit is not None else None + ) + + def run_moe(): + return torch.ops.trtllm.fused_moe( + x, + selected, + scores, + w1, + None, + w2, + None, + dtype, + [], + activation_type=int(activation_type), + swiglu_alpha=alpha, + swiglu_beta=beta, + swiglu_limit=limit, + )[0] + + if capture_first: + from tensorrt_llm._torch.custom_ops.torch_custom_ops import MoERunner + + # Isolate the native runner from earlier parameterized cases. Construct + # it before capture, with no eager forward or cached neutral defaults. + key = (dtype, dtype, dtype, False, False, False, False, False) + runner = torch.classes.trtllm.FusedMoeRunner(*key, True) + monkeypatch.setattr(MoERunner, "runner_dict", {key: runner}) + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.graph(graph, stream=stream): + captured = run_moe() + # Same stream, before replay: the eager call must not read defaults + # whose initialization was only recorded in the captured graph. + with torch.cuda.stream(stream): + actual = run_moe() + torch.cuda.current_stream().wait_stream(stream) + else: + actual = run_moe() + + expected = [] + for row, expert in zip(x, selected[:, 0]): + fc1 = torch.nn.functional.linear(row.float(), w1[expert].float()) + if gated: + up, gate = fc1.chunk(2) + if clamp_limit is not None: + up = up.clamp(min=-clamp_limit, max=clamp_limit) + gate = gate.clamp(max=clamp_limit) + activated = ( + torch.nn.functional.gelu(gate, approximate="tanh") + if activation_type == ActivationType.Geglu + else torch.nn.functional.silu(gate) + ) * up + elif activation_type == ActivationType.Gelu: + activated = torch.nn.functional.gelu(fc1, approximate="tanh") + elif activation_type == ActivationType.Relu: + activated = torch.nn.functional.relu(fc1) + elif activation_type == ActivationType.Silu: + activated = torch.nn.functional.silu(fc1) + else: + activated = fc1 + expected.append( + torch.nn.functional.linear(activated.to(dtype).float(), w2[expert].float()).to(dtype) + ) + + torch.testing.assert_close(actual, torch.stack(expected), atol=0.01, rtol=0.05) + if clamp_limit is not None: + if not capture_first: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run_moe() + graph.replay() + torch.testing.assert_close(captured, torch.stack(expected), atol=0.01, rtol=0.05) + + _MEGAMOE_BACKEND_TYPES = { MoeBackendType.MEGAMOE_DEEPGEMM, MoeBackendType.MEGAMOE_CUTEDSL, diff --git a/tests/unittest/_torch/moe/test_moe_impl.py b/tests/unittest/_torch/moe/test_moe_impl.py index b751cb0b5306..418e9c2ce3d4 100644 --- a/tests/unittest/_torch/moe/test_moe_impl.py +++ b/tests/unittest/_torch/moe/test_moe_impl.py @@ -1320,3 +1320,295 @@ class _SelfCollidingImpl: with pytest.raises(ValueError, match=r"token 'torch' is already a value of field 'provider'"): registry.register(_SelfCollidingImpl) + + +@pytest.mark.parametrize("is_oss", [False, True], ids=["internal", "oss"]) +@pytest.mark.parametrize( + "activation,quant_algo,sm,reason", + [ + pytest.param("SiTu", None, 90, MoERejectReason.ACTIVATION_UNSUPPORTED, id="situ"), + pytest.param("Relu2", None, 90, MoERejectReason.ACTIVATION_UNSUPPORTED, id="relu2"), + pytest.param( + "Swiglu", QuantAlgo.W4A16_MXFP4, 90, MoERejectReason.QUANT_UNSUPPORTED, id="w4a16_mxfp4" + ), + pytest.param( + "Swiglu", QuantAlgo.W4A8_AWQ, 90, MoERejectReason.QUANT_UNSUPPORTED, id="w4a8_awq" + ), + pytest.param("Swiglu", QuantAlgo.MXFP8, 100, MoERejectReason.QUANT_UNSUPPORTED, id="mxfp8"), + pytest.param( + "Swiglu", + QuantAlgo.W4A8_MXFP4_MXFP8, + 100, + MoERejectReason.QUANT_UNSUPPORTED, + id="w4a8_mxfp4_mxfp8", + ), + ], +) +def test_cutlass_build_mode_resolution(is_oss, activation, quant_algo, sm, reason): + report = resolve_moe_impl( + ModelConfig(moe_backend="CUTLASS"), + problem=MoEProblem( + quant=quant_algo.value if quant_algo is not None else None, + dtype_act=torch.bfloat16, + hidden_size=512, + intermediate_size=512, + num_experts=8, + top_k=2, + activation=activation, + ), + deployment=MoEDeployment( + ep_size=1, + tp_size=1, + parallel_size=1, + use_dp=False, + num_slots=8, + env=MoEEnvironment(sm=sm, oss_cutlass_moe=is_oss), + ), + ) + if is_oss: + assert impl_class_for(report) is CutlassFusedMoE + assert not report.rejected + else: + assert report.winner is None + assert len(report.rejected) == 1 + assert report.rejected[0].reason is reason + with pytest.raises(ValueError, match="USING_OSS_CUTLASS_MOE_GEMM=ON"): + impl_class_for(report) + + +@pytest.mark.parametrize("is_oss", [False, True]) +@pytest.mark.parametrize("dynamic", [False, True]) +def test_cutlass_dynamic_scale_resolution(is_oss: bool, dynamic: bool) -> None: + from tensorrt_llm._torch.moe.fused_moe.moe_resolution import build_moe_deployment + + deployment = build_moe_deployment( + ModelConfig(force_dynamic_quantization=dynamic), + num_experts=8, + environment=MoEEnvironment(sm=100, oss_cutlass_moe=is_oss), + ) + assert deployment.force_dynamic_quantization is dynamic + report = resolve_moe_impl( + ModelConfig(moe_backend="CUTLASS"), + problem=MoEProblem( + quant=QuantAlgo.NVFP4.value, + dtype_act=torch.bfloat16, + hidden_size=512, + intermediate_size=512, + num_experts=8, + top_k=2, + ), + deployment=deployment, + ) + if dynamic and not is_oss: + assert report.winner is None + assert report.rejected[0].reason is MoERejectReason.QUANT_UNSUPPORTED + else: + assert impl_class_for(report) is CutlassFusedMoE + + +@pytest.mark.parametrize("is_oss", [False, True]) +@pytest.mark.parametrize("alltoall", [False, True]) +def test_cutlass_communication_validation(is_oss: bool, alltoall: bool) -> None: + from types import SimpleNamespace + + backend = CutlassFusedMoE.__new__(CutlassFusedMoE) + moe = SimpleNamespace(enable_alltoall=alltoall) + with override_moe_environment(MoEEnvironment(sm=90, oss_cutlass_moe=is_oss)): + if alltoall and not is_oss: + with pytest.raises(ValueError, match="all-to-all.*USING_OSS_CUTLASS_MOE_GEMM=ON"): + backend.validate_configurable_moe(moe) + else: + backend.validate_configurable_moe(moe) + + +@pytest.mark.parametrize("mode", [None, False, True]) +def test_cutlass_build_mode_probe(monkeypatch: pytest.MonkeyPatch, mode: bool | None) -> None: + from types import SimpleNamespace + + from tensorrt_llm._torch.moe.fused_moe import impl_environment + + query = MagicMock(return_value=mode) + debug = MagicMock() + monkeypatch.setattr(impl_environment.logger, "debug", debug) + namespace = SimpleNamespace() if mode is None else SimpleNamespace(is_oss_cutlass_moe=query) + monkeypatch.setattr(torch.ops, "trtllm", namespace) + impl_environment.reset_moe_environment_cache() + try: + expected = True if mode is None else mode + assert impl_environment._is_oss_cutlass_moe() is expected + assert impl_environment._is_oss_cutlass_moe() is expected + assert query.call_count == (0 if mode is None else 1) + if mode is None: + debug.assert_called_once_with( + "CUTLASS MoE build-mode query is unavailable; assuming OSS kernels" + ) + else: + debug.assert_not_called() + finally: + impl_environment.reset_moe_environment_cache() + + +def test_cutlass_build_mode_fingerprint() -> None: + assert ( + MoEEnvironment(sm=90).fingerprint() + != MoEEnvironment(sm=90, oss_cutlass_moe=False).fingerprint() + ) + + +@pytest.mark.parametrize("is_oss", [False, True]) +@pytest.mark.parametrize("method", [None, "ALLGATHER", "NVLINK_ONE_SIDED"]) +@pytest.mark.parametrize("force_source", ["model", "env"]) +def test_cutlass_communication_selection(monkeypatch, is_oss, method, force_source): + from types import SimpleNamespace + + from tensorrt_llm._torch.moe.fused_moe.communication import communication_factory + from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE + + monkeypatch.delenv("TRTLLM_FORCE_COMM_METHOD", raising=False) + if force_source == "env" and method is not None: + monkeypatch.setenv("TRTLLM_FORCE_COMM_METHOD", method) + mapping = SimpleNamespace( + enable_attention_dp=True, + dp_size=2, + moe_tp_size=1, + moe_ep_size=2, + has_cp_helix=lambda: False, + ) + config = SimpleNamespace( + mapping=mapping, + pretrained_config=SimpleNamespace(hidden_size=128), + torch_dtype=torch.bfloat16, + quant_config=None, + max_num_tokens=8, + moe_max_num_tokens=8, + use_cuda_graph=False, + use_low_precision_moe_combine=False, + moe_load_balancer=None, + ) + constructors = {} + for name in ( + "AllGatherReduceScatter", + "NVLinkOneSided", + "NVLinkTwoSided", + "NVLinkTwoSidedFlashinfer", + "NcclEP", + "DeepEP", + "DeepEPLowLatency", + ): + constructors[name] = MagicMock(name=name) + monkeypatch.setattr(communication_factory, name, constructors[name]) + monkeypatch.setattr(communication_factory, "get_wide_ep_ft_options", lambda _: (None, 60, 1)) + moe = SimpleNamespace( + backend=CutlassFusedMoE.__new__(CutlassFusedMoE), + enable_dwdp=False, + model_config=config, + num_experts=4, + num_slots=4, + routing_method=SimpleNamespace(experts_per_token=1), + expert_size_per_partition=2, + use_flashinfer=False, + hidden_size=128, + communication_method=method if force_source == "model" else None, + ) + with override_moe_environment(MoEEnvironment(sm=90, oss_cutlass_moe=is_oss)): + if not is_oss and method == "NVLINK_ONE_SIDED": + with pytest.raises(ValueError, match="requires all-to-all.*use ALLGATHER"): + ConfigurableMoE._create_comm_strategy_auto(moe) + assert all(constructor.call_count == 0 for constructor in constructors.values()) + else: + strategy = ConfigurableMoE._create_comm_strategy_auto(moe) + selected = ( + "AllGatherReduceScatter" + if not is_oss or method == "ALLGATHER" + else "NVLinkOneSided" + ) + assert strategy is constructors[selected].return_value + constructors[selected].assert_called_once() + assert all( + constructor.call_count == 0 + for name, constructor in constructors.items() + if name != selected + ) + + +@pytest.mark.parametrize("is_oss", [False, True]) +@pytest.mark.parametrize("sm", [120, 121]) +@pytest.mark.parametrize("quant_algo", [QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4]) +@pytest.mark.parametrize("activation", ["Relu2", "Swiglu"]) +@pytest.mark.parametrize("dynamic", [False, True]) +def test_b12x_cutlass_build_mode_resolution(is_oss, sm, quant_algo, activation, dynamic): + from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE + + report = resolve_moe_impl( + ModelConfig(moe_backend="CUTEDSL"), + problem=MoEProblem( + quant=canonical_quant(quant_algo), + dtype_act=torch.bfloat16, + hidden_size=512, + intermediate_size=512, + num_experts=8, + top_k=2, + activation=activation, + ), + deployment=MoEDeployment( + ep_size=1, + tp_size=1, + parallel_size=1, + use_dp=False, + num_slots=8, + force_dynamic_quantization=dynamic, + env=MoEEnvironment(sm=sm, oss_cutlass_moe=is_oss, available_deps=(MoEDep.FLASHINFER,)), + ), + ) + rejection = not is_oss and quant_algo == QuantAlgo.NVFP4 and (activation == "Relu2" or dynamic) + if rejection: + assert report.winner is None + b12x_rejection = next( + r for r in report.rejected if r.legacy_backend == "CuteDslB12xFusedMoE" + ) + expected = ( + MoERejectReason.ACTIVATION_UNSUPPORTED + if activation == "Relu2" + else MoERejectReason.QUANT_UNSUPPORTED + ) + assert b12x_rejection.reason is expected + assert "USING_OSS_CUTLASS_MOE_GEMM=ON" in b12x_rejection.detail + assert "CuteDslB12xFusedMoE" in b12x_rejection.detail + assert b12x_rejection.detail.startswith("NVFP4 prefill routes through CUTLASS:") + else: + assert impl_class_for(report) is CuteDslB12xFusedMoE + + +@pytest.mark.parametrize("is_oss", [False, True]) +def test_b12x_lora_rejection_is_build_independent(is_oss): + from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE + + verdict = CuteDslB12xFusedMoE.can_implement( + MoEProblem(quant=canonical_quant(QuantAlgo.NVFP4), dtype_act=torch.bfloat16), + MoEDeployment( + ep_size=1, + tp_size=1, + parallel_size=1, + use_dp=False, + num_slots=8, + moe_lora_enabled=True, + env=MoEEnvironment(sm=120, oss_cutlass_moe=is_oss, available_deps=(MoEDep.FLASHINFER,)), + ), + ) + assert not verdict.eligible + assert verdict.reject_reason is MoERejectReason.LORA_UNSUPPORTED + assert verdict.detail == "CuteDslB12xFusedMoE does not fuse routed-expert LoRA" + + +@pytest.mark.parametrize("quant_algo", [QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4, None]) +@pytest.mark.parametrize("num_tokens", [63, 64, 65]) +def test_b12x_cutlass_prefill_routing(quant_algo, num_tokens): + from types import SimpleNamespace + + from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE + + backend = CuteDslB12xFusedMoE.__new__(CuteDslB12xFusedMoE) + torch.nn.Module.__init__(backend) + backend.quant_config = SimpleNamespace(quant_algo=quant_algo) if quant_algo else None + expected = quant_algo == QuantAlgo.NVFP4 and num_tokens >= 64 + assert backend._route_to_cutlass(torch.empty(num_tokens, 1)) is expected