[https://nvbugs/6432948][fix] Add per-engine flag `_trtllm_gen_jit_warmup_done`; skip… by trtllm-agent · Pull Request #16205 · NVIDIA/TensorRT-LLM

trtllm-agent · 2026-07-09T20:51:05Z

Summary

Root cause: create_py_executor's two-phase KV-cache-estimation flow re-invokes model_engine.warmup on the same model_engine, so _run_attention_warmup runs the TRTLLM-Gen FMHA JIT warmup grid twice; the second (Phase-2) run compiles nothing new (JIT cache is process-global) yet — under memory pressure after autotuner + CUDA-graph capture — occasionally faults with an illegal memory access.
Fix: Add per-engine flag _trtllm_gen_jit_warmup_done; skip _run_attention_warmup's JIT warmup body on the second and subsequent invocations with an informational log — Phase-1 semantics unchanged, redundant Phase-2 grid enumeration eliminated.
Automated fix generated by repair-bot

Test plan

Verify fix on the same GPU type as the original failure
Check for regressions in related tests

Links

Bug: https://nvbugs/6432948

Dev Engineer Review

Updated fp8BlockScaleMoe.cpp to exclude tiles below the configured minimum for all tactic-selection callers.
Updated the fallback path to select only explicitly validated configurations.
Added an error path when no valid configuration exists.
Removed the previous non-fused default-index path.
No public or exported declarations changed.
No configuration files or test-list files changed.
The change improves tactic-selection consistency and prevents invalid small-tile selection.
The new error path requires callers to handle the no-valid-configuration case correctly.

QA Engineer Review

No test changes.

coderabbitai · 2026-07-09T20:54:22Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: cb05ece1-3787-4861-9fd9-c9137860f7a4

📥 Commits

Reviewing files that changed from the base of the PR and between 92e7c88 and d6e0cc8.

📒 Files selected for processing (1)

cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp

🚧 Files skipped from review as they are similar to previous changes (1)

cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp

Walkthrough

FP8 block-scale MoE tactic selection now excludes undersized tiles for fused and unfused workloads. Fallback selection validates configurations, skips undersized tiles, prioritizes the heuristic tile, and errors when no valid configuration exists.

Changes

FP8 MoE tactic selection

Layer / File(s)	Summary
Universal tile filtering and fallback validation `cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp`	The minimum tile-size filter now applies to fused and unfused workloads. Fallback selection uses validated configurations, skips undersized tiles, prioritizes the heuristic tile, and reports an error when no valid configuration exists.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: asfiyab-nvidia, brnguyen2, barry-delaney

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Title check	⚠️ Warning	The title describes a TRTLLM-Gen JIT warmup flag, but the changeset fixes FP8 block-scale MoE tactic selection.	Update the title to describe the FP8 block-scale MoE tactic-selection fix, while retaining the valid NVBugs identifier and fix type.
Description check	⚠️ Warning	The description explains a different JIT warmup change and omits the required Description, Test Coverage, and PR Checklist sections.	Rewrite the description to match the FP8 block-scale MoE changes and add the required template sections, including relevant tests and checklist confirmation.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

tburt-nv · 2026-07-15T22:44:19Z

 perf/test_perf.py::test_perf[whisper_large_v3-bench-float16-input_output_len:128,20] SKIP
 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_1k1k] SKIP (https://nvbugs/6445332)
-perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp8_blackwell-r1_fp8_tp8_mtp3_8k1k] SKIP (https://nvbugs/6432948)


The code change looks like a good improvement, but I'm not seeing how it addresses the root cause of the flaky failure. The error was cudaStreamDestroy(stream): an illegal memory access was encountered, not an out of memory issue.

If you want to move forward with this as a fix, please run the test several times to ensure it passes consistently.

coderabbitai · 2026-08-06T16:54:42Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

coderabbitai

🧹 Nitpick comments (1)

cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp (1)

480-504: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Use braces and const in the fallback candidate loop.

Add braces to the if body at Lines 482-483. Declare unmodified t, env, and valid values as const.

Proposed fix

-            for (auto t : mSupportedTileN)
+            for (auto const t : mSupportedTileN)
             {
                 if (t != tileN)
+                {
                     tileN_candidates.push_back(t);
+                }
             }
             static int const fusedMinTileNFallback = []()
             {
-                char const* env = std::getenv("TLLM_MOE_FUSED_MIN_TILEN");
+                char const* 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(

As per coding guidelines, “always brace if/else, loop, and switch bodies” and “declare unmodified variables as const.”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp` around lines 480 - 504, Update
the fallback candidate loop to brace the if body that skips candidates below
fusedMinTileNFallback, and declare the unmodified loop variable t, environment
pointer env, and valid configuration result as const. Preserve the existing
candidate filtering and selection behavior.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp`:
- Around line 480-504: Update the fallback candidate loop to brace the if body
that skips candidates below fusedMinTileNFallback, and declare the unmodified
loop variable t, environment pointer env, and valid configuration result as
const. Preserve the existing candidate filtering and selection behavior.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c17ff890-8c14-4a9a-8ea0-6dae2406b3fc

📥 Commits

Reviewing files that changed from the base of the PR and between f12c5e5 and acb9a9d.

📒 Files selected for processing (1)

cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp

coderabbitai · 2026-08-06T22:08:28Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

leslie-fang25

The title doesn't match what changes in this PR.

leslie-fang25 · 2026-08-06T22:28:57Z

        for (auto& [tileN, runner] : mRunners)
        {
-            if (numFusedSharedExpert.value_or(0) > 0 && tileN < fusedMinTileN)
+            if (tileN < fusedMinTileN)


I think the IMA issues we met only happens when numFusedSharedExpert large than 0. Please explain why apply it for all the fp8 tactic.

coderabbitai · 2026-08-07T06:18:08Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

…ase-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>

Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>

…-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 NVIDIA#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>

trtllm-agent requested a review from a team as a code owner July 9, 2026 20:51

trtllm-agent requested a review from Tabrizian July 9, 2026 20:51

trtllm-agent assigned QiJune Jul 9, 2026

github-actions Bot assigned trtllm-agent Jul 9, 2026

trtllm-agent force-pushed the repair-bot-bug6432948 branch 3 times, most recently from 98dcf31 to 39b1dd5 Compare July 15, 2026 01:53

trtllm-agent requested review from a team as code owners July 15, 2026 01:53

trtllm-agent requested review from BowenFu, asfiyab-nvidia, brnguyen2, nv-xtf, tburt-nv, yingguo-trt, yuanjingx87 and yufeiwu-nv July 15, 2026 01:53

BowenFu approved these changes Jul 15, 2026

View reviewed changes

tburt-nv approved these changes Jul 15, 2026

View reviewed changes

trtllm-agent force-pushed the repair-bot-bug6432948 branch from 39b1dd5 to 2995242 Compare July 16, 2026 00:19

yufeiwu-nv removed their request for review July 16, 2026 05:39

yingguo-trt approved these changes Jul 16, 2026

View reviewed changes

trtllm-agent force-pushed the repair-bot-bug6432948 branch from 2995242 to 4cb6611 Compare July 18, 2026 08:36

mikeiovine removed the request for review from Tabrizian July 20, 2026 15:44

trtllm-agent force-pushed the repair-bot-bug6432948 branch from 4cb6611 to acb9a9d Compare August 6, 2026 16:54

trtllm-agent requested a review from a team as a code owner August 6, 2026 16:54

trtllm-agent requested review from Barry-Delaney and leslie-fang25 August 6, 2026 16:54

coderabbitai Bot reviewed Aug 6, 2026

View reviewed changes

trtllm-agent force-pushed the repair-bot-bug6432948 branch from acb9a9d to 8946584 Compare August 6, 2026 22:08

leslie-fang25 reviewed Aug 6, 2026

View reviewed changes

trtllm-agent force-pushed the repair-bot-bug6432948 branch from 8946584 to d6e0cc8 Compare August 7, 2026 06:18

trtllm-agent added 3 commits August 7, 2026 07:23

[nvbugs/6432948][chore] Remove stale waiver after fix

4153e85

Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>

trtllm-agent force-pushed the repair-bot-bug6432948 branch from d6e0cc8 to cac9664 Compare August 7, 2026 14:24

Navigation Menu

Uh oh!

[https://nvbugs/6432948][fix] Add per-engine flag `_trtllm_gen_jit_warmup_done`; skip… - #16205

[https://nvbugs/6432948][fix] Add per-engine flag `_trtllm_gen_jit_warmup_done`; skip…#16205
trtllm-agent wants to merge 3 commits into
NVIDIA:mainfrom
tensorrt-cicd:repair-bot-bug6432948

trtllm-agent commented Jul 9, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Jul 9, 2026 •

edited

Loading

❌ Failed checks (2 warnings)

Uh oh!

tburt-nv Jul 15, 2026

Uh oh!

coderabbitai Bot commented Aug 6, 2026

Uh oh!

coderabbitai Bot left a comment

Uh oh!

coderabbitai Bot commented Aug 6, 2026

Uh oh!

leslie-fang25 left a comment

Uh oh!

leslie-fang25 Aug 6, 2026

Uh oh!

coderabbitai Bot commented Aug 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

trtllm-agent commented Jul 9, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Links

Dev Engineer Review

QA Engineer Review

Uh oh!

coderabbitai Bot commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

❌ Failed checks (2 warnings)

Uh oh!

tburt-nv Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot commented Aug 6, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot commented Aug 6, 2026

Uh oh!

leslie-fang25 left a comment

Choose a reason for hiding this comment

Uh oh!

leslie-fang25 Aug 6, 2026

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot commented Aug 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

trtllm-agent commented Jul 9, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 9, 2026 •

edited

Loading