[https://nvbugs/6422343][fix] Retain the CPU source tensor as self._flash_mla_src_block_ids_cpu on the…

[trtllm-agent]

Summary

• Root cause: prepare_flash_mla issues non_blocking=True H2D copies from a locally-scoped pinned CPU tensor that Python can reclaim before the DMA completes, leaving stale/garbage entries in the FlashMLA block-ID buffers and triggering an OOB read in the FlashMLA kernel under MTP+ADP+cuda_graph+torch_compile+chunked_prefill warmup.
• Fix: Retain the CPU source tensor as self._flash_mla_src_block_ids_cpu on the metadata object so its lifetime spans the async DMA. Also removed the H20 waivers for this bug.
• 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/6422343

Summary by CodeRabbit

• Bug Fixes

  • Improved stability during attention preparation, preventing rare crashes or invalid reads during back-to-back graph-captured runs.
  • Better preserves correct results when using FlashMLA-based attention in demanding inference setups.

• Tests

  • Removed outdated test waivers for a previously unstable scenario, reflecting improved reliability.

[coderabbitai]

📝 Walkthrough

Walkthrough

Modifies prepare_flash_mla() to retain a persistent reference to the pinned CPU block_ids_per_seq tensor before non-blocking H2D copies, preventing premature deallocation during CUDA-graph capture. Also removes two now-unneeded waiver entries for a DeepSeekV3Lite accuracy test from the integration test waivers list.

Changes

FlashMLA Metadata Fix and Test Waiver Cleanup

Layer / File(s)	Summary
Retain pinned CPU source buffer tensorrt_llm/_torch/attention_backend/trtllm.py	Saves the pinned CPU block_ids_per_seq tensor to self._flash_mla_src_block_ids_cpu before issuing non-blocking device copies, preventing the source tensor from being freed during CUDA-graph capture and avoiding out-of-bounds reads.
Remove obsolete test waivers tests/integration/test_lists/waives.txt	Removes two SKIP/waiver entries for TestDeepSeekV3Lite::test_bfloat16 (mtp_nextn=2) covering both v2_kv_cache=False and v2_kv_cache=True variants.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#15251: Increases/multiplies prepare_flash_mla() calls during speculative draft KV cache switching, directly interacting with the buffer lifetime fix.
• NVIDIA/TensorRT-LLM#16030: Updates the same H20 TestDeepSeekV3Lite waiver entries in tests/integration/test_lists/waives.txt.
• NVIDIA/TensorRT-LLM#16034: Also modifies TestDeepSeekV3Lite SKIP/waiver entries in the same waives.txt file.

Suggested reviewers: zheyuf, ziyixiong-nv, pengbowang-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the bug fix and correctly identifies the main change, including the NVBugs reference and fix type.
Description check	✅ Passed	The description covers the summary, test plan, and bug link, and is mostly complete despite omitting the PR checklist section.
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.}

[coderabbitai]

🧹 Nitpick comments (1)

tensorrt_llm/_torch/attention_backend/trtllm.py (1)

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

Declare _flash_mla_src_block_ids_cpu as a typed dataclass field for consistency.

Sibling private state on this dataclass (e.g. _flash_mla_metadata_valid) is declared via field(default=..., init=False, repr=False) with a type annotation. _flash_mla_src_block_ids_cpu is instead created ad hoc via plain attribute assignment, which is inconsistent with the class's own convention and less friendly to static type checkers.

♻️ Suggested fix

     _flash_mla_metadata_valid: bool = field(default=False,
                                             init=False,
                                             repr=False)
+    # Retains the pinned CPU source tensor for the FlashMLA H2D copies so it
+    # outlives async, non-blocking copies issued during CUDA-graph capture.
+    _flash_mla_src_block_ids_cpu: Optional[torch.Tensor] = field(
+        default=None, init=False, repr=False)

As per coding guidelines, "Annotate class members and variables when necessary, especially for dataclasses and NamedTuple."

🤖 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 `@tensorrt_llm/_torch/attention_backend/trtllm.py` at line 690, Declare
`_flash_mla_src_block_ids_cpu` as an explicit typed dataclass field on the class
instead of creating it ad hoc in the `trtllm.py` attention backend. Update the
dataclass definition for `TRTLLM`/the owning class to include a type annotation
and a `field(default=..., init=False, repr=False)` entry, matching the existing
convention used by `_flash_mla_metadata_valid` and related private state. Then
remove the plain assignment in the initialization path and rely on the declared
field so the member is consistent and static type checkers can see it.

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 `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Line 690: Declare `_flash_mla_src_block_ids_cpu` as an explicit typed
dataclass field on the class instead of creating it ad hoc in the `trtllm.py`
attention backend. Update the dataclass definition for `TRTLLM`/the owning class
to include a type annotation and a `field(default=..., init=False, repr=False)`
entry, matching the existing convention used by `_flash_mla_metadata_valid` and
related private state. Then remove the plain assignment in the initialization
path and rely on the declared field so the member is consistent and static type
checkers can see it.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 73136071-c4c0-4817-8e5d-25e9aee47569

📥 Commits

Reviewing files that changed from the base of the PR and between 0457051 and c8b9154.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/attention_backend/trtllm.py
• tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt