[TRTLLM-12732][fix] Fence V2 _batched_migrate behind execution_stream

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -2116,6 +2116,10 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
                                              event_manager=self.event_manager)
             else:
                 raise
+        # Pass execution_stream to StorageManager so _batched_migrate
+        # can record a fresh event (V1-style syncWithBufferManager)
+        # to order copy streams after the latest forward-pass work.
+        self.impl._storage._execution_stream = self._stream.cuda_stream
 
         if self.event_manager is not None:
             self.event_manager.set_layer_group_window_sizes(

tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py

@@ -612,8 +612,6 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo
                     slot = slots[lc].pop()
                     self._scratch_slots[lc].append(ScratchSlotLock(slot, self, lc, skip_wait=True))
 
-            normal_slots = tuple(chain.from_iterable(slots))
-            new_slot_ready_events = tuple(slot.ready_event for slot in normal_slots)
             for beam_indices in self._base_page_indices:
                 for indices in beam_indices:
                     if type(indices) is array.array:
@@ -622,23 +620,9 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo
                     else:
                         if len(indices) < new_num_blocks:
                             raise ValueError("User-provided base page indices is too short")
-            if any(
-                type(indices) is memoryview
-                for beam_indices in self._base_page_indices
-                for indices in beam_indices
-            ):
-                # User-provided base page indices are host-visible. Do not publish
-                # a recycled slot id there until the slot is actually ready; a
-                # stream wait only orders GPU work and does not protect CPU-side
-                # metadata readers.
-                pending_ready_events = [
-                    event for event in set(new_slot_ready_events) if not event.query_complete()
-                ]
-                for event in pending_ready_events:
-                    event.synchronize()
-                for slot in normal_slots:
-                    slot.ready_event = CachedCudaEvent.NULL
-            stream_wait_events(self.cuda_stream, new_slot_ready_events)
+            stream_wait_events(
+                self.cuda_stream, (s.ready_event for s in chain.from_iterable(slots))
+            )
             for ordinal in typed_range(old_num_blocks, new_num_blocks):
                 block = make_typed(
                     lambda _: filled_list(cast(BlockPage, None), num_life_cycles), beam_width
@@ -656,9 +640,7 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo
                             if stale_beg <= ordinal < stale_end:
                                 continue
                         slot = slots[lc].pop()
-                        # Slot ready events are already handled above: either
-                        # synchronized for host-visible page-index buffers, or
-                        # enqueued as stream waits for internal buffers.
+                        # We have already waited for ready_event of the slots.
                         block[beam_index][lc] = UncommittedPage(
                             self, ordinal, lc, GPU_LEVEL, slot, beam_index
                         ).lock(self, beam_index, ordinal, lc, skip_wait=True)

tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py

@@ -29,6 +29,7 @@
     BlockOrdinal,
     CacheLevel,
     CacheTier,
+    CudaStream,
     LayerId,
     MemAddress,
     PageStatus,
@@ -175,6 +176,7 @@ class StorageManager:
         "_levels",
         "_min_slots",
         "_event_manager",
+        "_execution_stream",
         "__rawref__",
     )
     _life_cycles: LifeCycleRegistry
@@ -188,6 +190,7 @@ class StorageManager:
     _levels: TypedIndexList[CacheLevel, CacheLevelManager]
     _min_slots: TypedIndexList[PoolGroupIndex, int]
     _event_manager: "KVCacheEventManager | None"
+    _execution_stream: CudaStream | None
     __rawref__: rawref.ref["StorageManager"]
 
     def __init__(
@@ -202,6 +205,7 @@ def __init__(
     ) -> None:
         self.__rawref__ = rawref.NULL
         self._event_manager = event_manager
+        self._execution_stream: CudaStream | None = None
         assert config.cache_tiers[GPU_LEVEL].tier == CacheTier.GPU_MEM, (
             "The first cache tier must be GPU memory"
         )
@@ -510,6 +514,14 @@ def _batched_migrate(
         try:
             assert len(dst_slots) == num_slots
             prior_events: set[CachedCudaEvent] = set()
+            # Like V1's syncWithBufferManager(): record a fresh event on
+            # execution_stream so the copy stream waits for all prior
+            # forward-pass work, not just the (potentially stale) page
+            # ready_events.  This is the V2 equivalent of the explicit
+            # two-way sync that V1 performs between its BufferManager
+            # stream and offload/onboard streams.
+            if self._execution_stream is not None and not defrag:
+                prior_events.add(CachedCudaEvent(self._execution_stream))
             tasks_per_pool: TypedIndexList[PoolIndex, list[CopyTask]] = make_typed(
                 lambda _: list[CopyTask](), num_pools
             )

tests/integration/test_lists/waives.txt

@@ -87,8 +87,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backe
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6162115)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6095421)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6084824)
-accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_4gpus_static_eplb[moe_backend=MEGAMOE_DEEPGEMM] SKIP (https://nvbugs/6183549)
-accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_4gpus_static_eplb[moe_backend=TRTLLM] SKIP (worker init CUDA error masked by kv_cache_manager_v2/_exceptions.py:49 cuGetErrorString type bug; real failure unreadable until that wrapper is fixed)
 accuracy/test_llm_api_pytorch.py::TestEXAONE4::test_auto_dtype SKIP (https://nvbugs/6164924)
 accuracy/test_llm_api_pytorch.py::TestGLM4_5Air::test_nvfp4_2_model_mtp[2model] SKIP (https://nvbugs/5981293)
 accuracy/test_llm_api_pytorch.py::TestGLM4_5Air::test_nvfp4_2_model_mtp[2model_trtllm] SKIP (https://nvbugs/5981293)