[https://nvbugs/6095851][fix] Refresh FlashMLA metadata for draft KV cache

[yizhang-nv]

Summary by CodeRabbit

• Bug Fixes

  • Fixed state management for speculative decoding operations with KV-cache transitions.

• Tests

  • Additional test configurations now pass without waivers, improving test suite reliability.

Description

MTP draft replay temporarily switches attention metadata from the target KV cache manager to the draft KV cache manager. For FlashMLA, the separate block id table also needs to be rebuilt after that switch and after restoring the target KV cache manager. Without refreshing it, FlashMLA can use block ids from the wrong V2 KV cache pool and hit an illegal memory access.

This change refreshes FlashMLA metadata around the draft KV cache context and unwaives the two DeepSeekV3Lite H100 coverage cases for https://nvbugs/6084720 and https://nvbugs/6095851.

Test Coverage

Relevant tests:

• accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True]
• accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True]

Checks run locally:

• pre-commit hooks during git commit
• git diff --check
• python3 -m py_compile tensorrt_llm/_torch/speculative/interface.py

CI requested on this PR:

• DGX_H100-PyTorch-Post-Merge-1
• DGX_H100-PyTorch-Post-Merge-2

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[yizhang-nv]

/bot run --stage-list "DGX_H100-PyTorch-Post-Merge-1,DGX_H100-PyTorch-Post-Merge-2"

[coderabbitai]

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: 9ba8934b-6026-4a41-a3da-b1815bb875e8

📥 Commits

Reviewing files that changed from the base of the PR and between a622e30 and c85c30e.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/speculative/interface.py
• tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

---

📝 Walkthrough

Walkthrough

This PR re-initializes FlashMLA state at critical points during speculative execution: draft replay preparation and restoration, and KV-cache layout switching. It removes two related test waivers for DeepSeekV3Lite that now pass.

Changes

FlashMLA State Management in Speculative Execution

Layer / File(s)	Summary
FlashMLA state re-initialization in speculative operations tensorrt_llm/_torch/speculative/interface.py	Four locations now call attn_metadata.prepare_flash_mla() when enabled: in draft replay prepare/restore functions (lines 105–106, 160–161) and in the draft_kv_cache_context() manager after layout switches (lines 1285–1286, 1295–1296).
Test waiver removal for DeepSeekV3Lite configurations tests/integration/test_lists/waives.txt	Two SKIP waivers for TestDeepSeekV3Lite::test_bfloat16 with mtp_nextn=2 and attention_dp=False are removed, as the FlashMLA fixes enable these configurations to pass.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#15061: Both PRs modify tests/integration/test_lists/waives.txt to remove SKIP entries for the same TestDeepSeekV3Lite::test_bfloat16 configuration (un-waiving previously failing cases).

Suggested reviewers

• yechank-nvidia
• dpitman-nvda
• tburt-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main change: refreshing FlashMLA metadata for draft KV cache handling, which is the primary purpose of the PR.
Description check	✅ Passed	The PR description covers all required sections: explanation of the issue and solution, relevant test cases, local checks performed, CI requirements, and a completed checklist.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.}

[tensorrt-cicd]

PR_Github #53509 [ run ] triggered by Bot. Commit: c85c30e Link to invocation

[tensorrt-cicd]

PR_Github #53509 [ run ] completed with state SUCCESS. Commit: c85c30e
/LLM/main/L0_MergeRequest_PR pipeline #42667 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53796 [ run ] triggered by Bot. Commit: c85c30e Link to invocation

[tensorrt-cicd]

PR_Github #53796 [ run ] completed with state SUCCESS. Commit: c85c30e
/LLM/main/L0_MergeRequest_PR pipeline #42914 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54184 [ run ] triggered by Bot. Commit: c85c30e Link to invocation

[tensorrt-cicd]

PR_Github #54184 [ run ] completed with state FAILURE. Commit: c85c30e
/LLM/main/L0_MergeRequest_PR pipeline #43264 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54277 [ run ] triggered by Bot. Commit: eb69b34 Link to invocation

[tensorrt-cicd]

PR_Github #54277 [ run ] completed with state FAILURE. Commit: eb69b34
/LLM/main/L0_MergeRequest_PR pipeline #43350 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54435 [ run ] triggered by Bot. Commit: 23da8f5 Link to invocation

[tensorrt-cicd]

PR_Github #54435 [ run ] completed with state SUCCESS. Commit: 23da8f5
/LLM/main/L0_MergeRequest_PR pipeline #43499 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54723 [ run ] triggered by Bot. Commit: 23da8f5 Link to invocation

[tensorrt-cicd]

PR_Github #54723 [ run ] completed with state SUCCESS. Commit: 23da8f5
/LLM/main/L0_MergeRequest_PR pipeline #43750 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

[yizhang-nv]

/bot run

[tensorrt-cicd]

PR_Github #56272 [ run ] triggered by Bot. Commit: 0740bb9 Link to invocation

[tensorrt-cicd]

PR_Github #56272 [ run ] completed with state FAILURE. Commit: 0740bb9
/LLM/main/L0_MergeRequest_PR pipeline #45128 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #56867 [ run ] triggered by Bot. Commit: ad9a75a Link to invocation

[tensorrt-cicd]

PR_Github #56867 [ run ] completed with state SUCCESS. Commit: ad9a75a
/LLM/main/L0_MergeRequest_PR pipeline #45678 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #56956 [ run ] triggered by Bot. Commit: 39cfba0 Link to invocation

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #56986 [ run ] triggered by Bot. Commit: 39cfba0 Link to invocation

[tensorrt-cicd]

PR_Github #56956 [ run ] completed with state ABORTED. Commit: 39cfba0

Link to invocation

[tensorrt-cicd]

PR_Github #56986 [ run ] completed with state SUCCESS. Commit: 39cfba0
/LLM/main/L0_MergeRequest_PR pipeline #45787 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

[yizhang-nv]

/bot run --disable-fail-fast

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57142 [ run ] triggered by Bot. Commit: 39cfba0 Link to invocation

[tensorrt-cicd]

PR_Github #57142 [ run ] completed with state SUCCESS. Commit: 39cfba0
/LLM/main/L0_MergeRequest_PR pipeline #45925 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57364 [ run ] triggered by Bot. Commit: 39cfba0 Link to invocation

[tensorrt-cicd]

PR_Github #57364 [ run ] completed with state SUCCESS. Commit: 39cfba0
/LLM/main/L0_MergeRequest_PR pipeline #46115 completed with status: 'SUCCESS'

CI Report

Link to invocation