From 9dabe409a16471595ac84606582855054b1cf33c Mon Sep 17 00:00:00 2001
From: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
Date: Wed, 14 Oct 2026 03:04:03 -0700
Subject: [PATCH 1/7] Reuse PyExecutor across profiling resource transitions

Signed-off-by: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py       | 101 +++++--
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 251 ++++++++++++++----
 .../_torch/pyexecutor/py_executor_creator.py  | 129 ++++-----
 .../kv_cache/test_kv_cache_estimation.py      |  34 ++-
 .../executor/test_executor_resource_rebind.py | 248 +++++++++++++++++
 ...y_executor_creator_mla_cache_reuse_sync.py |  62 ++++-
 6 files changed, 679 insertions(+), 146 deletions(-)
 create mode 100644 tests/unittest/_torch/executor/test_executor_resource_rebind.py

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 1e18bf19e6f3..5539224245c4 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -73,7 +73,7 @@
                                            use_py_mamba_cache_manager)
 from .llm_request import ExecutorResponse, LlmRequestState
 from .model_engine import PyTorchModelEngine
-from .py_executor import PyExecutor
+from .py_executor import PyExecutor, PyExecutorResources
 from .resource_manager import (KVCacheCompressionManager, KVCacheManager,
                                PeftCacheManager, ResourceManager,
                                ResourceManagerType)
@@ -1830,7 +1830,7 @@ def configure_kv_cache_capacity(self,
                     kv_stats_draft.allocated_bytes
                     if kv_stats_draft is not None else 0)
                 py_executor.is_warmup = False
-                py_executor.shutdown_for_kv_cache_estimation()
+                py_executor.release_resources()
                 py_executor.enable_iter_perf_stats = origin_iter_stats
                 py_executor.set_gather_responses(False)
 
@@ -4123,7 +4123,7 @@ def validate_seq_slot_pool_covers_admission(max_num_sequences: int,
         "shortfall means one of them was re-derived from max_batch_size.")
 
 
-def create_py_executor_instance(
+def create_py_executor_resources(
     *,
     dist,
     resources,
@@ -4131,14 +4131,8 @@ def create_py_executor_instance(
     llm_args,
     ctx_chunk_config,
     model_engine,
-    start_worker,
-    sampler,
-    drafter,
-    guided_decoder: Optional[GuidedDecoder] = None,
     lora_config: Optional[LoraConfig] = None,
-    garbage_collection_gen0_threshold: Optional[int] = None,
     kv_connector_manager: Optional[KvCacheConnectorManager] = None,
-    resource_governor_queue=None,
     max_seq_len: Optional[int] = None,
     max_batch_size: Optional[int] = None,
     max_beam_width: Optional[int] = None,
@@ -4150,14 +4144,13 @@ def create_py_executor_instance(
     execution_stream: Optional[torch.cuda.Stream] = None,
     dwdp_manager: Optional[DwdpManager] = None,
     max_num_sequences: Optional[int] = None,
-) -> PyExecutor:
+) -> PyExecutorResources:
+    """Build the resource graph shared by initial execution and rebinding."""
     set_low_latency_dispatch(
         getattr(llm_args, 'enable_low_latency_host_dispatch', False))
 
     kv_cache_manager = resources.get(ResourceManagerType.KV_CACHE_MANAGER, None)
 
-    spec_config = model_engine.spec_config
-
     is_disagg = is_disagg_enabled(cache_transceiver_config)
 
     max_num_sequences = resolve_max_num_sequences(
@@ -4525,14 +4518,78 @@ def create_py_executor_instance(
             executor_max_seq_len = min(executor_max_seq_len,
                                        int(decoder_position_limit))
 
+    return PyExecutorResources(
+        resource_manager=resource_manager,
+        scheduler=scheduler,
+        kv_cache_transceiver=kv_cache_transceiver,
+        peft_cache_config=peft_cache_config,
+        virtual_memory_pools=virtual_memory_pools,
+        kv_connector_manager=kv_connector_manager,
+        dwdp_manager=dwdp_manager,
+        max_seq_len=executor_max_seq_len,
+        max_num_sequences=max_num_sequences,
+        waiting_queue_policy=waiting_queue_policy,
+    )
+
+
+def create_py_executor_instance(
+    *,
+    dist,
+    resources,
+    mapping,
+    llm_args,
+    ctx_chunk_config,
+    model_engine,
+    start_worker,
+    sampler,
+    drafter,
+    guided_decoder: Optional[GuidedDecoder] = None,
+    lora_config: Optional[LoraConfig] = None,
+    garbage_collection_gen0_threshold: Optional[int] = None,
+    kv_connector_manager: Optional[KvCacheConnectorManager] = None,
+    resource_governor_queue=None,
+    max_seq_len: Optional[int] = None,
+    max_batch_size: Optional[int] = None,
+    max_beam_width: Optional[int] = None,
+    max_num_tokens: Optional[int] = None,
+    peft_cache_config: Optional[PeftCacheConfig] = None,
+    scheduler_config: Optional[SchedulerConfig] = None,
+    cache_transceiver_config: Optional[CacheTransceiverConfig] = None,
+    virtual_memory_pools: Optional[dict] = None,
+    execution_stream: Optional[torch.cuda.Stream] = None,
+    dwdp_manager: Optional[DwdpManager] = None,
+    max_num_sequences: Optional[int] = None,
+) -> PyExecutor:
+    execution_resources = create_py_executor_resources(
+        dist=dist,
+        resources=resources,
+        mapping=mapping,
+        llm_args=llm_args,
+        ctx_chunk_config=ctx_chunk_config,
+        model_engine=model_engine,
+        lora_config=lora_config,
+        kv_connector_manager=kv_connector_manager,
+        max_seq_len=max_seq_len,
+        max_batch_size=max_batch_size,
+        max_beam_width=max_beam_width,
+        max_num_tokens=max_num_tokens,
+        peft_cache_config=peft_cache_config,
+        scheduler_config=scheduler_config,
+        cache_transceiver_config=cache_transceiver_config,
+        virtual_memory_pools=virtual_memory_pools,
+        execution_stream=execution_stream,
+        dwdp_manager=dwdp_manager,
+        max_num_sequences=max_num_sequences,
+    )
+    spec_config = model_engine.spec_config
     return PyExecutor(
-        resource_manager,
-        scheduler,
+        execution_resources.resource_manager,
+        execution_resources.scheduler,
         model_engine=model_engine,
         sampler=sampler,
         drafter=drafter,
         dist=dist,
-        max_num_sequences=max_num_sequences,
+        max_num_sequences=execution_resources.max_num_sequences,
         disable_overlap_scheduler=llm_args.disable_overlap_scheduler,
         enable_early_first_token_response=llm_args.
         enable_early_first_token_response,
@@ -4542,18 +4599,18 @@ def create_py_executor_instance(
         if spec_config is not None else 0,
         max_total_draft_tokens=(spec_config.tokens_per_gen_step -
                                 1) if spec_config is not None else 0,
-        kv_cache_transceiver=kv_cache_transceiver,
+        kv_cache_transceiver=execution_resources.kv_cache_transceiver,
         guided_decoder=guided_decoder,
         start_worker=start_worker,
         garbage_collection_gen0_threshold=garbage_collection_gen0_threshold,
-        kv_connector_manager=kv_connector_manager,
+        kv_connector_manager=execution_resources.kv_connector_manager,
         resource_governor_queue=resource_governor_queue,
-        max_seq_len=executor_max_seq_len,
-        peft_cache_config=peft_cache_config,
-        virtual_memory_pools=virtual_memory_pools,
+        max_seq_len=execution_resources.max_seq_len,
+        peft_cache_config=execution_resources.peft_cache_config,
+        virtual_memory_pools=execution_resources.virtual_memory_pools,
         execution_stream=execution_stream,
-        waiting_queue_policy=waiting_queue_policy,
-        dwdp_manager=dwdp_manager,
+        waiting_queue_policy=execution_resources.waiting_queue_policy,
+        dwdp_manager=execution_resources.dwdp_manager,
         enable_kv_pool_rebalance=llm_args.kv_cache_config.
         enable_kv_pool_rebalance,
         kv_pool_rebalance_check_interval=llm_args.kv_cache_config.
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 0b36518c142c..08febdd1f8a4 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -408,6 +408,41 @@ class PendingEncoderStep:
     result: Optional[EncoderStepResult] = None
 
 
+@dataclasses.dataclass
+class PyExecutorResources:
+    """Resources and capacities replaced between profiling and serving.
+
+    The executor owns their execution lifetime. Model engines and weights
+    remain owned by the executor across resource transitions.
+    """
+
+    resource_manager: ResourceManager
+    scheduler: RequestScheduler
+    kv_cache_transceiver: Optional[KvCacheTransceiver]
+    peft_cache_config: Optional[PeftCacheConfig]
+    virtual_memory_pools: Optional[dict]
+    kv_connector_manager: Optional[KvCacheConnectorManager]
+    dwdp_manager: Optional[DwdpManager]
+    max_seq_len: Optional[int]
+    max_num_sequences: int
+    waiting_queue_policy: WaitingQueuePolicy
+
+
+@dataclasses.dataclass(frozen=True)
+class _ExecutorExecutionConfig:
+    disable_overlap_scheduler: bool
+    enable_early_first_token_response: bool
+    max_input_len: int
+    max_batch_size: int
+    max_beam_width: int
+    max_draft_len: int
+    max_total_draft_tokens: int
+    garbage_collection_gen0_threshold: Optional[int]
+    resource_governor_queue: object
+    hang_detection_timeout: Optional[int]
+    enable_kv_pool_rebalance: bool
+
+
 class PyExecutor:
     # Minimum number of async micro batches for async PP execution.
     # This is a trade-off between memory usage and performance.
@@ -523,9 +558,79 @@ def __init__(
         # the manager operates on those fields via a weakref.
         self._profile_manager = PyExecutorProfileManager(self)
 
+        self.model_engine = model_engine
+        self.sampler = sampler
+        self.drafter = drafter
+        self.guided_decoder = guided_decoder
+        self.dist = dist
+        self._execution_config = _ExecutorExecutionConfig(
+            disable_overlap_scheduler=disable_overlap_scheduler,
+            enable_early_first_token_response=enable_early_first_token_response,
+            max_input_len=max_input_len,
+            max_batch_size=max_batch_size,
+            max_beam_width=max_beam_width,
+            max_draft_len=max_draft_len,
+            max_total_draft_tokens=max_total_draft_tokens,
+            garbage_collection_gen0_threshold=garbage_collection_gen0_threshold,
+            resource_governor_queue=resource_governor_queue,
+            hang_detection_timeout=hang_detection_timeout,
+            enable_kv_pool_rebalance=enable_kv_pool_rebalance,
+        )
+        self._resources_released = False
+        self._closed = False
+        self._initialize_execution(
+            PyExecutorResources(
+                resource_manager=resource_manager,
+                scheduler=scheduler,
+                kv_cache_transceiver=kv_cache_transceiver,
+                peft_cache_config=peft_cache_config,
+                virtual_memory_pools=virtual_memory_pools,
+                kv_connector_manager=kv_connector_manager,
+                dwdp_manager=dwdp_manager,
+                max_seq_len=max_seq_len,
+                max_num_sequences=max_num_sequences,
+                waiting_queue_policy=waiting_queue_policy,
+            ))
+        if start_worker:
+            self.start_worker()
+
+    def _initialize_execution(self, resources: PyExecutorResources) -> None:
+        """Initialize a stopped execution phase and warm its bound resources."""
+        model_engine = self.model_engine
+        sampler = self.sampler
+        drafter = self.drafter
+        guided_decoder = self.guided_decoder
+        dist = self.dist
+        resource_manager = resources.resource_manager
+        scheduler = resources.scheduler
+        kv_cache_transceiver = resources.kv_cache_transceiver
+        peft_cache_config = resources.peft_cache_config
+        virtual_memory_pools = resources.virtual_memory_pools
+        kv_connector_manager = resources.kv_connector_manager
+        dwdp_manager = resources.dwdp_manager
+        max_seq_len = resources.max_seq_len
+        max_num_sequences = resources.max_num_sequences
+        waiting_queue_policy = resources.waiting_queue_policy
+        disable_overlap_scheduler = self._execution_config.disable_overlap_scheduler
+        enable_early_first_token_response = self._execution_config.enable_early_first_token_response
+        max_input_len = self._execution_config.max_input_len
+        max_batch_size = self._execution_config.max_batch_size
+        max_beam_width = self._execution_config.max_beam_width
+        max_draft_len = self._execution_config.max_draft_len
+        max_total_draft_tokens = self._execution_config.max_total_draft_tokens
+        garbage_collection_gen0_threshold = self._execution_config.garbage_collection_gen0_threshold
+        resource_governor_queue = self._execution_config.resource_governor_queue
+        hang_detection_timeout = self._execution_config.hang_detection_timeout
+        enable_kv_pool_rebalance = self._execution_config.enable_kv_pool_rebalance
+        self.peft_cache_config = peft_cache_config
+        self.iter_counter = 0
+        self._metrics = {}
+        self._disagg_coordinator = None
+        self._ctx_mla_kv_len_cap = "unset"
         # related modules
         self.resource_manager = resource_manager
-        self.model_engine = model_engine
+        self.kv_connector_manager = kv_connector_manager
+        self.dwdp_manager = dwdp_manager
         self._enable_adp_dummy_fixes = getattr(model_engine,
                                                "_enable_adp_dummy_fixes", False)
         self._enable_scheduler_aware_adp_dummy = getattr(
@@ -542,8 +647,6 @@ def __init__(
         # the hang-detector status provider are initialized.
         self.kv_cache_transceiver = kv_cache_transceiver
         self.enable_attention_dp = model_engine.enable_attention_dp
-        self.dist = dist
-        self.sampler = sampler
         # Opt-in in-graph sampling: let the engine sample inside its forward CUDA
         # graph. The sampler resolves each batch to a tier and stages the
         # buffers the in-graph step reads; the engine keys its graph cache on
@@ -557,10 +660,8 @@ def __init__(
                 sampler.stage_in_graph_sampling)
             model_engine.register_sample_in_graph_callable(
                 self._sample_in_graph)
-        self.drafter = drafter
         self.draft_model_engine = getattr(self.drafter, "draft_model_engine",
                                           None)
-        self.guided_decoder = guided_decoder
         self.disable_overlap_scheduler = disable_overlap_scheduler
         self.enable_kv_pool_rebalance = enable_kv_pool_rebalance
         # Iteration throttle for the KV pool rebalance check.  See
@@ -1160,16 +1261,29 @@ def on_detected():
         # and race on the shared communicator, consuming the wrong ACKs.
         self._sleep_wakeup_lock = threading.Lock()
 
-        self.kv_connector_manager = kv_connector_manager
-
         self._maybe_init_kv_connector_manager()
 
-        self.dwdp_manager = dwdp_manager
-
         self._validate_moe_a2a_checkpoint_support()
 
-        if start_worker:
-            self.start_worker()
+    def rebind_resources(self, resources: PyExecutorResources) -> None:
+        """Bind fresh resources after release_resources(), without starting workers.
+
+        Call collectively on the startup thread. The previous phase must have
+        exited cleanly; failed and finally closed executors cannot be reused.
+        Warmup captures graphs against the new resource addresses.
+        """
+        if self._closed or not self._resources_released:
+            raise RuntimeError(
+                "Executor resources must be released before rebinding")
+        if self._event_loop_error is not None or self._fatal_error is not None:
+            raise RuntimeError(
+                "Cannot rebind an executor after an execution failure")
+        self._resources_released = False
+        try:
+            self._initialize_execution(resources)
+        except BaseException:
+            self.shutdown()
+            raise
 
     @property
     def metrics(self) -> dict[str, float | dict[str, float]]:
@@ -1961,33 +2075,81 @@ def _stop_mpi_progress_thread(self) -> None:
             self._mpi_progress_thread = None
             self._mpi_progress_thread_exit_hook = None
 
-    def shutdown(self):
-        """Shut down the terminal executor and all CUDA graphs it owns."""
-        self._shutdown(release_model_owned_graphs=True)
+    def release_resources(self) -> None:
+        """Stop a profiling phase and detach resources before cache replacement.
 
-    def shutdown_for_kv_cache_estimation(self):
-        """Shut down a temporary profiling executor whose model is reused."""
-        self._shutdown(release_model_owned_graphs=False)
+        Call collectively after consuming profiling responses. Graph destruction
+        precedes workspace release, and transfer registrations precede cache
+        release. Model engines, weights and the execution stream remain valid.
+        """
+        if self._closed or self._resources_released:
+            raise RuntimeError("Executor resources are already released")
+        if not self._release_execution_resources():
+            raise RuntimeError(
+                "Cannot release resources after an executor hang")
+        for engine in (self.model_engine, self.draft_model_engine):
+            if engine is not None:
+                engine.attn_metadata = None
+        self.resource_manager = None
+        self.scheduler = None
+        self.kv_cache_manager = None
+        self.draft_kv_cache_manager = None
+        self.kv_cache_transceiver = None
+        self.kv_connector_manager = None
+        self.async_transfer_manager = None
+        self.adp_router = None
+        self._kv_load_tracker = None
+        self.hang_detector = None
+        self._sleep_wakeup_comm = None
+        self._sleep_wakeup_listener_thread = None
+        self.request_broadcaster = None
+        self._disagg_coordinator = None
+        self._disagg_pp_termination_handler = None
+        self.previous_batch = None
+        self.micro_batches = []
+        self.pending_encoder_steps.clear()
+        self.active_requests.clear()
+        self.responses.clear()
+        self.result_wait_queues.clear()
+        self._pending_transfer_responses.clear()
+        self._pending_response_terminations.clear()
+        self.request_accumulated.clear()
+        self.control_requests.clear()
+        self.waiting_queue = None
+        self._pending_adp_dummy_request = None
+        if self.dist.pp_size > 1:
+            self.executed_batch_queue = None
+            self.executed_batch_response_queue = None
+        self._resources_released = True
 
-    def _shutdown(self, *, release_model_owned_graphs: bool):
-        """Shared shutdown implementation for explicit executor lifecycles.
+    def shutdown(self) -> None:
+        """Finally close the executor; closed executors cannot be rebound."""
+        if self._closed:
+            return
+        if not self._resources_released:
+            if not self._release_execution_resources():
+                return
+        self._closed = True
+        del self.model_engine
+        if self.draft_model_engine is not None:
+            del self.draft_model_engine
 
-        Model-owned graphs are terminal model resources, unlike executor-owned
-        graphs that may reference the resource managers torn down below.
-        """
+    def _release_execution_resources(self) -> bool:
+        """Stop execution and release GPU resources; retain model ownership."""
         node_server = getattr(self, "_openengine_node_server", None)
         if node_server is not None:
             node_server.close()
             self._openengine_node_server = None
-        self.executor_request_queue.enqueue_shutdown_request()
-        self.shutdown_event.wait()
+        if self.worker_started:
+            self.executor_request_queue.enqueue_shutdown_request()
+            self.shutdown_event.wait()
         # Tear down any profile window an HTTP caller left open (i.e.
         # /start_profile without a matching /stop_profile).  Runs before the
         # hang-detector early return below so the leaked window is closed on
         # every shutdown path.  A no-op when the executor loop already
         # unwound its own profiling state.
         self._profile_manager.cleanup()
-        if self.hang_detector.detected():
+        if self.hang_detector is not None and self.hang_detector.detected():
             # Early return here to avoid waiting for hanging threads.
             # Since `on_detected` has sent the error message as response,
             # this worker will be asked to shutdown immediately.
@@ -1999,11 +2161,12 @@ def _shutdown(self, *, release_model_owned_graphs: bool):
             # already in flight is not interrupted.
             self._pp_mpi_progress_stop.set()
             self._stop_mpi_progress_thread()
-            return
-        self.worker_thread.join()
+            return False
+        if self.worker_started:
+            self.worker_thread.join()
         if self.kv_connector_manager is not None:
             self.kv_connector_manager.shutdown()
-        if self.dist.pp_size > 1:
+        if self.dist.pp_size > 1 and self.worker_started:
             # The worker thread has exited, so no PP traffic needs progress.
             self._pp_mpi_progress_stop.set()
             self.executed_batch_queue.put(None)
@@ -2023,6 +2186,11 @@ def _shutdown(self, *, release_model_owned_graphs: bool):
             encoder_launch_executor.shutdown(wait=True)
             self.encoder_launch_executor = None
 
+        # Stop the sampler's async worker, if it was used
+        if (isinstance(self.sampler, AsyncWorkerMixin)
+                and self.sampler.async_worker_enabled()):
+            self.sampler.async_worker_stop()
+
         self.worker_started = False
         # Release CUDA graphs before resource managers free their GPU memory.
         # Resource managers (e.g. SuffixAutomatonManager) allocate GPU workspace
@@ -2043,33 +2211,20 @@ def _shutdown(self, *, release_model_owned_graphs: bool):
         # resource managers start freeing GPU-backed workspaces.
         if torch.cuda.is_available():
             torch.cuda.synchronize()
-        for manager in self.resource_manager.resource_managers.values():
-            if manager:
-                manager.shutdown()
-        # Note: do NOT call engine.cleanup() here. PyExecutor.shutdown() is
-        # paired with shutdown_for_kv_cache_estimation() mid-init in
-        # configure_kv_cache_capacity(); that path immediately reads
-        # model_engine.model.model_config
-        # to compute kv_cache_max_memory. cleanup() would set
-        # model_engine.model = None, breaking that read with
-        # `'NoneType' object has no attribute 'model_config'`.
-        # The engine's __del__ still calls cleanup() at terminal teardown
-        # (when the executor's reference is dropped), which is sufficient for
-        # the GMS daemon registry eviction the cleanup hook was added for.
-        del self.model_engine
-        if self.draft_model_engine is not None:
-            del self.draft_model_engine
+        if self.kv_cache_transceiver is not None:
+            self.kv_cache_transceiver.shutdown()
+        if self.resource_manager is not None:
+            for manager in self.resource_manager.resource_managers.values():
+                if manager:
+                    manager.shutdown()
         if self.virtual_memory_pools is not None:
             keys = list(self.virtual_memory_pools.keys())
             for key in keys:
                 del self.virtual_memory_pools[key]
-        # Stop the sampler's async worker, if it was used
-        if (isinstance(self.sampler, AsyncWorkerMixin)
-                and self.sampler.async_worker_enabled()):
-            self.sampler.async_worker_stop()
         if self.dwdp_manager is not None:
             self.dwdp_manager.__exit__(None, None, None)
             self.dwdp_manager = None
+        return True
 
     def can_enqueue_requests(self) -> bool:
         """
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index dbfff0093d7f..d6faa2029acb 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -43,8 +43,8 @@
 from ..virtual_memory import scope as virtual_memory_scope
 from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
                     compute_max_num_sequences, create_py_executor_instance,
-                    instantiate_sampler, is_disagg_enabled, is_mla,
-                    validate_feature_combination)
+                    create_py_executor_resources, instantiate_sampler,
+                    is_disagg_enabled, is_mla, validate_feature_combination)
 from .config_utils import (is_hybrid_linear, is_minimax_m3,
                            resolve_cache_transceiver_config,
                            uses_fp4_mla_attention, uses_vswa_kv_cache_layout)
@@ -1156,7 +1156,7 @@ def allocation_scope(current_stage: ExecutorMemoryType):
     with allocation_scope(
             ExecutorMemoryType.INIT_EXTRA_RESOURCES
             if estimating_kv_cache else ExecutorMemoryType.EXTRA_RESOURCES):
-        # run gc.collect() to free memory of the previous py_executor, avoid cudaFree overlap with cuda graph capture
+        # Reclaim unreachable allocations before warmup.
         gc.collect()
         py_executor = create_py_executor_instance(
             dist=dist,
@@ -1191,77 +1191,58 @@ def allocation_scope(current_stage: ExecutorMemoryType):
         peft_cache_config = py_executor.peft_cache_config
 
     if estimating_kv_cache:
-        assert kv_cache_creator is not None
-        _move_model_engine_metrics(py_executor, model_engine, "initial")
-        # record initial metrics as py_executor will be deleted later
-        initial_model_engine_metrics = dict(py_executor.metrics)
-        with allocation_scope(ExecutorMemoryType.MODEL_EXTRA):
-            kv_cache_creator.configure_kv_cache_capacity(py_executor)
-
-        with _startup_timer.phase("profiling_resource_teardown"):
-            # Shut down the transceiver before tearing down KV cache managers so
-            # that NIXL-registered (pinned) GPU memory is deregistered first;
-            # otherwise the old KV cache memory stays pinned and the subsequent
-            # KV cache allocation will OOM.
-            try:
-                if hasattr(py_executor, 'kv_cache_transceiver'
-                           ) and py_executor.kv_cache_transceiver is not None:
-                    with _startup_timer.phase("profiling_transceiver_shutdown"):
-                        py_executor.kv_cache_transceiver.shutdown()
-            finally:
-                with _startup_timer.phase("profiling_kv_cache_teardown"):
-                    kv_cache_creator.teardown_managers(resources)
-
-            # configure_kv_cache_capacity shuts down the Phase-1 executor, which
-            # releases its CUDA graphs before its resource managers. Only the
-            # profiling attention metadata remains to be discarded here.
-            model_engine.attn_metadata = None
-
-            del py_executor  # free before constructing new
-            gc.collect()
-            torch.cuda.empty_cache()
-
-        with allocation_scope(ExecutorMemoryType.KV_CACHE):
-            # Before estimating KV cache size, a minimal KV cache has been allocated using
-            # create_kv_cache_manager above, which caps kv_cache_creator.max_seq_len. Restoring
-            # the original value before creating the final KV cache.
-            kv_cache_creator._max_seq_len = model_engine_max_seq_len
-            kv_cache_creator.build_managers(resources, False)
-
-        model_engine._warmup_timer.purpose = "final_executor"
-
-        with allocation_scope(ExecutorMemoryType.EXTRA_RESOURCES):
-
-            # run gc.collect() to free memory of the previous py_executor, avoid cudaFree overlap with cuda graph capture
-            gc.collect()
-            py_executor = create_py_executor_instance(
-                dist=dist,
-                resources=resources,
-                mapping=mapping,
-                llm_args=llm_args,
-                ctx_chunk_config=ctx_chunk_config,
-                model_engine=model_engine,
-                start_worker=False,
-                sampler=sampler,
-                drafter=drafter,
-                guided_decoder=guided_decoder,
-                lora_config=lora_config,
-                garbage_collection_gen0_threshold=
-                garbage_collection_gen0_threshold,
-                kv_connector_manager=kv_connector_manager,
-                resource_governor_queue=resource_governor_queue,
-                max_seq_len=net_max_seq_len,
-                max_batch_size=max_batch_size,
-                max_beam_width=max_beam_width,
-                max_num_tokens=max_num_tokens,
-                peft_cache_config=peft_cache_config,
-                scheduler_config=scheduler_config,
-                cache_transceiver_config=cache_transceiver_config,
-                virtual_memory_pools=vm_pools,
-                execution_stream=execution_stream,
-                dwdp_manager=dwdp_manager,
-                max_num_sequences=max_num_seq_slots,
-            )
+        try:
+            assert kv_cache_creator is not None
+            _move_model_engine_metrics(py_executor, model_engine, "initial")
+            # Preserve profiling metrics across execution-state reinitialization.
+            initial_model_engine_metrics = dict(py_executor.metrics)
+            with allocation_scope(ExecutorMemoryType.MODEL_EXTRA):
+                kv_cache_creator.configure_kv_cache_capacity(py_executor)
+
+            with _startup_timer.phase("profiling_resource_teardown"):
+                # Profiling has released graphs, registrations and manager resources.
+                kv_cache_creator.teardown_managers(resources)
+                gc.collect()
+                torch.cuda.empty_cache()
+
+            with allocation_scope(ExecutorMemoryType.KV_CACHE):
+                # Before estimating KV cache size, a minimal KV cache has been allocated using
+                # create_kv_cache_manager above, which caps kv_cache_creator.max_seq_len. Restoring
+                # the original value before creating the final KV cache.
+                kv_cache_creator._max_seq_len = model_engine_max_seq_len
+                kv_cache_creator.build_managers(resources, False)
+
+            model_engine._warmup_timer.purpose = "final_executor"
+
+            with allocation_scope(ExecutorMemoryType.EXTRA_RESOURCES):
+
+                # Reclaim detached profiling state before capturing final graphs.
+                gc.collect()
+                execution_resources = create_py_executor_resources(
+                    dist=dist,
+                    resources=resources,
+                    mapping=mapping,
+                    llm_args=llm_args,
+                    ctx_chunk_config=ctx_chunk_config,
+                    model_engine=model_engine,
+                    lora_config=lora_config,
+                    kv_connector_manager=kv_connector_manager,
+                    max_seq_len=net_max_seq_len,
+                    max_batch_size=max_batch_size,
+                    max_beam_width=max_beam_width,
+                    max_num_tokens=max_num_tokens,
+                    peft_cache_config=peft_cache_config,
+                    scheduler_config=scheduler_config,
+                    cache_transceiver_config=cache_transceiver_config,
+                    virtual_memory_pools=vm_pools,
+                    execution_stream=execution_stream,
+                    dwdp_manager=dwdp_manager,
+                    max_num_sequences=max_num_seq_slots,
+                )
+                py_executor.rebind_resources(execution_resources)
+        except BaseException:
+            py_executor.shutdown()
+            raise
 
     _adjust_torch_mem_fraction()
 
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py
index 387f82e64804..ae03f8120558 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py
@@ -1544,7 +1544,7 @@ def test_estimation_temporarily_uses_inferred_pool_sizing(
     assert kv_cache_config.max_gpu_total_bytes == expected_budget
     assert creator._fp8_ctx_mla_kv_len_cap == expected_cap
     py_executor.start_worker.assert_called_once()
-    py_executor.shutdown_for_kv_cache_estimation.assert_called_once()
+    py_executor.release_resources.assert_called_once()
     py_executor.shutdown.assert_not_called()
     assert kv_cache_config.max_tokens == user_max_tokens
     assert kv_cache_config.pool_ratio == pool_ratio
@@ -1964,3 +1964,35 @@ def test_create_kv_cache_manager_layer_specs_preserve_window_guards(guard: str)
 
     assert manager_config is kv_cache_config
     assert manager_config.max_attention_window == ([1024] if guard == "explicit_window" else None)
+
+
+def test_profiling_failure_releases_resources_and_restores_response_state():
+    creator = object.__new__(KvCacheCreator)
+    creator._mapping = SimpleNamespace(cp_config={})
+    creator._kv_cache_config = SimpleNamespace(free_gpu_memory_fraction=0.9)
+    creator._skip_est = False
+    creator._dummy_reqs = []
+    creator._encode_dummy_inputs = Mock(return_value=None)
+    executor = Mock()
+    executor.dist.mapping.rank = 0
+    executor.dist.broadcast.side_effect = lambda value, root: value
+    executor.enable_iter_perf_stats = True
+    executor.await_responses.side_effect = RuntimeError("profiling request failed")
+    kv_manager = Mock()
+    kv_manager.get_kv_cache_stats.return_value = SimpleNamespace(allocated_bytes=0)
+    executor.resource_manager.resource_managers.get.side_effect = [kv_manager, None]
+
+    with (
+        patch.object(torch.cuda, "empty_cache"),
+        patch.object(torch.cuda, "reset_peak_memory_stats"),
+        patch.object(torch.cuda, "mem_get_info", return_value=(768, 1024)),
+        patch.object(torch.cuda, "memory_stats", return_value={"allocated_bytes.all.current": 128}),
+        pytest.raises(RuntimeError, match="profiling request failed"),
+    ):
+        creator.configure_kv_cache_capacity(executor)
+
+    executor.release_resources.assert_called_once_with()
+    executor.shutdown.assert_not_called()
+    assert executor.is_warmup is False
+    assert executor.enable_iter_perf_stats is True
+    assert executor.set_gather_responses.call_args.args == (False,)
diff --git a/tests/unittest/_torch/executor/test_executor_resource_rebind.py b/tests/unittest/_torch/executor/test_executor_resource_rebind.py
new file mode 100644
index 000000000000..3ffa33f6d232
--- /dev/null
+++ b/tests/unittest/_torch/executor/test_executor_resource_rebind.py
@@ -0,0 +1,248 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Contracts for switching a stopped executor from profiling to serving resources."""
+
+import threading
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor, PyExecutorResources
+from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType
+from tensorrt_llm._torch.pyexecutor.sampler import AsyncWorkerMixin
+
+pytestmark = pytest.mark.cpu_only
+
+
+@pytest.fixture
+def executor(monkeypatch):
+    events = []
+    engine = SimpleNamespace(
+        attn_metadata=object(),
+        _release_cuda_graphs=lambda: events.append("graphs"),
+    )
+    cache = SimpleNamespace(shutdown=lambda: events.append("cache"))
+    transceiver = SimpleNamespace(shutdown=lambda: events.append("unregister"))
+    ex = object.__new__(PyExecutor)
+    ex._closed = False
+    ex._resources_released = False
+    ex._event_loop_error = None
+    ex._fatal_error = None
+    ex.model_engine = engine
+    ex.draft_model_engine = None
+    ex.execution_stream = object()
+    ex.dist = SimpleNamespace(pp_size=1)
+    ex.worker_started = True
+    ex._mpi_progress_thread = None
+    ex._mpi_progress_thread_exit_hook = None
+    ex.worker_thread = Mock()
+    ex.worker_thread.join.side_effect = lambda: events.append("join")
+    ex.executor_request_queue = Mock()
+    ex.executor_request_queue.enqueue_shutdown_request.side_effect = lambda: events.append("stop")
+    ex.shutdown_event = threading.Event()
+    ex.shutdown_event.set()
+    ex._profile_manager = Mock()
+    ex.hang_detector = Mock()
+    ex.hang_detector.detected.return_value = False
+    ex._shutdown_sleep_wakeup_listeners = Mock()
+    ex.encoder_launch_executor = None
+    ex.kv_connector_manager = None
+    ex.kv_cache_transceiver = transceiver
+    ex.resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: cache})
+    ex.virtual_memory_pools = None
+    ex.sampler = object()
+    ex.dwdp_manager = None
+    ex.pending_encoder_steps = [object()]
+    ex.active_requests = [object()]
+    ex.responses = {17: object()}
+    ex.result_wait_queues = {17: object()}
+    ex._pending_transfer_responses = [object()]
+    ex._pending_response_terminations = [object()]
+    ex.request_accumulated = [object()]
+    ex.control_requests = [object()]
+    ex.waiting_queue = [object()]
+    ex._pending_adp_dummy_request = object()
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "synchronize", lambda: events.append("synchronize"))
+    return ex, events
+
+
+def test_release_preserves_model_and_orders_gpu_teardown(executor):
+    ex, events = executor
+    engine, stream, sampler = ex.model_engine, ex.execution_stream, ex.sampler
+    ex.release_resources()
+
+    assert events == ["stop", "join", "graphs", "synchronize", "unregister", "cache"]
+    assert ex.model_engine is engine
+    assert ex.execution_stream is stream
+    assert ex.sampler is sampler
+    assert engine.attn_metadata is None
+    assert ex.resource_manager is None
+    assert ex.scheduler is None
+    assert ex.kv_cache_transceiver is None
+    assert ex.hang_detector is None  # Drops the old transceiver's status-provider reference.
+    assert not ex.responses and not ex.active_requests and not ex.result_wait_queues
+    assert not ex.pending_encoder_steps
+    assert not ex.request_accumulated and not ex.control_requests
+    assert ex.waiting_queue is None and ex._pending_adp_dummy_request is None
+    assert ex._resources_released and not ex._closed
+
+
+def test_release_before_worker_start_does_not_wait_for_shutdown(executor):
+    ex, events = executor
+    ex.worker_started = False
+    ex.shutdown_event.clear()
+    ex.release_resources()
+    assert events == ["graphs", "synchronize", "unregister", "cache"]
+    ex.worker_thread.join.assert_not_called()
+
+
+def test_final_shutdown_after_release_does_not_release_resources_twice(executor):
+    ex, events = executor
+    ex.release_resources()
+    before = events.copy()
+    ex.shutdown()
+    ex.shutdown()
+    assert events == before
+    assert ex._closed
+    assert not hasattr(ex, "model_engine")
+
+
+@pytest.mark.parametrize("closed,released", [(False, False), (True, False), (True, True)])
+def test_rebind_requires_released_open_executor(executor, closed, released):
+    ex, _ = executor
+    ex._closed = closed
+    ex._resources_released = released
+    ex._initialize_execution = Mock()
+    with pytest.raises(RuntimeError, match="released before rebinding"):
+        ex.rebind_resources(Mock(spec=PyExecutorResources))
+    ex._initialize_execution.assert_not_called()
+
+
+@pytest.mark.parametrize("error_field", ["_event_loop_error", "_fatal_error"])
+def test_failed_executor_cannot_resume_with_new_resources(executor, error_field):
+    ex, _ = executor
+    ex.release_resources()
+    setattr(ex, error_field, RuntimeError("profiling failed"))
+    ex._initialize_execution = Mock()
+    with pytest.raises(RuntimeError, match="execution failure"):
+        ex.rebind_resources(Mock(spec=PyExecutorResources))
+    ex._initialize_execution.assert_not_called()
+
+
+def test_rebind_uses_same_executor_and_does_not_start_worker(executor):
+    ex, _ = executor
+    engine, stream = ex.model_engine, ex.execution_stream
+    ex.release_resources()
+    resources = Mock(spec=PyExecutorResources)
+    ex._initialize_execution = Mock()
+    ex.start_worker = Mock()
+    assert ex.rebind_resources(resources) is None
+    ex._initialize_execution.assert_called_once_with(resources)
+    ex.start_worker.assert_not_called()
+    assert ex.model_engine is engine and ex.execution_stream is stream
+    assert not ex._resources_released
+
+
+def test_failed_rebind_cannot_be_retried_over_partial_resources(executor):
+    ex, _ = executor
+    ex.release_resources()
+    cache = Mock()
+    resources = Mock(spec=PyExecutorResources)
+    resources.resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: cache})
+
+    def fail_warmup(bound_resources):
+        ex.resource_manager = bound_resources.resource_manager
+        raise RuntimeError("warmup failed")
+
+    ex._initialize_execution = Mock(side_effect=fail_warmup)
+    with pytest.raises(RuntimeError, match="warmup failed"):
+        ex.rebind_resources(resources)
+    with pytest.raises(RuntimeError, match="released before rebinding"):
+        ex.rebind_resources(resources)
+    ex._initialize_execution.assert_called_once()
+    assert ex._closed
+    cache.shutdown.assert_called_once_with()
+
+
+def test_release_destroys_draft_graphs_before_cache(executor):
+    ex, events = executor
+    draft_engine = SimpleNamespace(
+        attn_metadata=object(),
+        _release_cuda_graphs=lambda: events.append("draft_graphs"),
+    )
+    ex.draft_model_engine = draft_engine
+    ex.release_resources()
+    assert events.index("draft_graphs") < events.index("synchronize") < events.index("cache")
+    assert ex.draft_model_engine is draft_engine
+    assert draft_engine.attn_metadata is None
+
+
+def test_release_cannot_be_repeated(executor):
+    ex, events = executor
+    ex.release_resources()
+    before = events.copy()
+    with pytest.raises(RuntimeError, match="already released"):
+        ex.release_resources()
+    assert events == before
+
+
+def test_hung_worker_does_not_release_live_resources(executor):
+    ex, events = executor
+    ex.hang_detector.detected.return_value = True
+    ex._pp_mpi_progress_stop = threading.Event()
+    progress_thread = Mock()
+    ex._mpi_progress_thread = progress_thread
+    ex._mpi_progress_thread_exit_hook = progress_thread.stop
+    with pytest.raises(RuntimeError, match="executor hang"):
+        ex.release_resources()
+    assert events == ["stop"]
+    progress_thread.stop.assert_called_once_with()
+    assert ex._mpi_progress_thread is None
+    assert ex.resource_manager is not None
+    assert not ex._resources_released
+
+
+def test_async_sampler_stops_before_gpu_resource_release(executor):
+    ex, events = executor
+    ex.sampler = Mock(spec=AsyncWorkerMixin)
+    ex.sampler.async_worker_enabled.return_value = True
+    ex.sampler.async_worker_stop.side_effect = lambda: events.append("sampler_stop")
+    ex.release_resources()
+    assert events.index("join") < events.index("sampler_stop") < events.index("graphs")
+
+
+def test_transfer_shutdown_failure_keeps_cache_owned(executor):
+    ex, events = executor
+    manager = ex.resource_manager
+    ex.kv_cache_transceiver.shutdown = Mock(side_effect=RuntimeError("DMA still active"))
+    with pytest.raises(RuntimeError, match="DMA still active"):
+        ex.release_resources()
+    assert "cache" not in events
+    assert ex.resource_manager is manager
+    assert not ex._resources_released
+
+
+def test_release_stops_pipeline_progress_before_gpu_teardown(executor):
+    ex, events = executor
+    ex.dist.pp_size = 2
+    ex._pp_mpi_progress_stop = threading.Event()
+    ex._pp_mpi_progress_quiesced = threading.Event()
+    ex.executed_batch_queue = Mock()
+    ex.broadcast_sample_state_handler = Mock()
+    ex.broadcast_sample_state_handler.join.side_effect = lambda: events.append("relay_join")
+    progress_thread = Mock()
+    progress_thread.stop.side_effect = lambda: events.append("progress_stop")
+    ex._mpi_progress_thread = progress_thread
+    ex._mpi_progress_thread_exit_hook = progress_thread.stop
+
+    ex.release_resources()
+
+    assert events.index("join") < events.index("relay_join")
+    assert events.index("relay_join") < events.index("progress_stop") < events.index("graphs")
+    assert ex._pp_mpi_progress_stop.is_set()
+    assert ex._pp_mpi_progress_quiesced.is_set()
+    assert ex._mpi_progress_thread is None
+    assert ex._mpi_progress_thread_exit_hook is None
diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
index a51f978c5d39..7fdb6c488dc5 100644
--- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
+++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
@@ -110,6 +110,13 @@ def __init__(self, resources, model_engine, peft_cache_config, execution_stream)
         self.peft_cache_config = peft_cache_config
         self.execution_stream = execution_stream
         self.started = False
+        self.closed = False
+
+    def shutdown(self):
+        self.closed = True
+
+    def rebind_resources(self, resources):
+        self.resource_manager = resources
 
     def start_worker(self):
         """Mark executor as started."""
@@ -359,24 +366,37 @@ def _create_model_engine(**kwargs):
 
     monkeypatch.setattr(py_executor_creator, "PyTorchModelEngine", _create_model_engine)
 
+    created_executors = []
+
     def _create_py_executor_instance(**kwargs):
         if ctx_chunk_configs is not None:
             ctx_chunk_configs.append(kwargs["ctx_chunk_config"])
-        return _DummyPyExecutor(
+        executor = _DummyPyExecutor(
             resources=kwargs["resources"],
             model_engine=kwargs["model_engine"],
             peft_cache_config=kwargs["peft_cache_config"],
             execution_stream=kwargs["execution_stream"],
         )
+        created_executors.append(executor)
+        return executor
+
+    def _create_py_executor_resources(**kwargs):
+        if ctx_chunk_configs is not None:
+            ctx_chunk_configs.append(kwargs["ctx_chunk_config"])
+        return _DummyResourceManager(kwargs["resources"])
 
     monkeypatch.setattr(
         py_executor_creator, "create_py_executor_instance", _create_py_executor_instance
     )
+    monkeypatch.setattr(
+        py_executor_creator, "create_py_executor_resources", _create_py_executor_resources
+    )
 
     py_executor = py_executor_creator.create_py_executor(
         llm_args=llm_args,
         checkpoint_dir=None,
     )
+    assert created_executors == [py_executor]
     kv_cache_manager = py_executor.resource_manager.get_resource_manager(
         ResourceManagerType.KV_CACHE_MANAGER
     )
@@ -725,3 +745,43 @@ def teardown_managers(self, resources):
         assert allocation_phases == ["final_kv_cache_allocation", "final_executor_creation"]
     assert phases[-1] == "executor_start_worker"
     assert timers[0].depth == 0
+
+
+@pytest.mark.parametrize("failure_stage", ["profiling", "allocation", "rebind"])
+def test_failed_resource_transition_closes_executor(monkeypatch, failure_stage):
+    executors = []
+
+    class FailingCacheCreator(_DummyKvCacheCreator):
+        def try_prepare_estimation(self):
+            return True
+
+        def configure_kv_cache_capacity(self, executor):
+            executors.append(executor)
+            if failure_stage == "profiling":
+                raise RuntimeError("startup failure")
+
+        def teardown_managers(self, resources):
+            resources.pop(ResourceManagerType.KV_CACHE_MANAGER)
+
+        def build_managers(self, resources, estimating_kv_cache):
+            if not estimating_kv_cache and failure_stage == "allocation":
+                raise RuntimeError("startup failure")
+            super().build_managers(resources, estimating_kv_cache)
+
+    if failure_stage == "rebind":
+
+        def fail_rebind(self, resources):
+            raise RuntimeError("startup failure")
+
+        monkeypatch.setattr(_DummyPyExecutor, "rebind_resources", fail_rebind)
+
+    with pytest.raises(RuntimeError, match="startup failure"):
+        _run_create_py_executor(
+            monkeypatch,
+            sm_version=100,
+            kv_cache_quant_algo=QuantAlgo.NO_QUANT,
+            kv_cache_creator_cls=FailingCacheCreator,
+        )
+    assert len(executors) == 1
+    assert executors[0].closed
+    assert not executors[0].started

From 5901356091d49f7f528054a61b133d9243fb77cb Mon Sep 17 00:00:00 2001
From: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
Date: Wed, 14 Oct 2026 19:14:30 -0700
Subject: [PATCH 2/7] Consolidate executor resource lifecycle regression tests

Signed-off-by: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
---
 .../executor/test_executor_resource_rebind.py | 55 ++++++++-----------
 1 file changed, 23 insertions(+), 32 deletions(-)

diff --git a/tests/unittest/_torch/executor/test_executor_resource_rebind.py b/tests/unittest/_torch/executor/test_executor_resource_rebind.py
index 3ffa33f6d232..2dd9fb555aa7 100644
--- a/tests/unittest/_torch/executor/test_executor_resource_rebind.py
+++ b/tests/unittest/_torch/executor/test_executor_resource_rebind.py
@@ -71,10 +71,29 @@ def executor(monkeypatch):
 
 def test_release_preserves_model_and_orders_gpu_teardown(executor):
     ex, events = executor
+    draft_engine = SimpleNamespace(
+        attn_metadata=object(),
+        _release_cuda_graphs=lambda: events.append("draft_graphs"),
+    )
+    ex.draft_model_engine = draft_engine
+    ex.sampler = Mock(spec=AsyncWorkerMixin)
+    ex.sampler.async_worker_enabled.return_value = True
+    ex.sampler.async_worker_stop.side_effect = lambda: events.append("sampler_stop")
     engine, stream, sampler = ex.model_engine, ex.execution_stream, ex.sampler
     ex.release_resources()
 
-    assert events == ["stop", "join", "graphs", "synchronize", "unregister", "cache"]
+    assert events == [
+        "stop",
+        "join",
+        "sampler_stop",
+        "graphs",
+        "draft_graphs",
+        "synchronize",
+        "unregister",
+        "cache",
+    ]
+    assert ex.draft_model_engine is draft_engine
+    assert draft_engine.attn_metadata is None
     assert ex.model_engine is engine
     assert ex.execution_stream is stream
     assert ex.sampler is sampler
@@ -103,6 +122,9 @@ def test_final_shutdown_after_release_does_not_release_resources_twice(executor)
     ex, events = executor
     ex.release_resources()
     before = events.copy()
+    with pytest.raises(RuntimeError, match="already released"):
+        ex.release_resources()
+    assert events == before
     ex.shutdown()
     ex.shutdown()
     assert events == before
@@ -167,28 +189,6 @@ def fail_warmup(bound_resources):
     cache.shutdown.assert_called_once_with()
 
 
-def test_release_destroys_draft_graphs_before_cache(executor):
-    ex, events = executor
-    draft_engine = SimpleNamespace(
-        attn_metadata=object(),
-        _release_cuda_graphs=lambda: events.append("draft_graphs"),
-    )
-    ex.draft_model_engine = draft_engine
-    ex.release_resources()
-    assert events.index("draft_graphs") < events.index("synchronize") < events.index("cache")
-    assert ex.draft_model_engine is draft_engine
-    assert draft_engine.attn_metadata is None
-
-
-def test_release_cannot_be_repeated(executor):
-    ex, events = executor
-    ex.release_resources()
-    before = events.copy()
-    with pytest.raises(RuntimeError, match="already released"):
-        ex.release_resources()
-    assert events == before
-
-
 def test_hung_worker_does_not_release_live_resources(executor):
     ex, events = executor
     ex.hang_detector.detected.return_value = True
@@ -205,15 +205,6 @@ def test_hung_worker_does_not_release_live_resources(executor):
     assert not ex._resources_released
 
 
-def test_async_sampler_stops_before_gpu_resource_release(executor):
-    ex, events = executor
-    ex.sampler = Mock(spec=AsyncWorkerMixin)
-    ex.sampler.async_worker_enabled.return_value = True
-    ex.sampler.async_worker_stop.side_effect = lambda: events.append("sampler_stop")
-    ex.release_resources()
-    assert events.index("join") < events.index("sampler_stop") < events.index("graphs")
-
-
 def test_transfer_shutdown_failure_keeps_cache_owned(executor):
     ex, events = executor
     manager = ex.resource_manager

From 2f0cb23e705cfd4ddc60c7e4c07ef6fdfdd6b1c6 Mon Sep 17 00:00:00 2001
From: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
Date: Wed, 14 Oct 2026 23:30:42 -0700
Subject: [PATCH 3/7] Reset profiling windows when rebinding executor resources

Signed-off-by: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  19 ++-
 .../executor/test_executor_resource_rebind.py | 121 ++++++++++++++++--
 2 files changed, 121 insertions(+), 19 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 08febdd1f8a4..a58395a8ec0c 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -509,16 +509,6 @@ def __init__(
 
         self.peft_cache_config = peft_cache_config
 
-        self.iter_counter = 0
-        # profile config
-        _env_start_iters, _env_stop_iters = _load_iteration_indexes(
-            PROFILE_START_STOP_ENV_VAR_NAME)
-        # Convert to mutable sets so start_profile()/stop_profile() can
-        # .add() runtime-configured iteration indices.
-        # _load_iteration_indexes returns frozensets.
-        self.profile_start_iters = set(_env_start_iters)
-        self.profile_stop_iters = set(_env_stop_iters)
-
         # Runtime-configurable profiling state. Populated by start_profile()
         # when triggered via HTTP endpoints (e.g. trtllm-serve /start_profile)
         # or API calls. When these are None/False the _profiler() context
@@ -624,6 +614,15 @@ def _initialize_execution(self, resources: PyExecutorResources) -> None:
         enable_kv_pool_rebalance = self._execution_config.enable_kv_pool_rebalance
         self.peft_cache_config = peft_cache_config
         self.iter_counter = 0
+        # profile config
+        _env_start_iters, _env_stop_iters = _load_iteration_indexes(
+            PROFILE_START_STOP_ENV_VAR_NAME)
+        # Convert to mutable sets so start_profile()/stop_profile() can
+        # .add() runtime-configured iteration indices.
+        # _load_iteration_indexes returns frozensets.
+        self.profile_start_iters = set(_env_start_iters)
+        self.profile_stop_iters = set(_env_stop_iters)
+
         self._metrics = {}
         self._disagg_coordinator = None
         self._ctx_mla_kv_len_cap = "unset"
diff --git a/tests/unittest/_torch/executor/test_executor_resource_rebind.py b/tests/unittest/_torch/executor/test_executor_resource_rebind.py
index 2dd9fb555aa7..598e4d2b9888 100644
--- a/tests/unittest/_torch/executor/test_executor_resource_rebind.py
+++ b/tests/unittest/_torch/executor/test_executor_resource_rebind.py
@@ -3,13 +3,19 @@
 """Contracts for switching a stopped executor from profiling to serving resources."""
 
 import threading
+from contextlib import nullcontext
 from types import SimpleNamespace
 from unittest.mock import Mock
 
 import pytest
 import torch
 
-from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor, PyExecutorResources
+from tensorrt_llm._torch.pyexecutor import py_executor as executor_module
+from tensorrt_llm._torch.pyexecutor.py_executor import (
+    PyExecutor,
+    PyExecutorResources,
+    _ExecutorExecutionConfig,
+)
 from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType
 from tensorrt_llm._torch.pyexecutor.sampler import AsyncWorkerMixin
 
@@ -154,18 +160,115 @@ def test_failed_executor_cannot_resume_with_new_resources(executor, error_field)
     ex._initialize_execution.assert_not_called()
 
 
-def test_rebind_uses_same_executor_and_does_not_start_worker(executor):
+def test_rebind_resets_phase_state_and_profile_windows(executor, monkeypatch):
     ex, _ = executor
+    monkeypatch.setenv("TLLM_PROFILE_START_STOP", "2-4")
+    executor_module._load_iteration_indexes.cache_clear()
+    monkeypatch.setattr(torch.cuda, "Stream", Mock())
+    monkeypatch.setattr(torch.cuda, "Event", Mock())
+    monkeypatch.setattr(torch.cuda, "current_stream", Mock())
+    monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext())
+    monkeypatch.setattr(executor_module, "_distributed_warmup_guard", lambda *args: nullcontext())
+    monkeypatch.setattr(executor_module.ADPRouter, "create", Mock())
+    monkeypatch.setattr(executor_module, "HangDetector", Mock())
+    monkeypatch.setattr(executor_module, "ExecutorRequestQueue", Mock())
+    monkeypatch.setattr(executor_module, "RequestBroadcaster", Mock())
+    monkeypatch.setattr(executor_module, "mpi_disabled", lambda: True)
+    monkeypatch.setattr(executor_module.GpuKeepalive, "create_from_env", Mock(return_value=None))
+    monkeypatch.setattr(
+        executor_module, "create_executor_hang_diagnostics", Mock(return_value=(None, None))
+    )
+    for name in (
+        "_emit_initial_stats",
+        "_set_global_steady_clock_offset",
+        "_initialize_remote_tail_adp",
+        "_maybe_init_kv_connector_manager",
+        "_validate_moe_a2a_checkpoint_support",
+        "_warn_if_kv_block_budget_unchecked",
+    ):
+        monkeypatch.setattr(ex, name, Mock())
+    ex.device_id = 0
+    ex.global_rank = 0
+    ex.dist = SimpleNamespace(pp_size=1, rank=0, mapping=Mock(), barrier=Mock())
+    ex.execution_stream = Mock()
+    ex.drafter = None
+    ex.guided_decoder = None
+    ex._execution_config = _ExecutorExecutionConfig(
+        disable_overlap_scheduler=True,
+        enable_early_first_token_response=False,
+        max_input_len=128,
+        max_batch_size=2,
+        max_beam_width=1,
+        max_draft_len=0,
+        max_total_draft_tokens=0,
+        garbage_collection_gen0_threshold=None,
+        resource_governor_queue=None,
+        hang_detection_timeout=None,
+        enable_kv_pool_rebalance=False,
+    )
+    ex.model_engine.enable_attention_dp = False
+    ex.model_engine.model = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False))
+    ex.model_engine.get_max_num_sequences = Mock(return_value=2)
+    ex.model_engine.warmup = Mock()
+    ex.model_engine.llm_args = SimpleNamespace(
+        max_stats_len=0,
+        max_num_tokens=128,
+        print_iter_log=False,
+        enable_iter_perf_stats=False,
+        enable_iter_req_stats=False,
+        iter_perf_stats_interval=1,
+        stream_interval=1,
+        attention_dp_config=None,
+        batch_wait_timeout_ms=0,
+        batch_wait_timeout_iters=0,
+        batch_wait_max_tokens_ratio=0,
+    )
+    cache = Mock(
+        event_buffer_max_size=0,
+        streaming_kv_events_enabled=False,
+        enable_block_reuse=False,
+        tokens_per_block=16,
+    )
+    cache.get_kv_cache_stats.return_value = SimpleNamespace(max_num_blocks=8, tokens_per_block=16)
+    resources = PyExecutorResources(
+        resource_manager=ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: cache}),
+        scheduler=Mock(),
+        kv_cache_transceiver=None,
+        peft_cache_config=None,
+        virtual_memory_pools=None,
+        kv_connector_manager=None,
+        dwdp_manager=None,
+        max_seq_len=128,
+        max_num_sequences=2,
+        waiting_queue_policy=executor_module.WaitingQueuePolicy.FCFS,
+    )
+    ex.iter_counter = 42
+    ex._ctx_mla_kv_len_cap = None
+    ex.profile_start_iters = {99}
+    ex.profile_stop_iters = set()
+    old_shutdown_event = ex.shutdown_event
     engine, stream = ex.model_engine, ex.execution_stream
     ex.release_resources()
-    resources = Mock(spec=PyExecutorResources)
-    ex._initialize_execution = Mock()
+    ex._disagg_coordinator = object()
     ex.start_worker = Mock()
-    assert ex.rebind_resources(resources) is None
-    ex._initialize_execution.assert_called_once_with(resources)
-    ex.start_worker.assert_not_called()
-    assert ex.model_engine is engine and ex.execution_stream is stream
-    assert not ex._resources_released
+    try:
+        ex.rebind_resources(resources)
+        assert ex.iter_counter == 0
+        assert ex.profile_start_iters == {2}
+        assert ex.profile_stop_iters == {4}
+        assert ex._ctx_mla_kv_len_cap == "unset"
+        assert ex._disagg_coordinator is None
+        assert not ex.worker_started
+        assert ex.shutdown_event is not old_shutdown_event
+        assert not ex.shutdown_event.is_set()
+        assert ex.resource_manager is resources.resource_manager
+        assert ex.kv_cache_manager is cache
+        assert ex.model_engine is engine and ex.execution_stream is stream
+        assert not ex._resources_released
+        engine.warmup.assert_called_once_with(resources.resource_manager)
+        ex.start_worker.assert_not_called()
+    finally:
+        executor_module._load_iteration_indexes.cache_clear()
 
 
 def test_failed_rebind_cannot_be_retried_over_partial_resources(executor):

From 000654b2855ee9112ad03e0bf4344b22a40b59ab Mon Sep 17 00:00:00 2001
From: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
Date: Mon, 19 Oct 2026 03:47:33 -0700
Subject: [PATCH 4/7] Document phase resource ownership and guard rebind
 coverage

Signed-off-by: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  8 ++++++
 .../executor/test_executor_resource_rebind.py | 26 ++++++++++++++++++-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index a58395a8ec0c..08f8f2c3ee20 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -414,6 +414,14 @@ class PyExecutorResources:
 
     The executor owns their execution lifetime. Model engines and weights
     remain owned by the executor across resource transitions.
+
+    Add externally constructed, phase-dependent resources here and bind them
+    in ``_initialize_execution``. Create derived phase state in that method,
+    not in ``__init__``. Keep cross-phase objects on their existing owners.
+    Resources requiring teardown must also be handled explicitly by
+    ``_release_execution_resources`` and detached by ``release_resources``;
+    this input bundle does not perform automatic cleanup. Extend the resource
+    rebind tests when adding fields, including any cleanup ordering constraints.
     """
 
     resource_manager: ResourceManager
diff --git a/tests/unittest/_torch/executor/test_executor_resource_rebind.py b/tests/unittest/_torch/executor/test_executor_resource_rebind.py
index 598e4d2b9888..8b34dc0ed526 100644
--- a/tests/unittest/_torch/executor/test_executor_resource_rebind.py
+++ b/tests/unittest/_torch/executor/test_executor_resource_rebind.py
@@ -4,6 +4,7 @@
 
 import threading
 from contextlib import nullcontext
+from dataclasses import fields
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -162,6 +163,8 @@ def test_failed_executor_cannot_resume_with_new_resources(executor, error_field)
 
 def test_rebind_resets_phase_state_and_profile_windows(executor, monkeypatch):
     ex, _ = executor
+    create_waiting_queue = Mock(wraps=executor_module.create_waiting_queue)
+    monkeypatch.setattr(executor_module, "create_waiting_queue", create_waiting_queue)
     monkeypatch.setenv("TLLM_PROFILE_START_STOP", "2-4")
     executor_module._load_iteration_indexes.cache_clear()
     monkeypatch.setattr(torch.cuda, "Stream", Mock())
@@ -249,6 +252,24 @@ def test_rebind_resets_phase_state_and_profile_windows(executor, monkeypatch):
     old_shutdown_event = ex.shutdown_event
     engine, stream = ex.model_engine, ex.execution_stream
     ex.release_resources()
+    bound_fields = {
+        "resource_manager",
+        "scheduler",
+        "kv_cache_transceiver",
+        "peft_cache_config",
+        "virtual_memory_pools",
+        "kv_connector_manager",
+        "dwdp_manager",
+        "max_seq_len",
+    }
+    # Queue policy constructs a queue; sequence capacity comes from the engine.
+    assert {field.name for field in fields(PyExecutorResources)} == bound_fields | {
+        "waiting_queue_policy",
+        "max_num_sequences",
+    }, "New phase inputs need binding and lifecycle coverage in this test."
+    # Stale values also catch omitted assignments for optional resources set to None.
+    for name in bound_fields:
+        setattr(ex, name, object())
     ex._disagg_coordinator = object()
     ex.start_worker = Mock()
     try:
@@ -261,7 +282,10 @@ def test_rebind_resets_phase_state_and_profile_windows(executor, monkeypatch):
         assert not ex.worker_started
         assert ex.shutdown_event is not old_shutdown_event
         assert not ex.shutdown_event.is_set()
-        assert ex.resource_manager is resources.resource_manager
+        for name in bound_fields:
+            assert getattr(ex, name) is getattr(resources, name), name
+        create_waiting_queue.assert_called_once_with(resources.waiting_queue_policy)
+        assert ex.max_num_active_requests == engine.get_max_num_sequences.return_value
         assert ex.kv_cache_manager is cache
         assert ex.model_engine is engine and ex.execution_stream is stream
         assert not ex._resources_released

From 948e29402a2fcfcc59053ec0424fc571b626f8cb Mon Sep 17 00:00:00 2001
From: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
Date: Tue, 20 Oct 2026 04:10:13 -0700
Subject: [PATCH 5/7] Adapt executor test fixtures to resource rebinding
 lifecycle

Signed-off-by: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
---
 tests/unittest/_torch/executor/test_gpu_keepalive.py        | 6 ++++--
 .../_torch/executor/test_kv_connector_executor_lifetime.py  | 2 ++
 .../unittest/_torch/executor/test_pp_mpi_progress_thread.py | 3 +++
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/tests/unittest/_torch/executor/test_gpu_keepalive.py b/tests/unittest/_torch/executor/test_gpu_keepalive.py
index 5c3318c00690..6d1c0c7bfacd 100644
--- a/tests/unittest/_torch/executor/test_gpu_keepalive.py
+++ b/tests/unittest/_torch/executor/test_gpu_keepalive.py
@@ -298,10 +298,12 @@ class Stub:
         assert result[1] is False and order == []
         k.drain.assert_not_called()
 
-    # Exactly one `self._gpu_keepalive = GpuKeepalive.create_from_env(self.device_id)` in __init__.
+    # Each execution phase (initial construction or rebind) gets a fresh keepalive.
     tree = ast.parse(inspect.getsource(pe))
     cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "PyExecutor")
-    init = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "__init__")
+    init = next(
+        n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "_initialize_execution"
+    )
     constructions = [
         n
         for n in ast.walk(init)
diff --git a/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py b/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py
index a8263f4df93e..2fe2927761be 100644
--- a/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py
+++ b/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py
@@ -27,6 +27,8 @@
 
 def _executor() -> PyExecutor:
     executor = object.__new__(PyExecutor)
+    executor._closed = False
+    executor._resources_released = False
     executor.kv_connector_manager = Mock(spec=KvCacheConnectorManager)
     # Set explicitly: `capacity_only` is not on the spec, so a double that
     # leaves it unset fails loudly rather than silently lifting the executor
diff --git a/tests/unittest/_torch/executor/test_pp_mpi_progress_thread.py b/tests/unittest/_torch/executor/test_pp_mpi_progress_thread.py
index 06e92da33795..d75e57d1c68b 100644
--- a/tests/unittest/_torch/executor/test_pp_mpi_progress_thread.py
+++ b/tests/unittest/_torch/executor/test_pp_mpi_progress_thread.py
@@ -180,6 +180,9 @@ def _executor(monkeypatch, pp_size: int) -> PyExecutor:
     monkeypatch.setenv(MPI_PROGRESS_POLL_MS_ENV_VAR_NAME, "0")
     monkeypatch.delenv(PP_MPI_PROGRESS_INTERVAL_US_ENV_VAR_NAME, raising=False)
     executor = object.__new__(PyExecutor)
+    executor._closed = False
+    executor._resources_released = False
+    executor.kv_cache_transceiver = None
     executor.dist = SimpleNamespace(rank=0, world_size=1, pp_size=pp_size)
     executor.worker_lock = threading.Lock()
     executor.worker_started = False

From 1aa02c64d34d9570dce80e40360aa7b5c9660f01 Mon Sep 17 00:00:00 2001
From: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
Date: Wed, 21 Oct 2026 08:30:36 -0700
Subject: [PATCH 6/7] Preserve rebalance interval across executor resource
 rebinding

Signed-off-by: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 08f8f2c3ee20..a9b88b96f707 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -497,6 +497,9 @@ def __init__(
                 "kv_pool_rebalance_check_interval must be positive, got "
                 f"{kv_pool_rebalance_check_interval}")
         super(PyExecutor, self).__init__()
+        # Iteration throttle for the KV pool rebalance check.  See
+        # _can_pause_for_rebalance / _agreed_need_adjustment.
+        self._rebalance_check_interval = kv_pool_rebalance_check_interval
         self._metrics: dict[str, float | dict[str, float]] = {}
         self.device_id = torch.cuda.current_device()
         self.global_rank = dist.rank
@@ -671,9 +674,6 @@ def _initialize_execution(self, resources: PyExecutorResources) -> None:
                                           None)
         self.disable_overlap_scheduler = disable_overlap_scheduler
         self.enable_kv_pool_rebalance = enable_kv_pool_rebalance
-        # Iteration throttle for the KV pool rebalance check.  See
-        # _can_pause_for_rebalance / _agreed_need_adjustment.
-        self._rebalance_check_interval = kv_pool_rebalance_check_interval
         # Countdown of drain iterations remaining before a pipeline-parallel
         # rebalance can run; None when no rebalance is pending.  Only
         # _executor_loop_pp uses it -- see _start_pp_rebalance_drain.

From 0f5d7affb6e46f43b533c838081847e58a7c7eab Mon Sep 17 00:00:00 2001
From: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
Date: Thu, 22 Oct 2026 22:19:57 -0700
Subject: [PATCH 7/7] Preserve model-owned graph lifetime across executor
 resource release

Signed-off-by: ZhaoyangWang <265160349+zhaoyangwang-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py    | 16 ++++++++++++++--
 .../executor/test_executor_resource_rebind.py    |  5 ++++-
 .../unittest/_torch/executor/test_py_executor.py | 14 +++++++++++++-
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index a9b88b96f707..1c57cc29b746 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -2134,14 +2134,26 @@ def shutdown(self) -> None:
         if self._closed:
             return
         if not self._resources_released:
-            if not self._release_execution_resources():
+            if not self._release_execution_resources(
+                    release_model_owned_graphs=True):
                 return
+        else:
+            for engine in (self.model_engine, self.draft_model_engine):
+                release_model_graphs = getattr(
+                    engine, "_release_model_owned_cuda_graphs", None)
+                if release_model_graphs is not None:
+                    release_model_graphs()
+            if torch.cuda.is_available():
+                torch.cuda.synchronize()
         self._closed = True
         del self.model_engine
         if self.draft_model_engine is not None:
             del self.draft_model_engine
 
-    def _release_execution_resources(self) -> bool:
+    def _release_execution_resources(self,
+                                     *,
+                                     release_model_owned_graphs: bool = False
+                                     ) -> bool:
         """Stop execution and release GPU resources; retain model ownership."""
         node_server = getattr(self, "_openengine_node_server", None)
         if node_server is not None:
diff --git a/tests/unittest/_torch/executor/test_executor_resource_rebind.py b/tests/unittest/_torch/executor/test_executor_resource_rebind.py
index 8b34dc0ed526..2f248563bd3a 100644
--- a/tests/unittest/_torch/executor/test_executor_resource_rebind.py
+++ b/tests/unittest/_torch/executor/test_executor_resource_rebind.py
@@ -132,9 +132,12 @@ def test_final_shutdown_after_release_does_not_release_resources_twice(executor)
     with pytest.raises(RuntimeError, match="already released"):
         ex.release_resources()
     assert events == before
+    engine = ex.model_engine
+    engine._release_model_owned_cuda_graphs = Mock()
     ex.shutdown()
     ex.shutdown()
-    assert events == before
+    engine._release_model_owned_cuda_graphs.assert_called_once()
+    assert events == before + ["synchronize"]
     assert ex._closed
     assert not hasattr(ex, "model_engine")
 
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index 36fbcf622a7a..336f87975090 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -74,12 +74,24 @@
     ("shutdown_method", "releases_model_owned_graphs"),
     [
         ("shutdown", True),
-        ("shutdown_for_kv_cache_estimation", False),
+        ("release_resources", False),
     ],
 )
 def test_shutdown_graph_lifecycle(shutdown_method, releases_model_owned_graphs):
     engine = Mock()
     executor = object.__new__(PyExecutor)
+    executor._closed = False
+    executor._resources_released = False
+    executor.worker_started = True
+    executor.kv_cache_transceiver = None
+    executor.pending_encoder_steps = []
+    executor.active_requests = []
+    executor.responses = {}
+    executor.result_wait_queues = {}
+    executor._pending_transfer_responses = []
+    executor._pending_response_terminations = []
+    executor.request_accumulated = []
+    executor.control_requests = []
     executor.executor_request_queue = Mock()
     executor.shutdown_event = Mock()
     executor.hang_detector = Mock()