xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-20 14:57:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add threading for load-balancing to different workers

Browse Source
This commit is contained in:
Pravein Govindan Kannan 2025-06-30 14:40:18 +05:30
parent 5d8b665366
commit 17546dc79f
1 changed files with 12 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
View File

@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Optional
import msgspec import msgspec
import torch import torch
import zmq import zmq
from concurrent.futures import ThreadPoolExecutor, as_completed
from vllm import envs from vllm import envs
from vllm.attention.selector import backend_name_to_enum, get_attn_backend from vllm.attention.selector import backend_name_to_enum, get_attn_backend
@ -986,16 +987,23 @@ class NixlConnectorWorker:
# Prepare transfer with Nixl. # Prepare transfer with Nixl.
CHUNK_SIZE = 100 CHUNK_SIZE = 100
handles = [] handles = []
for i in range(0, len(local_block_descs_ids), CHUNK_SIZE): futures = []
handles.append( with ThreadPoolExecutor() as executor:
self.nixl_wrapper.make_prepped_xfer( for i in range(0, len(local_block_descs_ids), CHUNK_SIZE):
future = executor.submit(
self.nixl_wrapper.make_prepped_xfer,
"READ", "READ",
local_xfer_side_handle, local_xfer_side_handle,
local_block_descs_ids[i:i + CHUNK_SIZE], local_block_descs_ids[i:i + CHUNK_SIZE],
remote_xfer_side_handle, remote_xfer_side_handle,
remote_block_descs_ids[i:i + CHUNK_SIZE], remote_block_descs_ids[i:i + CHUNK_SIZE],
skip_desc_merge=True, skip_desc_merge=True,
)) )
futures.append(future)
for future in futures:
handles.append(future.result())
# Begin async xfer. # Begin async xfer.
start = time.perf_counter() start = time.perf_counter()