Add missing sync before the forward call.

Signed-off-by: SimengLiu-nv <simengl@nvidia.com>

Author: SimengLiu-nv

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -261,6 +261,7 @@ def __init__(self,
         # KVCacheTransferManager's onboard/offload operations.
         self.is_warmup = True
 
+        self.execution_stream.wait_stream(torch.cuda.current_stream())
         with torch.cuda.stream(self.execution_stream):
             self.model_engine.warmup(self.resource_manager)
             if self.draft_model_engine is not None:
@@ -2228,6 +2229,7 @@ def forward(scheduled_requests, resource_manager, new_tensors_device,
 
             # Run model forward on the execution stream for proper synchronization
             # with KVCacheTransferManager's onboard/offload operations.
+            self.execution_stream.wait_stream(torch.cuda.current_stream())
             with torch.cuda.stream(self.execution_stream):
                 outputs = forward(scheduled_requests, self.resource_manager,
                                   new_tensors_device, gather_context_logits,