From 47e9b5e191a7b42c9e6e707b75178550b1c520fb Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Sun, 23 Aug 2026 18:59:42 -0700 Subject: [PATCH] [https://nvbugs/6432948][fix] Exclude TRTLLM-Gen small tileN for all FP8 block-scale MoE 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). tileN >= 32 is unaffected (10/10 clean vs minutes-to-crash baseline on B300 TP=4). The existing workaround excluded those tiles only when shared experts were fused into the grouped GEMM (num_fused_shared_experts > 0). That scoping was wrong: the defect is in the shared small-tile cubins and is not caused by expert fusion. DeepSeek-R1 FP8 TP=8 (unfused) faults identically during warmup, where the 1/2/8-token shapes are the only ones that can select tileN 8/16 (12288 tokens gets tileN 64/128 and always passes). The same was already observed in #15297, where the IMA reproduced with num_fused_shared_experts=0. Changes: - Hoist the threshold into a single moeMinTileN() accessor built on common::getIntEnv, replacing two independently parsed function-local statics that both used unchecked std::atoi. Rename the knob to TLLM_MOE_MIN_TILEN and keep TLLM_MOE_FUSED_MIN_TILEN as a deprecated alias, since the exclusion is no longer scoped to the fused path. - Precompute mEligibleTileN in the ctor and drive both tactic selection and the tileN heuristic from it. Running computeSelectedTileN on the eligible list rather than on mSupportedTileN keeps the excluded tiles from consuming the returned neighbourhood: a shape whose heuristic tile is 8 now gets {32, 64, 128} instead of being left with 32 as its only candidate. - Reject a threshold that excludes every supported tile at construction time instead of silently returning an empty tactic list. - Include the problem dimensions in the fused fallback's no-valid-config error so a report from an unchecked model is actionable without a repro. The unfused fallback keeps getDefaultValidConfigIndex. An earlier revision of this change routed it through getValidConfigIndices(...).front() as well, but that helper returns the first pair in raw cartesian order, whereas getDefaultValidConfigIndex returns the first entry of the list sorted by the perf heuristic in KernelRunner.cpp. Switching it would have changed the selected kernel config for all default-path traffic, which is unrelated to this fix. Signed-off-by: ZhaoyangWang --- cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp | 111 +++++++++++++-------- 1 file changed, 70 insertions(+), 41 deletions(-) diff --git a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp index e5debeda9e6f..4fece941b277 100644 --- a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp @@ -14,6 +14,7 @@ * limitations under the License. */ +#include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h" #include "tensorrt_llm/thop/thUtils.h" @@ -24,9 +25,10 @@ #include #include -#include +#include #include #include +#include TRTLLM_NAMESPACE_BEGIN @@ -38,6 +40,38 @@ using tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::RoutingMethodTy using MoeRunnerType = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::Runner; using tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::computeSelectedTileN; +namespace +{ + +// 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); tileN >= 32 is +// unaffected (10/10 clean vs minutes-to-crash baseline on B300 TP=4). The exclusion was originally +// scoped to the fused shared-expert path, but the defect is in the shared small-tile cubins and is +// not caused by expert fusion: DeepSeek-R1 FP8 TP=8 (unfused) faults identically during warmup, +// where the 1/2/8-token shapes are the only ones that can select tileN 8/16 (12288 tokens gets +// tileN 64/128 and always passes). The tiles stay in mSupportedTileN: the ctor builds one runner +// per tile and each asserts a non-empty passing-config list, so the exclusion has to happen at +// tactic-selection time. +// +// TLLM_MOE_MIN_TILEN overrides the threshold (0 disables the WAR) for A/B experiments. +// TLLM_MOE_FUSED_MIN_TILEN is the deprecated original name, still honoured as an alias because the +// exclusion used to apply only to the fused shared-expert path. +int32_t moeMinTileN() +{ + static int32_t const minTileN = [] + { + auto value = tensorrt_llm::common::getIntEnv("TLLM_MOE_MIN_TILEN"); + if (!value.has_value()) + { + value = tensorrt_llm::common::getIntEnv("TLLM_MOE_FUSED_MIN_TILEN"); + } + return value.value_or(32); + }(); + return minTileN; +} + +} // namespace + at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logits, 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, @@ -403,6 +437,13 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder { mRunners.emplace(tileN, std::make_unique(mDtypeElt, mUseDeepSeekFp8, tileN)); } + // Tiles the small-tile WAR (see moeMinTileN) allows tactic selection to pick from. Kept as + // a separate sorted list so the tileN heuristic only ever proposes an eligible tile. + std::copy_if(mSupportedTileN.begin(), mSupportedTileN.end(), std::back_inserter(mEligibleTileN), + [](int32_t tileN) { return tileN >= moeMinTileN(); }); + TORCH_CHECK(!mEligibleTileN.empty(), "The minimum tileN in force (", moeMinTileN(), + ") excludes every supported tileN (max ", mSupportedTileN.back(), + "). Lower TLLM_MOE_MIN_TILEN or set it to 0 to disable the small-tile WAR."); } [[nodiscard]] std::vector> getValidConfigs(int64_t topK, @@ -412,33 +453,22 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder 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); - // 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; - }(); + // Only offer tiles the small-tile WAR allows (see moeMinTileN). The heuristic runs on the + // eligible list rather than on mSupportedTileN so that excluded tiles do not consume the + // neighbourhood computeSelectedTileN returns -- otherwise a small shape whose heuristic + // tile is 8 would be left with tileN 32 as its only candidate instead of {32, 64, 128}. + auto const chosen = computeSelectedTileN(mEligibleTileN, numTokens, totalExpertsPerToken, numTotalLocalExperts); // returns (tileN, config) std::vector> tactics; - for (auto& [tileN, runner] : mRunners) + for (auto const tileN : mEligibleTileN) { - if (numFusedSharedExpert.value_or(0) > 0 && tileN < fusedMinTileN) - { - continue; - } - auto chosen = computeSelectedTileN(mSupportedTileN, numTokens, totalExpertsPerToken, numTotalLocalExperts); if (chosen.find(tileN) == chosen.end()) { continue; } - auto config_indices_per_runner = runner->getValidConfigIndices( + auto const config_indices_per_runner = mRunners.at(tileN)->getValidConfigIndices( totalExpertsPerToken, hiddenSize, intermediateSize, numTotalLocalExperts, numTokens); - for (auto cfg : config_indices_per_runner) + for (auto const cfg : config_indices_per_runner) { tactics.push_back({tileN, cfg}); } @@ -472,35 +502,29 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder 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()); + // Snap the heuristic tile up to the smallest eligible one: small warmup batches would + // otherwise land on tileN 8/16, which is exactly the path that reaches the defective + // small-tile cubins (see moeMinTileN). + auto const heuristicTileN + = std::lower_bound(mEligibleTileN.begin(), mEligibleTileN.end(), nextPowerOfTwo(avg_tokens_per_expert)); + tileN = heuristicTileN == mEligibleTileN.end() ? mEligibleTileN.back() : *heuristicTileN; 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. + // The inflated fused expert/topK counts can leave the heuristic tile without a + // valid config, so walk the remaining eligible tiles instead of failing outright. config = -1; std::vector tileN_candidates{static_cast(tileN)}; - for (auto t : mSupportedTileN) + for (auto const t : mEligibleTileN) { 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) + for (auto const t : tileN_candidates) { - if (t < fusedMinTileNFallback) - { - continue; - } - auto valid = mRunners.at(t)->getValidConfigIndices( + auto const valid = mRunners.at(t)->getValidConfigIndices( total_experts_per_token, hidden_size, intermediate_size, num_total_local_experts, num_tokens); if (!valid.empty()) { @@ -509,8 +533,12 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder break; } } - TLLM_CHECK_WITH_INFO( - config != -1, "No valid TRTLLM-Gen config found for fused shared-expert FP8 block-scale MoE."); + TLLM_CHECK_WITH_INFO(config != -1, + "No valid TRTLLM-Gen config found for fused shared-expert FP8 block-scale MoE with num_tokens=%ld, " + "hidden_size=%ld, intermediate_size=%ld, experts_per_token=%ld, local_experts=%ld, min_tileN=%d.", + static_cast(num_tokens), static_cast(hidden_size), + static_cast(intermediate_size), total_experts_per_token, num_total_local_experts, + moeMinTileN()); } else { @@ -529,6 +557,7 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder using RunnerType = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::Runner; std::vector const mSupportedTileN; + std::vector mEligibleTileN; std::unordered_map> mRunners; btg::Dtype mDtypeElt{btg::Dtype::E4m3}; // FP8 runner so hard-coded