Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
trtllm-agent wants to merge 1 commit into
base: main
Choose a base branch
from
+5 −0
Open

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

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions tensorrt_llm/_torch/attention/backends/trtllm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -944,6 +944,11 @@ def prepare_flash_mla(self) -> None:
block_ids_per_seq = maybe_pin_memory(
self.kv_cache_manager.get_block_ids_per_seq(self.request_ids))
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

@brnguyen2 brnguyen2 Aug 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

# garbage entries that trigger an OOB read in the FlashMLA kernel.
self._flash_mla_src_block_ids_cpu = block_ids_per_seq
self.kv_block_ids_per_seq.fill_(0)
self.kv_block_ids_per_seq[:self.num_seqs, :num_blocks].copy_(
block_ids_per_seq, non_blocking=True)
Expand Down
Loading