iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[2/N][KV Connector][NIXL] Support per-region transfer geometry (#53780)

Browse Source 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
Matthew Bonanni
2026-09-08 16:01:12 -04:00
committed by GitHub
co-authored by OpenAI Codex OpenAI Codex
parent 2c9d68f80b
commit e0aaef85f3
9 changed files with 894 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_nixl_connector.py
+125 -4
View File
@@ -518,8 +518,15 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
slot_size_bytes = 4096
self.slot_size_per_layer = [slot_size_bytes]
self.block_len_per_layer = [slot_size_bytes * self.block_size]
self.num_blocks = 1
self.num_regions = 1
self.block_stride_per_layer = list(self.block_len_per_layer)
self.region_num_blocks = [self.num_blocks]
self.region_group_ids = [0]
self.region_mem_types = [self.nixl_memory_type]
self.dst_num_blocks[self.engine_id] = self.num_blocks
self.dst_region_num_blocks[self.engine_id] = self.region_num_blocks
self.dst_region_group_ids[self.engine_id] = self.region_group_ids
self.dst_region_mem_types[self.engine_id] = self.region_mem_types
assert expected_engine_id == self.REMOTE_ENGINE_ID
@@ -548,7 +555,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
agent_metadata=FakeNixlWrapper.AGENT_METADATA,
kv_caches_base_addr=[0],
device_id=remote_tp_rank,
num_blocks=1,
num_blocks=self.num_blocks,
block_lens=remote_block_lens,
block_strides=remote_block_lens,
kv_cache_layout="LBHNC",
@@ -556,6 +563,9 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
ssm_sizes=(0, 0),
attn_backend_name=self.backend_name,
physical_blocks_per_logical_kv_block=1,
region_num_blocks=self.region_num_blocks,
region_group_ids=self.region_group_ids,
region_mem_types=self.region_mem_types,
),
remote_tp_rank=remote_tp_rank,
remote_tp_size=remote_tp_size,
@@ -640,7 +650,7 @@ class TestNixlHandshake:
request_id = "req_id"
# Test worker role in decode server.
kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=10)
connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
connector.connector_worker = FakeNixlConnectorWorker(
vllm_config,
@@ -1167,12 +1177,19 @@ class TestNixlHandshake:
source_ranks_per_group=((0,), (0,)),
rank_offset_factor=1,
)
meta = MagicMock(
meta = NixlAgentMetadata(
engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
agent_metadata=FakeNixlWrapper.AGENT_METADATA,
kv_caches_base_addr=[0x1000],
device_id=0,
num_blocks=1,
block_lens=[remote_block_len],
block_strides=[remote_block_len],
kv_cache_layout="HND",
block_size=worker.block_size,
ssm_sizes=(0, 0),
attn_backend_name=worker.backend_name,
physical_blocks_per_logical_kv_block=1,
)
assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [
@@ -1846,6 +1863,110 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
llm.llm_engine.engine_core.shutdown()
def test_mixed_memory_local_descriptors_split_by_memory_type():
worker = object.__new__(NixlConnectorWorker)
worker.transfer_topo = MagicMock()
worker.block_size = 16
worker.engine_id = "local"
worker.tp_rank = 0
worker.device_id = 3
worker.kv_caches_base_addr = {"local": {0: [100, 200]}}
worker._has_mamba = False
worker._mixed_mem_types = True
worker.region_mem_types = ["DRAM", "VRAM"]
worker.region_num_blocks = [2, 2]
worker._desc_is_dram_by_block_size = {}
worker._desc_pos_by_block_size = {}
worker._dram_src_handles_by_block_size = {}
worker.nixl_memory_type = "VRAM"
worker._build_fa_local = MagicMock( # type: ignore[method-assign]
return_value=np.array(
[
[100, 10, 3],
[110, 10, 3],
[200, 10, 3],
[210, 10, 3],
]
)
)
worker.nixl_wrapper = MagicMock()
worker.nixl_wrapper.get_xfer_descs.side_effect = lambda blocks, memory_type: (
memory_type,
blocks,
)
worker.nixl_wrapper.prep_xfer_dlist.side_effect = [11, 22]
handle, blocks = worker.register_local_xfer_handler(worker.block_size)
assert handle == 22
assert worker._dram_src_handles_by_block_size[worker.block_size] == 11
assert [block[2] for block in blocks] == [0, 0, 3, 3]
assert [
call.args[1] for call in worker.nixl_wrapper.get_xfer_descs.call_args_list
] == ["DRAM", "VRAM"]
def test_mixed_memory_read_notifies_after_both_transfers_finish():
worker = object.__new__(NixlConnectorWorker)
worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])}
worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])}
worker._dram_src_handles_by_block_size = {16: 10}
worker._recving_metadata = {"request": MagicMock()}
worker._recving_transfers = defaultdict(list)
worker._pending_recv_notifs = {}
worker.xfer_stats = MagicMock()
worker.nixl_wrapper = MagicMock()
worker.nixl_wrapper.make_prepped_xfer.side_effect = [101, 102]
worker.nixl_wrapper.check_xfer_state.return_value = "DONE"
worker._read_blocks_mixed(
request_id="request",
local_block_size_key=16,
local_device_handle=20,
local_dram_handle=10,
remote_xfer_side_handle=30,
local_block_descs_ids=np.array([0, 2]),
remote_block_descs_ids=np.array([5, 7]),
notif_agent="prefill",
notif_id=b"request:1",
)
dram_read, device_read = worker.nixl_wrapper.make_prepped_xfer.call_args_list
assert dram_read.args[:2] == ("READ", 10)
assert dram_read.args[3] == 30
np.testing.assert_array_equal(dram_read.args[2], [0])
np.testing.assert_array_equal(dram_read.args[4], [5])
assert device_read.args[:2] == ("READ", 20)
assert device_read.args[3] == 30
np.testing.assert_array_equal(device_read.args[2], [0])
np.testing.assert_array_equal(device_read.args[4], [7])
worker.nixl_wrapper.send_notif.assert_not_called()
assert worker._pop_done_transfers(worker._recving_transfers) == {"request"}
worker.nixl_wrapper.send_notif.assert_called_once_with(
"prefill", notif_msg=b"request:1"
)
def test_mixed_memory_read_failure_does_not_notify_producer():
"""A failed half of a split READ must suppress its success notification."""
worker = object.__new__(NixlConnectorWorker)
worker._recving_metadata = {"request": MagicMock()}
worker._recving_transfers = {"request": [101, 102]}
worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]}
worker._is_hma_required = True
worker._failed_recv_reqs = MagicMock()
worker._log_failure = MagicMock() # type: ignore[method-assign]
worker.xfer_stats = MagicMock()
worker.nixl_wrapper = MagicMock()
worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"]
worker._pop_done_transfers(worker._recving_transfers)
assert "request" not in worker._pending_recv_notifs
worker.nixl_wrapper.send_notif.assert_not_called()
def element_byte_addrs(view: torch.Tensor) -> list[int]:
"""Absolute byte addresses of every element of a (possibly strided) view."""
offsets = torch.zeros(view.shape, dtype=torch.int64)
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+59
View File
@@ -861,6 +861,7 @@ def _make_mock_worker_for_desc_ids(
worker = MagicMock(spec=NixlConnectorWorker)
worker.num_regions = num_regions
worker._uses_region_group_mapping = False
worker._has_mamba = has_mamba
worker._group_spec_types = group_spec_types
worker.block_len_per_layer = block_len_per_layer or [100]
@@ -913,6 +914,64 @@ def test_get_block_descs_ids_hybrid_ssm():
assert list(result) == expected, f"Expected {expected}, got {list(result)}"
@pytest.mark.cpu_test
def test_get_block_descs_ids_selects_attention_regions_by_group():
"""Each attention group's blocks address only that group's regions."""
from vllm.v1.kv_cache_interface import FullAttentionSpec
worker = _make_mock_worker_for_desc_ids(
num_regions=3,
has_mamba=False,
group_spec_types=(FullAttentionSpec, FullAttentionSpec),
block_len_per_layer=[100, 100, 100],
)
worker.region_group_ids = [0, 0, 1]
worker._uses_region_group_mapping = True
result = worker._compute_desc_ids(
block_ids=([1, 2], [7]),
dst_num_blocks=10,
block_size_ratio=None,
physical_blocks_per_logical=1,
)
assert result.tolist() == [1, 2, 11, 12, 27]
remapped = worker._compute_desc_ids(
block_ids=([1, 2], [7]),
dst_num_blocks=10,
block_size_ratio=None,
physical_blocks_per_logical=1,
region_group_ids=[1, 1, 0],
)
assert remapped.tolist() == [21, 22, 7, 17]
@pytest.mark.cpu_test
def test_get_block_descs_ids_uses_per_region_pool_capacity():
"""Independent host and device pools use cumulative descriptor offsets."""
from vllm.v1.kv_cache_interface import FullAttentionSpec
worker = _make_mock_worker_for_desc_ids(
num_regions=2,
has_mamba=False,
group_spec_types=(FullAttentionSpec, FullAttentionSpec),
block_len_per_layer=[100, 100],
)
worker.region_group_ids = [0, 1]
worker._uses_region_group_mapping = True
result = worker._compute_desc_ids(
block_ids=([4], [8]),
dst_num_blocks=10,
block_size_ratio=None,
physical_blocks_per_logical=1,
region_num_blocks=[5, 10],
)
assert result.tolist() == [4, 13]
@pytest.mark.cpu_test
def test_get_block_descs_ids_kernel_block_mismatch():
"""Test _compute_desc_ids uses different strides for FA
tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+148 -2
View File
@@ -11,7 +11,8 @@ incoming transfer can overwrite a co-resident request's KV or mamba state
mid-decode (silent corruption of an unrelated request).
"""
from unittest.mock import patch
from collections import defaultdict
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
@@ -27,13 +28,14 @@ class _RecordingNixl:
def __init__(self, *args, **kwargs):
self.dlists: dict[int, np.ndarray] = {}
self.xfers: list[tuple] = []
self.registered: list[tuple[list[tuple], object]] = []
self._next_handle = 1
def get_reg_descs(self, caches_data, mem_type):
return caches_data
def register_memory(self, descs, backends=None):
pass
self.registered.append((descs, backends))
def deregister_memory(self, descs):
pass
@@ -98,6 +100,150 @@ class _RecordingNixl:
pass
@pytest.mark.cpu_test
def test_local_descriptors_follow_each_region_pool_capacity():
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
NixlConnectorWorker,
)
worker = object.__new__(NixlConnectorWorker)
worker.transfer_topo = MagicMock()
worker.device_id = 0
worker.block_len_per_layer = [16, 16]
worker.block_stride_per_layer = [16, 16]
worker.region_num_blocks = [2, 3]
descriptors = worker._build_fa_local([100, 1000], block_size_ratio=1)
assert descriptors[:, 0].tolist() == [100, 116, 1000, 1016, 1032]
@pytest.mark.cpu_test
def test_overlaid_transfer_groups_share_region_geometry():
"""Groups overlaid on one allocation share its transfer region."""
import msgspec
from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
NixlAgentMetadata,
)
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
NixlConnectorWorker,
)
from vllm.v1.kv_cache_interface import (
KVCacheConfig,
KVCacheGroupSpec,
KVCacheTensor,
MLAAttentionSpec,
)
num_blocks = 4
spec = MLAAttentionSpec(
block_size=4,
num_kv_heads=1,
head_size=8,
dtype=torch.uint8,
)
page_size = spec.page_size_bytes
block_stride = 2 * page_size
backing = torch.zeros(num_blocks, block_stride, dtype=torch.uint8)
caches = {
"layer.0": backing[:, :page_size],
"layer.1": backing[:, :page_size],
}
groups = [KVCacheGroupSpec([layer_name], spec) for layer_name in caches]
worker = object.__new__(NixlConnectorWorker)
worker.tp_rank = 0
worker.world_size = 1
worker.block_size = 4
worker.engine_id = "local-engine"
worker.use_mla = True
worker.model_config = MagicMock()
worker.model_config.get_total_num_kv_heads.return_value = 1
worker.attn_backends = []
worker._has_mamba = False
worker._is_csa_linear = False
worker.vllm_config = MagicMock()
worker.backend_name = "FLASHMLA"
worker.num_blocks = num_blocks
worker.nixl_memory_type = "VRAM"
worker.nixl_backends = None
worker.nixl_wrapper = _RecordingNixl()
worker._registered_descs = []
worker.dst_num_blocks = {}
worker.dst_region_num_blocks = {}
worker.dst_region_group_ids = {}
worker.dst_uses_region_group_mapping = {}
worker.dst_region_mem_types = {}
worker._desc_is_dram_by_block_size = {}
worker._desc_pos_by_block_size = {}
worker._dram_src_handles_by_block_size = {}
worker.src_xfer_handles_by_block_size = {}
worker.kv_caches_base_addr = defaultdict(dict)
worker._mamba_ssm_size = (0, 0)
worker.kv_cache_layout = "NHD"
worker.host_buffer_kv_cache_layout = "NHD"
worker._physical_blocks_per_logical_kv_block = 1
worker._logical_num_blocks = num_blocks
worker.region_group_ids = []
worker.region_mem_types = []
worker._mixed_mem_types = False
worker.region_names = []
worker.region_num_blocks = []
worker._region_is_mla = []
worker.block_len_per_layer = []
worker.block_stride_per_layer = []
worker.device_id = 0
worker.use_host_buffer = False
worker.host_xfer_buffers = {}
worker.device_kv_caches = {}
worker.pp_size = 1
worker.dcp_size = 1
worker.pcp_size = 1
worker.kv_buffer_device = "cuda"
worker._layer_specs = {name: spec for name in caches}
worker.kv_cache_config = KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=[
KVCacheTensor(
size=backing.nbytes,
layers=[name],
layer_stride=page_size,
block_stride=block_stride,
)
for name in caches
],
kv_cache_groups=groups,
)
transfer_topology = MagicMock()
with (
patch.object(bw, "TransferTopology", return_value=transfer_topology),
patch.object(bw, "compute_nixl_compatibility_hash", return_value="hash"),
):
worker.register_kv_caches(caches)
assert worker.region_group_ids == [-1]
assert worker.block_stride_per_layer == [block_stride]
assert worker.nixl_wrapper.registered[0][0] == [
(backing.data_ptr(), backing.nbytes, 0, "")
]
expected_addrs = [
backing.data_ptr() + block * block_stride for block in range(num_blocks)
]
assert worker.src_blocks_data[:, 0].tolist() == expected_addrs
metadata = msgspec.msgpack.decode(
worker.xfer_handshake_metadata.agent_metadata_bytes,
type=NixlAgentMetadata,
)
assert metadata.region_group_ids == [-1]
assert metadata.region_num_blocks == [num_blocks]
assert worker._block_ids_by_region(([0], [2]), worker.region_group_ids) == [[0, 2]]
def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blocks):
"""Build a real pull worker with a hybrid MLA + 2xKDA HMA layout."""
from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+15
View File
@@ -357,6 +357,12 @@ class _StubWriterWorker(NixlPushConnectorWorker):
w.engine_id = "test-decode-engine"
w._remote_agents = {}
w._physical_blocks_per_logical_kv_block = 1
w._uses_region_group_mapping = False
w.region_group_ids = [0]
w.dst_region_num_blocks = {}
w.dst_region_group_ids = {}
w.dst_uses_region_group_mapping = {}
w.dst_region_mem_types = {}
# Single non-hybrid attention group, matching the stub block id lists.
w._has_mamba = False
w._is_csa_linear = False
@@ -1314,6 +1320,15 @@ class TestPushPrefixCaching:
)
}
w.dst_num_blocks = {engine_id: 10_000, w.engine_id: 10_000}
w.dst_region_num_blocks = {
engine_id: [10_000],
w.engine_id: [10_000],
}
w.dst_region_group_ids = {engine_id: [0], w.engine_id: [0]}
w.dst_uses_region_group_mapping = {
engine_id: False,
w.engine_id: False,
}
w.dst_xfer_side_handles = {engine_id: {0: 5000}}
w.src_xfer_handles_by_block_size = {16: 2000}
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+320 -35
View File
@@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, cast
import msgspec
import numpy as np
import regex as re
import torch
import zmq
@@ -92,6 +93,16 @@ if TYPE_CHECKING:
logger = init_logger(__name__)
_SHARED_REGION_GROUP_ID = -1
def _region_sort_key(layer_name: str) -> tuple[tuple[int, int | str], ...]:
"""Sort transfer regions in model-layer order, then by cache name."""
return tuple(
(0, int(part)) if part.isdigit() else (1, part)
for part in re.split(r"(\d+)", layer_name)
)
def _share_storage_and_block_stride(caches: list[torch.Tensor]) -> bool:
"""Return whether all views share storage and a block stride."""
@@ -154,6 +165,9 @@ class NixlBaseConnectorWorker:
dst_num_blocks: int,
block_size_ratio: float | None,
physical_blocks_per_logical: int,
region_num_blocks: list[int] | None = None,
region_group_ids: list[int] | None = None,
uses_region_group_mapping: bool | None = None,
) -> np.ndarray:
"""Compute NIXL descriptor IDs for given block IDs."""
num_ssm_regions = 0
@@ -172,21 +186,51 @@ class NixlBaseConnectorWorker:
num_blocks = dst_num_blocks
if block_size_ratio is not None:
num_blocks = int(num_blocks * block_size_ratio)
num_fa_descs = self.num_regions * num_blocks
if region_num_blocks is None:
region_num_blocks = [num_blocks] * self.num_regions
elif block_size_ratio is not None:
region_num_blocks = [
int(count * block_size_ratio) for count in region_num_blocks
]
assert len(region_num_blocks) == self.num_regions
region_offsets = np.cumsum([0, *region_num_blocks[:-1]])
num_fa_descs = sum(region_num_blocks)
# All-attention fast path: single vectorized broadcast.
if num_ssm_regions == 0:
if region_group_ids is None:
region_group_ids = self.region_group_ids
if uses_region_group_mapping is None:
uses_region_group_mapping = self._uses_region_group_mapping
assert len(region_group_ids) == self.num_regions
if uses_region_group_mapping is None:
uses_region_group_mapping = len(set(region_group_ids)) > 1
region_group_ids_array = np.asarray(region_group_ids, dtype=np.int32)
# NOTE (NickLucche) With HMA, every kv group has the same number of layers
# and layers from different groups share the same kv tensor.
# eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be
# read across all regions, same for [3], but group0-group1 blocks will
# always differ (different areas). Therefore we can just flatten the
# block_ids and compute the descs ids for all groups at once.
block_arr = np.concatenate(
[np.asarray(g, dtype=np.int32) for g in block_ids]
)[None, :]
region_ids = np.arange(self.num_regions, dtype=np.int32)[:, None]
return (region_ids * num_blocks + block_arr).ravel()
if not uses_region_group_mapping:
block_arr = np.concatenate(
[np.asarray(group, dtype=np.int32) for group in block_ids]
)[None, :]
return (region_offsets[:, None] + block_arr).ravel()
desc_ids = []
for group_id, group in enumerate(block_ids):
if not group:
continue
region_ids = np.flatnonzero(
(region_group_ids_array == group_id)
| (region_group_ids_array == _SHARED_REGION_GROUP_ID)
)[:, None]
assert region_ids.size > 0
group_blocks = np.asarray(group, dtype=np.int32)[None, :]
desc_ids.append((region_offsets[region_ids] + group_blocks).ravel())
if not desc_ids:
return np.array([], dtype=np.int64)
return np.concatenate(desc_ids)
# Compute desc ids per group using the right stride: FA descs have
# num_blocks entries per region (kernel granularity, expanded by
@@ -556,6 +600,12 @@ class NixlBaseConnectorWorker:
# Number of NIXL regions. Currently one region per cache
# (so 1 per layer for MLA, otherwise 2 per layer)
self.num_regions = 0
self.region_group_ids: list[int] = []
self._uses_region_group_mapping = False
self.region_names: list[str] = []
self.region_num_blocks: list[int] = []
self.region_mem_types: list[str] = []
self._mixed_mem_types = False
# PP>1 (push mode): this worker holds a contiguous layer slice and
# transfers into the matching sub-range of a PP=1 remote's regions.
@@ -584,12 +634,20 @@ class NixlBaseConnectorWorker:
# Populated dynamically during handshake based on remote configuration.
# Per-source split handles, keyed by (tp_ratio, remote_block_size).
self.src_xfer_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {}
self._dram_src_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {}
# Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}.
self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict)
# Map of engine_id -> num_blocks. All ranks in the same deployment will
# have the same number of blocks.
self.dst_num_blocks: dict[EngineId, int] = {}
self.dst_region_num_blocks: dict[EngineId, list[int]] = {}
self.dst_region_group_ids: dict[EngineId, list[int]] = {}
self.dst_uses_region_group_mapping: dict[EngineId, bool] = {}
self.dst_region_mem_types: dict[EngineId, list[str]] = {}
self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {}
self._desc_pos_by_block_size: dict[int, np.ndarray] = {}
self._dram_src_handles_by_block_size: dict[int, int] = {}
self._registered_descs: list[Any] = []
# In progress transfers.
@@ -607,6 +665,7 @@ class NixlBaseConnectorWorker:
# Uses Queue for thread-safe cross-thread coordination with the
# background handshake thread, matching the _ready_requests pattern.
self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()
self._pending_recv_notifs: dict[ReqId, list[tuple[str, bytes]]] = {}
# Handshake metadata of this worker for NIXL transfers.
self.xfer_handshake_metadata: NixlHandshakePayload | None = None
@@ -1179,8 +1238,8 @@ class NixlBaseConnectorWorker:
self.use_host_buffer,
)
caches_data = []
seen_storage_addresses: set[int] = set()
registration_ranges: dict[tuple[int, str], tuple[int, int, int, str]] = {}
region_mem_types: list[str] = []
seen_base_addresses: list[int] = []
self._ssm_region_indices = []
self._scratch_region_indices = []
@@ -1219,7 +1278,13 @@ class NixlBaseConnectorWorker:
# K and V are packed into the content dim, so each attention layer is a
# single NIXL region whose block transfers as one unit. Mamba layers instead
# register separate conv/ssm sub-regions (see `_build_mamba_local`).
for layer_name, cache in xfer_buffers.items():
layer_names = (
xfer_buffers
if self._is_csa_linear
else sorted(xfer_buffers, key=_region_sort_key)
)
for layer_name in layer_names:
cache = xfer_buffers[layer_name]
# NOTE (NickLucche) Hybrid SSM mamba/FA physical page_size may differ when
# kernel requires a specific block size. This leads to SSM and FA layers
# having different num_blocks.
@@ -1240,6 +1305,7 @@ class NixlBaseConnectorWorker:
else layer_spec.page_size_bytes
// self._physical_blocks_per_logical_kv_block
)
group_id = self.kv_cache_config.transfer_group_index_by_layer[layer_name]
num_blocks = (
self._logical_num_blocks
if isinstance(layer_spec, MambaSpec)
@@ -1277,13 +1343,17 @@ class NixlBaseConnectorWorker:
)
continue
# Memory registration follows allocations, while transfer regions follow
# logical layers (or contiguous head segments). This keeps strided
# cross-layer views inside their registered allocation.
if storage_addr not in seen_storage_addresses:
seen_storage_addresses.add(storage_addr)
self.device_id = max(cache.get_device(), 0)
caches_data.append((storage_addr, storage.nbytes(), self.device_id, ""))
if cache.device.type == "cpu":
mem_type = "DRAM"
region_device_id = 0
else:
mem_type = self.nixl_memory_type
region_device_id = max(cache.get_device(), 0)
self.device_id = region_device_id
registration_ranges.setdefault(
(storage_addr, mem_type),
(storage_addr, storage.nbytes(), region_device_id, ""),
)
is_mla_region = isinstance(
layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
@@ -1366,12 +1436,19 @@ class NixlBaseConnectorWorker:
for base_addr, block_len, block_stride in region_specs:
if base_addr in seen_base_addresses:
region_index = seen_base_addresses.index(base_addr)
assert region_mem_types[region_index] == mem_type
self._region_is_mla[region_index] |= is_mla_region
if self.region_group_ids[region_index] != group_id:
self.region_group_ids[region_index] = _SHARED_REGION_GROUP_ID
else:
region_index = len(seen_base_addresses)
seen_base_addresses.append(base_addr)
self.block_len_per_layer.append(block_len)
self.block_stride_per_layer.append(block_stride)
self.region_group_ids.append(group_id)
self.region_names.append(layer_name)
self.region_num_blocks.append(num_blocks)
region_mem_types.append(mem_type)
self._region_is_mla.append(is_mla_region)
if isinstance(layer_spec, MambaSpec):
@@ -1391,10 +1468,10 @@ class NixlBaseConnectorWorker:
# caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
if (
self._physical_blocks_per_logical_kv_block == 1
and cache.shape[0] != num_blocks
and cache.shape[0] < num_blocks
):
raise AssertionError(
"All kv cache tensors must have the same number of "
"KV cache tensor has too few "
f"blocks; layer={layer_name}, "
f"expected_num_blocks={num_blocks}, "
f"cache_shape={tuple(cache.shape)}, "
@@ -1414,12 +1491,21 @@ class NixlBaseConnectorWorker:
== len(seen_base_addresses)
== len(self._region_is_mla)
== len(self.block_stride_per_layer)
== len(self.region_group_ids)
== len(self.region_names)
== len(self.region_num_blocks)
== len(region_mem_types)
)
# Descriptor ids must be region-ordered, matching the remote side.
self._scratch_region_indices.sort()
self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
self.num_regions = len(seen_base_addresses)
self.region_mem_types = region_mem_types
self._uses_region_group_mapping = len(set(self.region_group_ids)) > 1
self._mixed_mem_types = len(set(region_mem_types)) > 1
if self._has_mamba:
self.region_num_blocks = [self.num_blocks] * self.num_regions
if self.pp_size > 1:
start_layer, end_layer = self.model_config.get_layers_start_end_indices(
@@ -1431,16 +1517,29 @@ class NixlBaseConnectorWorker:
self._remote_region_offset = regions_per_layer * start_layer
# Total local FA descriptors (boundary between FA and mamba descs).
self.num_descs = self.num_regions * self.num_blocks
self.num_descs = sum(self.region_num_blocks)
descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)
logger.debug("Registering descs: %s", caches_data)
self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
logger.debug("Done registering descs")
self._registered_descs.append(descs)
for mem_type in sorted(set(region_mem_types)):
ranges = [
registration
for (
_,
registration_mem_type,
), registration in registration_ranges.items()
if registration_mem_type == mem_type
]
descs = self.nixl_wrapper.get_reg_descs(ranges, mem_type)
self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
self._registered_descs.append(descs)
self.device_kv_caches = kv_caches
self.dst_num_blocks[self.engine_id] = self.num_blocks
self.dst_region_num_blocks[self.engine_id] = self.region_num_blocks
self.dst_region_group_ids[self.engine_id] = self.region_group_ids
self.dst_uses_region_group_mapping[self.engine_id] = (
self._uses_region_group_mapping
)
self.dst_region_mem_types[self.engine_id] = self.region_mem_types
if self._has_mamba:
logger.info(
@@ -1479,6 +1578,10 @@ class NixlBaseConnectorWorker:
physical_blocks_per_logical_kv_block=(
self._physical_blocks_per_logical_kv_block
),
region_num_blocks=self.region_num_blocks,
region_group_ids=self.region_group_ids,
region_names=self.region_names,
region_mem_types=self.region_mem_types,
dcp_size=self.dcp_size,
pcp_size=self.pcp_size,
)
@@ -1642,14 +1745,23 @@ class NixlBaseConnectorWorker:
"""Build local FA descriptors for all layers as an Nx3 uint64 array."""
assert self.transfer_topo is not None
assert base_addresses, "Local KV cache base addresses must not be empty."
num_blocks = self.num_blocks * block_size_ratio
device_id = self.device_id
block_arange = np.arange(num_blocks, dtype=np.uint64)
parts: list[np.ndarray] = []
for i, base_addr in enumerate(base_addresses):
block_len = self.block_len_per_layer[i] // block_size_ratio
block_stride = self.block_stride_per_layer[i] // block_size_ratio
addrs = base_addr + block_arange * block_stride
logical_blocks = np.repeat(
np.arange(self.region_num_blocks[i], dtype=np.uint64),
block_size_ratio,
)
split_offsets = np.tile(
np.arange(block_size_ratio, dtype=np.uint64),
self.region_num_blocks[i],
)
addrs = (
base_addr
+ logical_blocks * self.block_stride_per_layer[i]
+ split_offsets * block_len
)
parts.append(self._stack_descs(addrs, block_len, device_id))
return np.concatenate(parts)
@@ -1670,9 +1782,11 @@ class NixlBaseConnectorWorker:
# SPLIT regions read their head slice from this many remote ranks at a
# per-rank offset; REPLICATE regions read the whole block once.
split_reads = len(plan.source_ranks_per_group[fa_group_idx])
num_blocks = nixl_agent_meta.num_blocks
region_num_blocks = nixl_agent_meta.region_num_blocks or [
nixl_agent_meta.num_blocks
] * len(nixl_agent_meta.kv_caches_base_addr)
device_id = nixl_agent_meta.device_id
block_arange = np.arange(num_blocks, dtype=np.uint64)
block_strides = nixl_agent_meta.block_strides
parts: list[np.ndarray] = []
for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
replicated = self._is_region_replicated(i)
@@ -1691,8 +1805,8 @@ class NixlBaseConnectorWorker:
)
local_block_len = local_block_len // num_reads
page_size = nixl_agent_meta.block_strides[i]
addrs = base_addr + rank_offset + block_arange * page_size
block_arange = np.arange(region_num_blocks[i], dtype=np.uint64)
addrs = base_addr + rank_offset + block_arange * block_strides[i]
parts.append(self._stack_descs(addrs, local_block_len, device_id))
return np.concatenate(parts)
@@ -1735,7 +1849,42 @@ class NixlBaseConnectorWorker:
mamba = self._build_mamba_local(local_base_addresses)
blocks_data = np.concatenate([blocks_data, mamba])
descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
if self._mixed_mem_types:
desc_is_dram = np.concatenate(
[
np.full(count * block_size_ratio, mem_type == "DRAM", dtype=bool)
for mem_type, count in zip(
self.region_mem_types,
self.region_num_blocks,
strict=True,
)
]
)
assert len(desc_is_dram) == len(blocks_data)
desc_pos = np.empty(len(desc_is_dram), dtype=np.int64)
dram_indices = np.flatnonzero(desc_is_dram)
device_indices = np.flatnonzero(~desc_is_dram)
desc_pos[dram_indices] = np.arange(len(dram_indices), dtype=np.int64)
desc_pos[device_indices] = np.arange(len(device_indices), dtype=np.int64)
self._desc_is_dram_by_block_size[block_size] = desc_is_dram
self._desc_pos_by_block_size[block_size] = desc_pos
blocks_data[desc_is_dram, 2] = 0
dram_descs = self.nixl_wrapper.get_xfer_descs(
blocks_data[desc_is_dram], "DRAM"
)
self._dram_src_handles_by_block_size[block_size] = (
self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs)
)
device_descs = self.nixl_wrapper.get_xfer_descs(
blocks_data[~desc_is_dram], self.nixl_memory_type
)
return (
self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", device_descs),
blocks_data,
)
mem_type = self.region_mem_types[0]
descs = self.nixl_wrapper.get_xfer_descs(blocks_data, mem_type)
# NIXL_INIT_AGENT to be used for preparations of local descs.
return self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs), blocks_data
@@ -1817,6 +1966,20 @@ class NixlBaseConnectorWorker:
]
nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end]
nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end]
if nixl_agent_meta.region_num_blocks is not None:
nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[
start:end
]
if nixl_agent_meta.region_group_ids is not None:
nixl_agent_meta.region_group_ids = nixl_agent_meta.region_group_ids[
start:end
]
if nixl_agent_meta.region_names is not None:
nixl_agent_meta.region_names = nixl_agent_meta.region_names[start:end]
if nixl_agent_meta.region_mem_types is not None:
nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[
start:end
]
### Register remote engine in TransferTopology (idempotent).
assert self.transfer_topo is not None
@@ -1855,7 +2018,24 @@ class NixlBaseConnectorWorker:
block_size_ratio = transfer_topo.block_size_ratio(nixl_agent_meta.block_size)
if engine_id not in self.dst_num_blocks:
num_remote_regions = len(nixl_agent_meta.kv_caches_base_addr)
self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks
self.dst_region_num_blocks[engine_id] = (
nixl_agent_meta.region_num_blocks
or [nixl_agent_meta.num_blocks] * num_remote_regions
)
self.dst_region_group_ids[engine_id] = nixl_agent_meta.region_group_ids or (
self.region_group_ids
if len(self.region_group_ids) == num_remote_regions
else [0] * num_remote_regions
)
self.dst_uses_region_group_mapping[engine_id] = (
len(set(self.dst_region_group_ids[engine_id])) > 1
)
self.dst_region_mem_types[engine_id] = (
nixl_agent_meta.region_mem_types
or [self.nixl_memory_type] * num_remote_regions
)
# Keep track of remote agent kv caches base addresses.
self.kv_caches_base_addr[engine_id][remote_tp_rank] = (
@@ -1902,6 +2082,8 @@ class NixlBaseConnectorWorker:
# while the SSM state is sharded across every remote TP rank.
# We only do this once per remote (tp_size, block_size).
self.src_xfer_handles_by_tp_ratio[split_key] = []
if self._mixed_mem_types:
self._dram_src_handles_by_tp_ratio[split_key] = []
for handle_data in self._build_local_splits_from_plan(
plan,
@@ -1909,6 +2091,17 @@ class NixlBaseConnectorWorker:
self.num_descs * block_size_ratio,
block_size_ratio,
):
if self._mixed_mem_types:
handle_data = np.asarray(handle_data, dtype=np.uint64)
desc_is_dram = self._desc_is_dram_by_block_size[remote_block_size]
dram_descs = self.nixl_wrapper.get_xfer_descs(
handle_data[desc_is_dram], "DRAM"
)
dram_handle = self.nixl_wrapper.prep_xfer_dlist(
"NIXL_INIT_AGENT", dram_descs
)
self._dram_src_handles_by_tp_ratio[split_key].append(dram_handle)
handle_data = handle_data[~desc_is_dram]
descs = self.nixl_wrapper.get_xfer_descs(
handle_data, self.nixl_memory_type
)
@@ -1943,8 +2136,14 @@ class NixlBaseConnectorWorker:
mamba = self._build_mamba_remote(nixl_agent_meta, tp_ratio, transfer_info)
blocks_data = np.concatenate([blocks_data, mamba])
# Register with NIXL.
descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
# A mixed local destination is split into multiple READs below. The
# producer remains a single descriptor list in this prerequisite.
remote_mem_types = set(self.dst_region_mem_types[engine_id])
if len(remote_mem_types) != 1:
raise NotImplementedError(
"NIXL pull does not yet support a mixed-memory producer"
)
descs = self.nixl_wrapper.get_xfer_descs(blocks_data, remote_mem_types.pop())
self.dst_xfer_side_handles[engine_id][remote_tp_rank] = (
self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs)
)
@@ -2133,6 +2332,27 @@ class NixlBaseConnectorWorker:
assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks
# Same number of regions/~layers.
assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer)
num_remote_regions = len(nixl_agent_meta.kv_caches_base_addr)
if nixl_agent_meta.region_num_blocks is not None:
assert len(nixl_agent_meta.region_num_blocks) == num_remote_regions
assert len(nixl_agent_meta.block_strides) == num_remote_regions
assert all(
stride >= block_len
for stride, block_len in zip(
nixl_agent_meta.block_strides,
nixl_agent_meta.block_lens,
strict=True,
)
)
if nixl_agent_meta.region_group_ids is not None:
assert len(nixl_agent_meta.region_group_ids) == num_remote_regions
if nixl_agent_meta.region_names is not None:
assert nixl_agent_meta.region_names == self.region_names, (
"NIXL transfer regions must name the same model caches in the "
"same order"
)
if nixl_agent_meta.region_mem_types is not None:
assert len(nixl_agent_meta.region_mem_types) == num_remote_regions
def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):
"""copy recved kv from host buffer to device."""
@@ -2526,6 +2746,7 @@ class NixlBaseConnectorWorker:
# don't report it again, just drop the remaining handles.
if req_id in self._recving_metadata:
done_req_ids.add(req_id)
self._send_pending_recv_notifs(req_id)
del transfers[req_id]
else:
transfers[req_id] = in_progress
@@ -2545,6 +2766,8 @@ class NixlBaseConnectorWorker:
# reported failed on the first one, which pops its metadata in
# get_finished(); on the later failures only the handle cleanup is left.
# TODO (NickLucche) handle failed transfer for HMA.
# A split READ's notification is sent only after every handle succeeds.
self._pending_recv_notifs.pop(req_id, None)
if (meta := self._recving_metadata.get(req_id)) is not None:
if not self._is_hma_required:
self._invalid_block_ids.put(set(meta.local_block_ids[0]))
@@ -2553,6 +2776,21 @@ class NixlBaseConnectorWorker:
self.nixl_wrapper.release_xfer_handle(handle)
self.xfer_stats.record_failed_transfer()
def _send_pending_recv_notifs(self, req_id: str) -> None:
for agent_name, notif_id in self._pending_recv_notifs.pop(req_id, []):
try:
self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)
except Exception as e:
self._log_failure(
failure_type="notification_failed",
msg="P worker blocks will be freed after timeout. "
"This may indicate network issues.",
req_id=req_id,
error=e,
remote_agent_name=agent_name,
)
self.xfer_stats.record_failed_notification()
def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None:
"""
Send heartbeat notifications to remote engines, extending lease on KV blocks.
@@ -2726,6 +2964,42 @@ class NixlBaseConnectorWorker:
matched_blocks = min(len(local_slice), len(remote_slice))
return local_slice[:matched_blocks], remote_slice[:matched_blocks]
@staticmethod
def _block_ids_by_region(
block_ids: BlockIds, region_group_ids: list[int]
) -> BlockIds:
"""Expand group block IDs into the corresponding region order."""
shared_block_ids = list(itertools.chain.from_iterable(block_ids))
block_ids_by_region = []
for group_id in region_group_ids:
if group_id == _SHARED_REGION_GROUP_ID:
block_ids_by_region.append(shared_block_ids.copy())
continue
block_ids_by_region.append(list(block_ids[group_id]))
return block_ids_by_region
@staticmethod
def _apply_prefix_caching_by_region(
decode_block_ids: BlockIds, prefill_block_ids: BlockIds
) -> tuple[BlockIds, BlockIds]:
"""Pair an uncached decode suffix with the same prefill regions."""
assert len(decode_block_ids) == len(prefill_block_ids)
if not any(decode_block_ids):
return [], prefill_block_ids
trimmed_prefill: list[list[int]] = []
for decode_region, prefill_region in zip(
decode_block_ids, prefill_block_ids, strict=True
):
if len(decode_region) > len(prefill_region):
raise ValueError(
"Decode allocated more KV pages than the prefill worker supplied"
)
trimmed_prefill.append(
list(prefill_region[-len(decode_region) :]) if decode_region else []
)
return decode_block_ids, trimmed_prefill
def _apply_prefix_caching(
self,
decode_block_ids: BlockIds,
@@ -2893,6 +3167,10 @@ class NixlBaseConnectorWorker:
self.kv_caches_base_addr.pop(engine_id, None)
self.dst_num_blocks.pop(engine_id, None)
self.dst_region_num_blocks.pop(engine_id, None)
self.dst_region_group_ids.pop(engine_id, None)
self.dst_uses_region_group_mapping.pop(engine_id, None)
self.dst_region_mem_types.pop(engine_id, None)
self.tp_mappings.pop(engine_id, None)
if self.transfer_topo is not None:
self.transfer_topo.unregister_remote_engine(engine_id)
@@ -2929,6 +3207,13 @@ class NixlBaseConnectorWorker:
for handle in handles:
self.nixl_wrapper.release_dlist_handle(handle)
self.src_xfer_handles_by_tp_ratio.clear()
for handles in self._dram_src_handles_by_tp_ratio.values():
for handle in handles:
self.nixl_wrapper.release_dlist_handle(handle)
self._dram_src_handles_by_tp_ratio.clear()
for handle in self._dram_src_handles_by_block_size.values():
self.nixl_wrapper.release_dlist_handle(handle)
self._dram_src_handles_by_block_size.clear()
for engine_id in list(self._remote_agents):
self._cleanup_remote_engine(engine_id, log_eviction=False)
for desc in self._registered_descs:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+6 -1
View File
@@ -46,8 +46,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"
# 8: Add dcp_size and pcp_size to NixlAgentMetadata
# 9: Add block_strides
# 10: Add dense virtual transfer pages for compressed MLA caches
# 11: Add per-region transfer geometry and memory types to NixlAgentMetadata
#
NIXL_CONNECTOR_VERSION: int = 10
NIXL_CONNECTOR_VERSION: int = 11
@dataclass
@@ -64,6 +65,10 @@ class NixlAgentMetadata:
ssm_sizes: tuple[int, int]
attn_backend_name: str
physical_blocks_per_logical_kv_block: int
region_num_blocks: list[int] | None = None
region_group_ids: list[int] | None = None
region_names: list[str] | None = None
region_mem_types: list[str] | None = None
dcp_size: int = 1
pcp_size: int = 1
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+210 -70
View File
@@ -5,6 +5,8 @@
import time
from typing import TYPE_CHECKING
import numpy as np
from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
NixlBaseConnectorWorker,
@@ -156,60 +158,86 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
remote_info = self.transfer_topo.get_engine_info(engine_id)
tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)
remote_logical_block_ids = meta.remote.block_ids
meta.remote.block_ids = self._logical_to_kernel_block_ids(
remote_logical_block_ids,
remote_info.remote_physical_blocks_per_logical,
)
num_groups = len(meta.local_block_ids)
dcp_active = self.dcp_size > 1 or remote_info.remote_dcp_size > 1
def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]:
return [
list(block_ids[g]) if rank in plan.source_ranks_per_group[g] else []
for g in range(num_groups)
]
read_specs = []
for rank in plan.all_source_ranks:
if dcp_active:
# DCP interleaves at block granularity, so slicing happens
# here on logical blocks, before kernel-block expansion.
local_ids = group_ids(meta.local_block_ids, rank)
remote_ids = group_ids(remote_logical_block_ids, rank)
for g in range(num_groups):
if not local_ids[g]:
continue
# Prefix cache hit may lead to skip some of the remote reads
# TODO (NickLucche) consider unifying prefix cache handling on
# logical blocks here for both dcp and non-dcp
local_ids[g], remote_ids[g] = self._apply_dcp_prefix_caching(
local_ids[g],
remote_ids[g],
remote_rank=rank,
local_dcp_size=self.dcp_size,
local_dcp_rank=self.dcp_rank,
remote_dcp_size=remote_info.remote_dcp_size,
local_num_computed_blocks=meta.local_num_computed_blocks[g],
)
local_physical_ids = self._logical_to_kernel_block_ids(
local_ids, self._physical_blocks_per_logical_kv_block
)
remote_physical_ids = self._logical_to_kernel_block_ids(
remote_ids, remote_info.remote_physical_blocks_per_logical
)
else:
# No DCP realignment needed: reuse the already-expanded full
# physical lists instead of re-deriving them from logical ids.
local_physical_ids = group_ids(meta.local_physical_block_ids, rank)
remote_physical_ids = group_ids(meta.remote.block_ids, rank)
read_specs.append(
ReadSpec(
remote_rank=rank,
local_block_ids=local_physical_ids,
remote_block_ids=remote_physical_ids,
local_block_ids = meta.local_physical_block_ids
local_region_groups = getattr(self, "region_group_ids", [])
remote_region_groups = getattr(self, "dst_region_group_ids", {}).get(
engine_id, local_region_groups
)
if local_region_groups != remote_region_groups:
if not self.use_mla or self._has_mamba:
raise NotImplementedError(
"Different NIXL cache-group layouts are only supported for "
"pure MLA models"
)
assert len(plan.all_source_ranks) == 1
remote_physical_block_ids = self._logical_to_kernel_block_ids(
meta.remote.block_ids,
remote_info.remote_physical_blocks_per_logical,
)
remote_by_region = self._block_ids_by_region(
remote_physical_block_ids, remote_region_groups
)
read_specs = [
ReadSpec(
remote_rank=plan.all_source_ranks[0],
local_block_ids=self._block_ids_by_region(
local_block_ids, local_region_groups
),
remote_block_ids=remote_by_region,
block_ids_by_region=True,
)
]
else:
remote_logical_block_ids = meta.remote.block_ids
meta.remote.block_ids = self._logical_to_kernel_block_ids(
remote_logical_block_ids,
remote_info.remote_physical_blocks_per_logical,
)
num_groups = len(meta.local_block_ids)
def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]:
return [
list(block_ids[g]) if rank in plan.source_ranks_per_group[g] else []
for g in range(num_groups)
]
read_specs = []
for rank in plan.all_source_ranks:
if dcp_active:
local_ids = group_ids(meta.local_block_ids, rank)
remote_ids = group_ids(remote_logical_block_ids, rank)
for g in range(num_groups):
if not local_ids[g]:
continue
local_ids[g], remote_ids[g] = self._apply_dcp_prefix_caching(
local_ids[g],
remote_ids[g],
remote_rank=rank,
local_dcp_size=self.dcp_size,
local_dcp_rank=self.dcp_rank,
remote_dcp_size=remote_info.remote_dcp_size,
local_num_computed_blocks=(
meta.local_num_computed_blocks[g]
),
)
local_physical_ids = self._logical_to_kernel_block_ids(
local_ids, self._physical_blocks_per_logical_kv_block
)
remote_physical_ids = self._logical_to_kernel_block_ids(
remote_ids,
remote_info.remote_physical_blocks_per_logical,
)
else:
local_physical_ids = group_ids(meta.local_physical_block_ids, rank)
remote_physical_ids = group_ids(meta.remote.block_ids, rank)
read_specs.append(
ReadSpec(
remote_rank=rank,
local_block_ids=local_physical_ids,
remote_block_ids=remote_physical_ids,
)
)
# D may have to perform multiple reads from different remote ranks.
# Pure MLA reads once because its cache is replicated. Hybrid
@@ -234,27 +262,39 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
# reads. Get the memory chunk onto which we will write to.
split_key = (tp_ratio, remote_block_size)
local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i]
local_dram_handle = (
self._dram_src_handles_by_tp_ratio[split_key][i]
if self._mixed_mem_types
else None
)
else:
# Single read from remote, we write to the whole memory region.
# Also handle remote block size different from local block size.
local_xfer_side_handle = self.src_xfer_handles_by_block_size[
remote_block_size
]
local_dram_handle = (
self._dram_src_handles_by_block_size[remote_block_size]
if self._mixed_mem_types
else None
)
# Destination handle: remote_engine_id -> remote_rank -> handle.
remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][
spec.remote_rank
]
self._read_blocks(
if not self._read_blocks(
read_spec=spec,
request_id=req_id,
dst_engine_id=meta.remote.engine_id,
remote_request_id=meta.remote.request_id,
local_xfer_side_handle=local_xfer_side_handle,
local_dram_handle=local_dram_handle,
remote_xfer_side_handle=remote_xfer_side_handle,
expected_consumers=plan.local_consumers,
)
):
return
if self.use_mla and tp_ratio < 0 and len(read_specs) == 1:
# ..but we still need to notify the other remote ranks that we
@@ -274,9 +314,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
request_id: str,
remote_request_id: str,
local_xfer_side_handle: int,
local_dram_handle: int | None,
remote_xfer_side_handle: int,
expected_consumers: int,
):
) -> bool:
"""
Post a READ point-to-point xfer request from a single local worker to
a single remote worker.
@@ -291,6 +332,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
remote_info.remote_block_size
)
if block_size_ratio > 1:
if read_spec.block_ids_by_region:
raise NotImplementedError(
"Region-mapped NIXL transfers require matching physical block sizes"
)
local_block_ids, remote_block_ids = (
self._map_block_ids_for_block_size_ratio(
local_block_ids, remote_block_ids, block_size_ratio
@@ -328,26 +373,30 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
remote_agent_name=agent_name,
)
self.xfer_stats.record_failed_notification()
return
return True
assert (
len(remote_block_ids)
== len(local_block_ids)
== len(self.kv_cache_config.transfer_groups)
)
if not (self.dcp_size > 1 or remote_info.remote_dcp_size > 1):
# DCP-active reads were already trimmed (and DCP-realigned, when
# sizes mismatch) in _read_blocks_for_req, at logical granularity.
local_block_ids, remote_block_ids = self._apply_prefix_caching(
if read_spec.block_ids_by_region:
local_block_ids, remote_block_ids = self._apply_prefix_caching_by_region(
decode_block_ids=local_block_ids,
prefill_block_ids=remote_block_ids,
decode_physical_per_logical=(
self._physical_blocks_per_logical_kv_block
),
prefill_physical_per_logical=(
remote_info.remote_physical_blocks_per_logical
),
)
else:
assert (
len(remote_block_ids)
== len(local_block_ids)
== len(self.kv_cache_config.transfer_groups)
)
if not (self.dcp_size > 1 or remote_info.remote_dcp_size > 1):
local_block_ids, remote_block_ids = self._apply_prefix_caching(
decode_block_ids=local_block_ids,
prefill_block_ids=remote_block_ids,
decode_physical_per_logical=(
self._physical_blocks_per_logical_kv_block
),
prefill_physical_per_logical=(
remote_info.remote_physical_blocks_per_logical
),
)
# NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
# corresponding rank. With heterogeneous TP, fixing D>P, the D tp
@@ -359,12 +408,34 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
dst_num_blocks=self.dst_num_blocks[dst_engine_id],
block_size_ratio=None,
physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
region_num_blocks=self.dst_region_num_blocks[dst_engine_id],
region_group_ids=(
list(range(self.num_regions))
if read_spec.block_ids_by_region
else self.dst_region_group_ids[dst_engine_id]
),
uses_region_group_mapping=(
self.num_regions > 1
if read_spec.block_ids_by_region
else self.dst_uses_region_group_mapping[dst_engine_id]
),
)
local_block_descs_ids = self._compute_desc_ids(
block_ids=local_block_ids,
dst_num_blocks=self.dst_num_blocks[self.engine_id],
block_size_ratio=block_size_ratio,
physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
region_num_blocks=self.dst_region_num_blocks[self.engine_id],
region_group_ids=(
list(range(self.num_regions))
if read_spec.block_ids_by_region
else self.region_group_ids
),
uses_region_group_mapping=(
self.num_regions > 1
if read_spec.block_ids_by_region
else self._uses_region_group_mapping
),
)
assert len(local_block_descs_ids) == len(remote_block_descs_ids)
@@ -372,6 +443,19 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
# Prepare transfer with Nixl.
handle = None
try:
if self._mixed_mem_types:
self._read_blocks_mixed(
request_id=request_id,
local_block_size_key=remote_info.remote_block_size,
local_device_handle=local_xfer_side_handle,
local_dram_handle=local_dram_handle,
remote_xfer_side_handle=remote_xfer_side_handle,
local_block_descs_ids=local_block_descs_ids,
remote_block_descs_ids=remote_block_descs_ids,
notif_agent=self._remote_agents[dst_engine_id][(0, remote_rank)],
notif_id=notif_id,
)
return True
handle = self.nixl_wrapper.make_prepped_xfer(
"READ",
local_xfer_side_handle,
@@ -386,6 +470,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
# Use handle to check completion in future step().
self._recving_transfers[request_id].append(handle)
return True
except Exception as e:
# mark all (logical) blocks for this request as invalid
self._log_failure(
@@ -397,6 +482,61 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
remote_rank=remote_rank,
)
self._handle_failed_transfer(request_id, handle)
return False
def _read_blocks_mixed(
self,
request_id: str,
local_block_size_key: int,
local_device_handle: int,
local_dram_handle: int | None,
remote_xfer_side_handle: int,
local_block_descs_ids: np.ndarray,
remote_block_descs_ids: np.ndarray,
notif_agent: str,
notif_id: bytes,
) -> None:
"""Split a READ across the local DRAM and device descriptor lists."""
desc_is_dram = self._desc_is_dram_by_block_size[local_block_size_key]
desc_pos = self._desc_pos_by_block_size[local_block_size_key]
local_ids = np.asarray(local_block_descs_ids)
remote_ids = np.asarray(remote_block_descs_ids)
is_dram = desc_is_dram[local_ids]
assert local_dram_handle is not None
reads = (
(is_dram, local_dram_handle),
(~is_dram, local_device_handle),
)
handles: list[int] = []
try:
for mask, local_handle in reads:
if mask.any():
handles.append(
self.nixl_wrapper.make_prepped_xfer(
"READ",
local_handle,
desc_pos[local_ids[mask]],
remote_xfer_side_handle,
remote_ids[mask],
)
)
except Exception:
for handle in handles:
self.nixl_wrapper.release_xfer_handle(handle)
raise
self._pending_recv_notifs.setdefault(request_id, []).append(
(notif_agent, notif_id)
)
for index, handle in enumerate(handles):
try:
self.nixl_wrapper.transfer(handle)
except Exception:
for unstarted in handles[index:]:
self.nixl_wrapper.release_xfer_handle(unstarted)
raise
self._recving_transfers[request_id].append(handle)
def _get_new_notifs(self) -> set[str]:
"""
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+10
View File
@@ -132,6 +132,10 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
def register_kv_caches(self, kv_caches: dict[str, "torch.Tensor"]):
super().register_kv_caches(kv_caches)
if self._mixed_mem_types:
raise NotImplementedError(
"NixlPushConnector does not support mixed-memory KV caches"
)
if self._push_writer_thread is None:
self._push_writer_thread = threading.Thread(
target=self._push_writer_loop,
@@ -661,12 +665,18 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
dst_num_blocks=self.dst_num_blocks[dst_engine_id],
block_size_ratio=None,
physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
region_num_blocks=self.dst_region_num_blocks[dst_engine_id],
region_group_ids=self.dst_region_group_ids[dst_engine_id],
uses_region_group_mapping=self.dst_uses_region_group_mapping[dst_engine_id],
)
local_block_descs_ids = self._compute_desc_ids(
block_ids=local_block_ids,
dst_num_blocks=self.dst_num_blocks[self.engine_id],
block_size_ratio=block_size_ratio,
physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
region_num_blocks=self.dst_region_num_blocks[self.engine_id],
region_group_ids=self.region_group_ids,
uses_region_group_mapping=self._uses_region_group_mapping,
)
assert len(local_block_descs_ids) == len(remote_block_descs_ids)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py
+1
View File
@@ -26,6 +26,7 @@ class ReadSpec:
remote_rank: int
local_block_ids: BlockIds
remote_block_ids: BlockIds
block_ids_by_region: bool = False
def _is_attention_spec(spec_type: type[KVCacheSpec]) -> bool: