[https://nvbugs/6191524][fix] In MLA.forward_context, also call the warmup when has_cached_kv_for_mla_context by tensorrt-cicd · Pull Request #14536 · NVIDIA/TensorRT-LLM

tensorrt-cicd · 2026-05-25T17:26:54Z

Summary

Root cause: maybe_compiled_cat at attention.py:2155 had no warmup on the cached-kv prefill path; torch.compile recompiled per num_ctx_tokens shape and one recompile stalled in inductor's compile worker, tripping the 300s hang detector.
Fix: In MLA.forward_context, also call the warmup when has_cached_kv_for_mla_context is true, so the cached-kv path uses a single mark_dynamic-based compile instead of per-shape recompiles.
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/6191524

Summary by CodeRabbit

Performance
- Enhanced attention mechanism warm-up logic to execute in additional scenarios, improving compatibility with more inference contexts.

coderabbitai · 2026-05-25T17:30:28Z

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: 6848444d-e970-4603-8f7f-897dafe1af73

📥 Commits

Reviewing files that changed from the base of the PR and between 63b1d8d and 042d661.

📒 Files selected for processing (1)

tensorrt_llm/_torch/modules/attention.py

📝 Walkthrough

Walkthrough

The change expands the warm-up logic in MLA.forward_context by broadening its condition. Previously, maybe_compiled_cat was warmed up only for the chunked-prefill path; now it warms up for both chunked-prefill and cached-KV paths, ensuring both context execution modes are prepared.

Changes

MLA Attention Warm-up Coverage

Layer / File(s)	Summary
Forward context warm-up condition expansion `tensorrt_llm/_torch/modules/attention.py`	The condition guarding warm-up of `maybe_compiled_cat` in `MLA.forward_context` is broadened from a single check to an OR of two checks, so the warm-up now triggers for both chunked-prefill and cached-KV context paths.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the specific fix: expanding the warmup logic in MLA.forward_context to also trigger for cached_kv_for_mla_context, which directly matches the core change in the PR.
Description check	✅ Passed	The PR description includes a clear root cause analysis, the specific fix applied, test verification (cached-kv path now uses single compile vs per-shape recompiles), and links to the bug tracker, covering all essential aspects of the change.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

zhhuang-nv · 2026-05-26T08:46:18Z

/bot run

tensorrt-cicd · 2026-05-26T08:52:06Z

PR_Github #50327 [ run ] triggered by Bot. Commit: 329de3a Link to invocation

tensorrt-cicd · 2026-05-26T13:31:34Z

PR_Github #50327 [ run ] completed with state FAILURE. Commit: 329de3a
/LLM/main/L0_MergeRequest_PR pipeline #39853 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

zhhuang-nv · 2026-05-28T01:54:56Z

/bot run

tensorrt-cicd · 2026-05-28T02:00:39Z

PR_Github #50654 [ run ] triggered by Bot. Commit: d5666d3 Link to invocation

tensorrt-cicd · 2026-05-28T02:38:46Z

PR_Github #50654 [ run ] completed with state FAILURE. Commit: d5666d3
/LLM/main/L0_MergeRequest_PR pipeline #40146 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

zhhuang-nv · 2026-05-28T03:19:31Z

/bot run

tensorrt-cicd · 2026-05-28T03:26:19Z

PR_Github #50680 [ run ] triggered by Bot. Commit: 8b52e64 Link to invocation

tensorrt-cicd · 2026-05-28T08:31:42Z

PR_Github #50680 [ run ] completed with state SUCCESS. Commit: 8b52e64
/LLM/main/L0_MergeRequest_PR pipeline #40168 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

In MLA prefill, forward_context_with_cached_kv calls maybe_compiled_cat on shapes that vary with num_ctx_tokens. The existing warmup (cached_warmup_forward_context_with_chunked_prefill) pre-compiles this op via torch._dynamo.maybe_mark_dynamic, but it was only invoked when runtime_features.chunked_prefill was set. Tests that enable block reuse without chunked prefill (e.g. DeepSeek-R1 throughput_mtp_trtllm) hit forward_context_with_cached_kv with no warmup, causing torch.compile to recompile maybe_compiled_cat for each new num_ctx_tokens shape; one of these compilations stalls inside inductor's compile worker, tripping PyExecutor's hang detector. Invoke the same warmup whenever has_cached_kv_for_mla_context is true, so the cached-kv prefill path also gets a single mark_dynamic-based compile rather than per-shape recompiles. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>

Signed-off-by: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com>

zhhuang-nv · 2026-06-01T06:19:58Z

/bot run

tensorrt-cicd · 2026-06-01T06:29:46Z

PR_Github #51322 [ run ] triggered by Bot. Commit: 843bbed Link to invocation

tensorrt-cicd · 2026-06-01T10:36:21Z

PR_Github #51322 [ run ] completed with state SUCCESS. Commit: 843bbed
/LLM/main/L0_MergeRequest_PR pipeline #40736 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

tensorrt-cicd requested a review from a team as a code owner May 25, 2026 17:26

tensorrt-cicd requested a review from PerkzZheng May 25, 2026 17:26

tensorrt-cicd assigned bobboli May 25, 2026

github-actions Bot assigned tensorrt-cicd May 25, 2026

PerkzZheng approved these changes May 26, 2026

View reviewed changes

zhhuang-nv force-pushed the repair-bot-bug6191524 branch from 85bc8b0 to 329de3a Compare May 26, 2026 08:41

zhhuang-nv force-pushed the repair-bot-bug6191524 branch from 329de3a to d5666d3 Compare May 28, 2026 01:54

zhhuang-nv force-pushed the repair-bot-bug6191524 branch from d5666d3 to 8b52e64 Compare May 28, 2026 03:17

tburt-nv mentioned this pull request May 28, 2026

[TRTLLMINF-67][infra] use pre-configured idle GPU exemption #14587

Merged

1 task

tensorrt-cicd and others added 2 commits June 1, 2026 14:19

unify code paths and rename functions

843bbed

Signed-off-by: Zhen Huang <145532724+zhhuang-nv@users.noreply.github.com>

zhhuang-nv force-pushed the repair-bot-bug6191524 branch from 8b52e64 to 843bbed Compare June 1, 2026 06:19

byshiue merged commit 209f371 into NVIDIA:main Jun 2, 2026
6 of 7 checks passed

Uh oh!

Conversation

tensorrt-cicd commented May 25, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Links

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented May 25, 2026

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

zhhuang-nv commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 26, 2026

Uh oh!

zhhuang-nv commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

zhhuang-nv commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

zhhuang-nv commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

tensorrt-cicd commented May 25, 2026 •

edited by coderabbitai Bot

Loading