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

Merge pull request #17 from praveingk/batching

Browse Source 
Load balance across multiple workers
This commit is contained in:
Robert Shaw 2025-06-30 08:21:03 -04:00 committed by GitHub
commit 39e6bd19fd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
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
@ -984,18 +985,25 @@ class NixlConnectorWorker:
assert len(local_block_descs_ids) == len(remote_block_descs_ids) assert len(local_block_descs_ids) == len(remote_block_descs_ids)
# Prepare transfer with Nixl. # Prepare transfer with Nixl.
CHUNK_SIZE = 100 CHUNK_SIZE = 1000
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()