From 0971d6803e57fd5bc1fc24e89d2dc0b214e30549 Mon Sep 17 00:00:00 2001 From: leslief Date: Thu, 11 Jun 2026 22:23:11 -0700 Subject: [PATCH 01/19] [None][feat] Cherry-pick #11143 to Main Signed-off-by: leslief --- .../blockScaleMoe/routing/RoutingDeepSeek.cu | 37 +++++- .../blockScaleMoe/routing/RoutingKernel.h | 16 +++ .../trtllmGenKernels/blockScaleMoe/runner.cu | 71 +++++++--- .../trtllmGenKernels/blockScaleMoe/runner.h | 15 ++- cpp/tensorrt_llm/thop/cuteDslMoeUtilsOp.cpp | 7 +- cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp | 5 +- cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp | 123 +++++++++++++----- .../thop/fp8PerTensorScaleMoe.cpp | 5 +- cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp | 5 +- .../custom_ops/trtllm_gen_custom_ops.py | 15 ++- .../_torch/models/modeling_deepseekv3.py | 96 +++++++++----- .../modules/fused_moe/configurable_moe.py | 11 ++ .../modules/fused_moe/fused_moe_trtllm_gen.py | 21 ++- .../modules/fused_moe/moe_op_backend.py | 7 + .../_torch/modules/fused_moe/quantization.py | 67 +++++++++- tests/unittest/_torch/thop/serial/test_moe.py | 116 ++++++++++++----- 16 files changed, 468 insertions(+), 149 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu index 476bba2bb317..55cfca91707e 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu @@ -380,17 +380,31 @@ __global__ void routingMainKernel(KernelParams params) auto finalScore = OutputT{scoreNorm * params.mRouteScale / redNorm}; // write expert idx out already - auto idxTopK = blockIdx.x * params.mTopK + laneIdx; + auto idxTopK = blockIdx.x * params.mTotalExpertsPerToken + laneIdx; + auto idxShared = blockIdx.x * params.mTotalExpertsPerToken + params.mTopK + laneIdx; if (laneIdx < params.mTopK && params.mPtrTopKPacked != nullptr) { PackedScoreIdx packedScore{static_cast(finalScore), static_cast(expertIdx)}; params.mPtrTopKPacked[idxTopK] = packedScore; } + if (laneIdx < params.mNumFusedSharedExperts && params.mPtrTopKPacked != nullptr) + { + PackedScoreIdx packedScore{ + static_cast(1.0F), static_cast(params.mNumExperts + laneIdx)}; + params.mPtrTopKPacked[idxShared] = packedScore; + } + if (laneIdx < params.mTopK && params.mPtrTopKWeights != nullptr && params.mPtrTopKIds == nullptr) { params.mPtrTopKWeights[idxTopK] = finalScore; } + + // Write score of 1.0 for shared expert if enabled + if (laneIdx < params.mNumFusedSharedExperts && params.mPtrTopKWeights != nullptr) + { + params.mPtrTopKWeights[idxShared] = static_cast(1.0F); + } } } } @@ -551,12 +565,29 @@ void run(Data& data, void* stream) } int const numBlocks = data.mNumTokens; - int const numThreadsHist = getMaxNumExperts(data.mNumExperts); + // Account for fused shared experts (appended after the routed experts) when sizing the histogram. + int const numThreadsHist = getMaxNumExperts(data.mNumExperts + data.mNumFusedSharedExperts); static int const smMajor = tensorrt_llm::common::getSMVersion() / 10; - // Step 1: Run DeepSeek-specific topK computation (writes to mPtrTopKPacked) + + TLLM_CHECK_WITH_INFO(data.mNumFusedSharedExperts <= WarpSize, + "Number of fused shared experts (%d) must be less than warp size (%d).", data.mNumFusedSharedExperts, WarpSize); + + // Step 1: Run DeepSeek-specific topK computation (writes to mPtrTopKPacked). + // When fused shared experts are enabled, the main kernel also appends them (index mNumExperts + i, weight 1.0), + // so it must run before mNumExperts/mTopK are expanded below. int const numThreadsMain = max(data.mNumExpertGroups * WarpSize, getMaxNumExperts(data.mNumExperts)); launchMainKernel(data, numBlocks, numThreadsMain, stream); + // Fused shared experts are appended after the routed experts; expand the expert/topK counts so the + // permutation pipeline below (which reads data.mNumExperts / data.mTopK) accounts for them. + if (data.mNumFusedSharedExperts > 0) + { + data.mNumExperts += data.mNumFusedSharedExperts; + data.mTopK += data.mNumFusedSharedExperts; + data.mNumLocalExperts += data.mNumFusedSharedExperts; + // data.mLocalExpertsStartIdx += data.mNumFusedSharedExperts; + } + // Step 2: Permutation pipeline (reads from mPtrTopKPacked written by step 1) if (data.mPtrPermutedIdxSize != nullptr) { diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h index bea23942e49d..e972b5809687 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h @@ -109,6 +109,12 @@ struct DataBase int32_t mLocalExpertsStartIdx; int32_t mLocalExpertsStrideLog2; int32_t mNumLocalExperts; + + /// For fused shared expert + int32_t mNumFusedSharedExperts; + int32_t mSharedExpertTokenOffset; + int32_t mSharedExpertNumTokens; + int32_t mTotalExpertsPerToken; }; template @@ -145,6 +151,11 @@ struct KernelParamsBase int32_t mLocalExpertsStrideLog2 = 0; int32_t mNumLocalExperts = 0; + int32_t mNumFusedSharedExperts; + int32_t mSharedExpertTokenOffset; + int32_t mSharedExpertNumTokens; + int32_t mTotalExpertsPerToken; + // Public initialization function - make it a template to accept different Data types template void setBaseParams(DataType const& data) @@ -171,6 +182,11 @@ struct KernelParamsBase mLocalExpertsStartIdx = data.mLocalExpertsStartIdx; mLocalExpertsStrideLog2 = data.mLocalExpertsStrideLog2; mNumLocalExperts = data.mNumLocalExperts; + + mNumFusedSharedExperts = data.mNumFusedSharedExperts; + mSharedExpertTokenOffset = data.mSharedExpertTokenOffset; + mSharedExpertNumTokens = data.mSharedExpertNumTokens; + mTotalExpertsPerToken = data.mTotalExpertsPerToken; } }; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu index 39021ce642b1..a86ecc49124b 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu @@ -61,8 +61,8 @@ Runner::Runner(int32_t tileTokensDim, int32_t clusterSizeInBatchDim) } void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int32_t numExperts, int32_t topK, - int32_t nGroup, int32_t topkGroup, int32_t localExpertOffset, int32_t localNumExperts, float routedScalingFactor, - int32_t* routingExpertIndexes, int32_t* expertCountHistogram, int32_t* permutedIdxSize, + int32_t numFusedSharedExpert, int32_t nGroup, int32_t topkGroup, int32_t localExpertOffset, int32_t localNumExperts, + float routedScalingFactor, int32_t* routingExpertIndexes, int32_t* expertCountHistogram, int32_t* permutedIdxSize, int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToExpandedIdx, int32_t* permutedIdxToTokenIdx, void* expertWeights, int32_t* expertIds, int32_t* numTokensPerExpert, int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, btg::Dtype dtypeElt, bool useRoutingScalesOnInput, @@ -71,7 +71,10 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 { if (routingMethodType == RoutingMethodType::DeepSeekV3 && nGroup <= 1) { - // DeepSeek no-groups case: use routingCustom with SigmoidBias preprocess + // DeepSeek no-groups case: use routingCustom with SigmoidBias preprocess. + // NOTE: routingCustom does not implement fused shared experts; when fusion is + // requested we fall through to the routingDeepSeek path below (which handles it + // for both grouped and non-grouped routing). // and ScaledSumNormalize postprocess. This is more efficient than the full DeepSeek // kernel because it uses the warp-level routingTopKExperts flow. moe::dev::routing::routingCustom::Data routingData; @@ -234,6 +237,8 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 routingData.mDtypeOutput = btg::Dtype::Bfloat16; routingData.mUsePdl = tensorrt_llm::common::getEnvEnablePDL(); + int32_t const totalExpertsPerToken = topK + numFusedSharedExpert; + // output: routingData.mPtrTopKPacked = routingExpertIndexes; routingData.mPtrExpertCounts = expertCountHistogram; @@ -255,9 +260,11 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 routingData.mPtrTopKIds = expertIds; routingData.mNumTokens = numTokens; routingData.mNumExperts = numExperts; + routingData.mNumFusedSharedExperts = numFusedSharedExpert; routingData.mNumExpertGroups = nGroup; routingData.mNumLimitedGroups = topkGroup; routingData.mTopK = topK; + routingData.mTotalExpertsPerToken = totalExpertsPerToken; routingData.mPaddingLog2 = computeLog2(mTileTokensDim); routingData.mTileTokensDim = mTileTokensDim; routingData.mLocalExpertsStartIdx = localExpertOffset; @@ -265,6 +272,23 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 routingData.mNumLocalExperts = localNumExperts; routingData.mRouteScale = routedScalingFactor; routingData.mUseRoutingSoftmax = false; + + // TODO Should these be passed directly instead? This does assume a constant number of experts per device + int32_t const numDevices = numExperts / localNumExperts; + int32_t const deviceIndex = localExpertOffset / localNumExperts; + int32_t const baseTokensPerDevice = numTokens / numDevices; + int32_t const remainingTokens = numTokens % numDevices; + + if (deviceIndex < remainingTokens) + { + routingData.mSharedExpertTokenOffset = (baseTokensPerDevice + 1) * deviceIndex; + routingData.mSharedExpertNumTokens = baseTokensPerDevice + 1; + } + else + { + routingData.mSharedExpertTokenOffset = remainingTokens + deviceIndex * baseTokensPerDevice; + routingData.mSharedExpertNumTokens = baseTokensPerDevice; + } moe::dev::routing::routingDeepSeek::run(routingData, stream); } else if (routingMethodType == RoutingMethodType::Llama4) @@ -274,6 +298,8 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 { TLLM_LOG_WARNING("For Llama routing method, nGroup/topkGroup is ignored, got %d/%d.", nGroup, topkGroup); } + TLLM_CHECK_WITH_INFO(numFusedSharedExpert == 0, "Llama routing method does not support fusing shared expert"); + moe::dev::routing::routingLlama4::Data routingData; routingData.mDtypeOutput = btg::Dtype::Bfloat16; routingData.mUsePdl = tensorrt_llm::common::getEnvEnablePDL(); @@ -318,6 +344,9 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 else if (routingMethodType == RoutingMethodType::Renormalize || routingMethodType == RoutingMethodType::RenormalizeNaive || routingMethodType == RoutingMethodType::Default) { + TLLM_CHECK_WITH_INFO( + numFusedSharedExpert == 0, "Renormalize routing method does not support fusing shared expert"); + moe::dev::routing::routingCustom::Data routingData; // @@ -649,6 +678,9 @@ void Runner::setOpsData(MoERunnerArgs const& args, MoEWorkspace const& workspace moe::dev::convertsf::Data& convertSfData, moe::dev::activation::Data& activationData, moe::dev::finalize::Data& finalizeData) { + int32_t const totalNumExperts = args.num_experts + args.num_fused_shared_experts; + int32_t const totalExpertsPerToken = args.top_k + args.num_fused_shared_experts; + // Setup sf conversion data if needed convertSfData.inSfPtr = args.hidden_states_scale; convertSfData.outSfPtr = workspace.hidden_states_scale_linear; @@ -667,7 +699,7 @@ void Runner::setOpsData(MoERunnerArgs const& args, MoEWorkspace const& workspace activationData.inDqSfsPtr = workspace.gemm1_output_scale; activationData.outDqSfsPtr = workspace.activation_output_scale; activationData.innerDim = args.intermediate_size * (mActType == ActType::SwiGlu ? 2 : 1); - activationData.topK = args.top_k; + activationData.topK = totalExpertsPerToken; // TODO Rename topK in activation data struct activationData.numTokens = args.num_tokens; activationData.expandedIdxToPermutedIdx = workspace.expanded_idx_to_permuted_idx; // For DeepSeek FP8 the activation runs as a separate kernel rather than @@ -699,8 +731,8 @@ void Runner::setOpsData(MoERunnerArgs const& args, MoEWorkspace const& workspace } finalizeData.expandedIdxToPermutedIdx = workspace.expanded_idx_to_permuted_idx; finalizeData.numTokens = args.num_tokens; - finalizeData.numExperts = args.num_experts; - finalizeData.topK = args.top_k; + finalizeData.numExperts = totalNumExperts; // TODO Is this used? + finalizeData.topK = totalExpertsPerToken; // TODO Rename topK in finalize data struct // We want to fuse unpadding into the finalize kernel, so we need to use the output hidden size. finalizeData.hiddenDim = args.valid_hidden_size.value_or(args.hidden_size); finalizeData.hiddenDimPadded = args.output_hidden_size.value_or(args.hidden_size); @@ -710,12 +742,15 @@ void Runner::setOpsData(MoERunnerArgs const& args, MoEWorkspace const& workspace std::tuple Runner::getWorkspaceSizeInBytes(MoERunnerArgs const& args, int64_t configIndex) const { + int32_t const totalLocalExperts = args.local_num_experts + args.num_fused_shared_experts; + int32_t const totalExpertsPerToken = args.top_k + args.num_fused_shared_experts; + auto const& config = mPassingConfigs[configIndex]; - auto workspace_size_fc1 = static_cast(mPermuteGemm1.getWorkspaceSizeInBytes(args.top_k, args.hidden_size, - args.intermediate_size, args.local_num_experts, args.num_tokens, config.gemm1Config)); - auto workspace_size_fc2 = static_cast(mGemm2.getWorkspaceSizeInBytes(args.top_k, args.hidden_size, - args.intermediate_size, args.local_num_experts, args.num_tokens, config.gemm2Config)); + auto workspace_size_fc1 = static_cast(mPermuteGemm1.getWorkspaceSizeInBytes(totalExpertsPerToken, + args.hidden_size, args.intermediate_size, totalLocalExperts, args.num_tokens, config.gemm1Config)); + auto workspace_size_fc2 = static_cast(mGemm2.getWorkspaceSizeInBytes(totalExpertsPerToken, + args.hidden_size, args.intermediate_size, totalLocalExperts, args.num_tokens, config.gemm2Config)); return std::make_tuple(workspace_size_fc1, workspace_size_fc2); } @@ -750,7 +785,6 @@ std::vector Runner::getValidConfigIndices(int32_t topK, int32_t hiddenS int64_t Runner::getDefaultValidConfigIndex(int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numLocalExperts, int32_t numTokens, int32_t validHiddenSize, int32_t validIntermediateSize) const { - int32_t indexGemm1 = mPermuteGemm1.getDefaultValidConfigIndex( topK, hiddenSize, intermediateSize, numLocalExperts, numTokens, validHiddenSize, validIntermediateSize); int32_t indexGemm2 = mGemm2.getDefaultValidConfigIndex( @@ -773,6 +807,9 @@ void Runner::run( sync_check_cuda_error(stream); setOpsData(args, workspace, convertSfData, activationData, finalizeData); + int32_t const totalLocalExperts = args.local_num_experts + args.num_fused_shared_experts; + int32_t const totalExpertsPerToken = args.top_k + args.num_fused_shared_experts; + void* hidden_states_scale_linear{args.hidden_states_scale}; auto const& config = mPassingConfigs[configIndex]; @@ -780,7 +817,7 @@ void Runner::run( mPermuteGemm1.run(args.hidden_states, hidden_states_scale_linear, args.gemm1_weights, args.gemm1_weights_scale, workspace.expert_weights, args.output1_scales_scalar, args.output1_scales_gate_scalar, args.gemm1_bias, args.gemm1_alpha, args.gemm1_beta, args.gemm1_clamp_limit, workspace.gemm1_output, workspace.gemm1_output_scale, - args.top_k, args.hidden_size, args.intermediate_size, args.local_num_experts, args.num_tokens, + totalExpertsPerToken, args.hidden_size, args.intermediate_size, totalLocalExperts, args.num_tokens, workspace.permuted_idx_to_token_idx, workspace.num_non_exiting_ctas, workspace.total_num_padded_tokens, workspace.cta_idx_xy_to_batch_idx, workspace.cta_idx_xy_to_mn_limit, workspace.bmm1_workspace, args.mUseRoutingScalesOnInput, device, stream, config.gemm1Config, @@ -801,11 +838,11 @@ void Runner::run( // Run gemm2 mGemm2.run(gemm2_input, gemm2_input_scale, args.gemm2_weights, args.gemm2_weights_scale, args.output2_scales_scalar, - args.gemm2_bias, workspace.gemm2_output, workspace.gemm2_output_scale, args.top_k, - args.output_hidden_size.value_or(args.hidden_size), args.intermediate_size, args.local_num_experts, - args.num_tokens, workspace.num_non_exiting_ctas, workspace.total_num_padded_tokens, - workspace.cta_idx_xy_to_batch_idx, workspace.cta_idx_xy_to_mn_limit, workspace.bmm2_workspace, device, stream, - config.gemm2Config, args.valid_hidden_size.value_or(args.hidden_size), + args.gemm2_bias, workspace.gemm2_output, workspace.gemm2_output_scale, totalExpertsPerToken, + args.output_hidden_size.value_or(args.hidden_size), args.intermediate_size, totalLocalExperts, args.num_tokens, + workspace.num_non_exiting_ctas, workspace.total_num_padded_tokens, workspace.cta_idx_xy_to_batch_idx, + workspace.cta_idx_xy_to_mn_limit, workspace.bmm2_workspace, device, stream, config.gemm2Config, + args.valid_hidden_size.value_or(args.hidden_size), args.valid_intermediate_size.value_or(args.intermediate_size)); // Run finalize diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h index 97d77eddad0a..6525b2547343 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h @@ -165,13 +165,13 @@ class Runner explicit Runner(int32_t tileTokensDim, int32_t clusterSizeInBatchDim = 1); void run(void* routingLogits, void* routingBias, int32_t numTokens, int32_t numExperts, int32_t topK, - int32_t nGroups, int32_t topkGroups, int32_t localExpertOffset, int32_t localNumExperts, - float routedScalingFactor, int32_t* routingExpertIndexes, int32_t* expertCountHistogram, - int32_t* permutedIdxSize, int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToExpandedIdx, - int32_t* permutedIdxToTokenIdx, void* expertWeights, int32_t* expertIds, int32_t* numTokensPerExpert, - int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, - batchedGemm::trtllm::gen::Dtype dtypeElt, bool useRoutingScalesOnInput, bool useDeepSeekFp8, - RoutingMethodType routingMethodType, cudaStream_t stream, + int32_t numFusedSharedExpert, int32_t nGroups, int32_t topkGroups, int32_t localExpertOffset, + int32_t localNumExperts, float routedScalingFactor, int32_t* routingExpertIndexes, + int32_t* expertCountHistogram, int32_t* permutedIdxSize, int32_t* expandedIdxToPermutedIdx, + int32_t* permutedIdxToExpandedIdx, int32_t* permutedIdxToTokenIdx, void* expertWeights, int32_t* expertIds, + int32_t* numTokensPerExpert, int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, + int32_t* numNonExitingCtas, batchedGemm::trtllm::gen::Dtype dtypeElt, bool useRoutingScalesOnInput, + bool useDeepSeekFp8, RoutingMethodType routingMethodType, cudaStream_t stream, batchedGemm::trtllm::gen::Dtype dtypeRoutingLogits = batchedGemm::trtllm::gen::Dtype::Bfloat16, batchedGemm::trtllm::gen::Dtype dtypeRoutingBias = batchedGemm::trtllm::gen::Dtype::Bfloat16); @@ -297,6 +297,7 @@ struct MoERunnerArgs int32_t num_tokens{0}; int32_t num_experts{0}; + int32_t num_fused_shared_experts{0}; // Hidden dimension input of MoE block. It might be padded. int32_t hidden_size{0}; // Hidden dimension output of MoE block. It might be padded. diff --git a/cpp/tensorrt_llm/thop/cuteDslMoeUtilsOp.cpp b/cpp/tensorrt_llm/thop/cuteDslMoeUtilsOp.cpp index cb6765ac6a56..73d41ba6f4a1 100644 --- a/cpp/tensorrt_llm/thop/cuteDslMoeUtilsOp.cpp +++ b/cpp/tensorrt_llm/thop/cuteDslMoeUtilsOp.cpp @@ -79,9 +79,10 @@ std::vector moe_topk_sort_impl(torch::optional con auto const dtypeRoutingLogits = routing_logits.has_value() ? (routing_logits->scalar_type() == at::ScalarType::Float ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16) : btg::Dtype::Bfloat16; - routing_runner.run(routing_logits_ptr, routing_bias_ptr, num_tokens, num_experts, top_k, n_group.value_or(0), - topk_group.value_or(0), local_expert_offset, local_num_experts, routed_scaling_factor.value_or(1.0), - expert_indexes.data_ptr(), expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), + routing_runner.run(routing_logits_ptr, routing_bias_ptr, num_tokens, num_experts, top_k, + /* num_fused_shared_expert */ 0, n_group.value_or(0), topk_group.value_or(0), local_expert_offset, + local_num_experts, routed_scaling_factor.value_or(1.0), expert_indexes.data_ptr(), + expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), expanded_idx_to_permuted_idx.data_ptr(), permuted_idx_to_expanded_idx.data_ptr(), nullptr /*permuted_idx_to_token_idx.data_ptr()*/, token_final_scales_ptr, token_selected_experts_ptr, num_tokens_per_expert.data_ptr(), tile_idx_to_expert_idx.data_ptr(), diff --git a/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp index b90fbcf0bf70..74440d2f74de 100644 --- a/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp @@ -280,8 +280,9 @@ std::vector run_fp4_block_scale_moe_runner(torch::optional(), expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), + /* num_fused_shared_expert */ 0, args.n_group, args.topk_group, args.local_expert_offset, + args.local_num_experts, args.routed_scaling_factor, expert_indexes.data_ptr(), + expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), expanded_idx_to_permuted_idx.data_ptr(), nullptr, /*permuted_idx_to_expanded_idx.data_ptr(),*/ permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, args.topk_ids, num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), diff --git a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp index e3079b3e71ee..8dcf5875317a 100644 --- a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp @@ -40,11 +40,12 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit std::optional const& routing_bias, at::Tensor const& hidden_states, at::Tensor const& hidden_states_scale, at::Tensor const& gemm1_weights, at::Tensor const& gemm1_weights_scale, at::Tensor const& gemm2_weights, at::Tensor const& gemm2_weights_scale, int64_t const num_experts, - int64_t const top_k, std::optional const n_group, std::optional const topk_group, - int64_t const intermediate_size, int64_t const local_expert_offset, int64_t const local_num_experts, - std::optional const routed_scaling_factor, int64_t const tile_tokens_dim, int64_t const routing_method_type, - MoeRunnerType& moe_runner, int64_t moeConfigIndex, std::optional const& topk_weights, - std::optional const& topk_ids, std::optional const& gemm1_clamp_limit = std::nullopt, + int64_t const top_k, std::optional const num_fused_shared_experts, std::optional const n_group, + std::optional const topk_group, int64_t const intermediate_size, int64_t const local_expert_offset, + int64_t const local_num_experts, std::optional const routed_scaling_factor, int64_t const tile_tokens_dim, + int64_t const routing_method_type, MoeRunnerType& moe_runner, int64_t moeConfigIndex, + std::optional const& topk_weights, std::optional const& topk_ids, + std::optional const& gemm1_clamp_limit = std::nullopt, std::optional const& out_tensor = std::nullopt) { TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP8 block scale MOE"); @@ -138,6 +139,13 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::MoERunnerArgs args; tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::MoEWorkspace workspace; + int64_t const num_total_experts + = num_experts + (num_fused_shared_experts.value_or(0) > 0 ? num_fused_shared_experts.value() : 0); + int64_t const total_experts_per_token + = top_k + (num_fused_shared_experts.value_or(0) > 0 ? num_fused_shared_experts.value() : 0); + int64_t const num_total_local_experts + = local_num_experts + (num_fused_shared_experts.value_or(0) > 0 ? num_fused_shared_experts.value() : 0); + // setup args // note: the assumption is that output data type is always Bfloat16 (the default) args.mDtypeElt = btg::Dtype::E4m3; @@ -161,6 +169,7 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit args.num_experts = num_experts; args.hidden_size = hidden_states.sizes()[1]; args.top_k = top_k; + args.num_fused_shared_experts = num_fused_shared_experts.value_or(0); args.n_group = n_group.value_or(0); args.topk_group = topk_group.value_or(0); args.local_expert_offset = local_expert_offset; @@ -193,7 +202,7 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit = at::detail::empty_cuda({num_experts}, at::ScalarType::Int, routing_device, std::nullopt); int32_t max_num_padded_tokens = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::getMaxPermutedPaddedCount( - args.num_tokens, top_k, num_experts, tile_tokens_dim); + args.num_tokens, total_experts_per_token, num_total_experts, tile_tokens_dim); int32_t max_num_padded_tokens_gemm1 = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::maybeGetMinTokenCount( max_num_padded_tokens, 2 * args.intermediate_size, btg::dtypeGetNumBits(args.mDtypeElt)); @@ -202,8 +211,8 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit max_num_padded_tokens, args.hidden_size, btg::dtypeGetNumBits(args.mDtypeOut)); at::Tensor total_num_padded_tokens = at::empty({}, at::TensorOptions().device(routing_device).dtype(at::ScalarType::Int)); - at::Tensor expanded_idx_to_permuted_idx - = at::detail::empty_cuda({args.num_tokens * args.top_k}, at::ScalarType::Int, routing_device, std::nullopt); + at::Tensor expanded_idx_to_permuted_idx = at::detail::empty_cuda( + {args.num_tokens * total_experts_per_token}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor permuted_idx_to_token_idx = at::detail::empty_cuda({max_num_padded_tokens}, at::ScalarType::Int, routing_device, std::nullopt); // expert_weights is the routing kernel's topk-weights output and is consumed by moe_finalize, @@ -224,9 +233,9 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit } }(); at::Tensor expert_weights = at::detail::empty_cuda( - {args.num_tokens, args.top_k}, expert_weights_scalar_type, routing_device, std::nullopt); - at::Tensor expert_indexes - = at::detail::empty_cuda({args.num_tokens, args.top_k}, at::ScalarType::Int, routing_device, std::nullopt); + {args.num_tokens, total_experts_per_token}, expert_weights_scalar_type, routing_device, std::nullopt); + at::Tensor expert_indexes = at::detail::empty_cuda( + {args.num_tokens, total_experts_per_token}, at::ScalarType::Int, routing_device, std::nullopt); int64_t const size_of_expert_count_histogram = std::max(num_experts * 2, int64_t(256 * 2)); at::Tensor expert_count_histogram = at::detail::empty_cuda({size_of_expert_count_histogram}, at::ScalarType::Int, routing_device, std::nullopt); @@ -244,7 +253,7 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit {max_num_padded_tokens_gemm2, args.hidden_size}, at::ScalarType::BFloat16, routing_device, std::nullopt); int32_t max_num_ctas = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::getMaxNumCtasInBatchDim( - args.num_tokens, args.top_k, args.num_experts, tile_tokens_dim); + args.num_tokens, total_experts_per_token, num_total_experts, tile_tokens_dim); at::Tensor cta_idx_xy_to_batch_idx = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor cta_idx_xy_to_mn_limit @@ -262,12 +271,13 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit ? (routing_logits.value().scalar_type() == at::ScalarType::Float ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16) : btg::Dtype::Bfloat16; routing_runner.run(args.routing_logits, args.routing_bias, args.num_tokens, args.num_experts, args.top_k, - args.n_group, args.topk_group, args.local_expert_offset, args.local_num_experts, args.routed_scaling_factor, - expert_indexes.data_ptr(), expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), - expanded_idx_to_permuted_idx.data_ptr(), nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, - permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, args.topk_ids, - num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), - cta_idx_xy_to_mn_limit.data_ptr(), num_non_exiting_ctas.data_ptr(), args.mDtypeElt, false, true, + args.num_fused_shared_experts, args.n_group, args.topk_group, args.local_expert_offset, args.local_num_experts, + args.routed_scaling_factor, expert_indexes.data_ptr(), expert_count_histogram.data_ptr(), + total_num_padded_tokens.data_ptr(), expanded_idx_to_permuted_idx.data_ptr(), + nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, permuted_idx_to_token_idx.data_ptr(), + expert_weights_ptr, args.topk_ids, num_tokens_per_expert.data_ptr(), + cta_idx_xy_to_batch_idx.data_ptr(), cta_idx_xy_to_mn_limit.data_ptr(), + num_non_exiting_ctas.data_ptr(), args.mDtypeElt, false, true, static_cast(routing_method_type), stream, dtypeRoutingLogits, args.mDtypeBias); // MoE kernel except routing @@ -279,25 +289,26 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit TORCH_CHECK(hidden_states_scale.sizes()[1] == args.num_tokens, "hidden_states_scale dim1 must match num_tokens."); TORCH_CHECK(gemm1_weights.scalar_type() == at::ScalarType::Float8_e4m3fn, "gemm1_weights must be fp8."); TORCH_CHECK(gemm1_weights.dim() == 3, "gemm1_weights must be 3D."); + TORCH_CHECK(gemm1_weights.sizes()[0] == num_total_local_experts, "gemm1_weights has incorrect shape."); TORCH_CHECK(gemm1_weights.sizes()[1] % 2 == 0, "the second dimension of weights must be even."); TORCH_CHECK(intermediate_size == gemm1_weights.sizes()[1] / 2, "intermediate_size has incorrect shape."); TORCH_CHECK(gemm1_weights.sizes()[2] == hidden_states.sizes()[1], "the third dimension of weights must be equal to hidden_size."); TORCH_CHECK(gemm1_weights_scale.scalar_type() == at::ScalarType::Float, "gemm1_weights_scale must be float."); TORCH_CHECK(gemm1_weights_scale.dim() == 3, "gemm1_weights_scale must be 3D."); - - TORCH_CHECK(gemm1_weights_scale.sizes()[0] == local_num_experts, "gemm1_weights_scale has incorrect shape."); + TORCH_CHECK(gemm1_weights_scale.sizes()[0] == num_total_local_experts, "gemm1_weights_scale has incorrect dim 0."); TORCH_CHECK(intermediate_size % 128 == 0, "the second dimension of weights must be a multiple of 128."); TORCH_CHECK( gemm1_weights_scale.sizes()[1] == 2 * intermediate_size / 128, "gemm1_weights_scale has incorrect shape."); TORCH_CHECK(gemm1_weights_scale.sizes()[2] == args.hidden_size / 128, "gemm1_weights_scale has incorrect shape."); TORCH_CHECK(gemm2_weights.scalar_type() == at::ScalarType::Float8_e4m3fn, "gemm2_weights must be fp8."); TORCH_CHECK(gemm2_weights.dim() == 3, "gemm2_weights must be 3D."); + TORCH_CHECK(gemm2_weights.sizes()[0] == num_total_local_experts, "gemm2_weights has incorrect shape."); TORCH_CHECK(gemm2_weights.sizes()[2] == intermediate_size, "the third dimension of weights must be equal to intermediate_size."); TORCH_CHECK(gemm2_weights_scale.scalar_type() == at::ScalarType::Float, "gemm2_weights_scale must be float."); TORCH_CHECK(gemm2_weights_scale.dim() == 3, "gemm2_weights_scale must be 3D."); - TORCH_CHECK(gemm2_weights_scale.sizes()[0] == local_num_experts, "gemm2_weights_scale has incorrect shape."); + TORCH_CHECK(gemm2_weights_scale.sizes()[0] == num_total_local_experts, "gemm2_weights_scale has incorrect dim 0."); TORCH_CHECK(gemm2_weights_scale.sizes()[1] == args.hidden_size / 128, "gemm2_weights_scale has incorrect shape."); TORCH_CHECK(gemm2_weights_scale.sizes()[2] == intermediate_size / 128, "gemm2_weights_scale has incorrect shape."); @@ -372,20 +383,25 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder } } - [[nodiscard]] std::vector> getValidConfigs( - int64_t topK, int64_t hiddenSize, int64_t intermediateSize, int64_t numLocalExperts, int64_t numTokens) const + [[nodiscard]] std::vector> getValidConfigs(int64_t topK, + std::optional const numFusedSharedExpert, int64_t hiddenSize, int64_t intermediateSize, + int64_t numLocalExperts, int64_t numTokens) const { + int64_t const totalExpertsPerToken + = topK + (numFusedSharedExpert.value_or(0) > 0 ? numFusedSharedExpert.value() : 0); + int64_t const numTotalLocalExperts + = numLocalExperts + (numFusedSharedExpert.value_or(0) > 0 ? numFusedSharedExpert.value() : 0); // returns (tileN, config) std::vector> tactics; for (auto& [tileN, runner] : mRunners) { - auto chosen = computeSelectedTileN(mSupportedTileN, numTokens, topK, numLocalExperts); + auto chosen = computeSelectedTileN(mSupportedTileN, numTokens, totalExpertsPerToken, numTotalLocalExperts); if (chosen.find(tileN) == chosen.end()) { continue; } - auto config_indices_per_runner - = runner->getValidConfigIndices(topK, hiddenSize, intermediateSize, numLocalExperts, numTokens); + auto config_indices_per_runner = runner->getValidConfigIndices( + totalExpertsPerToken, hiddenSize, intermediateSize, numTotalLocalExperts, numTokens); for (auto cfg : config_indices_per_runner) { tactics.push_back({tileN, cfg}); @@ -398,9 +414,9 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder std::optional const& routing_bias, at::Tensor const& hidden_states, at::Tensor const& hidden_states_scale, at::Tensor const& gemm1_weights, at::Tensor const& gemm1_weights_scale, at::Tensor const& gemm2_weights, at::Tensor const& gemm2_weights_scale, int64_t num_experts, int64_t top_k, - std::optional const n_group, std::optional const topk_group, int64_t const intermediate_size, - int64_t const local_expert_offset, int64_t const local_num_experts, - std::optional const routed_scaling_factor, int64_t routing_method_type, + std::optional const num_fused_shared_experts, std::optional const n_group, + std::optional const topk_group, int64_t const intermediate_size, int64_t const local_expert_offset, + int64_t const local_num_experts, std::optional const routed_scaling_factor, int64_t routing_method_type, std::vector tile_config_pair, std::optional const& topk_weights, std::optional const& topk_ids, std::optional const& gemm1_clamp_limit = std::nullopt, std::optional const& output = std::nullopt) @@ -411,20 +427,57 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder // Autotuner has requested a default or 'fallback' config index if (tileN == -1 || config == -1) { + int64_t const total_experts_per_token + = top_k + (num_fused_shared_experts.value_or(0) > 0 ? num_fused_shared_experts.value() : 0); + int64_t const num_total_local_experts + = local_num_experts + (num_fused_shared_experts.value_or(0) > 0 ? num_fused_shared_experts.value() : 0); + auto const num_tokens = hidden_states.sizes()[0]; auto const hidden_size = hidden_states.sizes()[1]; - float const avg_tokens_per_expert = static_cast(num_tokens * top_k) / local_num_experts; + float const avg_tokens_per_expert + = static_cast(num_tokens * total_experts_per_token) / num_total_local_experts; tileN = std::clamp(nextPowerOfTwo(avg_tokens_per_expert), mSupportedTileN.front(), mSupportedTileN.back()); - config = mRunners.at(tileN)->getDefaultValidConfigIndex( - top_k, hidden_size, intermediate_size, local_num_experts, num_tokens); + if (num_fused_shared_experts.value_or(0) > 0) + { + // getDefaultValidConfigIndex only pairs the per-GEMM "default" indices without + // re-validating them against the actual problem size. For the inflated fused + // expert/topK counts that can return a config whose kernel is absent (illegal + // memory access at launch). Pick an explicitly-validated config instead -- the + // same set the autotuner draws from -- searching the heuristic tileN first. + config = -1; + std::vector tileN_candidates{static_cast(tileN)}; + for (auto t : mSupportedTileN) + { + if (t != tileN) + tileN_candidates.push_back(t); + } + for (auto t : tileN_candidates) + { + auto valid = mRunners.at(t)->getValidConfigIndices( + total_experts_per_token, hidden_size, intermediate_size, num_total_local_experts, num_tokens); + if (!valid.empty()) + { + tileN = t; + config = valid.front(); + break; + } + } + TLLM_CHECK_WITH_INFO( + config != -1, "No valid TRTLLM-Gen config found for fused shared-expert FP8 block-scale MoE."); + } + else + { + config = mRunners.at(tileN)->getDefaultValidConfigIndex( + total_experts_per_token, hidden_size, intermediate_size, num_total_local_experts, num_tokens); + } } return run_fp8_block_scale_moe(routing_logits, routing_bias, hidden_states, hidden_states_scale, gemm1_weights, - gemm1_weights_scale, gemm2_weights, gemm2_weights_scale, num_experts, top_k, n_group, topk_group, - intermediate_size, local_expert_offset, local_num_experts, routed_scaling_factor, tileN, - routing_method_type, *mRunners.at(tileN), config, topk_weights, topk_ids, gemm1_clamp_limit, output); + gemm1_weights_scale, gemm2_weights, gemm2_weights_scale, num_experts, top_k, num_fused_shared_experts, + n_group, topk_group, intermediate_size, local_expert_offset, local_num_experts, routed_scaling_factor, + tileN, routing_method_type, *mRunners.at(tileN), config, topk_weights, topk_ids, gemm1_clamp_limit, output); } private: diff --git a/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp index 5a2674c5c735..433a477b5b4c 100644 --- a/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp @@ -240,8 +240,9 @@ torch::Tensor fp8_per_tensor_scale_moe_runner(torch::optional con ? (routing_logits.value().scalar_type() == at::ScalarType::Float ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16) : btg::Dtype::Bfloat16; routing_runner.run(args.routing_logits, args.routing_bias, args.num_tokens, args.num_experts, args.top_k, - args.n_group, args.topk_group, args.local_expert_offset, args.local_num_experts, args.routed_scaling_factor, - expert_indexes.data_ptr(), expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), + /* num_fused_shared_expert */ 0, args.n_group, args.topk_group, args.local_expert_offset, + args.local_num_experts, args.routed_scaling_factor, expert_indexes.data_ptr(), + expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), expanded_idx_to_permuted_idx.data_ptr(), nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, args.topk_ids, num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), diff --git a/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp index e641b6546c92..e3d48e6350ac 100644 --- a/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp @@ -293,8 +293,9 @@ torch::Tensor dtype_mxe2m1_block_scale_moe_runner(torch::optional ? (routing_logits.value().scalar_type() == at::ScalarType::Float ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16) : btg::Dtype::Bfloat16; routing_runner.run(args.routing_logits, args.routing_bias, args.num_tokens, args.num_experts, args.top_k, - args.n_group, args.topk_group, args.local_expert_offset, args.local_num_experts, args.routed_scaling_factor, - expert_indexes.data_ptr(), expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), + /* num_fused_shared_expert */ 0, args.n_group, args.topk_group, args.local_expert_offset, + args.local_num_experts, args.routed_scaling_factor, expert_indexes.data_ptr(), + expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), expanded_idx_to_permuted_idx.data_ptr(), nullptr, /*permuted_idx_to_expanded_idx.data_ptr(),*/ permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, args.topk_ids, num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index a8bbb4cb2f09..023fc0c32098 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -820,6 +820,7 @@ def __init__( self, num_experts: int, top_k: int, + num_fused_shared_experts: Optional[int], n_group: Optional[int], topk_group: Optional[int], intermediate_size: int, @@ -835,6 +836,7 @@ def __init__( self.num_experts = num_experts self.top_k = top_k + self.num_fused_shared_experts = num_fused_shared_experts self.n_group = n_group self.topk_group = topk_group self.intermediate_size = intermediate_size @@ -882,10 +884,11 @@ def forward( args.hidden_states_scale, args.gemm1_weights, args.gemm1_weights_scale, args.gemm2_weights, args.gemm2_weights_scale, self.num_experts, self.top_k, - self.n_group, self.topk_group, self.intermediate_size, - self.local_expert_offset, self.local_num_experts, - self.routed_scaling_factor, self.routing_method_type, tactic, - args.topk_weights, args.topk_ids, self.gemm1_clamp_limit_value, + self.num_fused_shared_experts, self.n_group, self.topk_group, + self.intermediate_size, self.local_expert_offset, + self.local_num_experts, self.routed_scaling_factor, + self.routing_method_type, tactic, args.topk_weights, args.topk_ids, + self.gemm1_clamp_limit_value, output) def get_valid_tactics(self, inputs: List[torch.Tensor], @@ -901,6 +904,7 @@ def get_valid_tactics(self, inputs: List[torch.Tensor], tactics = kernel_runner.get_valid_configs( self.top_k, + self.num_fused_shared_experts, hidden_size, self.intermediate_size, self.local_num_experts, @@ -1007,6 +1011,7 @@ def fp8_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], gemm2_weights_scale: torch.Tensor, num_experts: int, top_k: int, + num_fused_shared_experts: Optional[int], n_group: Optional[int], topk_group: Optional[int], intermediate_size: int, @@ -1026,6 +1031,7 @@ def fp8_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], kernel_runner = FP8BlockScaleMoERunner( num_experts, top_k, + num_fused_shared_experts, n_group, topk_group, intermediate_size, @@ -1108,6 +1114,7 @@ def _(routing_logits: torch.Tensor, gemm2_weights_scale: torch.Tensor, num_experts: int, top_k: int, + num_fused_shared_experts: Optional[int], n_group: Optional[int], topk_group: Optional[int], intermediate_size: int, diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 2ccc34a3cda4..4011c868054e 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -1016,7 +1016,7 @@ def __init__(self, and shared_quant_config.group_size is not None): block_size = shared_quant_config.group_size - shared_tp_size, self.shared_output_scale = self._compute_shared_expert_tp_size( + self.shared_tp_size, self.shared_output_scale = self._compute_shared_expert_tp_size( shared_expert_intermediate_size, block_size) self.shared_experts = GatedMLP( @@ -1025,7 +1025,7 @@ def __init__(self, bias=False, dtype=dtype, config=shared_model_config, - overridden_tp_size=shared_tp_size, + overridden_tp_size=self.shared_tp_size, reduce_output=False, use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, ) @@ -1069,6 +1069,9 @@ def _compute_shared_expert_tp_size( if self.use_dp: # If using attention DP, the shared experts also use DP instead of TP. shared_tp_size = 1 + elif hasattr(self.experts, 'num_fused_shared_expert' + ) and self.experts.num_fused_shared_expert > 0: + shared_tp_size = self.mapping.moe_tp_size else: # Due to the restriction of block scale size (i.e., 128), the supported TP sizes only include 1, 2, 4, 8, and 16. # The math.gcd operation ensures that shared_tp_size falls in the supported TP sizes. @@ -1181,53 +1184,61 @@ def _compute_routed_output(): # NOTE: define compiled helpers at module scope to avoid defining decorators inside compiled frames - routed_output, shared_output = maybe_execute_in_parallel( - _compute_routed_output, - _compute_shared_output, - self.event_dict[EventType.Main], - self.event_dict[EventType.MoeShared], - self.aux_stream, - disable_on_compile=True) + if self.shared_experts is not None: + routed_output, shared_output = maybe_execute_in_parallel( + _compute_routed_output, + _compute_shared_output, + self.event_dict[EventType.Main], + self.event_dict[EventType.MoeShared], + self.aux_stream, + disable_on_compile=True) + else: + # Shared experts have been fused into the routed experts (see post_load_weights); + # routed_output already contains the shared expert contribution. + shared_output = None + routed_output = _compute_routed_output() if not do_finalize: return [shared_output, *routed_output] else: - if not isinstance(shared_output, torch.Tensor): + if shared_output is None: + final_hidden_states = routed_output + elif not isinstance(shared_output, torch.Tensor): final_hidden_states = shared_output + routed_output if not self.use_dp and self.mapping.tp_size > 1: final_hidden_states = self.allreduce( final_hidden_states, all_reduce_params=final_all_reduce_params) return final_hidden_states - output_tensor = None - if not self.use_dp and self.mapping.tp_size > 1: - w, actual_kind = torch.ops.trtllm.allocate_output( - shared_output, self.allreduce.output_buffer_kind, - self.mapping.tp_group) - if actual_kind == int(BufferKind.NCCL_WINDOW): - output_tensor = w - if routed_output.dim() == 3: - assert shared_output.numel( - ) * self.top_k == routed_output.numel( - ), 'unmatched tensor shape' - final_hidden_states = moe_reduce_add_shared_output( - routed_output, shared_output, out=output_tensor) else: - assert shared_output.size() == routed_output.size( - ), 'unmatched tensor shape' - if output_tensor is not None: - final_hidden_states = torch.add(shared_output, - routed_output, - out=output_tensor) + output_tensor = None + if not self.use_dp and self.mapping.tp_size > 1: + w, actual_kind = torch.ops.trtllm.allocate_output( + shared_output, self.allreduce.output_buffer_kind, + self.mapping.tp_group) + if actual_kind == int(BufferKind.NCCL_WINDOW): + output_tensor = w + if routed_output.dim() == 3: + assert shared_output.numel( + ) * self.top_k == routed_output.numel( + ), 'unmatched tensor shape' + final_hidden_states = moe_reduce_add_shared_output( + routed_output, shared_output, out=output_tensor) else: - # In-place add to avoid allocating a temporary tensor, reducing peak memory - final_hidden_states = shared_output.add_(routed_output) + assert shared_output.size() == routed_output.size( + ), 'unmatched tensor shape' + if output_tensor is not None: + final_hidden_states = torch.add(shared_output, + routed_output, + out=output_tensor) + else: + # In-place add to avoid allocating a temporary tensor, reducing peak memory + final_hidden_states = shared_output.add_(routed_output) if not self.use_dp and self.mapping.tp_size > 1: final_hidden_states = self.allreduce( final_hidden_states, all_reduce_params=final_all_reduce_params) - return final_hidden_states @@ -2035,3 +2046,24 @@ def setup_aliases(self) -> None: layer.post_attention_layernorm.nvfp4_scale = ( _static_nvfp4_input_scale(getattr(mlp, "gate_up_proj", None))) + + # Note: merge shared expert into FusedMoe module + if idx >= self.config.first_k_dense_replace and idx % self.config.moe_layer_freq == 0: + if hasattr(layer.mlp.experts, 'num_fused_shared_expert' + ) and layer.mlp.experts.num_fused_shared_expert > 0: + layer.mlp.experts.fuse_shared_expert( + layer.mlp.shared_experts) + layer.mlp.shared_experts = None + + # Also process MTP layers if present + if self.draft_model is not None and hasattr(self.draft_model, + 'mtp_layers'): + for layer in self.model.layers[self.config.num_hidden_layers:]: + # MTP layers also have MoE, need to fuse shared experts + if hasattr(layer, 'mlp') and hasattr(layer.mlp, 'experts'): + if hasattr( + layer.mlp.experts, 'num_fused_shared_expert' + ) and layer.mlp.experts.num_fused_shared_expert > 0: + layer.mlp.experts.fuse_shared_expert( + layer.mlp.shared_experts) + layer.mlp.shared_experts = None diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index 1c183310c83a..9042c9c05be3 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -362,6 +362,17 @@ def _supports_load_balancer(self) -> bool: return self.use_dp and self.parallel_size > 1 return self.backend._supports_load_balancer() + @property + def num_fused_shared_expert(self) -> int: + """Expose the backend's fused-shared-expert count so model code (e.g. + DeepseekV3 post_load_weights / shared-expert TP sizing) sees it through + this wrapper. Returns 0 when the backend does not support fusion.""" + return getattr(self.backend, "num_fused_shared_expert", 0) + + def fuse_shared_expert(self, shared_experts): + """Delegate shared-expert fusion to the backend (e.g. TRTLLMGenFusedMoE).""" + return self.backend.fuse_shared_expert(shared_experts) + def validate_config(self): """ Validate configuration parameters diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index e6b192a3f5ed..f47359e80882 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -34,6 +34,7 @@ from ...expert_statistic import ExpertStatistic from ...model_config import ModelConfig from ...utils import ActivationType, AuxStreamType, Fp4QuantizedTensor +from ..gated_mlp import GatedMLP from .interface import AlltoallMethodType, MoE, MoEWeightLoadingMode from .moe_op_backend import MoEOpBackend, get_op_backend from .wide_ep_ft import get_wide_ep_ft_options @@ -325,8 +326,17 @@ def __init__( self.moe_a2a = None self._weights_created = False + self.num_fused_shared_expert = 0 if not model_config.skip_create_weights_in_init: self.create_weights() + self.layer_idx = layer_idx + + if model_config.mapping.dp_size == 1 and self.quant_config.layer_quant_mode.has_fp8_block_scales( + ): + # Not all models that use this backend define shared experts (e.g. non-DeepSeek + # MoEs), so fall back to 0 when the config has no `n_shared_experts`. + self.num_fused_shared_expert = getattr( + model_config.pretrained_config, "n_shared_experts", 0) or 0 def _to_trtllm_gen_activation_type(self, activation_type: ActivationType) -> int: @@ -525,7 +535,10 @@ def create_weights(self): return self.quant_method = self._get_quant_method() - self.quant_method.create_weights(self) + if self.quant_config.layer_quant_mode.has_fp8_block_scales(): + self.quant_method.create_weights(self, self.num_fused_shared_expert) + else: + self.quant_method.create_weights(self) self._weights_created = True self._check_configs() @@ -674,6 +687,11 @@ def _extract_routing_params(self) -> RoutingParams: routed_scaling_factor=None, ) + def fuse_shared_expert(self, shared_experts: GatedMLP): + assert self._weights_created + self.quant_method.fuse_shared_expert(self, shared_experts, + self.num_fused_shared_expert) + def run_moe( self, x: torch.Tensor, @@ -775,6 +793,7 @@ def run_moe( self.w2_weight_scaling_factor, self.num_slots, top_k, + self.num_fused_shared_expert, n_group, topk_group, self.intermediate_size_per_partition, diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py b/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py index cf0e1ad3c6b1..fbab22d78b81 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py @@ -105,6 +105,7 @@ def run_fp8_block_scale_moe( gemm2_weights_scale: torch.Tensor, num_experts: int, top_k: int, + num_fused_shared_experts: Optional[int], n_group: Optional[int], topk_group: Optional[int], intermediate_size: int, @@ -247,6 +248,7 @@ def run_fp8_block_scale_moe( gemm2_weights_scale, num_experts, top_k, + num_fused_shared_experts, n_group, topk_group, intermediate_size, @@ -276,6 +278,7 @@ def run_fp8_block_scale_moe( gemm2_weights_scale, num_experts, top_k, + num_fused_shared_experts, n_group, topk_group, intermediate_size, @@ -592,6 +595,7 @@ def run_fp8_block_scale_moe( gemm2_weights_scale, num_experts, top_k, + num_fused_shared_experts, n_group, topk_group, intermediate_size, @@ -610,6 +614,9 @@ def run_fp8_block_scale_moe( tune_max_num_tokens=8192, use_dp=False, ): + assert not num_fused_shared_experts, ( + "Flashinfer backend does not support fusing shared experts" + ) if gemm1_clamp_limit is not None: raise NotImplementedError( "FlashinferOpBackend.run_fp8_block_scale_moe does not yet " diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index a39b53dcdbdb..281d28a8601b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -37,6 +37,7 @@ from ...mmap_utils import advise_tensor_pageout from ...utils import (ActivationType, replace_parameter_and_save_metadata, swizzle_sf, unswizzle_sf) +from ..gated_mlp import GatedMLP from ..linear import TensorParallelMode, load_weight_shard from .interface import MoEWeightLoadingMode @@ -1067,14 +1068,15 @@ class DeepSeekFP8BlockScalesFusedMoEMethod(FusedMoEMethodBase): eplb_support_status = EplbSupportStatus.NOT_VERIFIED FP8_QUANT_BLOCK_SIZE = 128 - def create_weights(self, module: torch.nn.Module): + def create_weights(self, module: torch.nn.Module, n_shared_experts=0): weight_dtype = torch.float8_e4m3fn - w3_w1_weight_shape = (module.expert_size_per_partition, + w3_w1_weight_shape = (module.expert_size_per_partition + + n_shared_experts, module.intermediate_size_per_partition * 2, module.hidden_size) w2_weight_shape = ( - module.expert_size_per_partition, + module.expert_size_per_partition + n_shared_experts, module.hidden_size, module.intermediate_size_per_partition, ) @@ -1083,7 +1085,7 @@ def create_weights(self, module: torch.nn.Module): cell_div = lambda x, y: (x + y - 1) // y w3_w1_weight_scaling_factor = nn.Parameter(torch.empty( - (module.expert_size_per_partition, + (module.expert_size_per_partition + n_shared_experts, cell_div(module.intermediate_size_per_partition, self.FP8_QUANT_BLOCK_SIZE) * 2, cell_div(w3_w1_weight_shape[2], self.FP8_QUANT_BLOCK_SIZE)), @@ -1093,7 +1095,7 @@ def create_weights(self, module: torch.nn.Module): w3_w1_weight_scaling_factor) w2_weight_scaling_factor = nn.Parameter(torch.empty( - (module.expert_size_per_partition, + (module.expert_size_per_partition + n_shared_experts, cell_div(w2_weight_shape[1], self.FP8_QUANT_BLOCK_SIZE), cell_div(w2_weight_shape[2], self.FP8_QUANT_BLOCK_SIZE)), dtype=torch.float32), @@ -1113,6 +1115,61 @@ def load_weights(self, super().load_weights(module, weights, weight_loading_mode, allow_partial_loading) + def fuse_shared_expert(self, module: torch.nn.Module, + shared_experts: GatedMLP, n_shared_experts: int): + # Fuse the shared expert(s) into the trailing routed-expert slots + # (module.expert_size_per_partition + i). On this trtllm-gen FP8 block-scale path the + # routed-expert weights are stored in plain (non-shuffled) layout, so the shared expert + # weights are reshaped and copied without any extra layout transform. + # gate_up_proj stores [gate(w1); up(w3)]; the routed expert tensor stores [w3; w1]. + w1_weight, w3_weight = shared_experts.gate_up_proj.weight.data.chunk( + 2, dim=0) + w1_weight = w1_weight.view(n_shared_experts, + module.w3_w1_weight.shape[1] // 2, + module.w3_w1_weight.shape[2]) + w3_weight = w3_weight.view(n_shared_experts, + module.w3_w1_weight.shape[1] // 2, + module.w3_w1_weight.shape[2]) + w2_weight = shared_experts.down_proj.weight.view( + module.w2_weight.shape[1], n_shared_experts, + module.w2_weight.shape[2]).permute(1, 0, 2).contiguous() + + w1_w3_weight_scale = shared_experts.gate_up_proj.weight_scale.data + w1_weight_scale, w3_weight_scale = w1_w3_weight_scale.chunk(2, dim=0) + w1_weight_scale = w1_weight_scale.view( + n_shared_experts, module.w3_w1_weight_scaling_factor.shape[1] // 2, + module.w3_w1_weight_scaling_factor.shape[2]) + w3_weight_scale = w3_weight_scale.view( + n_shared_experts, module.w3_w1_weight_scaling_factor.shape[1] // 2, + module.w3_w1_weight_scaling_factor.shape[2]) + # down_proj weight_scale is (hidden_blocks, n_shared * intermediate_blocks); + # the per-expert intermediate blocks are the trailing dim, so split there and + # move the expert axis to the front (mirrors the w2 weight reshape above). + w2_weight_scale = shared_experts.down_proj.weight_scale.data.view( + module.w2_weight_scaling_factor.shape[1], n_shared_experts, + module.w2_weight_scaling_factor.shape[2]).permute(1, 0, + 2).contiguous() + + for i in range(n_shared_experts): + slot = module.expert_size_per_partition + i + # Routed-expert layout is [w3; w1] along dim 0 (see load_expert_w3_w1_weight). + dst_w3_weight, dst_w1_weight = module.w3_w1_weight[slot].chunk( + 2, dim=0) + dst_w3_weight.copy_(w3_weight[i].view(dst_w3_weight.dtype), + non_blocking=True) + dst_w1_weight.copy_(w1_weight[i].view(dst_w1_weight.dtype), + non_blocking=True) + module.w2_weight[slot].copy_(w2_weight[i].view( + module.w2_weight.dtype), + non_blocking=True) + + dst_w3_scale, dst_w1_scale = module.w3_w1_weight_scaling_factor[ + slot].chunk(2, dim=0) + dst_w3_scale.copy_(w3_weight_scale[i].view(dst_w3_scale.dtype)) + dst_w1_scale.copy_(w1_weight_scale[i].view(dst_w1_scale.dtype)) + module.w2_weight_scaling_factor[slot].copy_(w2_weight_scale[i].view( + module.w2_weight_scaling_factor.dtype)) + def setup_quant_scales(self, module: torch.nn.Module): module.quant_scales = FusedMoEQuantScalesDeepSeekFP8BlockScales( fc_weight_scales=module.w3_w1_weight_scaling_factor, diff --git a/tests/unittest/_torch/thop/serial/test_moe.py b/tests/unittest/_torch/thop/serial/test_moe.py index 53c70ee21c04..514a1f9913e9 100644 --- a/tests/unittest/_torch/thop/serial/test_moe.py +++ b/tests/unittest/_torch/thop/serial/test_moe.py @@ -139,30 +139,52 @@ def __init__(self, self.act_type = act_type -def routing_reference(expertLogits, topK, padding): +def routing_reference(expertLogits, topK, padding, num_fused_shared_experts=0): originalDevice = expertLogits.device expertLogits = expertLogits.cpu() numTokens, numExperts = expertLogits.shape assert topK <= numExperts - numTokensPerExpert = torch.zeros(numExperts, dtype=torch.int64) - expandedTokenIdxToExpert = -torch.ones(numTokens * topK, dtype=torch.int64) - expandedTokenIdxToIdxInExpert = -torch.ones(numTokens * topK, - dtype=torch.int64) + numTotalExperts = numExperts + num_fused_shared_experts + totalExpertsPerToken = topK + num_fused_shared_experts + + numTokensPerExpert = torch.zeros(numTotalExperts, dtype=torch.int64) + expandedTokenIdxToExpert = -torch.ones(numTokens * totalExpertsPerToken, + dtype=torch.int64) + expandedTokenIdxToIdxInExpert = -torch.ones( + numTokens * totalExpertsPerToken, dtype=torch.int64) topKLogits, topKIndices = torch.topk(expertLogits, topK, dim=1) + if num_fused_shared_experts > 0: + # Shared experts will use weight of 1.0 + sharedLogits = torch.ones(numTokens, + num_fused_shared_experts, + dtype=topKLogits.dtype) + topKLogits = torch.cat((topKLogits, sharedLogits), dim=1) + assert numTokens == topKLogits.shape[0] + assert totalExpertsPerToken == topKLogits.shape[1] + + # Shared experts will have index starting at number of local experts + sharedIndices = torch.range(numExperts, + numExperts + num_fused_shared_experts - 1, + dtype=topKIndices.dtype) + sharedIndices = torch.unsqueeze(sharedIndices, 0) + sharedIndices = sharedIndices.repeat(numTokens, 1) + topKIndices = torch.cat((topKIndices, sharedIndices), dim=1) + assert numTokens == topKIndices.shape[0] + assert totalExpertsPerToken == topKIndices.shape[1] for tokenIdx in range(numTokens): - for k in range(topK): - expandedIdx = tokenIdx * topK + k + for k in range(totalExpertsPerToken): + expandedIdx = tokenIdx * totalExpertsPerToken + k expertIndex = topKIndices[tokenIdx, k] expandedTokenIdxToExpert[expandedIdx] = expertIndex expandedTokenIdxToIdxInExpert[expandedIdx] = numTokensPerExpert[ expertIndex] numTokensPerExpert[expertIndex] += 1 - paddedTokensPerExpertPrefixSum = torch.zeros(numExperts + 1, + paddedTokensPerExpertPrefixSum = torch.zeros(numTotalExperts + 1, dtype=torch.int64) - for ii in range(numExperts): + for ii in range(numTotalExperts): def divUpMul(a, b): return (a + b - 1) // b * b @@ -170,16 +192,16 @@ def divUpMul(a, b): paddedTokensPerExpertPrefixSum[ ii + 1] = paddedTokensPerExpertPrefixSum[ii] + divUpMul( numTokensPerExpert[ii], padding) - permutedBufferSize = paddedTokensPerExpertPrefixSum[numExperts] + permutedBufferSize = paddedTokensPerExpertPrefixSum[numTotalExperts] - expandedTokenIdxToPermutedIdx = -torch.ones(numTokens * topK, - dtype=torch.int64) + expandedTokenIdxToPermutedIdx = -torch.ones( + numTokens * totalExpertsPerToken, dtype=torch.int64) permutedIdxToExpandedIdx = -torch.ones(permutedBufferSize, dtype=torch.int64) permutedIdxToTokenIdx = -torch.ones(permutedBufferSize, dtype=torch.int64) for tokenIdx in range(numTokens): - for k in range(topK): - expandedIdx = tokenIdx * topK + k + for k in range(totalExpertsPerToken): + expandedIdx = tokenIdx * totalExpertsPerToken + k expert = expandedTokenIdxToExpert[expandedIdx] offsetWithinExpert = expandedTokenIdxToIdxInExpert[expandedIdx] offsetForExpert = paddedTokensPerExpertPrefixSum[expert] @@ -188,6 +210,7 @@ def divUpMul(a, b): expandedTokenIdxToPermutedIdx[expandedIdx] = permutedIdx permutedIdxToExpandedIdx[permutedIdx] = expandedIdx permutedIdxToTokenIdx[permutedIdx] = tokenIdx + return { "paddedTokensPerExpertPrefixSum": paddedTokensPerExpertPrefixSum.to(originalDevice), @@ -257,7 +280,8 @@ def routing_reference_no_aux(expert_logits, top_k_groups, routed_scaling, padding, - use_routing_scales_on_input=False): + use_routing_scales_on_input=False, + num_fused_shared_experts=0): routing_logits = expert_logits.to(dtype=torch.float, device='cuda') if use_routing_scales_on_input: # if using routing scales on input, topK == 1 and the score is a plain sigmoid @@ -265,7 +289,8 @@ def routing_reference_no_aux(expert_logits, else: scores = noaux_tc_ref(routing_logits, routing_bias, n_groups, top_k_groups, top_k, routed_scaling) - permute_info = routing_reference(scores, top_k, padding) + permute_info = routing_reference(scores, top_k, padding, + num_fused_shared_experts) return permute_info, scores @@ -897,8 +922,10 @@ class TestMoeFP8: """ @pytest.mark.parametrize("num_tokens", [16, 64, 1024]) - @pytest.mark.parametrize("expert_info", [(32, 8, 4, 8), (72, 1, 1, 6), - (256, 8, 4, 8)]) + @pytest.mark.parametrize("expert_info", + [(32, 8, 4, 8, 0), (72, 1, 1, 6, 0), + (256, 8, 4, 8, 0), (32, 8, 4, 8, 1), + (256, 8, 4, 8, 2), (72, 1, 1, 6, 2)]) @pytest.mark.parametrize("hidden_size", [512]) @pytest.mark.parametrize("intermediate_size", [512]) def test_autotune(self, num_tokens: int, expert_info: Tuple[int, int, int, @@ -913,7 +940,9 @@ def test_autotune(self, num_tokens: int, expert_info: Tuple[int, int, int, use_topk_as_input=False) @pytest.mark.parametrize("num_tokens", [16]) - @pytest.mark.parametrize("expert_info", [(32, 8, 4, 8), (384, 1, 1, 8)]) + @pytest.mark.parametrize("expert_info", [(32, 8, 4, 8, 0), (72, 1, 1, 6, 0), + (72, 1, 1, 6, 1), + (72, 1, 1, 6, 2)]) @pytest.mark.parametrize("hidden_size", [512]) @pytest.mark.parametrize("intermediate_size", [512]) @pytest.mark.parametrize("use_topk_as_input", [False, True], @@ -939,7 +968,7 @@ def run_moe_fp8_test(self, num_tokens: int, expert_info: Tuple[int, int, # # Data Generation # - num_experts, n_groups, top_k_groups, top_k = expert_info + num_experts, n_groups, top_k_groups, top_k, num_fused_shared_experts = expert_info padding = 8 routed_scaling = 2.5 routing_method_type = RoutingMethodType.DeepSeekV3 @@ -948,6 +977,14 @@ def run_moe_fp8_test(self, num_tokens: int, expert_info: Tuple[int, int, assert top_k <= 8 assert num_experts % 4 == 0 + total_experts_per_token = top_k + num_fused_shared_experts + num_experts_total = num_experts + num_fused_shared_experts + + if use_topk_as_input and num_fused_shared_experts > 0: + pytest.skip( + "use_topk_as_input is tested only with num_fused_shared_experts=0" + ) + if are_groups_valid(top_k_groups, n_groups): assert top_k_groups <= 4 assert num_experts > n_groups @@ -966,27 +1003,34 @@ def run_moe_fp8_test(self, num_tokens: int, expert_info: Tuple[int, int, (hidden_size // 128, num_tokens), device='cuda').to(torch.float) gemm1_weights = torch.randn( - (num_experts, 2 * intermediate_size, hidden_size), + (num_experts_total, 2 * intermediate_size, hidden_size), device='cuda').to(torch.float8_e4m3fn) gemm1_scales = 2 * torch.rand( - (num_experts, 2 * intermediate_size // 128, hidden_size // 128), + (num_experts_total, 2 * intermediate_size // 128, + hidden_size // 128), device='cuda').to(torch.float) gemm2_weights = torch.randn( - (num_experts, hidden_size, intermediate_size), + (num_experts_total, hidden_size, intermediate_size), device='cuda').to(torch.float8_e4m3fn) gemm2_scales = 2 * torch.rand( - (num_experts, hidden_size // 128, intermediate_size // 128), + (num_experts_total, hidden_size // 128, intermediate_size // 128), device='cuda').to(torch.float) - permute_info, scores = routing_reference_no_aux(expert_logits, - routing_bias, top_k, - n_groups, top_k_groups, - routed_scaling, padding) - - args = moe_args(num_tokens, num_experts, hidden_size, intermediate_size, - top_k, padding, hidden_states, hidden_states_scale, - None, scores, gemm1_weights, gemm1_scales, None, - gemm2_weights, gemm2_scales, None, permute_info, False) + permute_info, scores = routing_reference_no_aux( + expert_logits, + routing_bias, + top_k, + n_groups, + top_k_groups, + routed_scaling, + padding, + num_fused_shared_experts=num_fused_shared_experts) + + args = moe_args(num_tokens, num_experts_total, hidden_size, + intermediate_size, total_experts_per_token, padding, + hidden_states, hidden_states_scale, None, scores, + gemm1_weights, gemm1_scales, None, gemm2_weights, + gemm2_scales, None, permute_info, False) if use_topk_as_input: topk_ids = permute_info["topKIndices"].to(torch.int32) @@ -1001,9 +1045,9 @@ def run_moe_fp8_test(self, num_tokens: int, expert_info: Tuple[int, int, output = torch.ops.trtllm.fp8_block_scale_moe_runner( expert_logits, routing_bias, hidden_states, hidden_states_scale, gemm1_weights, gemm1_scales, gemm2_weights, gemm2_scales, - num_experts, top_k, n_groups, top_k_groups, intermediate_size, - 0, num_experts, routed_scaling, routing_method_type, - topk_weights, topk_ids) + num_experts, top_k, num_fused_shared_experts, n_groups, + top_k_groups, intermediate_size, 0, num_experts, routed_scaling, + routing_method_type, topk_weights, topk_ids) torch.cuda.synchronize() output_dequant_actual = output.to(torch.float) # From 289fa1fced5b29fa94247a8e1f1c5d579526736a Mon Sep 17 00:00:00 2001 From: leslief Date: Thu, 11 Jun 2026 22:27:18 -0700 Subject: [PATCH 02/19] Add E2E testing and fix Signed-off-by: leslief --- .../trtllmGenKernels/blockScaleMoe/runner.cu | 2 +- .../defs/accuracy/test_llm_api_pytorch.py | 39 +++++++++++++++++++ tests/unittest/_torch/thop/serial/test_moe.py | 8 ++-- 3 files changed, 45 insertions(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu index a86ecc49124b..bde059e9170f 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu @@ -69,7 +69,7 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 bool useDeepSeekFp8, RoutingMethodType routingMethodType, cudaStream_t stream, btg::Dtype dtypeRoutingLogits, btg::Dtype dtypeRoutingBias) { - if (routingMethodType == RoutingMethodType::DeepSeekV3 && nGroup <= 1) + if (routingMethodType == RoutingMethodType::DeepSeekV3 && nGroup <= 1 && numFusedSharedExpert == 0) { // DeepSeek no-groups case: use routingCustom with SigmoidBias preprocess. // NOTE: routingCustom does not implement fused shared experts; when fusion is diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 2305c0852664..1250483d45f6 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1910,6 +1910,45 @@ def test_fp8_block_scales(self, mtp, fp8kv, attention_dp, cuda_graph, task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + def test_fp8_block_scales_fusion_tmp(self): + # TEMP (PR#11143 validation, NOT for commit): force the TRT-LLM Gen MoE + # backend + disable attention DP so the shared-expert -> sparse-expert + # fusion path (dp_size==1 + fp8_block_scales gate) is actually exercised + # end-to-end. The stock test uses DEEPGEMM (no fusion) on SM>=100. + # TP from env TRTLLM_VALIDATE_TP (default 4). Run from tests/integration/defs: + # LLM_MODELS_ROOT=/llm-models TRTLLM_VALIDATE_TP=4 python3 -m pytest -s -v \ + # accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_fusion_tmp + tp = int(os.environ.get("TRTLLM_VALIDATE_TP", "4")) + # TRTLLM_VALIDATE_BACKEND: TRTLLM (fusion) | DEEPGEMM (stock, no fusion). + # DEEPGEMM is used to prove the TP>1 tunable_allreduce hang is independent + # of the shared-expert fusion (fusion is not active with DEEPGEMM). + moe_backend = os.environ.get("TRTLLM_VALIDATE_BACKEND", "TRTLLM") + # TRTLLM_VALIDATE_AR: override allreduce_strategy (AUTO|NCCL|...). Used to + # work around the AUTO-strategy collective deadlock seen at TP>1 on B300. + ar_strategy = os.environ.get("TRTLLM_VALIDATE_AR", "AUTO") + # TRTLLM_VALIDATE_ADP: enable_attention_dp (0|1). Diagnostic only: with 1, + # attention uses DP (no pure-TP o_proj allreduce) and fusion is OFF; used to + # bound whether the TP=4 hang is specific to the attention_dp=False TP path. + attention_dp = os.environ.get("TRTLLM_VALIDATE_ADP", "0") == "1" + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75) + pytorch_config = dict( + disable_overlap_scheduler=True, + cuda_graph_config=None, + moe_config=MoeConfig(backend=moe_backend), + allreduce_strategy=ar_strategy, + ) + with LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", + tensor_parallel_size=tp, + moe_expert_parallel_size=1, + kv_cache_config=kv_cache_config, + **pytorch_config, + enable_attention_dp=attention_dp) as llm: + + assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + @skip_pre_blackwell @parametrize_with_ids("torch_compile", [False]) @parametrize_with_ids( diff --git a/tests/unittest/_torch/thop/serial/test_moe.py b/tests/unittest/_torch/thop/serial/test_moe.py index 514a1f9913e9..bfbe34ba7358 100644 --- a/tests/unittest/_torch/thop/serial/test_moe.py +++ b/tests/unittest/_torch/thop/serial/test_moe.py @@ -940,9 +940,11 @@ def test_autotune(self, num_tokens: int, expert_info: Tuple[int, int, int, use_topk_as_input=False) @pytest.mark.parametrize("num_tokens", [16]) - @pytest.mark.parametrize("expert_info", [(32, 8, 4, 8, 0), (72, 1, 1, 6, 0), - (72, 1, 1, 6, 1), - (72, 1, 1, 6, 2)]) + @pytest.mark.parametrize("expert_info", + [(32, 8, 4, 8, 0), (384, 1, 1, 8, 0), + (72, 1, 1, 6, 0), (72, 1, 1, 6, 1), + (72, 1, 1, 6, 2), (384, 1, 1, 8, 1), + (384, 1, 1, 8, 2)]) @pytest.mark.parametrize("hidden_size", [512]) @pytest.mark.parametrize("intermediate_size", [512]) @pytest.mark.parametrize("use_topk_as_input", [False, True], From 4ff0d64028b3105316c960523ceb18074406c5fd Mon Sep 17 00:00:00 2001 From: leslief Date: Fri, 12 Jun 2026 06:53:28 +0000 Subject: [PATCH 03/19] [None][feat] Enhance the check to enable this fusion path Enhance the check to enable this fusion path Signed-off-by: leslief --- .../_torch/modules/fused_moe/fused_moe_trtllm_gen.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index f47359e80882..bb208d5bc839 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -36,7 +36,7 @@ from ...utils import ActivationType, AuxStreamType, Fp4QuantizedTensor from ..gated_mlp import GatedMLP from .interface import AlltoallMethodType, MoE, MoEWeightLoadingMode -from .moe_op_backend import MoEOpBackend, get_op_backend +from .moe_op_backend import MoEOpBackend, TRTLLMOpBackend, get_op_backend from .wide_ep_ft import get_wide_ep_ft_options # isort: off @@ -331,7 +331,13 @@ def __init__( self.create_weights() self.layer_idx = layer_idx - if model_config.mapping.dp_size == 1 and self.quant_config.layer_quant_mode.has_fp8_block_scales( + # Set TLLM_MOE_DISABLE_SHARED_EXPERT_FUSION=1 to disable fusing the + # shared experts into the routed-expert grouped GEMM. + fusion_disabled = os.environ.get( + "TLLM_MOE_DISABLE_SHARED_EXPERT_FUSION", "0") == "1" + # Only the trtllm op backend implements fused shared experts + on_trtllm_backend = isinstance(self.op_backend, TRTLLMOpBackend) + if not fusion_disabled and on_trtllm_backend and model_config.mapping.dp_size == 1 and self.quant_config.layer_quant_mode.has_fp8_block_scales( ): # Not all models that use this backend define shared experts (e.g. non-DeepSeek # MoEs), so fall back to 0 when the config has no `n_shared_experts`. From 9753441f33c8a9fb011e395cde4e66d30465afa3 Mon Sep 17 00:00:00 2001 From: leslief Date: Wed, 17 Jun 2026 05:07:28 +0000 Subject: [PATCH 04/19] [None][test] bench_moe: shared-expert fusion support Signed-off-by: leslief --- tests/microbenchmarks/bench_moe/build.py | 89 +++++++++++++++++++ .../microbenchmarks/bench_moe/case_runner.py | 22 +++-- tests/microbenchmarks/bench_moe/cli.py | 24 +++++ tests/microbenchmarks/bench_moe/mapping.py | 19 +++- tests/microbenchmarks/bench_moe/specs.py | 6 ++ 5 files changed, 153 insertions(+), 7 deletions(-) diff --git a/tests/microbenchmarks/bench_moe/build.py b/tests/microbenchmarks/bench_moe/build.py index 0715bc783b23..a5ad3dccb435 100644 --- a/tests/microbenchmarks/bench_moe/build.py +++ b/tests/microbenchmarks/bench_moe/build.py @@ -103,6 +103,63 @@ def _create_moe_for_benchmark(**kwargs): return create_moe(**kwargs) +class _UnfusedSharedMoE(torch.nn.Module): + """Pre-fusion baseline: routed MoE plus a separate shared-expert GatedMLP. + + The routed MoE is built without fused shared experts; a standalone shared-expert + ``GatedMLP`` is summed onto its output, mirroring the pre-#11143 model path. The + routed MoE and the shared GatedMLP are run with ``maybe_execute_in_parallel`` on + a dedicated aux stream (same mechanism DeepseekV3MoE uses for the unfused path), + so the routed grouped GEMM and the shared MLP can overlap. Multi-stream is + enabled only under CUDA graph (``use_cuda_graph``), matching the real engine + (cuda_graph_runner sets with_multi_stream(True) during capture); in eager the + two run serially, since stream-switch host overhead makes overlap a net loss. + Exposes the same ``forward(x, router_logits, **kwargs)`` signature as the MoE; + introspection attributes (backend / routing_method / ...) delegate to the MoE. + """ + + def __init__( + self, + moe: torch.nn.Module, + shared_mlp: torch.nn.Module, + use_cuda_graph: bool = False, + ): + super().__init__() + self.moe = moe + self.shared_mlp = shared_mlp + # Overlap only when the case is timed under CUDA graph (mirrors the real + # model, which only enables multi-stream during graph capture). + self._multi_stream = bool(use_cuda_graph) + self.aux_stream = torch.cuda.Stream() + self.event_main = torch.cuda.Event() + self.event_shared = torch.cuda.Event() + + def __getattr__(self, name): + # nn.Module.__getattr__ resolves registered submodules/params/buffers + # (incl. ``moe``/``shared_mlp``); anything else (routing_method, backend, + # comm, scheduler, ...) is delegated to the wrapped MoE for introspection. + try: + return super().__getattr__(name) + except AttributeError: + return getattr(super().__getattr__("moe"), name) + + def forward(self, x, router_logits, **kwargs): + from tensorrt_llm._torch.modules.multi_stream_utils import ( + maybe_execute_in_parallel, + with_multi_stream, + ) + + with with_multi_stream(self._multi_stream): + routed, shared = maybe_execute_in_parallel( + lambda: self.moe.forward(x, router_logits, **kwargs), + lambda: self.shared_mlp(x), + self.event_main, + self.event_shared, + self.aux_stream, + ) + return routed + shared + + def _build_moe_module( *, model: ModelSpec, @@ -121,6 +178,20 @@ def _build_moe_module( Returns ``(moe_module, routing_logits_dtype)``. """ + # Shared-expert support (PR #11143) currently only works in TTP: attention TP + # (dp_size==1, so the fusion gate fires) + MoE TP (moe_ep_size==1; the EP fused + # path is unimplemented). Reject other multi-GPU layouts (DEP/TEP/DTP) up front + # instead of silently dropping the shared experts or hitting the dead EP path. + # (Single GPU resolves to dp_size==1/moe_ep_size==1 for every mode, so it passes.) + if model.n_shared_experts > 0 and (mapping.dp_size != 1 or mapping.moe_ep_size != 1): + raise NotImplementedError( + f"bench_moe shared-expert support (n_shared_experts={model.n_shared_experts}) " + f"only supports TTP (attention TP + MoE TP: dp_size==1 and moe_ep_size==1). " + f"Got parallel_mode={config.parallel_mode!r} -> dp_size={mapping.dp_size}, " + f"moe_ep_size={mapping.moe_ep_size}. Re-run with --parallel_mode TTP " + f"(or drop --n_shared_experts)." + ) + if enable_perfect_router: os.environ["ENABLE_PERFECT_ROUTER"] = "1" else: @@ -206,4 +277,22 @@ def _build_moe_module( moe.post_load_weights() moe.cuda(f"cuda:{torch.cuda.current_device()}") + # "unfused" baseline: the routed MoE above was built without fused shared + # experts (mapping.py forces pretrained_config.n_shared_experts=0 in this mode); + # add a standalone shared-expert GatedMLP and sum it, mirroring the pre-fusion + # model path so the benchmark can compare against the fused path. + if model.n_shared_experts > 0 and model.shared_expert_mode == "unfused": + from tensorrt_llm._torch.modules.gated_mlp import GatedMLP + + shared_mlp = GatedMLP( + hidden_size=mc.hidden_size, + intermediate_size=model.n_shared_experts * mc.intermediate_size, + bias=False, + dtype=dtype, + config=model_config, + reduce_output=False, + ) + shared_mlp.cuda(f"cuda:{torch.cuda.current_device()}") + moe = _UnfusedSharedMoE(moe, shared_mlp, use_cuda_graph=use_cuda_graph) + return moe, routing_logits_dtype diff --git a/tests/microbenchmarks/bench_moe/case_runner.py b/tests/microbenchmarks/bench_moe/case_runner.py index 4b7acf8e35fb..3bcc8c30a93b 100644 --- a/tests/microbenchmarks/bench_moe/case_runner.py +++ b/tests/microbenchmarks/bench_moe/case_runner.py @@ -693,11 +693,23 @@ def _run_one_candidate( mapping=mapping, moe_backend=config.backend, use_cuda_graph=bool(config.cuda_graph), - # Symmetric-memory comm backends (e.g. NVLINK_ONE_SIDED) size their - # workspace from max_num_tokens and require every rank to allocate the - # same size, so use the global per-rank maximum rather than this rank's - # local token count (which differs under uneven attention-DP shards). - max_num_tokens=max(int(max(per_rank)) if per_rank else 0, 1), + # Under CUDA graph, capture cannot cross the multi-chunk MoE scheduler + # path (`_forward_multiple_chunks` does a host `if` on a CUDA + # `chunked_used` tensor -> D2H sync, and overlaps chunks on an aux stream + # -- both illegal during capture). Production only captures the + # single-chunk path (num_chunks==1); calculate_num_chunks sums all ranks + # (num_rows=sum(all_rank_num_tokens)), so size max_num_tokens to the + # global token count under CUDA graph to force num_chunks==1. + # Otherwise use the global per-rank maximum: symmetric-memory comm + # backends (e.g. NVLINK_ONE_SIDED) size their workspace from + # max_num_tokens and require every rank to allocate the same size, so + # this rank's local token count (which differs under uneven + # attention-DP shards) must not be used. + max_num_tokens=( + max(sum(all_rank_num_tokens), 1) + if config.cuda_graph + else max(int(max(per_rank)) if per_rank else 0, 1) + ), use_low_precision_moe_combine=bool(config.use_low_precision_moe_combine), enable_perfect_router=enable_perfect_router, dtype=act_dtype, diff --git a/tests/microbenchmarks/bench_moe/cli.py b/tests/microbenchmarks/bench_moe/cli.py index 729eb7a10694..0b516983ada1 100644 --- a/tests/microbenchmarks/bench_moe/cli.py +++ b/tests/microbenchmarks/bench_moe/cli.py @@ -195,6 +195,24 @@ def parse_args() -> argparse.Namespace: default=None, help="DeepSeek-style number of routing groups kept per token.", ) + model_group.add_argument( + "--n_shared_experts", + type=int, + default=None, + help="Number of shared experts to fuse into the routed-expert grouped " + "GEMM (DeepSeek-style). Only takes effect on the TRTLLM backend with " + "FP8_BLOCK_SCALES and dp_size==1; ignored by other backends. Default 0.", + ) + model_group.add_argument( + "--shared_expert_mode", + type=lambda s: str(s).lower(), + default="fused", + choices=["fused", "unfused"], + help="How shared experts (n_shared_experts>0) are realized: 'fused' folds " + "them into the routed grouped GEMM (PR #11143); 'unfused' runs the routed " + "MoE plus a separate shared GatedMLP and sums them (pre-fusion baseline, " + "for measuring fusion's net benefit). Default fused.", + ) model_group.add_argument( "--quant", type=lambda s: QuantAlgo[str(s).upper()] if s is not None else None, @@ -526,6 +544,8 @@ def _resolve_model_from_args(args: argparse.Namespace) -> ModelSpec: routing_method=routing, n_group=args.n_group, topk_group=args.topk_group, + n_shared_experts=int(args.n_shared_experts) if args.n_shared_experts is not None else 0, + shared_expert_mode=args.shared_expert_mode, ) # Built-in model with optional per-field overrides. @@ -549,6 +569,10 @@ def _resolve_model_from_args(args: argparse.Namespace) -> ModelSpec: routing_method=routing, n_group=args.n_group if args.n_group is not None else base.n_group, topk_group=args.topk_group if args.topk_group is not None else base.topk_group, + n_shared_experts=int(args.n_shared_experts) + if args.n_shared_experts is not None + else base.n_shared_experts, + shared_expert_mode=args.shared_expert_mode, swiglu_alpha=base.swiglu_alpha, swiglu_beta=base.swiglu_beta, swiglu_limit=base.swiglu_limit, diff --git a/tests/microbenchmarks/bench_moe/mapping.py b/tests/microbenchmarks/bench_moe/mapping.py index 545b9109edcd..0d18b0f3069e 100644 --- a/tests/microbenchmarks/bench_moe/mapping.py +++ b/tests/microbenchmarks/bench_moe/mapping.py @@ -147,7 +147,11 @@ def _create_routing_method( def _build_pretrained_config( - num_experts: int, hidden_size: int, intermediate_size: int, dtype: torch.dtype + num_experts: int, + hidden_size: int, + intermediate_size: int, + dtype: torch.dtype, + n_shared_experts: int = 0, ) -> PretrainedConfig: """Construct a HF-style ``PretrainedConfig`` for ``ConfigurableMoE``.""" pc = PretrainedConfig() @@ -155,6 +159,10 @@ def _build_pretrained_config( pc.hidden_size = hidden_size pc.intermediate_size = intermediate_size pc.torch_dtype = dtype + # TRTLLMGenFusedMoE reads ``n_shared_experts`` off the pretrained config to + # decide num_fused_shared_expert (shared-expert fusion). Keep it at 0 unless + # explicitly requested so non-fusion cases are unaffected. + pc.n_shared_experts = n_shared_experts return pc @@ -169,8 +177,15 @@ def _build_model_config( dtype: torch.dtype, ) -> ModelConfig: """Build ``ModelConfig`` plumbed into ``create_moe``.""" + # In "unfused" mode the routed MoE must NOT fuse the shared experts (a separate + # GatedMLP is built and summed in build.py instead), so the fused count is 0. + fused_n_shared = model.n_shared_experts if model.shared_expert_mode != "unfused" else 0 pretrained_config = _build_pretrained_config( - model.num_experts, model.hidden_size, model.intermediate_size, dtype + model.num_experts, + model.hidden_size, + model.intermediate_size, + dtype, + n_shared_experts=fused_n_shared, ) quant_algo = model.quant_algo_enum diff --git a/tests/microbenchmarks/bench_moe/specs.py b/tests/microbenchmarks/bench_moe/specs.py index 36134c9b429a..6b2043d69c22 100644 --- a/tests/microbenchmarks/bench_moe/specs.py +++ b/tests/microbenchmarks/bench_moe/specs.py @@ -108,6 +108,12 @@ class ModelSpec: routing_method: str n_group: Optional[int] = None topk_group: Optional[int] = None + n_shared_experts: int = 0 + # How shared experts are realized when n_shared_experts > 0: + # "fused" -> fold them into the routed-expert grouped GEMM (PR #11143). + # "unfused" -> routed MoE (no fusion) + a separate shared GatedMLP, summed + # (the pre-fusion baseline, for measuring fusion's net benefit). + shared_expert_mode: str = "fused" swiglu_alpha: float = 1.0 swiglu_beta: float = 0.0 swiglu_limit: float = float("inf") From a9b2d25be47b7dbccc7566b8f4e6264b3992bf0a Mon Sep 17 00:00:00 2001 From: leslief Date: Mon, 22 Jun 2026 10:20:29 +0000 Subject: [PATCH 05/19] [None][test] bench_moe: route unfused shared GatedMLP FP8 GEMM via cute_dsl to avoid IMA Signed-off-by: leslief --- tests/microbenchmarks/bench_moe/build.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/microbenchmarks/bench_moe/build.py b/tests/microbenchmarks/bench_moe/build.py index a5ad3dccb435..90c366be3011 100644 --- a/tests/microbenchmarks/bench_moe/build.py +++ b/tests/microbenchmarks/bench_moe/build.py @@ -291,6 +291,12 @@ def _build_moe_module( dtype=dtype, config=model_config, reduce_output=False, + # Route the shared GatedMLP's FP8 block-scale GEMM through cute_dsl + # (cute_dsl_fp8_gemm_blackwell) instead of the SM100f default DeepGEMM + # `fp8_swap_ab_gemm`. The DeepGEMM path hits an intermittent + # cudaErrorIllegalAddress when the routed TRTLLM-Gen MoE and this shared + # GatedMLP run in the same process across growing token counts. + use_cute_dsl_blockscaling_mm=True, ) shared_mlp.cuda(f"cuda:{torch.cuda.current_device()}") moe = _UnfusedSharedMoE(moe, shared_mlp, use_cuda_graph=use_cuda_graph) From a8cd40d1dc986578f8aa7217fc1d2313432cd360 Mon Sep 17 00:00:00 2001 From: leslief Date: Sat, 11 Jul 2026 10:43:38 +0000 Subject: [PATCH 06/19] [None][fix] Restrict fused shared-expert MoE to tileN>=32 to avoid small-tile dynB IMA The small-tile (tileN 8/16) dynB TRTLLM-Gen batched-GEMM cubins flakily hit an illegal memory access (garbage TMA-descriptor pointer, MMU fault in the gemm2 K-loop) when shared experts are fused into the grouped GEMM. tileN >= 32 is unaffected (10/10 clean vs minutes-to-crash baseline on B300 TP=4). The bug reproduces unfused on stock main @ 24d598862d with the standalone capture-replay kit, so this is a kernel-side defect; keep the fused path on tileN >= 32 until the cubin fix lands (nvbug TBD). TLLM_MOE_FUSED_MIN_TILEN overrides the threshold (0 disables) for A/B experiments. Signed-off-by: leslief --- cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp | 26 ++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp index 8dcf5875317a..2b7d25102ba0 100644 --- a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp @@ -23,6 +23,7 @@ #include #include +#include #include #include @@ -391,10 +392,25 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder = topK + (numFusedSharedExpert.value_or(0) > 0 ? numFusedSharedExpert.value() : 0); int64_t const numTotalLocalExperts = numLocalExperts + (numFusedSharedExpert.value_or(0) > 0 ? numFusedSharedExpert.value() : 0); + // WAR: the small-tile (tileN 8/16) dynB TRTLLM-Gen batched-GEMM cubins flakily hit an + // illegal memory access (garbage TMA-descriptor pointer, MMU fault in the gemm2 K-loop) + // when shared experts are fused into the grouped GEMM (num_fused_shared_experts > 0); + // tileN >= 32 is unaffected (10/10 clean vs minutes-to-crash baseline on B300 TP=4). + // Restrict the fused path to tileN >= 32 until the kernel-side fix lands (nvbug TBD). + // TLLM_MOE_FUSED_MIN_TILEN overrides the threshold (0 disables) for A/B experiments. + static int const fusedMinTileN = []() + { + char const* env = std::getenv("TLLM_MOE_FUSED_MIN_TILEN"); + return env != nullptr ? std::atoi(env) : 32; + }(); // returns (tileN, config) std::vector> tactics; for (auto& [tileN, runner] : mRunners) { + if (numFusedSharedExpert.value_or(0) > 0 && tileN < fusedMinTileN) + { + continue; + } auto chosen = computeSelectedTileN(mSupportedTileN, numTokens, totalExpertsPerToken, numTotalLocalExperts); if (chosen.find(tileN) == chosen.end()) { @@ -453,8 +469,18 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder if (t != tileN) tileN_candidates.push_back(t); } + // Same small-tile exclusion as getValidConfigs (see the WAR comment there). + static int const fusedMinTileNFallback = []() + { + char const* env = std::getenv("TLLM_MOE_FUSED_MIN_TILEN"); + return env != nullptr ? std::atoi(env) : 32; + }(); for (auto t : tileN_candidates) { + if (t < fusedMinTileNFallback) + { + continue; + } auto valid = mRunners.at(t)->getValidConfigIndices( total_experts_per_token, hidden_size, intermediate_size, num_total_local_experts, num_tokens); if (!valid.empty()) From 9b0a2080ba5e486ac2885c957a4a925b6dc6c740 Mon Sep 17 00:00:00 2001 From: leslief Date: Sat, 11 Jul 2026 11:31:44 +0000 Subject: [PATCH 07/19] [None][chore] Fix yapf formatting Signed-off-by: leslief --- tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index 023fc0c32098..400067ea6c59 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -888,8 +888,7 @@ def forward( self.intermediate_size, self.local_expert_offset, self.local_num_experts, self.routed_scaling_factor, self.routing_method_type, tactic, args.topk_weights, args.topk_ids, - self.gemm1_clamp_limit_value, - output) + self.gemm1_clamp_limit_value, output) def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, From c7345fae4fdf0ad6711012cec3fe388430433665 Mon Sep 17 00:00:00 2001 From: leslief Date: Mon, 13 Jul 2026 06:47:53 +0000 Subject: [PATCH 08/19] [None][feat] Make TRTLLM-Gen shared-expert fusion opt-in via TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION Signed-off-by: leslief --- .../modules/fused_moe/fused_moe_trtllm_gen.py | 19 ++++++++++++++----- .../defs/accuracy/test_llm_api_pytorch.py | 4 ++++ tests/microbenchmarks/bench_moe/build.py | 7 +++++++ 3 files changed, 25 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index bb208d5bc839..5d16e7ac26cc 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -331,18 +331,27 @@ def __init__( self.create_weights() self.layer_idx = layer_idx - # Set TLLM_MOE_DISABLE_SHARED_EXPERT_FUSION=1 to disable fusing the - # shared experts into the routed-expert grouped GEMM. - fusion_disabled = os.environ.get( - "TLLM_MOE_DISABLE_SHARED_EXPERT_FUSION", "0") == "1" + # Fusing the shared experts into the routed-expert grouped GEMM is opt-in: + # set TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION=1 to enable it. The benefit is + # workload-dependent (small decode batches gain, large prefill chunks lose the + # aux-stream overlap of the unfused path), and the fused path additionally + # restricts tactics to tileN>=32 to avoid a small-tile dynB kernel defect. + fusion_enabled = os.environ.get("TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION", + "0") == "1" # Only the trtllm op backend implements fused shared experts on_trtllm_backend = isinstance(self.op_backend, TRTLLMOpBackend) - if not fusion_disabled and on_trtllm_backend and model_config.mapping.dp_size == 1 and self.quant_config.layer_quant_mode.has_fp8_block_scales( + if fusion_enabled and on_trtllm_backend and model_config.mapping.dp_size == 1 and self.quant_config.layer_quant_mode.has_fp8_block_scales( ): # Not all models that use this backend define shared experts (e.g. non-DeepSeek # MoEs), so fall back to 0 when the config has no `n_shared_experts`. self.num_fused_shared_expert = getattr( model_config.pretrained_config, "n_shared_experts", 0) or 0 + if self.num_fused_shared_expert > 0: + logger.info_once( + f"Shared-expert fusion enabled: folding " + f"{self.num_fused_shared_expert} shared expert(s) into the " + f"routed-expert grouped GEMM.", + key="trtllm_gen_shared_expert_fusion") def _to_trtllm_gen_activation_type(self, activation_type: ActivationType) -> int: diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 1250483d45f6..bdb5eb2b3309 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1918,6 +1918,10 @@ def test_fp8_block_scales_fusion_tmp(self): # TP from env TRTLLM_VALIDATE_TP (default 4). Run from tests/integration/defs: # LLM_MODELS_ROOT=/llm-models TRTLLM_VALIDATE_TP=4 python3 -m pytest -s -v \ # accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_fusion_tmp + # Shared-expert fusion is opt-in; this test exists to exercise the fused + # path, so enable it by default (spawned TP workers inherit the env). + # Export TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION=0 to A/B the unfused path. + os.environ.setdefault("TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION", "1") tp = int(os.environ.get("TRTLLM_VALIDATE_TP", "4")) # TRTLLM_VALIDATE_BACKEND: TRTLLM (fusion) | DEEPGEMM (stock, no fusion). # DEEPGEMM is used to prove the TP>1 tunable_allreduce hang is independent diff --git a/tests/microbenchmarks/bench_moe/build.py b/tests/microbenchmarks/bench_moe/build.py index 90c366be3011..a6215c30da93 100644 --- a/tests/microbenchmarks/bench_moe/build.py +++ b/tests/microbenchmarks/bench_moe/build.py @@ -197,6 +197,13 @@ def _build_moe_module( else: os.environ.pop("ENABLE_PERFECT_ROUTER", None) + # Shared-expert fusion is opt-in in TRTLLMGenFusedMoE; "fused" cases must + # enable it explicitly or they would silently benchmark the unfused path. + if model.n_shared_experts > 0 and model.shared_expert_mode != "unfused": + os.environ["TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION"] = "1" + else: + os.environ.pop("TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION", None) + mc = model.to_moe_model_config() swiglu_gptoss_style = model.swiglu_gptoss_style From f4d6f104ad71e4870b348d87011f9695c220b2ca Mon Sep 17 00:00:00 2001 From: leslief Date: Mon, 13 Jul 2026 08:13:40 +0000 Subject: [PATCH 09/19] [None][test] Add TRTLLM-Gen shared-expert fusion coverage to modules/moe backend tests Signed-off-by: leslief --- .../test_lists/test-db/l0_b200.yml | 2 + .../test_lists/test-db/l0_b300.yml | 2 + .../_torch/modules/moe/test_moe_backend.py | 353 +++++++++++++++++- 3 files changed, 355 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 464c3e553f91..fc87f04793d3 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -146,6 +146,8 @@ l0_b200: - unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py # ------------- MoE: test_moe_backend (by backend) --------------- - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe + - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts + - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout # ------------- MoE: test_single_gpu (by backend) --------------- - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and not None" - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "TRTLLM" diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 5d420918073e..a79ee846eb0e 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -36,6 +36,8 @@ l0_b300: - unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py # ------------- MoE: test_moe_backend (by backend) --------------- - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe + - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts + - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout # ------------- MoE: test_single_gpu (specific quant per backend) --------------- # CUTLASS backend: FP8, NVFP4, W4A8_MXFP4_MXFP8, W8A16 - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=FP8-routing=Renormalize] diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 060472b5b8be..ef91c892148b 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -67,7 +67,7 @@ W4A8MXFP4MXFP8MegaMoEDeepGemmMethod, ) from tensorrt_llm._torch.utils import ActivationType, is_gated_activation -from tensorrt_llm._utils import mpi_rank +from tensorrt_llm._utils import get_sm_version, mpi_rank from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo @@ -149,6 +149,7 @@ def create_test_backend( swiglu_limit: Optional[torch.Tensor] = None, weight_loading_mode: MoEWeightLoadingMode = MoEWeightLoadingMode.VANILLA, activation_type: ActivationType = ActivationType.Swiglu, + n_shared_experts: int = 0, ) -> MoE: """Create a MoE backend for testing.""" backend_cls = get_backend_class(backend_type) @@ -158,6 +159,11 @@ def create_test_backend( pretrained_config.hidden_size = hidden_size pretrained_config.intermediate_size = intermediate_size pretrained_config.torch_dtype = dtype + if n_shared_experts > 0: + # TRTLLMGenFusedMoE reads n_shared_experts off the pretrained config to + # decide num_fused_shared_expert (shared-expert fusion, opt-in via + # TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION=1). + pretrained_config.n_shared_experts = n_shared_experts # CUTE_DSL_B12X is internal-only: the user-facing API selects it on the # CUTEDSL path when SM120/121 + NVFP4 + flashinfer is importable. Route @@ -171,8 +177,13 @@ def create_test_backend( mapping=mapping, moe_backend=moe_backend_value, ) + if n_shared_experts > 0: + # The shared-expert-fusion gate runs after the eager create_weights() + # in __init__, so weight creation must be deferred (as the real model + # engine does) for the fused trailing slots to be allocated. + model_config.skip_create_weights_in_init = True - return create_moe_backend( + backend = create_moe_backend( moe_cls=backend_cls, routing_method=routing_method, num_experts=num_experts, @@ -189,6 +200,9 @@ def create_test_backend( weight_loading_mode=weight_loading_mode, activation_type=activation_type, ) + if n_shared_experts > 0: + backend.create_weights() + return backend # ============================================================================ @@ -1138,3 +1152,338 @@ def run_moe(): with torch.inference_mode(): output = run_moe() ref_fused_moe.check_accuracy(output, ref_output) + + +# ============================================================================ +# TRTLLM-Gen shared-expert fusion (migrated from deprecated +# tests/unittest/_torch/thop/serial/test_moe.py::TestMoeFP8 fusion coverage) +# ============================================================================ +# TRTLLMGenFusedMoE can fold n_shared_experts into the routed grouped GEMM as +# always-selected experts (opt-in via TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION=1; +# requires FP8_BLOCK_SCALES + dp_size==1 + DeepSeekV3 routing). The fused +# experts occupy trailing weight slots [num_experts, num_experts+n_fused) and +# receive routing weight 1.0, applied after routed_scaling_factor. + +# (num_experts, n_group, topk_group, top_k, n_fused) — shapes taken from the +# deprecated thop test's expert_info parametrization (n_fused > 0 variants). +FUSED_SHARED_EXPERT_INFOS = [ + (32, 8, 4, 8, 1), + (256, 8, 4, 8, 2), + (72, 1, 1, 6, 1), + (72, 1, 1, 6, 2), +] + + +class _AppendSharedExpertsRouting: + """Reference-side routing wrapper: appends the fused shared experts as + always-selected entries (ids [num_experts, num_experts+n_fused), weight + 1.0 after routed scaling) to the wrapped routing method's output.""" + + def __init__(self, routing_method, num_experts: int, n_fused: int): + self._routing_method = routing_method + self._num_experts = num_experts + self._n_fused = n_fused + + def apply(self, router_logits): + ids, weights = self._routing_method.apply(router_logits) + num_tokens = ids.shape[0] + shared_ids = torch.arange( + self._num_experts, + self._num_experts + self._n_fused, + dtype=ids.dtype, + device=ids.device, + ).expand(num_tokens, -1) + shared_weights = torch.ones( + (num_tokens, self._n_fused), dtype=weights.dtype, device=weights.device + ) + return torch.cat([ids, shared_ids], dim=1), torch.cat([weights, shared_weights], dim=1) + + +def _write_fused_shared_expert_slots(backend, weights, num_experts: int, n_fused: int): + """Copy the shared experts' quantized tensors into the trailing weight + slots, mirroring DeepSeekFP8BlockScalesFusedMoEMethod's routed-slot layout + (w3_w1 slot = [w3; w1] along dim 0; scales likewise).""" + for i in range(n_fused): + expert_id = num_experts + i + slot = num_experts + i + dst_w3, dst_w1 = backend.w3_w1_weight.data[slot].chunk(2, dim=0) + dst_w3.copy_(weights[f"{expert_id}.w3.weight"].view(dst_w3.dtype)) + dst_w1.copy_(weights[f"{expert_id}.w1.weight"].view(dst_w1.dtype)) + backend.w2_weight.data[slot].copy_( + weights[f"{expert_id}.w2.weight"].view(backend.w2_weight.dtype) + ) + dst_w3_scale, dst_w1_scale = backend.w3_w1_weight_scaling_factor.data[slot].chunk(2, dim=0) + dst_w3_scale.copy_(weights[f"{expert_id}.w3.weight_scale"]) + dst_w1_scale.copy_(weights[f"{expert_id}.w1.weight_scale"]) + backend.w2_weight_scaling_factor.data[slot].copy_(weights[f"{expert_id}.w2.weight_scale"]) + + +@pytest.mark.skipif( + get_sm_version() not in (100, 103), + reason="TRTLLM-Gen FP8 block scales requires SM100/103", +) +@pytest.mark.parametrize("seq_len", [16, 1024]) +@pytest.mark.parametrize( + "expert_info", + FUSED_SHARED_EXPERT_INFOS, + ids=lambda info: f"e{info[0]}g{info[1]}tg{info[2]}k{info[3]}fused{info[4]}", +) +def test_trtllm_fp8_block_scales_fused_shared_experts( + expert_info, + seq_len: int, + monkeypatch: pytest.MonkeyPatch, +): + """Fused-shared-expert accuracy for the TRTLLM backend (FP8 block scales). + + Loads routed experts through the regular path, writes the shared experts + into the trailing fused slots, and checks the fused kernel against a + reference that treats the shared experts as always-selected with weight + 1.0. Replays every captured tactic; also asserts the small-tile WAR + (fused path restricted to tileN >= 32) holds in the autotuner candidates. + """ + num_experts, n_group, topk_group, top_k, n_fused = expert_info + hidden_size = 512 + intermediate_size = 512 + dtype = torch.bfloat16 + backend_type = MoeBackendType.TRTLLM + + monkeypatch.setenv("TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION", "1") + + mapping = Mapping() + mapping.rank = mpi_rank() + + with torch.device(f"cuda:{mapping.rank}"): + torch.manual_seed(0) + torch.cuda.manual_seed(0) + AutoTuner.get().setup_distributed_state(mapping) + + e_score_correction_bias = torch.randn(num_experts, dtype=torch.bfloat16, device="cuda") + routing_method = DeepSeekV3MoeRoutingMethod( + top_k=top_k, + n_group=n_group, + topk_group=topk_group, + routed_scaling_factor=2.5, + callable_e_score_correction_bias=lambda: e_score_correction_bias, + ) + + router_logits = torch.randn((seq_len, num_experts), dtype=torch.float32, device="cuda") + + placeholder_x = torch.randn((seq_len, hidden_size), dtype=dtype, device="cuda") + quantize_util_cls, quant_config, quant_kwargs = get_test_quant_params( + QuantAlgo.FP8_BLOCK_SCALES, placeholder_x, backend_type + ) + # The quantize util covers routed + fused shared experts so create_weights + # emits per-expert tensors for all of them and the reference module gets + # a GatedMLP per fused slot as well. + quantize_util = quantize_util_cls( + num_experts=num_experts + n_fused, + dtype=dtype, + intermediate_size=intermediate_size, + hidden_size=hidden_size, + quant_config=quant_config, + ) + weights = quantize_util.create_weights(**quant_kwargs) + x = quantize_util.create_input(seq_len) + + backend = create_test_backend( + backend_type=backend_type, + routing_method=routing_method, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + quant_config=quant_config, + mapping=mapping, + n_shared_experts=n_fused, + ) + assert backend.num_fused_shared_expert == n_fused, ( + "shared-expert fusion gate did not activate" + ) + assert backend.w3_w1_weight.shape[0] == num_experts + n_fused + + # Mirror the real flow (modeling_deepseekv3): load routed experts, + # post-process, then fill the trailing fused slots. + routed_weights = {k: v for k, v in weights.items() if int(k.split(".")[0]) < num_experts} + backend.load_weights([routed_weights]) + backend.post_load_weights() + backend.cuda() + _write_fused_shared_expert_slots(backend, weights, num_experts, n_fused) + + ref_routing = _AppendSharedExpertsRouting(routing_method, num_experts, n_fused) + ref_fused_moe = quantize_util.create_ref_module(ref_routing) + ref_fused_moe.load_weights([weights]) + ref_fused_moe.cuda() + + AutoTuner.get().clear_cache() + + with torch.inference_mode(): + ref_output = ref_fused_moe.forward(x, router_logits) + + def run_moe(): + token_selected_experts, token_final_scales = routing_method.apply(router_logits) + x_quantized, x_sf = backend.quantize_input(x, post_quant_comm=False) + return run_backend_moe( + backend, + backend_type, + x_quantized, + x_sf, + token_selected_experts, + token_final_scales, + dtype, + router_logits=router_logits, + ) + + autotuner = AutoTuner.get() + autotuner.warmup = 0 + autotuner.repeat = 1 + autotuner.stream_delay_micro_secs = 10 + + with ( + torch.inference_mode(), + autotune(cache_path="/tmp/moe_autotuner_cache_fused_shared.json"), + ): + _ = run_moe() + + with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): + _ = run_moe() + + # WAR regression check: with num_fused_shared_experts > 0 the C++ side + # (fp8BlockScaleMoe.cpp fusedMinTileN, default 32) must exclude the + # small-tile (tileN 8/16) dynB cubins from the candidate tactics. + moe_tile_ns = [] + for combo in all_tactics: + for _, tactic in combo: + if isinstance(tactic, (list, tuple)) and len(tactic) == 2: + moe_tile_ns.append(int(tactic[0])) + assert moe_tile_ns, "expected [tileN, config] tactics from the fused MoE runner" + assert all(tile_n >= 32 for tile_n in moe_tile_ns), ( + f"fused path produced small-tile tactics (tileN < 32): {sorted(set(moe_tile_ns))}" + ) + + replay_tactics_and_check( + all_tactics=all_tactics, + run_moe_fn=run_moe, + check_accuracy_fn=ref_fused_moe.check_accuracy, + ref_output=ref_output, + backend_type=backend_type, + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, + ) + + +@pytest.mark.skipif( + get_sm_version() not in (100, 103), + reason="TRTLLM-Gen FP8 block scales requires SM100/103", +) +def test_trtllm_fp8_block_scales_fuse_shared_expert_layout(monkeypatch: pytest.MonkeyPatch): + """Layout-only check of fuse_shared_expert (no kernel launch). + + Builds a shared GatedMLP of intermediate n_fused*I, fuses it, and asserts + each trailing slot holds the expected per-expert sub-tensors by + definition: slot i's [w3; w1] = gate_up rows [i*I, (i+1)*I) of the up/gate + halves, slot i's w2 = down_proj columns [i*I, (i+1)*I) — including the + block-scale tensors (the w2 scale split is a historically bug-prone spot). + """ + from tensorrt_llm._torch.modules.gated_mlp import GatedMLP + + num_experts = 8 + top_k = 2 + hidden_size = 256 + intermediate_size = 256 + n_fused = 2 + dtype = torch.bfloat16 + scale_block = 128 + + monkeypatch.setenv("TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION", "1") + + mapping = Mapping() + mapping.rank = mpi_rank() + + with torch.device(f"cuda:{mapping.rank}"): + torch.manual_seed(0) + + routing_method = DeepSeekV3MoeRoutingMethod( + top_k=top_k, + n_group=1, + topk_group=1, + routed_scaling_factor=1.0, + callable_e_score_correction_bias=lambda: torch.zeros( + num_experts, dtype=dtype, device="cuda" + ), + ) + placeholder_x = torch.randn((1, hidden_size), dtype=dtype, device="cuda") + _, quant_config, _ = get_test_quant_params( + QuantAlgo.FP8_BLOCK_SCALES, placeholder_x, MoeBackendType.TRTLLM + ) + + backend = create_test_backend( + backend_type=MoeBackendType.TRTLLM, + routing_method=routing_method, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + quant_config=quant_config, + mapping=mapping, + n_shared_experts=n_fused, + ) + assert backend.num_fused_shared_expert == n_fused + + shared_mlp = GatedMLP( + hidden_size=hidden_size, + intermediate_size=n_fused * intermediate_size, + bias=False, + dtype=dtype, + config=ModelConfig(quant_config=quant_config), + ) + shared_mlp.cuda() + # Fill weights/scales with distinct random data (fp8 weights are filled + # via a bf16 sample viewed as fp8 bytes; values only need to be unique). + gate_up = shared_mlp.gate_up_proj + down = shared_mlp.down_proj + gate_up.weight.data.copy_( + torch.randn(gate_up.weight.shape, dtype=torch.bfloat16, device="cuda") + .to(torch.float8_e4m3fn) + .view(gate_up.weight.dtype) + ) + down.weight.data.copy_( + torch.randn(down.weight.shape, dtype=torch.bfloat16, device="cuda") + .to(torch.float8_e4m3fn) + .view(down.weight.dtype) + ) + gate_up.weight_scale.data.copy_( + torch.rand(gate_up.weight_scale.shape, dtype=torch.float32, device="cuda") + ) + down.weight_scale.data.copy_( + torch.rand(down.weight_scale.shape, dtype=torch.float32, device="cuda") + ) + + backend.fuse_shared_expert(shared_mlp) + torch.cuda.synchronize() + + # gate_up_proj rows are [gate(w1); up(w3)]; per-expert i owns rows + # [i*I, (i+1)*I) within each half. down_proj columns [i*I, (i+1)*I). + w1_all, w3_all = gate_up.weight.data.chunk(2, dim=0) + w1_scale_all, w3_scale_all = gate_up.weight_scale.data.chunk(2, dim=0) + inter_blocks = intermediate_size // scale_block + for i in range(n_fused): + slot = num_experts + i + rows = slice(i * intermediate_size, (i + 1) * intermediate_size) + scale_rows = slice(i * inter_blocks, (i + 1) * inter_blocks) + + dst_w3, dst_w1 = backend.w3_w1_weight.data[slot].chunk(2, dim=0) + torch.testing.assert_close(dst_w3, w3_all[rows].view(dst_w3.dtype)) + torch.testing.assert_close(dst_w1, w1_all[rows].view(dst_w1.dtype)) + torch.testing.assert_close( + backend.w2_weight.data[slot], + down.weight.data[:, rows].view(backend.w2_weight.dtype), + ) + + dst_w3_scale, dst_w1_scale = backend.w3_w1_weight_scaling_factor.data[slot].chunk( + 2, dim=0 + ) + torch.testing.assert_close(dst_w3_scale, w3_scale_all[scale_rows]) + torch.testing.assert_close(dst_w1_scale, w1_scale_all[scale_rows]) + torch.testing.assert_close( + backend.w2_weight_scaling_factor.data[slot], + down.weight_scale.data[:, scale_rows], + ) From 9432727a6dc89f37a2565639b4214f679681c7af Mon Sep 17 00:00:00 2001 From: leslief Date: Mon, 13 Jul 2026 08:16:00 +0000 Subject: [PATCH 10/19] [None][test] Drop fusion additions to deprecated thop test_moe.py (coverage moved to modules/moe) Signed-off-by: leslief --- tests/unittest/_torch/thop/serial/test_moe.py | 118 ++++++------------ 1 file changed, 36 insertions(+), 82 deletions(-) diff --git a/tests/unittest/_torch/thop/serial/test_moe.py b/tests/unittest/_torch/thop/serial/test_moe.py index bfbe34ba7358..53c70ee21c04 100644 --- a/tests/unittest/_torch/thop/serial/test_moe.py +++ b/tests/unittest/_torch/thop/serial/test_moe.py @@ -139,52 +139,30 @@ def __init__(self, self.act_type = act_type -def routing_reference(expertLogits, topK, padding, num_fused_shared_experts=0): +def routing_reference(expertLogits, topK, padding): originalDevice = expertLogits.device expertLogits = expertLogits.cpu() numTokens, numExperts = expertLogits.shape assert topK <= numExperts - numTotalExperts = numExperts + num_fused_shared_experts - totalExpertsPerToken = topK + num_fused_shared_experts - - numTokensPerExpert = torch.zeros(numTotalExperts, dtype=torch.int64) - expandedTokenIdxToExpert = -torch.ones(numTokens * totalExpertsPerToken, - dtype=torch.int64) - expandedTokenIdxToIdxInExpert = -torch.ones( - numTokens * totalExpertsPerToken, dtype=torch.int64) + numTokensPerExpert = torch.zeros(numExperts, dtype=torch.int64) + expandedTokenIdxToExpert = -torch.ones(numTokens * topK, dtype=torch.int64) + expandedTokenIdxToIdxInExpert = -torch.ones(numTokens * topK, + dtype=torch.int64) topKLogits, topKIndices = torch.topk(expertLogits, topK, dim=1) - if num_fused_shared_experts > 0: - # Shared experts will use weight of 1.0 - sharedLogits = torch.ones(numTokens, - num_fused_shared_experts, - dtype=topKLogits.dtype) - topKLogits = torch.cat((topKLogits, sharedLogits), dim=1) - assert numTokens == topKLogits.shape[0] - assert totalExpertsPerToken == topKLogits.shape[1] - - # Shared experts will have index starting at number of local experts - sharedIndices = torch.range(numExperts, - numExperts + num_fused_shared_experts - 1, - dtype=topKIndices.dtype) - sharedIndices = torch.unsqueeze(sharedIndices, 0) - sharedIndices = sharedIndices.repeat(numTokens, 1) - topKIndices = torch.cat((topKIndices, sharedIndices), dim=1) - assert numTokens == topKIndices.shape[0] - assert totalExpertsPerToken == topKIndices.shape[1] for tokenIdx in range(numTokens): - for k in range(totalExpertsPerToken): - expandedIdx = tokenIdx * totalExpertsPerToken + k + for k in range(topK): + expandedIdx = tokenIdx * topK + k expertIndex = topKIndices[tokenIdx, k] expandedTokenIdxToExpert[expandedIdx] = expertIndex expandedTokenIdxToIdxInExpert[expandedIdx] = numTokensPerExpert[ expertIndex] numTokensPerExpert[expertIndex] += 1 - paddedTokensPerExpertPrefixSum = torch.zeros(numTotalExperts + 1, + paddedTokensPerExpertPrefixSum = torch.zeros(numExperts + 1, dtype=torch.int64) - for ii in range(numTotalExperts): + for ii in range(numExperts): def divUpMul(a, b): return (a + b - 1) // b * b @@ -192,16 +170,16 @@ def divUpMul(a, b): paddedTokensPerExpertPrefixSum[ ii + 1] = paddedTokensPerExpertPrefixSum[ii] + divUpMul( numTokensPerExpert[ii], padding) - permutedBufferSize = paddedTokensPerExpertPrefixSum[numTotalExperts] + permutedBufferSize = paddedTokensPerExpertPrefixSum[numExperts] - expandedTokenIdxToPermutedIdx = -torch.ones( - numTokens * totalExpertsPerToken, dtype=torch.int64) + expandedTokenIdxToPermutedIdx = -torch.ones(numTokens * topK, + dtype=torch.int64) permutedIdxToExpandedIdx = -torch.ones(permutedBufferSize, dtype=torch.int64) permutedIdxToTokenIdx = -torch.ones(permutedBufferSize, dtype=torch.int64) for tokenIdx in range(numTokens): - for k in range(totalExpertsPerToken): - expandedIdx = tokenIdx * totalExpertsPerToken + k + for k in range(topK): + expandedIdx = tokenIdx * topK + k expert = expandedTokenIdxToExpert[expandedIdx] offsetWithinExpert = expandedTokenIdxToIdxInExpert[expandedIdx] offsetForExpert = paddedTokensPerExpertPrefixSum[expert] @@ -210,7 +188,6 @@ def divUpMul(a, b): expandedTokenIdxToPermutedIdx[expandedIdx] = permutedIdx permutedIdxToExpandedIdx[permutedIdx] = expandedIdx permutedIdxToTokenIdx[permutedIdx] = tokenIdx - return { "paddedTokensPerExpertPrefixSum": paddedTokensPerExpertPrefixSum.to(originalDevice), @@ -280,8 +257,7 @@ def routing_reference_no_aux(expert_logits, top_k_groups, routed_scaling, padding, - use_routing_scales_on_input=False, - num_fused_shared_experts=0): + use_routing_scales_on_input=False): routing_logits = expert_logits.to(dtype=torch.float, device='cuda') if use_routing_scales_on_input: # if using routing scales on input, topK == 1 and the score is a plain sigmoid @@ -289,8 +265,7 @@ def routing_reference_no_aux(expert_logits, else: scores = noaux_tc_ref(routing_logits, routing_bias, n_groups, top_k_groups, top_k, routed_scaling) - permute_info = routing_reference(scores, top_k, padding, - num_fused_shared_experts) + permute_info = routing_reference(scores, top_k, padding) return permute_info, scores @@ -922,10 +897,8 @@ class TestMoeFP8: """ @pytest.mark.parametrize("num_tokens", [16, 64, 1024]) - @pytest.mark.parametrize("expert_info", - [(32, 8, 4, 8, 0), (72, 1, 1, 6, 0), - (256, 8, 4, 8, 0), (32, 8, 4, 8, 1), - (256, 8, 4, 8, 2), (72, 1, 1, 6, 2)]) + @pytest.mark.parametrize("expert_info", [(32, 8, 4, 8), (72, 1, 1, 6), + (256, 8, 4, 8)]) @pytest.mark.parametrize("hidden_size", [512]) @pytest.mark.parametrize("intermediate_size", [512]) def test_autotune(self, num_tokens: int, expert_info: Tuple[int, int, int, @@ -940,11 +913,7 @@ def test_autotune(self, num_tokens: int, expert_info: Tuple[int, int, int, use_topk_as_input=False) @pytest.mark.parametrize("num_tokens", [16]) - @pytest.mark.parametrize("expert_info", - [(32, 8, 4, 8, 0), (384, 1, 1, 8, 0), - (72, 1, 1, 6, 0), (72, 1, 1, 6, 1), - (72, 1, 1, 6, 2), (384, 1, 1, 8, 1), - (384, 1, 1, 8, 2)]) + @pytest.mark.parametrize("expert_info", [(32, 8, 4, 8), (384, 1, 1, 8)]) @pytest.mark.parametrize("hidden_size", [512]) @pytest.mark.parametrize("intermediate_size", [512]) @pytest.mark.parametrize("use_topk_as_input", [False, True], @@ -970,7 +939,7 @@ def run_moe_fp8_test(self, num_tokens: int, expert_info: Tuple[int, int, # # Data Generation # - num_experts, n_groups, top_k_groups, top_k, num_fused_shared_experts = expert_info + num_experts, n_groups, top_k_groups, top_k = expert_info padding = 8 routed_scaling = 2.5 routing_method_type = RoutingMethodType.DeepSeekV3 @@ -979,14 +948,6 @@ def run_moe_fp8_test(self, num_tokens: int, expert_info: Tuple[int, int, assert top_k <= 8 assert num_experts % 4 == 0 - total_experts_per_token = top_k + num_fused_shared_experts - num_experts_total = num_experts + num_fused_shared_experts - - if use_topk_as_input and num_fused_shared_experts > 0: - pytest.skip( - "use_topk_as_input is tested only with num_fused_shared_experts=0" - ) - if are_groups_valid(top_k_groups, n_groups): assert top_k_groups <= 4 assert num_experts > n_groups @@ -1005,34 +966,27 @@ def run_moe_fp8_test(self, num_tokens: int, expert_info: Tuple[int, int, (hidden_size // 128, num_tokens), device='cuda').to(torch.float) gemm1_weights = torch.randn( - (num_experts_total, 2 * intermediate_size, hidden_size), + (num_experts, 2 * intermediate_size, hidden_size), device='cuda').to(torch.float8_e4m3fn) gemm1_scales = 2 * torch.rand( - (num_experts_total, 2 * intermediate_size // 128, - hidden_size // 128), + (num_experts, 2 * intermediate_size // 128, hidden_size // 128), device='cuda').to(torch.float) gemm2_weights = torch.randn( - (num_experts_total, hidden_size, intermediate_size), + (num_experts, hidden_size, intermediate_size), device='cuda').to(torch.float8_e4m3fn) gemm2_scales = 2 * torch.rand( - (num_experts_total, hidden_size // 128, intermediate_size // 128), + (num_experts, hidden_size // 128, intermediate_size // 128), device='cuda').to(torch.float) - permute_info, scores = routing_reference_no_aux( - expert_logits, - routing_bias, - top_k, - n_groups, - top_k_groups, - routed_scaling, - padding, - num_fused_shared_experts=num_fused_shared_experts) - - args = moe_args(num_tokens, num_experts_total, hidden_size, - intermediate_size, total_experts_per_token, padding, - hidden_states, hidden_states_scale, None, scores, - gemm1_weights, gemm1_scales, None, gemm2_weights, - gemm2_scales, None, permute_info, False) + permute_info, scores = routing_reference_no_aux(expert_logits, + routing_bias, top_k, + n_groups, top_k_groups, + routed_scaling, padding) + + args = moe_args(num_tokens, num_experts, hidden_size, intermediate_size, + top_k, padding, hidden_states, hidden_states_scale, + None, scores, gemm1_weights, gemm1_scales, None, + gemm2_weights, gemm2_scales, None, permute_info, False) if use_topk_as_input: topk_ids = permute_info["topKIndices"].to(torch.int32) @@ -1047,9 +1001,9 @@ def run_moe_fp8_test(self, num_tokens: int, expert_info: Tuple[int, int, output = torch.ops.trtllm.fp8_block_scale_moe_runner( expert_logits, routing_bias, hidden_states, hidden_states_scale, gemm1_weights, gemm1_scales, gemm2_weights, gemm2_scales, - num_experts, top_k, num_fused_shared_experts, n_groups, - top_k_groups, intermediate_size, 0, num_experts, routed_scaling, - routing_method_type, topk_weights, topk_ids) + num_experts, top_k, n_groups, top_k_groups, intermediate_size, + 0, num_experts, routed_scaling, routing_method_type, + topk_weights, topk_ids) torch.cuda.synchronize() output_dequant_actual = output.to(torch.float) # From 73b1a1130a0cd1d6caf4a069be00da73f015beb1 Mon Sep 17 00:00:00 2001 From: leslief Date: Mon, 13 Jul 2026 08:20:49 +0000 Subject: [PATCH 11/19] [None][test] Drop temporary in-tree fusion E2E test (kept as standalone local validation script) Signed-off-by: leslief --- .../defs/accuracy/test_llm_api_pytorch.py | 43 ------------------- 1 file changed, 43 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index bdb5eb2b3309..2305c0852664 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -1910,49 +1910,6 @@ def test_fp8_block_scales(self, mtp, fp8kv, attention_dp, cuda_graph, task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - def test_fp8_block_scales_fusion_tmp(self): - # TEMP (PR#11143 validation, NOT for commit): force the TRT-LLM Gen MoE - # backend + disable attention DP so the shared-expert -> sparse-expert - # fusion path (dp_size==1 + fp8_block_scales gate) is actually exercised - # end-to-end. The stock test uses DEEPGEMM (no fusion) on SM>=100. - # TP from env TRTLLM_VALIDATE_TP (default 4). Run from tests/integration/defs: - # LLM_MODELS_ROOT=/llm-models TRTLLM_VALIDATE_TP=4 python3 -m pytest -s -v \ - # accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_fusion_tmp - # Shared-expert fusion is opt-in; this test exists to exercise the fused - # path, so enable it by default (spawned TP workers inherit the env). - # Export TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION=0 to A/B the unfused path. - os.environ.setdefault("TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION", "1") - tp = int(os.environ.get("TRTLLM_VALIDATE_TP", "4")) - # TRTLLM_VALIDATE_BACKEND: TRTLLM (fusion) | DEEPGEMM (stock, no fusion). - # DEEPGEMM is used to prove the TP>1 tunable_allreduce hang is independent - # of the shared-expert fusion (fusion is not active with DEEPGEMM). - moe_backend = os.environ.get("TRTLLM_VALIDATE_BACKEND", "TRTLLM") - # TRTLLM_VALIDATE_AR: override allreduce_strategy (AUTO|NCCL|...). Used to - # work around the AUTO-strategy collective deadlock seen at TP>1 on B300. - ar_strategy = os.environ.get("TRTLLM_VALIDATE_AR", "AUTO") - # TRTLLM_VALIDATE_ADP: enable_attention_dp (0|1). Diagnostic only: with 1, - # attention uses DP (no pure-TP o_proj allreduce) and fusion is OFF; used to - # bound whether the TP=4 hang is specific to the attention_dp=False TP path. - attention_dp = os.environ.get("TRTLLM_VALIDATE_ADP", "0") == "1" - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75) - pytorch_config = dict( - disable_overlap_scheduler=True, - cuda_graph_config=None, - moe_config=MoeConfig(backend=moe_backend), - allreduce_strategy=ar_strategy, - ) - with LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", - tensor_parallel_size=tp, - moe_expert_parallel_size=1, - kv_cache_config=kv_cache_config, - **pytorch_config, - enable_attention_dp=attention_dp) as llm: - - assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES - - task = GSM8K(self.MODEL_NAME) - task.evaluate(llm) - @skip_pre_blackwell @parametrize_with_ids("torch_compile", [False]) @parametrize_with_ids( From dbc1c1dbe842687127e924294d342b1228db72b1 Mon Sep 17 00:00:00 2001 From: leslief Date: Wed, 15 Jul 2026 03:41:39 +0000 Subject: [PATCH 12/19] [None][chore] Drop unconsumed EP token-sharding plumbing for fused shared experts The shared-expert token offset/count computed in the DeepSeek routing branch was never consumed by any kernel (the fields only flowed into KernelParams and stopped there), and its EP derivation carried a known-bad assumption (TODO). Replace it with an explicit check rejecting fused shared experts under expert parallelism; re-add the sharding together with a kernel-side consumer when EP support is actually implemented. Signed-off-by: leslief --- .../blockScaleMoe/routing/RoutingDeepSeek.cu | 1 - .../blockScaleMoe/routing/RoutingKernel.h | 6 ----- .../trtllmGenKernels/blockScaleMoe/runner.cu | 23 +++++-------------- 3 files changed, 6 insertions(+), 24 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu index 55cfca91707e..feea69d48ab4 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu @@ -585,7 +585,6 @@ void run(Data& data, void* stream) data.mNumExperts += data.mNumFusedSharedExperts; data.mTopK += data.mNumFusedSharedExperts; data.mNumLocalExperts += data.mNumFusedSharedExperts; - // data.mLocalExpertsStartIdx += data.mNumFusedSharedExperts; } // Step 2: Permutation pipeline (reads from mPtrTopKPacked written by step 1) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h index e972b5809687..02387783fbcb 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h @@ -112,8 +112,6 @@ struct DataBase /// For fused shared expert int32_t mNumFusedSharedExperts; - int32_t mSharedExpertTokenOffset; - int32_t mSharedExpertNumTokens; int32_t mTotalExpertsPerToken; }; @@ -152,8 +150,6 @@ struct KernelParamsBase int32_t mNumLocalExperts = 0; int32_t mNumFusedSharedExperts; - int32_t mSharedExpertTokenOffset; - int32_t mSharedExpertNumTokens; int32_t mTotalExpertsPerToken; // Public initialization function - make it a template to accept different Data types @@ -184,8 +180,6 @@ struct KernelParamsBase mNumLocalExperts = data.mNumLocalExperts; mNumFusedSharedExperts = data.mNumFusedSharedExperts; - mSharedExpertTokenOffset = data.mSharedExpertTokenOffset; - mSharedExpertNumTokens = data.mSharedExpertNumTokens; mTotalExpertsPerToken = data.mTotalExpertsPerToken; } }; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu index bde059e9170f..29a7cdff0dda 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu @@ -233,6 +233,12 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 { TLLM_CHECK_WITH_INFO(topK <= 22, "For DeepSeek routing method, must have topK <= 22"); TLLM_CHECK_WITH_INFO(topkGroup <= 4, "For DeepSeek routing method, must have topkGroup <= 4"); + // Fused shared experts assume the full expert set is resident on this rank + // (no expert parallelism): the appended experts live at global indices + // [numExperts, numExperts + numFusedSharedExpert) on every rank. + TLLM_CHECK_WITH_INFO(numFusedSharedExpert == 0 || (localExpertOffset == 0 && localNumExperts == numExperts), + "Fused shared experts do not support expert parallelism yet."); + moe::dev::routing::routingDeepSeek::Data routingData; routingData.mDtypeOutput = btg::Dtype::Bfloat16; routingData.mUsePdl = tensorrt_llm::common::getEnvEnablePDL(); @@ -272,23 +278,6 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 routingData.mNumLocalExperts = localNumExperts; routingData.mRouteScale = routedScalingFactor; routingData.mUseRoutingSoftmax = false; - - // TODO Should these be passed directly instead? This does assume a constant number of experts per device - int32_t const numDevices = numExperts / localNumExperts; - int32_t const deviceIndex = localExpertOffset / localNumExperts; - int32_t const baseTokensPerDevice = numTokens / numDevices; - int32_t const remainingTokens = numTokens % numDevices; - - if (deviceIndex < remainingTokens) - { - routingData.mSharedExpertTokenOffset = (baseTokensPerDevice + 1) * deviceIndex; - routingData.mSharedExpertNumTokens = baseTokensPerDevice + 1; - } - else - { - routingData.mSharedExpertTokenOffset = remainingTokens + deviceIndex * baseTokensPerDevice; - routingData.mSharedExpertNumTokens = baseTokensPerDevice; - } moe::dev::routing::routingDeepSeek::run(routingData, stream); } else if (routingMethodType == RoutingMethodType::Llama4) From 2e059e9b0a846ba03de8cfe8e4b8ed4e976e269e Mon Sep 17 00:00:00 2001 From: leslief Date: Wed, 15 Jul 2026 06:29:43 +0000 Subject: [PATCH 13/19] [None][fix] bench_moe: drop stale cuda-graph max_num_tokens special-case Upstream #15397 made the non-DP (TEP/TTP) bench path pass all_rank_num_tokens=None (the MoE scheduler now derives chunking from this rank's rows), so the cuda-graph branch that computed sum(all_rank_num_tokens) crashed with TypeError on every fused TTP case, and the num_chunks==1 forcing it existed for is now guaranteed by the scheduler itself. Restore the plain per-rank-maximum expression. Signed-off-by: leslief --- .../microbenchmarks/bench_moe/case_runner.py | 26 +++++++------------ 1 file changed, 9 insertions(+), 17 deletions(-) diff --git a/tests/microbenchmarks/bench_moe/case_runner.py b/tests/microbenchmarks/bench_moe/case_runner.py index 3bcc8c30a93b..4a08c0de113b 100644 --- a/tests/microbenchmarks/bench_moe/case_runner.py +++ b/tests/microbenchmarks/bench_moe/case_runner.py @@ -693,23 +693,15 @@ def _run_one_candidate( mapping=mapping, moe_backend=config.backend, use_cuda_graph=bool(config.cuda_graph), - # Under CUDA graph, capture cannot cross the multi-chunk MoE scheduler - # path (`_forward_multiple_chunks` does a host `if` on a CUDA - # `chunked_used` tensor -> D2H sync, and overlaps chunks on an aux stream - # -- both illegal during capture). Production only captures the - # single-chunk path (num_chunks==1); calculate_num_chunks sums all ranks - # (num_rows=sum(all_rank_num_tokens)), so size max_num_tokens to the - # global token count under CUDA graph to force num_chunks==1. - # Otherwise use the global per-rank maximum: symmetric-memory comm - # backends (e.g. NVLINK_ONE_SIDED) size their workspace from - # max_num_tokens and require every rank to allocate the same size, so - # this rank's local token count (which differs under uneven - # attention-DP shards) must not be used. - max_num_tokens=( - max(sum(all_rank_num_tokens), 1) - if config.cuda_graph - else max(int(max(per_rank)) if per_rank else 0, 1) - ), + # Symmetric-memory comm backends (e.g. NVLINK_ONE_SIDED) size their + # workspace from max_num_tokens and require every rank to allocate the + # same size, so use the global per-rank maximum rather than this rank's + # local token count (which differs under uneven attention-DP shards). + # No CUDA-graph special-casing is needed anymore: since #15397 the + # non-DP (TEP/TTP) path passes all_rank_num_tokens=None and the MoE + # scheduler derives chunking from this rank's rows, so max(per_rank) + # already yields num_chunks==1 under graph capture. + max_num_tokens=max(int(max(per_rank)) if per_rank else 0, 1), use_low_precision_moe_combine=bool(config.use_low_precision_moe_combine), enable_perfect_router=enable_perfect_router, dtype=act_dtype, From ff3227e72acfc8c3a95d2c067d40117ae688bfac Mon Sep 17 00:00:00 2001 From: leslief Date: Wed, 15 Jul 2026 06:30:06 +0000 Subject: [PATCH 14/19] [None][chore] bench_moe: print traceback on build failures The build-error handler only printed a one-line summary while the timed-phase handler already prints the full traceback; align them so build failures are diagnosable from the log. Signed-off-by: leslief --- tests/microbenchmarks/bench_moe/case_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/microbenchmarks/bench_moe/case_runner.py b/tests/microbenchmarks/bench_moe/case_runner.py index 4a08c0de113b..53844bbc25a3 100644 --- a/tests/microbenchmarks/bench_moe/case_runner.py +++ b/tests/microbenchmarks/bench_moe/case_runner.py @@ -710,7 +710,7 @@ def _run_one_candidate( ) except Exception as exc: reason = f"build error: {type(exc).__name__}: {exc}" - _maybe_print_rank0(f"[bench_moe] build failed: {reason}") + _maybe_print_rank0(f"[bench_moe] build failed: {reason}\n{traceback.format_exc()}") return _short_circuit(result, "failed", reason) result.actual_backend = _backend_name_from_module(moe) From c7b31f0aeef48d187175b58cf1d32e242e29488b Mon Sep 17 00:00:00 2001 From: leslief Date: Fri, 24 Jul 2026 03:19:54 +0000 Subject: [PATCH 15/19] [None][fix] Guard quant_config None in TRTLLM-Gen MoE create_weights/gate The fused-shared-expert branch dereferenced self.quant_config.layer_quant_mode unconditionally to decide whether to pass num_fused_shared_experts. On the BF16 unquantized path quant_config is None, so create_weights crashed with AttributeError (test_trtllm_bf16_unquantized_moe, 16/16). Add the same 'quant_config is not None' guard already used elsewhere in this file; also guard the __init__ fusion gate (protected today only by env short-circuit). Signed-off-by: leslief --- .../_torch/modules/fused_moe/fused_moe_trtllm_gen.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 5d16e7ac26cc..0d255cf46c75 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -340,7 +340,7 @@ def __init__( "0") == "1" # Only the trtllm op backend implements fused shared experts on_trtllm_backend = isinstance(self.op_backend, TRTLLMOpBackend) - if fusion_enabled and on_trtllm_backend and model_config.mapping.dp_size == 1 and self.quant_config.layer_quant_mode.has_fp8_block_scales( + if fusion_enabled and on_trtllm_backend and model_config.mapping.dp_size == 1 and self.quant_config is not None and self.quant_config.layer_quant_mode.has_fp8_block_scales( ): # Not all models that use this backend define shared experts (e.g. non-DeepSeek # MoEs), so fall back to 0 when the config has no `n_shared_experts`. @@ -550,7 +550,8 @@ def create_weights(self): return self.quant_method = self._get_quant_method() - if self.quant_config.layer_quant_mode.has_fp8_block_scales(): + if self.quant_config is not None and self.quant_config.layer_quant_mode.has_fp8_block_scales( + ): self.quant_method.create_weights(self, self.num_fused_shared_expert) else: self.quant_method.create_weights(self) From 1756c1a64335d7313b325664b297da4d72d91084 Mon Sep 17 00:00:00 2001 From: leslief Date: Fri, 24 Jul 2026 03:27:49 +0000 Subject: [PATCH 16/19] [None][fix] Align AutoDeploy fp8_block_scale_moe_runner calls with new op schema The shared-expert-fusion op-schema change inserted num_fused_shared_experts after top_k, shifting the positional args of fp8_block_scale_moe_runner. The two AutoDeploy callers in trtllm_moe.py still passed args positionally against the old layout, so None (meant for routed_scaling_factor) landed on the required local_num_experts (test_trtllm_finegrained_fp8_moe_blackwell and test_trtllm_quant_finegrained_fp8_moe_fused_correctness). Pass num_fused_shared_experts=None and switch to keyword args from intermediate_size onward so the calls stay aligned across future schema additions. Signed-off-by: leslief --- .../custom_ops/fused_moe/trtllm_moe.py | 34 +++++++++++-------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py index 9ddbb9519aad..7987e5dba9ed 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py @@ -384,13 +384,16 @@ def _run_moe_with_alltoall( fc2_ws_f32, global_num_experts, top_k, - None, # n_group - None, # topk_group - intermediate_size, - local_expert_offset, - local_num_experts, - None, # routed_scaling_factor - RoutingMethodType.Renormalize, + # Keyword args from here on so the call stays aligned if the op + # schema gains new optional parameters (e.g. num_fused_shared_experts). + num_fused_shared_experts=None, + n_group=None, + topk_group=None, + intermediate_size=intermediate_size, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + routed_scaling_factor=None, + routing_method_type=RoutingMethodType.Renormalize, topk_weights=routing_weights_bf16, topk_ids=dispatched_selected, ) @@ -1043,13 +1046,16 @@ def trtllm_quant_finegrained_fp8_moe_fused( fc2_weight_scale_f32, num_experts, top_k, - None, # n_group - None, # topk_group - intermediate_size, - 0, # local_expert_offset - num_experts, # local_num_experts - None, # routed_scaling_factor - RoutingMethodType.Renormalize, + # Keyword args from here on so the call stays aligned if the op + # schema gains new optional parameters (e.g. num_fused_shared_experts). + num_fused_shared_experts=None, + n_group=None, + topk_group=None, + intermediate_size=intermediate_size, + local_expert_offset=0, + local_num_experts=num_experts, + routed_scaling_factor=None, + routing_method_type=RoutingMethodType.Renormalize, topk_weights=routing_weights_bf16, topk_ids=selected_experts, ) From c5a8e939d87894d7f4abfdb03d2306916dfec922 Mon Sep 17 00:00:00 2001 From: leslief Date: Fri, 24 Jul 2026 03:57:22 +0000 Subject: [PATCH 17/19] [None][fix] Default-init fused-shared-expert routing Data fields RoutingDeepSeekKernelTest constructs routingDeepSeek::Data directly (not via the thop runner), so the new mNumFusedSharedExperts/mTotalExpertsPerToken fields held uninitialized garbage: large values tripped the WarpSize assertion, small nonzero values corrupted the routing output or caused an illegal memory access. Default-init both to 0 in DataBase and derive mTotalExpertsPerToken = mTopK + mNumFusedSharedExperts inside routingDeepSeek::run() so any direct Data caller stays valid. Verified RoutingDeepSeekKernelTest 22/22 and the full routingKernelsTest 205/205. Signed-off-by: leslief --- .../blockScaleMoe/routing/RoutingDeepSeek.cu | 4 ++++ .../blockScaleMoe/routing/RoutingKernel.h | 8 +++++--- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu index feea69d48ab4..066ef26c1d68 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu @@ -565,6 +565,10 @@ void run(Data& data, void* stream) } int const numBlocks = data.mNumTokens; + // Derive the per-token entry count (topK routed + fused shared) here so callers + // that construct Data directly (unit tests) need not set it; the main kernel uses + // it as the output stride before mTopK is expanded below. + data.mTotalExpertsPerToken = data.mTopK + data.mNumFusedSharedExperts; // Account for fused shared experts (appended after the routed experts) when sizing the histogram. int const numThreadsHist = getMaxNumExperts(data.mNumExperts + data.mNumFusedSharedExperts); static int const smMajor = tensorrt_llm::common::getSMVersion() / 10; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h index 02387783fbcb..f31f94fa8cc5 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h @@ -110,9 +110,11 @@ struct DataBase int32_t mLocalExpertsStrideLog2; int32_t mNumLocalExperts; - /// For fused shared expert - int32_t mNumFusedSharedExperts; - int32_t mTotalExpertsPerToken; + /// For fused shared expert. Default to 0 (no fusion) so callers that + /// construct Data directly (e.g. the routing kernel unit tests) stay valid; + /// mTotalExpertsPerToken is derived in routingDeepSeek::run() from mTopK. + int32_t mNumFusedSharedExperts{0}; + int32_t mTotalExpertsPerToken{0}; }; template From b3791123cb7a05bc2867b248e10481b90610e90e Mon Sep 17 00:00:00 2001 From: leslief Date: Fri, 24 Jul 2026 06:14:52 +0000 Subject: [PATCH 18/19] [None][chore] Address shared-expert fusion review comments - Gate out fusion under expert parallelism (require moe_ep_size==1) so EP configs fall back to unfused instead of failing at the runtime EP check. - bench_moe: reject fused shared experts on unsupported backend/quant (fusion is TRTLLM-Gen + FP8_BLOCK_SCALES only) instead of silently dropping the shared experts. - bench_moe: validate ModelSpec shared-expert fields (Literal mode + non-negative n_shared_experts) so invalid inputs fail explicitly. - Annotate fusion bench/test helper signatures to match file conventions. Signed-off-by: leslief --- .../modules/fused_moe/fused_moe_trtllm_gen.py | 6 +++++- tests/microbenchmarks/bench_moe/build.py | 21 +++++++++++++++---- tests/microbenchmarks/bench_moe/specs.py | 12 +++++++++-- .../_torch/modules/moe/test_moe_backend.py | 10 +++++---- 4 files changed, 38 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 0d255cf46c75..b93066a1dac0 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -340,7 +340,11 @@ def __init__( "0") == "1" # Only the trtllm op backend implements fused shared experts on_trtllm_backend = isinstance(self.op_backend, TRTLLMOpBackend) - if fusion_enabled and on_trtllm_backend and model_config.mapping.dp_size == 1 and self.quant_config is not None and self.quant_config.layer_quant_mode.has_fp8_block_scales( + # Expert parallelism (moe_ep_size > 1) is not supported by the fused path yet + # (the routing kernel's shared-expert append assumes the full expert set is + # local); gate it out here so EP configs fall back to the unfused path instead + # of tripping the runtime EP check in the TRTLLM-Gen runner. + if fusion_enabled and on_trtllm_backend and model_config.mapping.dp_size == 1 and model_config.mapping.moe_ep_size == 1 and self.quant_config is not None and self.quant_config.layer_quant_mode.has_fp8_block_scales( ): # Not all models that use this backend define shared experts (e.g. non-DeepSeek # MoEs), so fall back to 0 when the config has no `n_shared_experts`. diff --git a/tests/microbenchmarks/bench_moe/build.py b/tests/microbenchmarks/bench_moe/build.py index a6215c30da93..10e4a0f6b2cd 100644 --- a/tests/microbenchmarks/bench_moe/build.py +++ b/tests/microbenchmarks/bench_moe/build.py @@ -25,7 +25,7 @@ from __future__ import annotations import os -from typing import Dict, List, Optional +from typing import Any, Dict, List, Optional import torch @@ -123,7 +123,7 @@ def __init__( moe: torch.nn.Module, shared_mlp: torch.nn.Module, use_cuda_graph: bool = False, - ): + ) -> None: super().__init__() self.moe = moe self.shared_mlp = shared_mlp @@ -134,7 +134,7 @@ def __init__( self.event_main = torch.cuda.Event() self.event_shared = torch.cuda.Event() - def __getattr__(self, name): + def __getattr__(self, name: str) -> Any: # nn.Module.__getattr__ resolves registered submodules/params/buffers # (incl. ``moe``/``shared_mlp``); anything else (routing_method, backend, # comm, scheduler, ...) is delegated to the wrapped MoE for introspection. @@ -143,7 +143,7 @@ def __getattr__(self, name): except AttributeError: return getattr(super().__getattr__("moe"), name) - def forward(self, x, router_logits, **kwargs): + def forward(self, x: torch.Tensor, router_logits: torch.Tensor, **kwargs) -> torch.Tensor: from tensorrt_llm._torch.modules.multi_stream_utils import ( maybe_execute_in_parallel, with_multi_stream, @@ -192,6 +192,19 @@ def _build_moe_module( f"(or drop --n_shared_experts)." ) + # Fusion is TRTLLM-Gen specific and only activates for FP8_BLOCK_SCALES; any other + # backend/quant would silently drop the shared experts and mislabel the benchmark. + # Reject fused candidates up front (unfused runs a separate GatedMLP, so it is fine + # on any backend and preserved here). + if model.n_shared_experts > 0 and model.shared_expert_mode != "unfused": + if moe_backend.upper() != "TRTLLM" or model.quant_algo_enum != QuantAlgo.FP8_BLOCK_SCALES: + raise NotImplementedError( + f"bench_moe fused shared experts require --backend TRTLLM with " + f"--quant FP8_BLOCK_SCALES (fusion is TRTLLM-Gen specific). Got " + f"backend={moe_backend!r}, quant={model.quant_algo!r}. Use " + f"--shared_expert_mode unfused to benchmark other backends." + ) + if enable_perfect_router: os.environ["ENABLE_PERFECT_ROUTER"] = "1" else: diff --git a/tests/microbenchmarks/bench_moe/specs.py b/tests/microbenchmarks/bench_moe/specs.py index 6b2043d69c22..9ec219dde43a 100644 --- a/tests/microbenchmarks/bench_moe/specs.py +++ b/tests/microbenchmarks/bench_moe/specs.py @@ -24,7 +24,7 @@ from __future__ import annotations from dataclasses import asdict, dataclass, field -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Literal, Optional, Tuple from tensorrt_llm._torch.modules.fused_moe.routing import ( DeepSeekV3MoeRoutingMethod, @@ -113,11 +113,19 @@ class ModelSpec: # "fused" -> fold them into the routed-expert grouped GEMM (PR #11143). # "unfused" -> routed MoE (no fusion) + a separate shared GatedMLP, summed # (the pre-fusion baseline, for measuring fusion's net benefit). - shared_expert_mode: str = "fused" + shared_expert_mode: Literal["fused", "unfused"] = "fused" swiglu_alpha: float = 1.0 swiglu_beta: float = 0.0 swiglu_limit: float = float("inf") + def __post_init__(self) -> None: + if self.n_shared_experts < 0: + raise ValueError(f"n_shared_experts must be >= 0, got {self.n_shared_experts}.") + if self.shared_expert_mode not in ("fused", "unfused"): + raise ValueError( + f"shared_expert_mode must be 'fused' or 'unfused', got {self.shared_expert_mode!r}." + ) + @property def routing_method_cls(self) -> type: return _ROUTING_METHODS[self.routing_method] diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index ef91c892148b..6c5d09c35949 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -30,7 +30,7 @@ import logging import os from types import SimpleNamespace -from typing import List, Optional +from typing import List, Optional, Tuple from unittest.mock import MagicMock import pytest @@ -1179,12 +1179,12 @@ class _AppendSharedExpertsRouting: always-selected entries (ids [num_experts, num_experts+n_fused), weight 1.0 after routed scaling) to the wrapped routing method's output.""" - def __init__(self, routing_method, num_experts: int, n_fused: int): + def __init__(self, routing_method, num_experts: int, n_fused: int) -> None: self._routing_method = routing_method self._num_experts = num_experts self._n_fused = n_fused - def apply(self, router_logits): + def apply(self, router_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: ids, weights = self._routing_method.apply(router_logits) num_tokens = ids.shape[0] shared_ids = torch.arange( @@ -1199,7 +1199,9 @@ def apply(self, router_logits): return torch.cat([ids, shared_ids], dim=1), torch.cat([weights, shared_weights], dim=1) -def _write_fused_shared_expert_slots(backend, weights, num_experts: int, n_fused: int): +def _write_fused_shared_expert_slots( + backend: MoE, weights: dict, num_experts: int, n_fused: int +) -> None: """Copy the shared experts' quantized tensors into the trailing weight slots, mirroring DeepSeekFP8BlockScalesFusedMoEMethod's routed-slot layout (w3_w1 slot = [w3; w1] along dim 0; scales likewise).""" From 5016db4fd3ce73ba695ca1b3925c12043e8fb40a Mon Sep 17 00:00:00 2001 From: leslief Date: Fri, 24 Jul 2026 06:52:39 +0000 Subject: [PATCH 19/19] [None][chore] Address shared-expert fusion review comments (C++) - Default-init KernelParamsBase fused-shared-expert fields (= 0) to match the other scalar members and stay safe if used before setBaseParams(). - fp8BlockScaleMoe: guard num_fused_shared_experts >= 0 and simplify the redundant '> 0 ? value() : 0' ternaries to value_or(0). - fp8BlockScaleMoe: reject external topk_ids/topk_weights together with fused shared experts (integrated routing appends the shared experts; the external-topk path would silently drop them). Signed-off-by: leslief --- .../blockScaleMoe/routing/RoutingKernel.h | 4 +-- cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp | 30 ++++++++++--------- 2 files changed, 18 insertions(+), 16 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h index f31f94fa8cc5..3ef3dd05cb43 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h @@ -151,8 +151,8 @@ struct KernelParamsBase int32_t mLocalExpertsStrideLog2 = 0; int32_t mNumLocalExperts = 0; - int32_t mNumFusedSharedExperts; - int32_t mTotalExpertsPerToken; + int32_t mNumFusedSharedExperts = 0; + int32_t mTotalExpertsPerToken = 0; // Public initialization function - make it a template to accept different Data types template diff --git a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp index 2b7d25102ba0..c89953f9d327 100644 --- a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp @@ -130,6 +130,13 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit TORCH_CHECK(num_experts > top_k, "num_experts must be greater than top_k"); + // Fused shared experts are appended by the integrated routing kernel (from routing_logits); + // that path is skipped when external topk_ids/topk_weights are supplied, which would silently + // drop the shared experts. Reject the unsupported combination up front. + TORCH_CHECK(num_fused_shared_experts.value_or(0) == 0 || (!topk_ids.has_value() && !topk_weights.has_value()), + "Fused shared experts require integrated routing; external topk_ids/topk_weights are not " + "supported when num_fused_shared_experts > 0."); + // If both routing inputs are provided, they must be on the same device if (routing_logits.has_value() && topk_ids.has_value()) { @@ -140,12 +147,10 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::MoERunnerArgs args; tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::MoEWorkspace workspace; - int64_t const num_total_experts - = num_experts + (num_fused_shared_experts.value_or(0) > 0 ? num_fused_shared_experts.value() : 0); - int64_t const total_experts_per_token - = top_k + (num_fused_shared_experts.value_or(0) > 0 ? num_fused_shared_experts.value() : 0); - int64_t const num_total_local_experts - = local_num_experts + (num_fused_shared_experts.value_or(0) > 0 ? num_fused_shared_experts.value() : 0); + TORCH_CHECK(num_fused_shared_experts.value_or(0) >= 0, "num_fused_shared_experts must be non-negative."); + int64_t const num_total_experts = num_experts + num_fused_shared_experts.value_or(0); + int64_t const total_experts_per_token = top_k + num_fused_shared_experts.value_or(0); + int64_t const num_total_local_experts = local_num_experts + num_fused_shared_experts.value_or(0); // setup args // note: the assumption is that output data type is always Bfloat16 (the default) @@ -388,10 +393,9 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder std::optional const numFusedSharedExpert, int64_t hiddenSize, int64_t intermediateSize, int64_t numLocalExperts, int64_t numTokens) const { - int64_t const totalExpertsPerToken - = topK + (numFusedSharedExpert.value_or(0) > 0 ? numFusedSharedExpert.value() : 0); - int64_t const numTotalLocalExperts - = numLocalExperts + (numFusedSharedExpert.value_or(0) > 0 ? numFusedSharedExpert.value() : 0); + TORCH_CHECK(numFusedSharedExpert.value_or(0) >= 0, "num_fused_shared_experts must be non-negative."); + int64_t const totalExpertsPerToken = topK + numFusedSharedExpert.value_or(0); + int64_t const numTotalLocalExperts = numLocalExperts + numFusedSharedExpert.value_or(0); // WAR: the small-tile (tileN 8/16) dynB TRTLLM-Gen batched-GEMM cubins flakily hit an // illegal memory access (garbage TMA-descriptor pointer, MMU fault in the gemm2 K-loop) // when shared experts are fused into the grouped GEMM (num_fused_shared_experts > 0); @@ -443,10 +447,8 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder // Autotuner has requested a default or 'fallback' config index if (tileN == -1 || config == -1) { - int64_t const total_experts_per_token - = top_k + (num_fused_shared_experts.value_or(0) > 0 ? num_fused_shared_experts.value() : 0); - int64_t const num_total_local_experts - = local_num_experts + (num_fused_shared_experts.value_or(0) > 0 ? num_fused_shared_experts.value() : 0); + int64_t const total_experts_per_token = top_k + num_fused_shared_experts.value_or(0); + int64_t const num_total_local_experts = local_num_experts + num_fused_shared_experts.value_or(0); auto const num_tokens = hidden_states.sizes()[0]; auto const hidden_size = hidden_states.sizes()[1];