From fc58e1d83d9267f5db6ab4a79d183902fd047b80 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Thu, 9 Jul 2026 13:42:04 -0700 Subject: [PATCH 1/3] [nvbugs/6432948][fix] Skip redundant TRTLLM-Gen FMHA JIT warmup in Phase-2 warmup create_py_executor's two-phase KV-cache-estimation flow instantiates PyExecutor twice against the same model_engine, so warmup - and therefore _run_attention_warmup - runs twice. The TRTLLM-Gen FMHA JIT kernel cache is process-global, so the second grid enumeration compiles nothing new. For DeepSeek-R1 FP8 TP=8 MTP3 with max_num_tokens=12288 on B200, running the grid a second time - after autotuner exploration and CUDA-graph capture have consumed most of GPU memory - occasionally triggers an illegal memory access asynchronously reported at torch.cuda.synchronize() in _run_attention_warmup. Add a per-engine _trtllm_gen_jit_warmup_done flag: first call runs the full grid and sets the flag; subsequent calls short-circuit with a logger.info. Correctness is preserved because the JIT cache from Phase-1 already covers every shape Phase-2 could request; any kernel not yet compiled would JIT-compile lazily on first request anyway. Verified: 1 passed in 230.48s on B200 tp8 (previously EXIT_CODE=1 at 1071s with CUDA IMA in _run_attention_warmup during Phase-2 restart). Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 92f79e2505ca..030175b66d85 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -825,6 +825,10 @@ def __init__( self.kv_cache_manager_key = ResourceManagerType.DRAFT_KV_CACHE_MANAGER if is_draft_model else ResourceManagerType.KV_CACHE_MANAGER self.lora_model_config: Optional[LoraModelConfig] = None self._trtllm_gen_jit_warmup = False + # KV-cache estimation re-instantiates PyExecutor and re-runs warmup on + # the same engine; the TRTLLM-Gen FMHA JIT cache is process-global, so + # skip subsequent passes. See nvbugs/6432948. + self._trtllm_gen_jit_warmup_done = False # Create config and runner cuda_graph_runner_config = CUDAGraphRunnerConfig( @@ -1601,6 +1605,12 @@ def _run_attention_warmup(self, if not issubclass(self.attn_backend.Metadata, TrtllmAttentionMetadata): return + if self._trtllm_gen_jit_warmup_done: + logger.info( + "Skipping TRTLLM-Gen FMHA JIT warmup: already populated by a prior warmup pass." + ) + return + @contextlib.contextmanager def trtllm_gen_fmha_jit_warmup(): previous = self._trtllm_gen_jit_warmup @@ -1666,6 +1676,8 @@ def trtllm_gen_fmha_jit_warmup(): resource_manager=resource_manager) torch.cuda.synchronize() + self._trtllm_gen_jit_warmup_done = True + @staticmethod def _release_megamoe_profiling_scratch(): # MegaMoE tuning resources are shared across layers, so only the engine From 6920fd940bd0c0458bc293177aa1a6468a35d273 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:50:35 -0700 Subject: [PATCH 2/3] [nvbugs/6432948][chore] Remove stale waiver after fix Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 6a003bf64f0a..6ac7b9caed69 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -342,7 +342,6 @@ perf/test_perf_sanity.py::test_e2e[aggr_upload-ctx_only-gb300_glm-5-fp4_8k1k_con perf/test_perf_sanity.py::test_e2e[aggr_upload-ctx_only-gb300_glm-5-fp4_8k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6517846) perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp4_v2_2_nodes_grace_blackwell-r1_fp4_v2_dep8_mtp1_8k1k] SKIP (https://nvbugs/6530213) perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp4_v2_grace_blackwell-r1_fp4_v2_dep4_mtp1_1k8k] SKIP (https://nvbugs/6422339) -perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp8_blackwell-r1_fp8_tp8_mtp3_8k1k] SKIP (https://nvbugs/6432948) perf/test_perf_sanity.py::test_e2e[aggr_upload-dynamo_gpt_oss_120b_fp4_blackwell-gpt_oss_fp4_tep4_adp_cutlass_8k1k] SKIP (https://nvbugs/6374910) perf/test_perf_sanity.py::test_e2e[aggr_upload-glm5_fp4_blackwell-glm5_fp4_tep8_mtp3_8k1k] SKIP (https://nvbugs/6329155) perf/test_perf_sanity.py::test_e2e[aggr_upload-k25_thinking_fp4_2_nodes_grace_blackwell-k25_thinking_fp4_tep8_32k8k] SKIP (https://nvbugs/6422339) From 1ac1ec8f308960b0b7c216fd1b7c63fa01d4fd20 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Thu, 6 Aug 2026 09:21:38 -0700 Subject: [PATCH 3/3] [nvbugs/6432948][fix] Exclude defective small tiles for all FP8 block-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR #15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp | 81 +++++++++---------- .../_torch/pyexecutor/model_engine.py | 12 --- tests/integration/test_lists/waives.txt | 1 + 3 files changed, 41 insertions(+), 53 deletions(-) diff --git a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp index c3007343bf68..b28b8242fb80 100644 --- a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp @@ -400,10 +400,16 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder 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); + // 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). - // Restrict the fused path to tileN >= 32 until the kernel-side fix lands (nvbug TBD). + // Originally scoped to the fused shared-expert path, but the defect is in the shared + // small-tile cubins and 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). Excluding + // the small tiles for every caller is safe because every FP8 block-scale MoE shape also + // offers a tileN >= 32 tactic, so the returned list is never emptied. 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_FUSED_MIN_TILEN overrides the threshold (0 disables) for A/B experiments. static int const fusedMinTileN = []() { @@ -414,7 +420,7 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder std::vector> tactics; for (auto& [tileN, runner] : mRunners) { - if (numFusedSharedExpert.value_or(0) > 0 && tileN < fusedMinTileN) + if (tileN < fusedMinTileN) { continue; } @@ -461,49 +467,42 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder = static_cast(num_tokens * total_experts_per_token) / num_total_local_experts; tileN = std::clamp(nextPowerOfTwo(avg_tokens_per_expert), mSupportedTileN.front(), mSupportedTileN.back()); - 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, which 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. Small warmup batches clamp the heuristic to + // mSupportedTileN.front(), so this fallback is the path that reaches the defective + // small-tile cubins; it needs the same tileN >= 32 exclusion as getValidConfigs + // (see the WAR comment there). + config = -1; + std::vector tileN_candidates{static_cast(tileN)}; + for (auto t : mSupportedTileN) { - // 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); + } + 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) { - if (t != tileN) - tileN_candidates.push_back(t); + continue; } - // 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) + auto valid = mRunners.at(t)->getValidConfigIndices( + total_experts_per_token, hidden_size, intermediate_size, num_total_local_experts, num_tokens); + if (!valid.empty()) { - 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()) - { - tileN = t; - config = valid.front(); - break; - } + 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); } + TLLM_CHECK_WITH_INFO(config != -1, "No valid TRTLLM-Gen config found for FP8 block-scale MoE."); } return run_fp8_block_scale_moe(routing_logits, routing_bias, hidden_states, hidden_states_scale, gemm1_weights, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 030175b66d85..92f79e2505ca 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -825,10 +825,6 @@ def __init__( self.kv_cache_manager_key = ResourceManagerType.DRAFT_KV_CACHE_MANAGER if is_draft_model else ResourceManagerType.KV_CACHE_MANAGER self.lora_model_config: Optional[LoraModelConfig] = None self._trtllm_gen_jit_warmup = False - # KV-cache estimation re-instantiates PyExecutor and re-runs warmup on - # the same engine; the TRTLLM-Gen FMHA JIT cache is process-global, so - # skip subsequent passes. See nvbugs/6432948. - self._trtllm_gen_jit_warmup_done = False # Create config and runner cuda_graph_runner_config = CUDAGraphRunnerConfig( @@ -1605,12 +1601,6 @@ def _run_attention_warmup(self, if not issubclass(self.attn_backend.Metadata, TrtllmAttentionMetadata): return - if self._trtllm_gen_jit_warmup_done: - logger.info( - "Skipping TRTLLM-Gen FMHA JIT warmup: already populated by a prior warmup pass." - ) - return - @contextlib.contextmanager def trtllm_gen_fmha_jit_warmup(): previous = self._trtllm_gen_jit_warmup @@ -1676,8 +1666,6 @@ def trtllm_gen_fmha_jit_warmup(): resource_manager=resource_manager) torch.cuda.synchronize() - self._trtllm_gen_jit_warmup_done = True - @staticmethod def _release_megamoe_profiling_scratch(): # MegaMoE tuning resources are shared across layers, so only the engine diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 6ac7b9caed69..6a003bf64f0a 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -342,6 +342,7 @@ perf/test_perf_sanity.py::test_e2e[aggr_upload-ctx_only-gb300_glm-5-fp4_8k1k_con perf/test_perf_sanity.py::test_e2e[aggr_upload-ctx_only-gb300_glm-5-fp4_8k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6517846) perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp4_v2_2_nodes_grace_blackwell-r1_fp4_v2_dep8_mtp1_8k1k] SKIP (https://nvbugs/6530213) perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp4_v2_grace_blackwell-r1_fp4_v2_dep4_mtp1_1k8k] SKIP (https://nvbugs/6422339) +perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp8_blackwell-r1_fp8_tp8_mtp3_8k1k] SKIP (https://nvbugs/6432948) perf/test_perf_sanity.py::test_e2e[aggr_upload-dynamo_gpt_oss_120b_fp4_blackwell-gpt_oss_fp4_tep4_adp_cutlass_8k1k] SKIP (https://nvbugs/6374910) perf/test_perf_sanity.py::test_e2e[aggr_upload-glm5_fp4_blackwell-glm5_fp4_tep8_mtp3_8k1k] SKIP (https://nvbugs/6329155) perf/test_perf_sanity.py::test_e2e[aggr_upload-k25_thinking_fp4_2_nodes_grace_blackwell-k25_thinking_fp4_tep8_32k8k] SKIP (https://nvbugs/6422339)