[None][fix] Fix disagg KV cache transfer crash on reuse budget skip

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -3018,8 +3018,16 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests):
                         resource_mgr_type].prepare_resources(
                             disagg_gen_init_to_prepare)
 
-            # Trigger KV cache exchange for new disagg_gen_init_requests
-            self._recv_disagg_gen_cache(fitting_disagg_gen_init_requests)
+            # Use the filtered batch after prepare_resources, not the original
+            # list.  prepare_resources may skip add_sequence for some requests
+            # (e.g. reuse budget exhaustion) and remove them from the batch via
+            # reset_context_requests.  Passing skipped requests to KV transfer
+            # causes unordered_map::at in getSequence because the sequence was
+            # never registered.  Skipped requests remain in active_requests
+            # with DISAGG_GENERATION_INIT state and will be rescheduled.
+            accepted_requests = list(
+                disagg_gen_init_to_prepare.context_requests)
+            self._recv_disagg_gen_cache(accepted_requests)
 
     @nvtx_range("_prepare_disagg_gen_transmission_complete")
     def _prepare_disagg_gen_transmission_complete(self, scheduled_batch):
@@ -3642,14 +3650,13 @@ def _handle_responses(self):
                                     f"Request {request.py_request_id} has no avg_decoded_tokens_per_iter"
                                 )
 
-                # If partial reuse is enabled, and the KV cache manager is not VSWA, and the PP size is 1,
-                # then we need to terminate the request. TODO: Remove this once disagg support from KVCache reuse
-                # path is fixed.
-                if self.enable_partial_reuse_for_disagg and not self.kv_cache_manager.is_vswa and self.dist.pp_size == 1:
+                # Never terminate a request while KV cache transfer is
+                # still in progress — the GEN worker may still be receiving
+                # data from the pinned blocks.  The request will be
+                # terminated by _end_transfer_and_maybe_terminate once the
+                # transfer completes (or times out).
+                if not request.is_disagg_context_transmission_state:
                     requests_to_terminate.append(request)
-                else:
-                    if not request.is_disagg_context_transmission_state:
-                        requests_to_terminate.append(request)
             else:
                 new_active_requests.append(request)