[https://nvbugs/6422343][fix] Retain the CPU source tensor as self._flash_mla_src_block_ids_cpu on the… - #16071
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review. Walkthrough
ChangesFlashMLA buffer lifetime
Estimated code review effort: 1 (Trivial) | ~5 minutes Merge Risk: ⚪ Minimal · up to The change retains the CPU source tensor long enough for the asynchronous transfer, preventing stale block IDs and the associated kernel out-of-bounds read; no actionable merge-blocking risk remains after normal checks. Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/attention_backend/trtllm.py (1)
690-690: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeclare
_flash_mla_src_block_ids_cpuas a typed dataclass field for consistency.Sibling private state on this dataclass (e.g.
_flash_mla_metadata_valid) is declared viafield(default=..., init=False, repr=False)with a type annotation._flash_mla_src_block_ids_cpuis 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
📒 Files selected for processing (2)
tensorrt_llm/_torch/attention_backend/trtllm.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
36f5286 to
6b8c6ed
Compare
6b8c6ed to
7613c09
Compare
ZhanruiSunCh
left a comment
ZhanruiSunCh
left a comment
There was a problem hiding this comment.
LGTM for infra part.
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Approved. Please ensure that the un-waived case get passed before merging.
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
LGTM — the change retains the pinned CPU block-ids buffer past the non-blocking H2D copy (fixes a use-after-free under CUDA-graph capture) and is gated behind enable_flash_mla, so non-FlashMLA paths are unaffected.
9915876 to
66f3e15
Compare
66f3e15 to
a09ceb5
Compare
55ff069 to
9e1ad8a
Compare
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Two things before this lands:
-
Evidence for the waiver removals. Eight waivers across five GPU types are dropped, but the PR body only mentions the H20 ones and "verified on the same GPU type" without a link or a repro count. This failure is intermittent (illegal memory access during attention warmup); a passing run isn't distinguishable from a lucky one. Please post the repro command, iteration count, and the pre-fix failure rate you reproduced against.
-
The stated root cause needs backing — see the inline comment.
| num_blocks = block_ids_per_seq.shape[1] | ||
| # Retain the source CPU buffer so it outlives the non-blocking H2D | ||
| # copies below; back-to-back prepare_flash_mla calls during CUDA-graph | ||
| # capture can otherwise free it before the DMA completes, leaving |
There was a problem hiding this comment.
This retains the buffer only until the next prepare_flash_mla() call, which rebinds the attribute and drops the last reference. If the described hazard is real — back-to-back calls during capture freeing the source before the DMA lands — this moves the window by exactly one call rather than closing it: call N's buffer is still released while its copy may be in flight, and the allocator can hand that block back to call N+1's pin_memory().
Also, tensors from pin_memory() come from PyTorch's caching host allocator, which records a stream event on free and won't reuse a block until the event completes. So the premise that a plain refcount drop can corrupt an in-flight H2D needs evidence, not just plausibility. Please point at the mechanism that defeats that guard here (e.g. behavior under graph capture), or state that the sequence was actually observed.
If you want a fix that's robust by construction rather than by lifetime accounting, allocate a persistent pinned staging buffer on the metadata once, copy_ into it, and issue the non-blocking H2D from that — no per-call allocation, no lifetime question.
9e1ad8a to
952a332
Compare
|
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. |
…pletes prepare_flash_mla uses non_blocking=True copies from a locally-scoped pinned CPU tensor. Under heavy warmup (MTP+ADP+cuda_graph+torch_compile+chunked_prefill), the tensor can be reclaimed by Python before the DMA finishes, leaving stale/garbage entries in the device block-ID buffers and later producing an illegal memory access from the FlashMLA kernel. Pin the source buffer to the metadata object so its lifetime spans the async copy. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
952a332 to
ad14d09
Compare
|
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. |
6 participants
Summary
self._flash_mla_src_block_ids_cpuon the metadata object so its lifetime spans the async DMA. Also removed the H20 waivers for this bug.Test plan
Links
Dev Engineer Review
prepare_flash_mla()retains the pinned CPUblock_ids_per_seqtensor inself._flash_mla_src_block_ids_cpu.non_blocking=TrueH2D copy.QA Engineer Review
No test changes.