[BugFix] Fix pipeline parallel (#24621)

njhill · web-flow · commit e2d8c27f6868 · 2025-09-10T23:05:30.000-07:00
Signed-off-by: Nick Hill &lt;nhill@redhat.com&gt;
diff --git a/vllm/executor/uniproc_executor.py b/vllm/executor/uniproc_executor.py
@@ -71,6 +71,10 @@ def reinitialize_distributed(
             self.shutdown()
         return
 
+    def shutdown(self) -> None:
+        if worker := self.driver_worker:
+            worker.shutdown()
+
 
 UniProcExecutorAsync = UniProcExecutor
 
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
@@ -2070,7 +2070,6 @@ def execute_model(
             sampler_output = self._sample(logits, spec_decode_metadata)
 
         with record_function_or_nullcontext("Bookkeep"):
-            assert isinstance(hidden_states, torch.Tensor)
             (
                 num_nans_in_logits,
                 logprobs_lists,
diff --git a/vllm/v1/worker/kv_connector_model_runner_mixin.py b/vllm/v1/worker/kv_connector_model_runner_mixin.py
@@ -45,7 +45,8 @@ def maybe_setup_kv_connector(scheduler_output: "SchedulerOutput"):
 
     @staticmethod
     def ensure_kv_transfer_shutdown() -> None:
-        if has_kv_transfer_group():
+        # has_kv_transfer_group can be None during interpreter shutdown.
+        if has_kv_transfer_group and has_kv_transfer_group():
             ensure_kv_transfer_shutdown()
 
     @staticmethod

Original file line number	Diff line number	Diff line change
`@@ -2070,7 +2070,6 @@ def execute_model(`
`2070`	`2070`	`sampler_output = self._sample(logits, spec_decode_metadata)`
`2071`	`2071`
`2072`	`2072`	`with record_function_or_nullcontext("Bookkeep"):`
`2073`		`- assert isinstance(hidden_states, torch.Tensor)`
`2074`	`2073`	`(`
`2075`	`2074`	`num_nans_in_logits,`
`2076`	`2075`	`logprobs_lists,`