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

clean up

Browse Source 
Signed-off-by: NickLucche <nlucches@redhat.com>
This commit is contained in:
NickLucche 2025-10-09 15:43:43 +00:00
parent 7bb3861faf
commit 9f38fed93c
1 changed files with 111 additions and 75 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -520,7 +520,9 @@ class NixlConnectorWorker:
Calculate the tensor parallel ratio between local and remote TP. Calculate the tensor parallel ratio between local and remote TP.
We can think of it as the number of local TP workers-per-remote TP We can think of it as the number of local TP workers-per-remote TP
workers. Local workers will read from the same remote TP worker in workers. Local workers will read from the same remote TP worker in
groups of size `tp_ratio`. groups of size `tp_ratio`. If remote tp_size > local tp_size, the
ratio is flipped (remote_size/local_size) and the returned value is
negative.
""" """
if remote_tp_size is None: if remote_tp_size is None:
assert remote_engine_id is not None assert remote_engine_id is not None
@ -539,7 +541,9 @@ class NixlConnectorWorker:
# P TP > D TP case, return the ratio as negative # P TP > D TP case, return the ratio as negative
return -remote_tp_size // self.tp_size return -remote_tp_size // self.tp_size
def is_kv_replicated(self, engine_id: Optional[EngineId] = None, tp_size: Optional[int] = None) -> bool: def is_kv_replicated(
self, engine_id: Optional[EngineId] = None, tp_size: Optional[int] = None
) -> bool:
""" """
Whether the KV cache is replicated across TP workers due to the Whether the KV cache is replicated across TP workers due to the
number of TP workers being greater than the number of KV heads. number of TP workers being greater than the number of KV heads.
@ -549,11 +553,14 @@ class NixlConnectorWorker:
tp_size = self.remote_tp_size[engine_id] tp_size = self.remote_tp_size[engine_id]
return tp_size // self.total_num_kv_heads >= 1 return tp_size // self.total_num_kv_heads >= 1
def replicates_kv_cache(self, remote_engine_id: Optional[EngineId] = None, remote_tp_size: Optional[int] = None) -> bool: def replicates_kv_cache(
self,
remote_engine_id: Optional[EngineId] = None,
remote_tp_size: Optional[int] = None,
) -> bool:
# MLA is always replicated as the hidden dim can't be split. # MLA is always replicated as the hidden dim can't be split.
return ( return self.is_mla or self.is_kv_replicated(
self.is_mla remote_engine_id, remote_tp_size
or self.is_kv_replicated(remote_engine_id, remote_tp_size)
) )
def get_target_remote_ranks( def get_target_remote_ranks(
@ -563,7 +570,8 @@ class NixlConnectorWorker:
) -> list[int]: ) -> list[int]:
""" """
Get the remote TP rank (on P) that the current local TP rank Get the remote TP rank (on P) that the current local TP rank
(on D) will read from. (on D) will read from. When remote tp_size > local tp_size, we
read from multiple remote ranks.
""" """
tp_ratio = self.tp_ratio(remote_engine_id, remote_tp_size) tp_ratio = self.tp_ratio(remote_engine_id, remote_tp_size)
if tp_ratio > 0: if tp_ratio > 0:
@ -573,8 +581,8 @@ class NixlConnectorWorker:
tp_ratio = -tp_ratio tp_ratio = -tp_ratio
if self.replicates_kv_cache(remote_engine_id, remote_tp_size): if self.replicates_kv_cache(remote_engine_id, remote_tp_size):
# When cache is replicated on remote, we only need to read # When cache is replicated on remote, we only need to read
# from one remote. # from one remote (they all have the same cache).
return [self.tp_rank*tp_ratio] return [self.tp_rank * tp_ratio]
return [self.tp_rank * tp_ratio + i for i in range(tp_ratio)] return [self.tp_rank * tp_ratio + i for i in range(tp_ratio)]
def __init__(self, vllm_config: VllmConfig, engine_id: str): def __init__(self, vllm_config: VllmConfig, engine_id: str):
@ -672,12 +680,10 @@ class NixlConnectorWorker:
# nixl_prepped_dlist_handle. # nixl_prepped_dlist_handle.
self.src_xfer_side_handle: int = 0 self.src_xfer_side_handle: int = 0
# TODO flexible enough to handle different P TP destinations? # Populated dynamically during handshake based on remote configuration.
# tp_ratio->handles # Keep track of regions at different tp_ratio values. tp_ratio->handles
# Only poulated during handshake when we read from multiple sources
self.src_xfer_side_chunked_handles: dict[int, list[int]] = {} self.src_xfer_side_chunked_handles: dict[int, list[int]] = {}
# Map of engine_id -> nixl_prepped_dlist_handle (int)]. # Map of engine_id -> nixl_prepped_dlist_handle (int)].
# TODO do I need tp_Ratio of this?
self.dst_xfer_side_handles: defaultdict[EngineId, dict[int, int]] = defaultdict( self.dst_xfer_side_handles: defaultdict[EngineId, dict[int, int]] = defaultdict(
dict dict
) )
@ -1033,7 +1039,9 @@ class NixlConnectorWorker:
self.tp_rank, self.tp_rank,
) )
descs = self.nixl_wrapper.get_xfer_descs(self.src_blocks_data, self.nixl_memory_type) descs = self.nixl_wrapper.get_xfer_descs(
self.src_blocks_data, self.nixl_memory_type
)
# NIXL_INIT_AGENT to be used for preparations of local descs. # NIXL_INIT_AGENT to be used for preparations of local descs.
self.src_xfer_side_handle = self.nixl_wrapper.prep_xfer_dlist( self.src_xfer_side_handle = self.nixl_wrapper.prep_xfer_dlist(
"NIXL_INIT_AGENT", descs "NIXL_INIT_AGENT", descs
@ -1093,10 +1101,12 @@ class NixlConnectorWorker:
In particular, handle both homogeneous and heterogeneous TP. The former In particular, handle both homogeneous and heterogeneous TP. The former
requires local rank_i to read from remote rank_i. requires local rank_i to read from remote rank_i.
The latter, assuming D.world_size > P.world_size, requires that two or The latter, in the case of D.world_size < P.world_size, requires that a
more local TP worker share the xfer from a single TP worker. local (D) TP worker reads from multiple remote (P) TP workers.
Conversely, assuming D.world_size > P.world_size, two or more local TP
workers will read from a single remote TP worker.
Here's an example (non-MLA case): Here's an example for the last case described above (non-MLA):
rank_offset p_remote_tp_rank rank_offset p_remote_tp_rank
(kv split no) (kv split no)
@ -1155,35 +1165,36 @@ class NixlConnectorWorker:
) )
self._validate_remote_agent_handshake(nixl_agent_meta, remote_tp_size) self._validate_remote_agent_handshake(nixl_agent_meta, remote_tp_size)
# Number of D TP workers reading from a single P TP worker. This is # This is 1 when P and D `--tensor-parallel-size` match. Otherwise,
# 1 when P and D `--tensor-parallel-size` match. If P TP > D TP, # this is the ratio between the two sizes.
# we don't need to use this for splitting the remote kv cache.
tp_ratio = self.kv_info.tp_ratio(engine_id) tp_ratio = self.kv_info.tp_ratio(engine_id)
# Handle tp_size>num_kv_heads: replicate KV cache. # Handle tp_size>num_kv_heads: replicate KV cache.
indexes_into_remote = (not self.kv_info.replicates_kv_cache(engine_id) \ indexes_into_remote = (
and tp_ratio < 0) not self.kv_info.replicates_kv_cache(engine_id) and tp_ratio < 0
)
# When you realize you're in P TP>DTP you have to split your regions ### (Optional) Register local agent memory regions
if tp_ratio < 0 and tp_ratio not in self.src_xfer_side_chunked_handles: if tp_ratio < 0 and tp_ratio not in self.src_xfer_side_chunked_handles:
# TODO use positive tp_ratio value? # Remote tp_size > local tp_size: read from multiple remote ranks.
# Logically "split" own regions into |tp_ratio| chunks. Mind that
# we only do this once per remote tp_size (replica-friendly).
self.src_xfer_side_chunked_handles[tp_ratio] = [] self.src_xfer_side_chunked_handles[tp_ratio] = []
# This is still needed even for MLA # MLA-optimization: only prepare one region.
# TODO actually only needs one!! # NOTE NickLucche: only a chunk of whole cache is used with MLA!
# Check if we have a split we can re-use, ie a remote P with same tp_ratio tp_ratio_opt = 1 if self.use_mla else -tp_ratio
for i in range(-tp_ratio): for i in range(tp_ratio_opt):
blocks_data = [] blocks_data = []
for memory_region in self.src_blocks_data: for memory_region in self.src_blocks_data:
addr, local_block_len, own_tp_rank = memory_region addr, local_block_len, own_tp_rank = memory_region
# Computing block len layer by layer allow for different # Computing block len layer by layer allows for different
# block sizes per layer # block sizes to be used.
# TODO this needs to be an assert when validating remote_block_len = local_block_len // (-tp_ratio)
# TODO is this the right dim we're splitting on? H?
remote_block_len = local_block_len//(-tp_ratio)
# Offset
addr = addr + i * remote_block_len addr = addr + i * remote_block_len
blocks_data.append((addr, remote_block_len, own_tp_rank)) # TODO same tp_rank? blocks_data.append((addr, remote_block_len, own_tp_rank))
descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) descs = self.nixl_wrapper.get_xfer_descs(
blocks_data, self.nixl_memory_type
)
handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs) handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)
self.src_xfer_side_chunked_handles[tp_ratio].append(handle) self.src_xfer_side_chunked_handles[tp_ratio].append(handle)
@ -1196,8 +1207,11 @@ class NixlConnectorWorker:
# Register all remote blocks, but only the corresponding kv heads. # Register all remote blocks, but only the corresponding kv heads.
for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
# TODO workaround # Read our whole local region size from remote.
kv_block_len = self.get_backend_aware_kv_block_len(layer_idx=i) // 2 kv_block_len = self.get_backend_aware_kv_block_len(layer_idx=i)
if tp_ratio < 0:
# Remote tp is bigger: read a chunk of local region from remote
kv_block_len = kv_block_len // (-tp_ratio)
rank_offset = ( rank_offset = (
self.tp_rank % tp_ratio * kv_block_len if indexes_into_remote else 0 self.tp_rank % tp_ratio * kv_block_len if indexes_into_remote else 0
) )
@ -1262,7 +1276,7 @@ class NixlConnectorWorker:
) )
remote_block_size = remote_block_len // (self.slot_size_per_layer[0]) remote_block_size = remote_block_len // (self.slot_size_per_layer[0])
else: else:
if tp_ratio > 1 and self.device_type == "xpu": if tp_ratio != 1 and self.device_type == "xpu":
# XPU uses NHD, hence it does not support splitting on H # XPU uses NHD, hence it does not support splitting on H
raise ValueError("Heterogeneous TP is not supported on XPU") raise ValueError("Heterogeneous TP is not supported on XPU")
# When MLA is not used, this is a list of the same block length # When MLA is not used, this is a list of the same block length
@ -1270,26 +1284,42 @@ class NixlConnectorWorker:
assert block_len == remote_block_len, ( assert block_len == remote_block_len, (
"All remote layers must have the same block size" "All remote layers must have the same block size"
) )
remote_block_size = remote_block_len // (
self.slot_size_per_layer[0] * tp_ratio if tp_ratio > 0:
) # Remote NHD/H'D*tp_ratio=N -page_size-
remote_block_size = remote_block_len // (
self.slot_size_per_layer[0] * tp_ratio
)
# Remote tp is smaller: remote block_len size is bigger
assert remote_block_len == self.block_len_per_layer[0] * tp_ratio, (
"Remote P worker KV layer cache must be of shape [2, N, "
"local_kv_heads*tp_ratio, page_size, head_dim] and same dtype."
) # noqa: E501
else:
# Remote NHD/(H'D/tp_ratio)=N -page_size-
remote_block_size = remote_block_len // (
self.slot_size_per_layer[0] // (-tp_ratio)
)
# Remote tp is bigger: remote block_len size is smaller
assert remote_block_len == self.block_len_per_layer[0] // (-tp_ratio), (
"Remote P worker KV layer cache must be of shape [2, N, "
"local_kv_heads/tp_ratio, page_size, head_dim] and same dtype."
) # noqa: E501
if self._use_flashinfer: if self._use_flashinfer:
# With flashinfer, KV are sent in the same message. # With flashinfer, KV are sent in the same message.
remote_block_size //= 2 remote_block_size //= 2
# TODO add asserts for P TP > D TP
# assert remote_block_len == self.block_len_per_layer[0] * tp_ratio, (
# "Remote P worker KV layer cache must be of shape [2, N, "
# "local_kv_heads*tp_ratio, block_size, head_dim] and same dtype."
# )
# assert self.block_size == remote_block_size, ( # We may allow it in the future with logical kvcache manager block_size
# "Remote P worker with different page/block size is not supported " assert self.block_size == remote_block_size, (
# f"{self.block_size=}, {remote_block_size=}" "Remote P worker with different page/block size is not supported "
# ) f"{self.block_size=}, {remote_block_size=}"
)
# # TP workers have same #blocks. # TP workers (handhshakes with same remote) have same #blocks.
# assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks
# assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer) # Same number of regions/~layers.
assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer)
def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta): def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):
"""copy recved kv from host buffer to device.""" """copy recved kv from host buffer to device."""
@ -1421,7 +1451,7 @@ class NixlConnectorWorker:
""" """
done_req_ids: set[str] = set() done_req_ids: set[str] = set()
for req_id, handles in list(transfers.items()): for req_id, handles in list(transfers.items()):
in_progress = False in_progress = []
for handle, _xfer_stime in handles: for handle, _xfer_stime in handles:
xfer_state = self.nixl_wrapper.check_xfer_state(handle) xfer_state = self.nixl_wrapper.check_xfer_state(handle)
if xfer_state == "DONE": if xfer_state == "DONE":
@ -1430,13 +1460,16 @@ class NixlConnectorWorker:
self.xfer_stats.record_transfer(res) self.xfer_stats.record_transfer(res)
self.nixl_wrapper.release_xfer_handle(handle) self.nixl_wrapper.release_xfer_handle(handle)
elif xfer_state == "PROC": elif xfer_state == "PROC":
in_progress = True in_progress.append((handle, _xfer_stime))
continue continue
else: else:
raise RuntimeError("Transfer failed with state %s", xfer_state) raise RuntimeError("Transfer failed with state %s", xfer_state)
if not in_progress: if not in_progress:
# Only report request as completed when all transfers are done.
done_req_ids.add(req_id) done_req_ids.add(req_id)
del transfers[req_id] del transfers[req_id]
else:
transfers[req_id] = in_progress
return done_req_ids return done_req_ids
def start_load_kv(self, metadata: NixlConnectorMetadata): def start_load_kv(self, metadata: NixlConnectorMetadata):
@ -1490,7 +1523,8 @@ class NixlConnectorWorker:
def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
remote_ranks = self.kv_info.get_target_remote_ranks(meta.remote_engine_id) remote_ranks = self.kv_info.get_target_remote_ranks(meta.remote_engine_id)
# D may perform multiple reads from different remote ranks. tp_ratio = self.kv_info.tp_ratio(meta.remote_engine_id)
# D may have to perform multiple reads from different remote ranks.
for i, remote_rank in enumerate(remote_ranks): for i, remote_rank in enumerate(remote_ranks):
logger.debug( logger.debug(
"Remote agent %s available, calling _read_blocks" "Remote agent %s available, calling _read_blocks"
@ -1499,13 +1533,17 @@ class NixlConnectorWorker:
remote_rank, remote_rank,
req_id, req_id,
) )
# TODO refactor properly and ONLY DO THIS FOR PTP>DTP if tp_ratio < 0:
tp_ratio = self.kv_info.tp_ratio(meta.remote_engine_id) # Remote tp_size > local tp_size: we must perform multiple
# Get nixl desc handles depending on whether we're reading from # reads. Get the memory chunk onto which we will write to.
# multiple sources or we're reading a chunk of local_xfer_side_handle = self.src_xfer_side_chunked_handles[tp_ratio][i]
local_xfer_side_handle = self.src_xfer_side_chunked_handles[tp_ratio][i] else:
remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote_engine_id][remote_rank] # Single read from remote, we write to the whole memory region.
# TODO multiread; notifs to all twice?? SPLIT LOCAL BLOCKS! local_xfer_side_handle = self.src_xfer_side_handle
# Destination handle: remote_engine_id -> remote_rank -> handle.
remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote_engine_id][
remote_rank
]
self._read_blocks( self._read_blocks(
request_id=req_id, request_id=req_id,
dst_engine_id=meta.remote_engine_id, dst_engine_id=meta.remote_engine_id,
@ -1526,6 +1564,10 @@ class NixlConnectorWorker:
local_xfer_side_handle: int, local_xfer_side_handle: int,
remote_xfer_side_handle: int, remote_xfer_side_handle: int,
): ):
"""
Post a READ xfer request from a single local worker to a single
remote worker.
"""
# NOTE(rob): having the staging blocks be on the READER side is # NOTE(rob): having the staging blocks be on the READER side is
# not going to work well (since we will have to call rearrange tensors). # not going to work well (since we will have to call rearrange tensors).
# after we detect the txn is complete (which means we cannot make the # after we detect the txn is complete (which means we cannot make the
@ -1543,7 +1585,7 @@ class NixlConnectorWorker:
notif_id = f"{request_id}:{tp_ratio}".encode() notif_id = f"{request_id}:{tp_ratio}".encode()
# Full prefix cache hit: do not need to read remote blocks, # Full prefix cache hit: do not need to read remote blocks,
# just notify P worker(s) that we have the blocks we need. # just notify P worker that we have the blocks we need.
num_local_blocks = len(local_block_ids) num_local_blocks = len(local_block_ids)
if num_local_blocks == 0: if num_local_blocks == 0:
agent_name = self._remote_agents[dst_engine_id][remote_rank] agent_name = self._remote_agents[dst_engine_id][remote_rank]
@ -1556,10 +1598,6 @@ class NixlConnectorWorker:
if num_local_blocks < num_remote_blocks: if num_local_blocks < num_remote_blocks:
remote_block_ids = remote_block_ids[-num_local_blocks:] remote_block_ids = remote_block_ids[-num_local_blocks:]
# Get side handles.
# local_xfer_side_handle = self.src_xfer_side_handle
# remote_xfer_side_handle = self.dst_xfer_side_handles[dst_engine_id][remote_rank]
# NOTE (nicolo) With homogeneous TP, each TP worker loads KV from # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
# corresponding rank. With heterogeneous TP, fixing D>P, the D tp # corresponding rank. With heterogeneous TP, fixing D>P, the D tp
# workers will issue xfers to parts of the P worker remote kv caches. # workers will issue xfers to parts of the P worker remote kv caches.
@ -1647,7 +1685,6 @@ class NixlConnectorWorker:
assert self.num_layers == self.num_regions assert self.num_layers == self.num_regions
region_ids = np.arange(layer_idx, layer_idx + 1) region_ids = np.arange(layer_idx, layer_idx + 1)
# TODO can this vary?
num_blocks = self.dst_num_blocks[engine_id] num_blocks = self.dst_num_blocks[engine_id]
# Compute the desc ids for each block. # Compute the desc ids for each block.
@ -1694,11 +1731,10 @@ class NixlConnectorWorker:
if self.src_xfer_side_handle: if self.src_xfer_side_handle:
self.nixl_wrapper.release_dlist_handle(self.src_xfer_side_handle) self.nixl_wrapper.release_dlist_handle(self.src_xfer_side_handle)
self.src_xfer_side_handle = 0 self.src_xfer_side_handle = 0
if self.src_xfer_side_chunked_handles: for handles in self.src_xfer_side_chunked_handles.values():
for handles in self.src_xfer_side_chunked_handles.values(): for handle in handles:
for handle in handles: self.nixl_wrapper.release_dlist_handle(handle)
self.nixl_wrapper.release_dlist_handle(handle) self.src_xfer_side_chunked_handles.clear()
self.src_xfer_side_chunked_handles.clear()
for dst_xfer_side_handles in self.dst_xfer_side_handles.values(): for dst_xfer_side_handles in self.dst_xfer_side_handles.values():
for dst_xfer_side_handle in dst_xfer_side_handles.values(): for dst_xfer_side_handle in dst_xfer_side_handles.values():
self.nixl_wrapper.release_dlist_handle(dst_xfer_side_handle) self.nixl_wrapper.release_dlist_handle(dst_xfer_side_handle)