[https://nvbugs/6402058][fix] [None][fix] Remove stale prepare_inputs event sync to unhang DS-V3 EP4 MTP

[trtllm-agent]

Summary

• Root cause: PyTorchModelEngine held a leftover _prepare_inputs_event recorded after _prepare_and_schedule_batch and synchronized via wait_for_input_copy(). In the DeepSeek-V3-Lite EP4 + MTP2 + attention_dp + CUDA-graph + overlap-scheduler + torch.compile configuration, this extra CUDA event synchronization interacted badly with the overlap scheduler on the 4-GPU post-merge stage, causing the executor to block indefinitely on IPC poll.
• Fix: Removed the unused _prepare_inputs_event field, its record site in the forward path, and the wait_for_input_copy() helper, since input-copy ordering is already guaranteed by the stream/graph semantics elsewhere. The corresponding waive entry for nvbugs/6402058 was dropped from waives.txt so the previously hanging DeepSeek-V3-Lite EP4 MTP2 test re-enters CI to confirm the fix.
• 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/6402058

Summary by CodeRabbit

• Bug Fixes
  • Simplified input handling by removing an extra synchronization step during request processing, reducing unnecessary waiting before inputs can be updated.
• Tests
  • Removed an outdated test waiver entry so the related integration test is now tracked normally.

[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: c6032043-f2f2-4073-8a19-36c3f890d94a

📥 Commits

Reviewing files that changed from the base of the PR and between 4cfcc64 and 43184ca.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/model_engine.py
• tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (2)

• tests/integration/test_lists/waives.txt
• tensorrt_llm/_torch/pyexecutor/model_engine.py

---

📝 Walkthrough

Walkthrough

This PR removes CUDA-event-based synchronization (_prepare_inputs_event) from PyTorchModelEngine, including its initialization, usage in the forward path, and the wait_for_input_copy() method. It also removes an unrelated test waiver entry from the waives list.

Changes

Prepare-Inputs Synchronization Removal

Layer / File(s)	Summary
Remove CUDA event synchronization tensorrt_llm/_torch/pyexecutor/model_engine.py	Removes _prepare_inputs_event field initialization, its creation/recording in the forward path, and the wait_for_input_copy() method that synchronized on it.
Remove obsolete test waiver tests/integration/test_lists/waives.txt	Deletes the SKIP waiver line for a specific TestDeepSeekV3Lite::test_bfloat16_4gpus parameterization.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: StanleySun639, LarryXFly, crazydemo

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific and matches the PR's main fix, despite a duplicate ticket/type prefix.
Description check	✅ Passed	It covers the bug, fix, and test plan, though the template's explicit Description, Test Coverage, and PR Checklist sections are not used.
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.}

[tongyuantongyu]

AI suggests this should wait #16088 to be merged first.

Detail

For the target DS-V3-Lite / V1 KV-cache path, yes: the original behavior is still correct because the race that _prepare_inputs_event papered over was later fixed at the actual mutable host-buffer source by #15546. But I would not claim this removal is generally safe for every V2/DSv4 path unless #16088 or an equivalent V2 staging fix is also present.

Why _prepare_inputs_event was needed

#14128 added _prepare_inputs_event to serialize the overlap scheduler against _prepare_inputs’ async host→device work:

• _prepare_inputs(...) builds attention/page-table inputs and enqueues non-blocking H2D copies.
• In the overlap scheduler, the CPU can start scheduling / updating the next iteration while the previous iteration’s H2D copies are still queued on execution_stream.
• Some H2D sources were persistent pinned host buffers, especially KV block/page-offset tables.
• The next iteration could overwrite those host buffers before the previous H2D copy actually read them, so the GPU could consume block offsets for the wrong batch.

That is why #14128 recorded an event immediately after _prepare_inputs and made the next executor iteration call wait_for_input_copy() before scheduler/resource mutations.

What made it unnecessary for the DS-V3-Lite target

The later fix is #15546: #15546

That PR moved the correctness boundary from a coarse global synchronization to safe per-call host staging in KVCacheManager.copy_batch_block_offsets.

In the current base, KVCacheManager.copy_batch_block_offsets no longer copies directly from the persistent host_kv_cache_block_offsets buffer. It first snapshots the rows needed for this batch into a fresh pinned buffer, then launches the non-blocking H2D from that private snapshot:

• resource_manager.py explains the exact race:
  • async H2D reads source at execution time
  • overlap scheduler can run one iteration ahead
  • persistent host buffer could be clobbered by the next iteration
• Then it stages through _stage_block_offsets_for_copy(...) before the H2D copy.

Relevant code/comments:

TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py

Lines 2186 to 2201 in 725d9a1

	# The H2D copies below are asynchronous, so their source is read at copy
	# execution time, not enqueue time. With the overlap scheduler the CPU
	# runs an iteration ahead, so copying straight from the persistent buffer
	# let the next iteration's in-place refill clobber the source of this
	# iteration's still-pending H2D -> the attention kernel indexed another
	# batch's blocks (nvbug 6293536; the window is widened when host
	# offloading stalls the execution stream in front of the H2D). Stage the
	# async copy through a freshly-allocated pinned buffer instead; the
	# caching host allocator holds it until the consuming copy completes,
	# matching the already-safe kv_lens / block_ids_per_seq staging. The
	# persistent buffer above is untouched by this and stays valid for the
	# synchronous CPU readers.
	host_block_offsets = self._stage_block_offsets_for_copy(num_seqs)
	for pool_idx in range(self.num_pools):
	dst_tensor[pool_idx, :num_seqs].copy_(host_block_offsets[pool_idx],
	non_blocking=True)

# The H2D copies below are asynchronous, so their source is read at copy
# execution time, not enqueue time. With the overlap scheduler the CPU
# runs an iteration ahead, so copying straight from the persistent buffer
# let the next iteration's in-place refill clobber the source of this
# iteration's still-pending H2D -> the attention kernel indexed another
# batch's blocks ...
host_block_offsets = self._stage_block_offsets_for_copy(num_seqs)
for pool_idx in range(self.num_pools):
    dst_tensor[pool_idx, :num_seqs].copy_(host_block_offsets[pool_idx],
                                          non_blocking=True)

That is the key distinction: after #15546, correctness no longer depends on preventing the next iteration from mutating the persistent host table. The in-flight copy reads a private snapshot, so it remains valid even if the next iteration reuses or rewrites the persistent table.

So for the DS-V3-Lite PR #15849 is targeting, the old event is redundant and harmful: it only adds a full CPU-side synchronization point in the overlap path, and that can interact badly with CUDA graph / overlap scheduler progress.

Why this is more than “it won’t crash”

The behavior remains correct because the ordering requirement changed:

• Before [https://nvbugs/6293536][fix] Stage KV block offsets through a fresh host buffer #15546: correctness required “do not mutate persistent host block-offset buffer until previous H2D has consumed it.”
• After [https://nvbugs/6293536][fix] Stage KV block offsets through a fresh host buffer #15546: the previous H2D no longer consumes the persistent buffer; it consumes a per-call snapshot. Therefore the next iteration may safely mutate the persistent buffer.

That preserves the semantic requirement — “the GPU sees the block offsets for the batch being executed” — without requiring a global event synchronize.

Important caveat: V2 / DSv4 is a separate case

The same broad race class exists in V2/DSv4 through a different mechanism. The current DSv4 manager still has paths that copy from persistent host staging into device staging asynchronously, for example:

TensorRT-LLM/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py

Lines 1261 to 1265 in 725d9a1

	device_copy_idx = self._copy_idx_to_device(copy_idx)
	self._device_kv_cache_block_offsets_input.copy_(
	self.host_kv_cache_block_offsets,
	non_blocking=True,
	)

self._device_kv_cache_block_offsets_input.copy_(
    self.host_kv_cache_block_offsets,
    non_blocking=True,
)

And #16088 explicitly calls out that V2 has the same class of race and needs a V2-specific staging/retirement fix:
#16088

So the strongest accurate statement is:

• Safe for the PR [https://nvbugs/6402058][fix] [None][fix] Remove stale prepare_inputs event sync to unhang DS-V3 EP4 MTP #15849 DS-V3-Lite/V1 target, because [https://nvbugs/6293536][fix] Stage KV block offsets through a fresh host buffer #15546 already fixed the underlying mutable-host-source race.
• Not proven as a general removal for all V2/DSv4 paths unless [https://nvbugs/6293536][fix] Stage v2 KV block offsets through fresh host buffers #16088 or equivalent V2 staging is also in the base.

Suggested PR explanation

I’d phrase the PR rationale like this:

_prepare_inputs_event was originally added as a coarse fence to prevent overlap-scheduler iterations from mutating persistent host input/page-table buffers before previous non-blocking H2D copies consumed them. For the DS-V3-Lite/V1 KV-cache path targeted here, that correctness requirement is now handled locally by KVCacheManager.copy_batch_block_offsets, which snapshots the persistent host block-offset table into a fresh pinned per-call buffer before launching the async H2D copy (#15546 / nvbug 6293536). Therefore the next iteration can safely mutate the persistent table without affecting the in-flight copy, and the global per-forward event synchronize is redundant. Removing it preserves correctness while avoiding an unnecessary overlap-scheduler synchronization that contributes to the nvbug/6402058 hang. Note that V2/DSv4 uses a different staging path and should rely on its own fix (#16088 or equivalent), not on this coarse model-engine event.

[jiaganc]

Even after #16088, we still need to insert a event wait here, because lots of buffers used by _prepare_inputs are not protected by double buffer.

But after #16088, we can remove the wait in executor loop, because only KVCM will modify CPU buffer out of _prepare_inputs.