A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 23:45 UTC
[BugFix] Fix KVConnector TP worker aggregation (#21473)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Jul 24, 2025 at 04:56 UTC
eec6942014c4408d8d9e4c3a37324f7ff35fc5aa
1 file changed
+10
-8
vllm/v1/worker/gpu_worker.py
+10
-8
@@ -16,7 +16,8 @@ from vllm.config import VllmConfig
16
from vllm.distributed import (ensure_model_parallel_initialized,
17
init_distributed_environment,
18
set_custom_all_reduce)
19
-from vllm.distributed.kv_transfer import ensure_kv_transfer_initialized
19
+from vllm.distributed.kv_transfer import (ensure_kv_transfer_initialized,
20
+ has_kv_transfer_group)
21
from vllm.distributed.parallel_state import get_pp_group, get_tp_group
22
from vllm.logger import init_logger
23
from vllm.lora.request import LoRARequest
@@ -342,19 +343,20 @@ class Worker(WorkerBase):
343
assert isinstance(output, IntermediateTensors)
344
get_pp_group().send_tensor_dict(output.tensors,
345
all_gather_group=get_tp_group())
346
+ if not has_kv_transfer_group():
347
+ return None
348
349
# In case of PP with kv transfer, we need to pass through the
350
# finished_sending and finished_recving buffers.
348
- empty_output = EMPTY_MODEL_RUNNER_OUTPUT
351
+ new_output = EMPTY_MODEL_RUNNER_OUTPUT
352
if output.finished_sending or output.finished_recving:
350
- empty_output = copy.copy(empty_output)
351
- empty_output.finished_sending = output.finished_sending
352
- empty_output.finished_recving = output.finished_recving
353
- output = empty_output
353
+ new_output = copy.copy(new_output)
354
+ new_output.finished_sending = output.finished_sending
355
+ new_output.finished_recving = output.finished_recving
356
+ output = new_output
357
358
assert isinstance(output, ModelRunnerOutput)
356
- # return output only from the driver worker
357
- return output if self.is_driver_worker else None
359
+ return output
360
361
def profile(self, is_start: bool = True):
362
if self.profiler is None: