From 601fa9a74e8cd2ae0a4fa127d68b91bbf363a24e Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 28 Jul 2026 06:45:02 -0700 Subject: [PATCH] [KV Connector] Support NIXL heterogeneous P/D block sizes for hybrid models (#49612) Signed-off-by: Nick Hill Co-authored-by: Claude Fable 5 --- .../kv_connector/unit/test_nixl_connector.py | 10 +- .../unit/test_nixl_connector_hma.py | 114 ++++ .../unit/test_nixl_desc_geometry.py | 619 ++++++++++++++++++ tests/v1/kv_connector/unit/test_tp_mapping.py | 54 ++ .../kv_connector/v1/nixl/base_worker.py | 271 ++++++-- .../kv_connector/v1/nixl/pull_worker.py | 35 +- .../kv_connector/v1/nixl/push_worker.py | 22 +- 7 files changed, 1014 insertions(+), 111 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_nixl_desc_geometry.py diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index d58f117fa31..c474077ab22 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -479,8 +479,9 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): super().__init__(*args, kv_cache_config=kv_cache_config, **kwargs) self._hand_shake_latency = hand_shake_latency self.kv_cache_layout = kv_cache_layout - # Mock register_kv_caches attribute needed for tests that do not call it. + # Mock register_kv_caches attributes needed for tests that do not call it. self.src_xfer_handles_by_block_size = {self.block_size: 1} + self.src_blocks_data = np.empty((0, 3), dtype=np.uint64) test_shape = self.attn_backends[0].get_kv_cache_shape( num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 ) @@ -765,8 +766,9 @@ class TestNixlHandshake: assert remote_info.remote_tp_size == remote_tp_size assert -tp_ratio == worker.transfer_topo.tp_ratio(remote_tp_size) # ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks - assert -tp_ratio in worker.src_xfer_handles_by_tp_ratio - assert len(worker.src_xfer_handles_by_tp_ratio[-tp_ratio]) == tp_ratio + split_key = (-tp_ratio, worker.block_size) + assert split_key in worker.src_xfer_handles_by_tp_ratio + assert len(worker.src_xfer_handles_by_tp_ratio[split_key]) == tp_ratio assert remote_engine_id in worker.dst_xfer_side_handles assert set(worker.dst_xfer_side_handles[remote_engine_id].keys()) == set( range(tp_ratio) @@ -2091,7 +2093,7 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init): # Mock register_kv_cache which registers local handle worker.src_xfer_handles_by_block_size = {worker.block_size: 455} # P TP = 2 * D TP case, we should register 2 local handles - worker.src_xfer_handles_by_tp_ratio = {-2: [456, 457]} + worker.src_xfer_handles_by_tp_ratio = {(-2, 16): [456, 457]} worker.dst_xfer_side_handles = {"engine1": {0: 789}} worker._remote_agents = {"engine1": {(0, 0): "agent1"}} # _cleanup_remote_engine (called by shutdown) also clears these: diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 4945942ba3a..1f7a62d2c9a 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -708,6 +708,120 @@ def test_get_block_descs_ids_kernel_block_mismatch(): assert list(result) == expected, f"Expected {expected}, got {list(result)}" +@pytest.mark.cpu_test +def test_get_block_descs_ids_hetero_block_size_hybrid(): + """With a block-size ratio, FA desc ids are ratio-expanded while SSM + desc ids keep the unexpanded logical stride (state blocks are never + sub-split).""" + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + worker = _make_mock_worker_for_desc_ids( + num_regions=2, + has_mamba=True, + group_spec_types=(FullAttentionSpec, MambaSpec), + block_len_per_layer=[100], + ) + + ratio = 4 + # FA ids are already remote-granularity (expanded) sub-block ids. + fa_sub_blocks = [3, 5] + ssm_blocks = [1] + result = worker._compute_desc_ids( + block_ids=(fa_sub_blocks, ssm_blocks), + dst_num_blocks=100, + block_size_ratio=ratio, + physical_blocks_per_logical=1, + ) + + # FA regions have 100*4 entries each; SSM regions (4 per layer) start at + # 2*400 and stride by the unexpanded 100 logical blocks. + expected = [3, 5, 403, 405, 801, 901, 1001, 1101] + assert list(result) == expected, f"Expected {expected}, got {list(result)}" + + +def _bind_worker_method(worker, name): + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + + method = getattr(NixlConnectorWorker, name) + setattr(worker, name, method.__get__(worker, NixlConnectorWorker)) + + +@pytest.mark.cpu_test +def test_map_block_ids_for_block_size_ratio_hybrid(): + """Attention groups expand to remote granularity and clip to the remote + coverage; mamba state blocks pass through 1:1.""" + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + worker = MagicMock(spec=NixlConnectorWorker) + worker._group_spec_types = (FullAttentionSpec, MambaSpec) + _bind_worker_method(worker, "get_mapped_blocks") + _bind_worker_method(worker, "_map_block_ids_for_block_size_ratio") + + local, remote = worker._map_block_ids_for_block_size_ratio( + [[1, 2, 3], [7]], + [list(range(30, 40)), [42]], + 4, + ) + # [1, 2, 3] expand to sub-blocks [4..15], clipped to the 10 remote blocks. + assert local == [list(range(4, 14)), [7]] + assert remote == [list(range(30, 40)), [42]] + + # Attention-only full prefix hit: empty local list is preserved. + worker._group_spec_types = (FullAttentionSpec,) + local, remote = worker._map_block_ids_for_block_size_ratio([[]], [[30, 31]], 4) + assert local == [] + + +@pytest.mark.cpu_test +def test_post_process_zeroes_untransferred_tail(): + """The untransferred sub-blocks of the last local block are zeroed on + receive; mamba state caches are untouched by the attention permute.""" + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + + ratio = 4 + block_tokens = 8 # 2 tokens per remote sub-block + + worker = MagicMock(spec=NixlConnectorWorker) + worker._group_spec_types = (FullAttentionSpec, MambaSpec) + worker.transfer_topo = MagicMock() + worker.device_type = "cpu" + worker.enable_permute_local_kv = False + attn_cache = torch.ones(6, block_tokens, 2, 4) + mamba_cache = torch.ones(6, 16) + worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache} + fa_group = MagicMock(layer_names=["attn.0"]) + ssm_group = MagicMock(layer_names=["mamba.0"]) + worker.kv_cache_config = MagicMock(kv_cache_groups=[fa_group, ssm_group]) + # The cached property filters mamba layers out of the permuted caches. + attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker) + assert len(attn_caches) == 1 and attn_caches[0] is attn_cache + worker._attention_kv_caches = attn_caches + _bind_worker_method(worker, "post_process_device_kv_on_receive") + + # Request occupies blocks [2, 3]; only 6 of 8 sub-blocks were received. + worker.post_process_device_kv_on_receive(ratio, [([2, 3], 6)]) + + # Block 2 fully covered; block 3 covered for 2 sub-blocks (4 tokens). + assert torch.all(attn_cache[2] == 1) + assert torch.all(attn_cache[3, :4] == 1) + assert torch.all(attn_cache[3, 4:] == 0) + # Untouched blocks and the mamba cache keep their content. + assert torch.all(attn_cache[4] == 1) + assert torch.all(mamba_cache == 1) + + @pytest.mark.cpu_test def test_nixl_metadata_hybrid_ssm_block_ids(): """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py new file mode 100644 index 00000000000..e2fc41a0229 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -0,0 +1,619 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end NIXL descriptor geometry invariants for hybrid MLA+SSM models +under heterogeneous P/D block geometry (TP-sharded KDA-style state, so +the mamba-aligned logical block size differs between P and D while the +kernel-granularity pages stay equal). + +The invariant under test: every LOCAL byte range a request's READ transfers +into must lie within that request's own blocks. A violation means an +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 + +import numpy as np +import pytest +import torch + +from .utils import create_vllm_config + + +class _RecordingNixl: + """Minimal NIXL wrapper stand-in that records descriptor lists and + prepared transfers so tests can resolve desc ids to byte ranges.""" + + def __init__(self, *args, **kwargs): + self.dlists: dict[int, np.ndarray] = {} + self.xfers: list[tuple] = [] + self._next_handle = 1 + + def get_reg_descs(self, caches_data, mem_type): + return caches_data + + def register_memory(self, descs, backends=None): + pass + + def deregister_memory(self, descs): + pass + + def get_agent_metadata(self): + return b"agent-meta" + + def get_xfer_descs(self, blocks_data, mem_type): + return blocks_data + + def prep_xfer_dlist(self, agent, descs): + handle = self._next_handle + self._next_handle += 1 + self.dlists[handle] = np.asarray(descs, dtype=np.uint64).reshape(-1, 3) + return handle + + def add_remote_agent(self, metadata): + return "remote-agent" + + def make_prepped_xfer( + self, op, local_handle, local_ids, remote_handle, remote_ids, notif_msg=None + ): + handle = self._next_handle + self._next_handle += 1 + self.xfers.append( + ( + op, + local_handle, + np.asarray(local_ids), + remote_handle, + np.asarray(remote_ids), + ) + ) + return handle + + def transfer(self, handle): + pass + + def check_xfer_state(self, handle): + return "DONE" + + def get_xfer_telemetry(self, handle): + from types import SimpleNamespace + + return SimpleNamespace( + xferDuration=1.0, postDuration=1.0, totalBytes=1, descCount=1 + ) + + def release_xfer_handle(self, handle): + pass + + def release_dlist_handle(self, handle): + pass + + def send_notif(self, agent, notif_msg=None): + pass + + def get_new_notifs(self): + return {} + + def remove_remote_agent(self, agent): + pass + + +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 ( + base_worker as bw, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum + from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + ) + + mla_spec = MLAAttentionSpec( + block_size=local_block_size, + num_kv_heads=1, + head_size=6, + dtype=torch.float16, + ) + unified_page = mla_spec.page_size_bytes + kda_spec = MambaSpec( + block_size=local_block_size, + shapes=((8, 3), (1, 4, 4)), + dtypes=(torch.float16, torch.float32), + page_size_padded=unified_page, + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + ) + kv_cache_config = KVCacheConfig( + num_blocks=num_logical_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=num_logical_blocks * unified_page, + shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"], + ) + for i in range(2) + ], + kv_cache_groups=[ + KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec), + KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec), + KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec), + ], + ) + + vllm_config = create_vllm_config(block_size=local_block_size) + vllm_config.cache_config.enable_prefix_caching = False + # kv_buffer_device defaults to the *real* platform's device type, which on + # a CPU-only test host would make this a host-buffer worker: host xfer + # buffers are per-layer, so the HMA shared-tensor regions this test builds + # would not be deduplicated. Pin it to the faked device type. + vllm_config.kv_transfer_config.kv_buffer_device = "cuda" + + from unittest.mock import MagicMock + + fake_backend = MagicMock() + fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size] + fake_backend.get_name.return_value = "FLASHMLA" + fake_backend.full_cls_name.return_value = "fake.FLASHMLA" + fake_platform = MagicMock() + fake_platform.device_type = "cuda" + fake_platform.get_nixl_memory_type.return_value = "VRAM" + + from vllm.config import set_current_vllm_config + + with ( + patch.object(bw, "NixlWrapper", _RecordingNixl), + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), + patch.object(bw, "current_platform", fake_platform), + patch( + "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout", + return_value="DS", + ), + set_current_vllm_config(vllm_config), + ): + worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config) + worker.use_mla = True + + # Attention caches are kernel-block granular on dim 0, as the + # receive post-process assumes. + ppl = local_block_size // kernel_block_size + tensors = [ + torch.zeros( + num_logical_blocks * ppl, unified_page // ppl, dtype=torch.uint8 + ) + for _ in range(2) + ] + worker.register_kv_caches( + { + "kda_a.0": tensors[0], + "mla.0": tensors[0], + "kda_b.0": tensors[0], + "kda_a.1": tensors[1], + "mla.1": tensors[1], + "kda_b.1": tensors[1], + } + ) + # Keep tensors alive alongside the worker; flat views for byte checks. + worker._test_tensors = [t.view(-1) for t in tensors] + worker._test_tensors_2d = tensors + worker._test_unified_page = unified_page + return worker + + +def _make_remote_meta( + worker, + remote_block_size, + remote_kernel_block_size, + remote_num_logical, + remote_ssm_sizes, +): + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlAgentMetadata, + ) + + remote_ppl = remote_block_size // remote_kernel_block_size + # Kernel-granularity pages are TP-independent for MLA hybrids and must + # match the local ones for the handshake to pass, scaled down by the + # block-size ratio when the remote's kernel block is smaller. + block_size_ratio = worker.block_size // remote_kernel_block_size + kernel_page = worker.block_len_per_layer[0] // block_size_ratio + return NixlAgentMetadata( + engine_id="remote-engine", + agent_metadata=b"remote-agent-meta", + device_id=0, + kv_caches_base_addr=[0x10_000_000, 0x20_000_000], + num_blocks=remote_num_logical * remote_ppl, + block_lens=[kernel_page, kernel_page], + kv_cache_layout=worker.kv_cache_layout, + block_size=remote_kernel_block_size, + ssm_sizes=remote_ssm_sizes, + attn_backend_name=worker.backend_name, + physical_blocks_per_logical_kv_block=remote_ppl, + ) + + +def _owned_byte_ranges(worker, group_logical_ids): + """Byte ranges owned by a request: for each HMA region tensor, every + logical block id of every group maps to one unified page.""" + unified_page = worker._test_unified_page + bases = [t.data_ptr() for t in worker._test_tensors] + owned = [] + for base in bases: + for ids in group_logical_ids: + for b in ids: + owned.append((base + b * unified_page, base + (b + 1) * unified_page)) + return owned + + +def _assert_local_writes_within(worker, owned_ranges): + nixl = worker.nixl_wrapper + assert nixl.xfers, "no transfers were posted" + violations = [] + total_descs = 0 + for op, local_handle, local_ids, _, remote_ids in nixl.xfers: + assert len(local_ids) == len(remote_ids) + desc_arr = nixl.dlists[local_handle] + for i in local_ids: + addr, length, _dev = desc_arr[int(i)] + addr, length = int(addr), int(length) + total_descs += 1 + if not any(lo <= addr and addr + length <= hi for lo, hi in owned_ranges): + violations.append((int(i), hex(addr), length)) + assert not violations, ( + f"{len(violations)}/{total_descs} local descriptors write outside " + f"the request's own blocks: {violations[:10]}" + ) + return total_descs + + +@pytest.mark.cpu_test +def test_hetero_ppl_multi_read_writes_stay_within_request_blocks(): + """MLA-hybrid hetero geometry: local (D, TP1) logical blocks of 12 tokens + (kernel 4, ppl=3) vs remote (P, TP2) logical blocks of 8 tokens (ppl=2), + equal kernel pages, tp_ratio=-2 multi-read with replicated MLA and + TP-sharded KDA state. Every local descriptor of the request's reads must + stay within its own blocks.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ) + + worker = _make_mla_hybrid_worker( + local_block_size=12, kernel_block_size=4, num_logical_blocks=8 + ) + assert worker._physical_blocks_per_logical_kv_block == 3 + + meta_r = _make_remote_meta( + worker, + remote_block_size=8, + remote_kernel_block_size=4, + remote_num_logical=12, + remote_ssm_sizes=(24, 32), + ) + for rank in (0, 1): + worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=2) + + # Request B: 17 matched tokens. Local: 2 logical blocks (24 tok + # capacity); remote: 16 prefilled tokens -> 2 remote logical blocks. + # Sparse, non-contiguous ids so neighbor blocks exist on all sides. + local_ids = ([2, 5], [1], [7]) + remote_ids = [[1, 4], [5], [2]] + + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="req-b", + local_block_ids=local_ids, + kv_transfer_params={ + "remote_block_ids": remote_ids, + "remote_engine_id": "remote-engine", + "remote_request_id": "prefill-req-b", + "remote_host": "localhost", + "remote_port": 1234, + "tp_size": 2, + }, + ) + meta = metadata.reqs_to_recv["req-b"] + meta.local_physical_block_ids = worker._logical_to_kernel_block_ids( + meta.local_block_ids, worker._physical_blocks_per_logical_kv_block + ) + worker._recving_metadata["req-b"] = meta + + worker._read_blocks_for_req("req-b", meta) + + owned = _owned_byte_ranges(worker, local_ids) + total = _assert_local_writes_within(worker, owned) + # Multi-read: rank 0 carries the replicated MLA + its SSM shard, + # rank 1 carries only its SSM shard. + assert len(worker.nixl_wrapper.xfers) == 2 + assert total > 0 + + +def _resolve( + desc_arr, + idx, + bases, + region_size, + unified_page, + desc_page, + logical_ids_attn, + block_tokens, +): + """Resolve a desc id to (region, kind, token_start) where kind is 'attn' + (desc-page sized, sub-block-aligned, in the request's attention blocks) + or 'mamba'. token_start is the request-relative token offset, so local + and remote are comparable even when their kernel blocks differ in size.""" + addr, length, _ = (int(x) for x in desc_arr[int(idx)]) + for region, base in enumerate(bases): + off = addr - base + if 0 <= off < region_size: + b = off // unified_page + rem = off % unified_page + if length == desc_page and rem % desc_page == 0 and b in logical_ids_attn: + pos = logical_ids_attn.index(b) + tokens_per_desc = block_tokens * desc_page // unified_page + sub = rem // desc_page + return (region, "attn", pos * block_tokens + sub * tokens_per_desc) + return (region, "mamba", None) + raise AssertionError(f"desc {idx} addr {addr:#x} not in any region") + + +def _run_hetero_case( + local_block, kernel, remote_block, num_tokens, tp_size=2, remote_kernel=None +): + """Full pull-path run for one geometry; returns pairing records. + + ``remote_kernel`` defaults to the local kernel block size; a smaller + value additionally exercises block_size_ratio > 1. + """ + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ) + + remote_kernel = remote_kernel or kernel + block_size_ratio = kernel // remote_kernel + remote_ppl = remote_block // remote_kernel + matched = num_tokens - 1 # mamba N-1 rule + n_local = -(-num_tokens // local_block) + n_remote = -(-matched // remote_block) + + worker = _make_mla_hybrid_worker( + local_block_size=local_block, + kernel_block_size=kernel, + num_logical_blocks=max(2 * n_local + 4, 8), + ) + # Local KDA state pages are (48, 64) bytes; the remote holds 1/tp_size + # shards of each. + meta_r = _make_remote_meta( + worker, + remote_block_size=remote_block, + remote_kernel_block_size=remote_kernel, + remote_num_logical=max(2 * n_remote + 4, 8), + remote_ssm_sizes=(48 // tp_size, 64 // tp_size), + ) + for rank in range(tp_size): + worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=tp_size) + + # Sparse ids so neighbors exist between the request's blocks. + local_attn = [2 * i + 1 for i in range(n_local)] + remote_attn = [2 * i + 2 for i in range(n_remote)] + local_ids = (local_attn, [0], [2 * n_local + 2]) + remote_ids = [remote_attn, [1], [0]] + + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="req-b", + local_block_ids=local_ids, + kv_transfer_params={ + "remote_block_ids": remote_ids, + "remote_engine_id": "remote-engine", + "remote_request_id": "prefill-req-b", + "remote_host": "localhost", + "remote_port": 1234, + "tp_size": tp_size, + }, + ) + meta = metadata.reqs_to_recv["req-b"] + meta.local_physical_block_ids = worker._logical_to_kernel_block_ids( + meta.local_block_ids, worker._physical_blocks_per_logical_kv_block + ) + worker._recving_metadata["req-b"] = meta + + # Sentinel-fill the local KV so untouched bytes are detectable. + for t in worker._test_tensors: + t.fill_(0xAA) + + worker._read_blocks_for_req("req-b", meta) + + # Invariant 1: all local writes within the request's own blocks. + owned = _owned_byte_ranges(worker, local_ids) + _assert_local_writes_within(worker, owned) + + # Invariant 2: local<->remote attention pairs are token-aligned. + nixl = worker.nixl_wrapper + local_bases = [t.data_ptr() for t in worker._test_tensors] + remote_bases = [0x10_000_000, 0x20_000_000] + local_unified = worker._test_unified_page + remote_unified = (local_unified // local_block) * remote_block + # With block_size_ratio > 1 the local page is split into ratio sub-descs, + # each the size of a whole remote kernel page. + desc_page = worker.block_len_per_layer[0] // block_size_ratio + meta_r_num_blocks_bytes = (meta_r.num_blocks // remote_ppl) * remote_unified + covered_tokens = set() + for op, lh, lids, rh, rids in nixl.xfers: + larr, rarr = nixl.dlists[lh], nixl.dlists[rh] + for li, ri in zip(lids, rids): + lreg, lkind, ltok = _resolve( + larr, + li, + local_bases, + len(worker._test_tensors[0]), + local_unified, + desc_page, + local_attn, + local_block, + ) + rreg, rkind, rtok = _resolve( + rarr, + ri, + remote_bases, + meta_r_num_blocks_bytes, + remote_unified, + desc_page, + remote_attn, + remote_block, + ) + assert lkind == rkind, ( + f"pair kind mismatch: local {lkind} vs remote {rkind} " + f"(local desc {li}, remote desc {ri})" + ) + assert lreg == rreg, ( + f"region mismatch: local {lreg} vs remote {rreg} for " + f"tokens {ltok} vs {rtok}" + ) + if lkind == "attn": + assert ltok == rtok, ( + f"TOKEN MISALIGNMENT: local sub-block holds tokens " + f"[{ltok}..) but receives remote tokens [{rtok}..) " + f"(geometry local_block={local_block}, " + f"remote_block={remote_block}, N={num_tokens})" + ) + covered_tokens.add(ltok) + + # Invariant 3: full coverage of the matched tokens, at the finest + # transfer granularity (the remote kernel block). + needed = {t for t in range(0, matched - matched % remote_kernel, remote_kernel)} + missing = needed - covered_tokens + assert not missing, ( + f"tokens never transferred: {sorted(missing)[:8]} " + f"(geometry local_block={local_block}, remote_block={remote_block}, " + f"N={num_tokens}, matched={matched})" + ) + + # Invariant 4: no stale bytes after receive completion. The scheduler + # excludes the blocks covering the matched tokens from alloc-time KV + # zeroing (the zeroing would race the RDMA write), so every byte of + # those blocks must be either written by the transfer or zeroed by the + # receive post-process. Stale bytes surface as mid-response garbage + # once decode grows into the untransferred tail. + for op, lh, lids, rh, rids in nixl.xfers: + larr = nixl.dlists[lh] + for li in lids: + addr, length, _ = (int(x) for x in larr[int(li)]) + for t in worker._test_tensors: + off = addr - t.data_ptr() + if 0 <= off < t.numel(): + t[off : off + length] = 0 # simulate the RDMA write + break + done_sending, done_recving = worker.get_finished() + assert "req-b" in done_recving + n_excluded = -(-matched // local_block) + stale = [] + for b in local_attn[:n_excluded]: + for region, t in enumerate(worker._test_tensors): + page = t[b * local_unified : (b + 1) * local_unified] + n_stale = int((page == 0xAA).sum()) + if n_stale: + stale.append((region, b, n_stale)) + assert not stale, ( + f"stale (unzeroed, untransferred) bytes in matched-range attention " + f"blocks (region, block, bytes): {stale} " + f"(geometry local_block={local_block}, remote_block={remote_block}, " + f"N={num_tokens}, matched={matched})" + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "local_block,remote_block", + [ + (12, 8), # ppl 3 vs 2 + (36, 8), # ppl 9 vs 2 (large ppl asymmetry, scaled) + (24, 4), # ppl 6 vs 1 + (16, 24), # remote larger than local (D_TP > P_TP direction) + ], +) +@pytest.mark.parametrize("num_tokens", list(range(2, 40))) +def test_hetero_ppl_token_alignment_sweep(local_block, remote_block, num_tokens): + """Sweep prompt lengths across block-boundary residues for several + hetero-ppl geometries; assert neighbor-safety, token alignment, and + coverage of every transferred kernel block.""" + _run_hetero_case( + local_block, kernel=4, remote_block=remote_block, num_tokens=num_tokens + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "num_tokens", + # Residues around the remote kernel block (4), the local kernel block + # (8), the remote logical block (8) and the local logical block (24). + [2, 5, 8, 9, 13, 16, 17, 21, 24, 25, 29, 32, 33, 41, 48, 49], +) +def test_hetero_ppl_with_block_size_ratio(num_tokens): + """Both hetero regimes at once: kernel blocks differ (local 8 / remote + 4, block_size_ratio=2) *and* physical_blocks_per_logical differs (3 vs + 2). The transfer is clipped at remote sub-block granularity by the + pairing and front-trimmed by _apply_prefix_caching, so the + untransferred tail can span both a partial block and whole blocks — + the case each of the two former zeroing paths handled only half of.""" + _run_hetero_case( + local_block=24, + kernel=8, + remote_block=8, + remote_kernel=4, + num_tokens=num_tokens, + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "num_tokens", + # Residues around every geometric boundary: kernel block (64), remote + # logical block (768), local logical block (5760), plus odd offsets. + [ + 2, + 63, + 64, + 65, + 127, + 128, + 300, + 640, + 767, + 768, + 769, + 831, + 832, + 1000, + 1535, + 1536, + 1537, + 2303, + 2304, + 2305, + 3001, + 5759, + 5760, + 5761, + 5824, + 6528, + 6529, + ], +) +def test_mla_hybrid_large_ppl_geometry(num_tokens): + """KimiLinear-scale MLA-hybrid geometry (TP8 prefill -> TP1 decode): + decode (local) logical block 5760 / kernel 64 (ppl=90), prefill + (remote) logical block 768 (ppl=12), tp_ratio=-8 multi-read with + replicated MLA and 8-way TP-sharded KDA state.""" + _run_hetero_case( + local_block=5760, + kernel=64, + remote_block=768, + num_tokens=num_tokens, + tp_size=8, + ) diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 7c735098230..d72bbc9ec2d 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -161,3 +161,57 @@ class TestMambaPlanSplitHandles: # FA: chunk=200//1=200, slot=0 (skip_fa) → (1000, 200, 0), (2000, 200, 0) # SSM: chunk=400//2=200, idx=1 → (3200, 200, 0) assert splits[1] == [(1000, 200, 0), (2000, 200, 0), (3200, 200, 0)] + + def test_hetero_block_size_splits(self): + """With a block-size ratio, single-source FA sub-block descs pass + through whole; SSM descs are unexpanded and split per source.""" + plan = TPMapping( + source_ranks_per_group=((0,), (0, 1)), + all_source_ranks=(0, 1), + rank_to_attention_slot={0: 0, 1: 0}, + rank_offset_factor=0, + ) + + worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec)) + # 2 FA blocks x ratio 2 sub-blocks + 1 SSM desc (never expanded). + src_blocks_data = np.array( + [ + (1000, 100, 0), + (1100, 100, 0), + (2000, 100, 0), + (2100, 100, 0), + (3000, 400, 0), + ], + dtype=np.uint64, + ) + + splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 4, 2)) + + assert len(splits) == 2 + fa_passthrough = [ + (1000, 100, 0), + (1100, 100, 0), + (2000, 100, 0), + (2100, 100, 0), + ] + assert splits[0] == fa_passthrough + [(3000, 200, 0)] + assert splits[1] == fa_passthrough + [(3200, 200, 0)] + + def test_hetero_block_size_head_sharded_asserts(self): + """Head-sharded FA reads (multiple FA sources) are incompatible with + a block-size mismatch and must fail loudly.""" + plan = TPMapping( + source_ranks_per_group=((0, 1), (0, 1)), + all_source_ranks=(0, 1), + rank_to_attention_slot={0: 0, 1: 1}, + rank_offset_factor=0, + ) + + worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec)) + src_blocks_data = np.array( + [(1000, 100, 0), (1100, 100, 0), (3000, 400, 0)], + dtype=np.uint64, + ) + + with pytest.raises(AssertionError, match="Head-sharded"): + list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2, 2)) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 9063308ce2e..480ac8937df 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -12,6 +12,7 @@ import uuid from collections import defaultdict from collections.abc import Iterator from concurrent.futures import Future, ThreadPoolExecutor +from functools import cached_property from typing import TYPE_CHECKING, Any, cast import msgspec @@ -67,6 +68,7 @@ from vllm.distributed.parallel_state import ( from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.network_utils import make_zmq_path +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -122,9 +124,11 @@ class NixlBaseConnectorWorker: return (region_ids * num_blocks + block_arr).flatten() # Compute desc ids per group using the right stride: FA descs have - # num_blocks entries per region (kernel granularity), SSM descs have - # logical_blocks entries per region (no kernel splitting). - logical_blocks = num_blocks // physical_blocks_per_logical + # num_blocks entries per region (kernel granularity, expanded by + # block_size_ratio for heterogeneous block sizes), SSM descs have + # logical_blocks entries per region (no kernel splitting, and never + # ratio-expanded since state blocks are indivisible). + logical_blocks = dst_num_blocks // physical_blocks_per_logical all_descs: list[np.ndarray] = [] for i, group in enumerate(block_ids): group_arr = np.asarray(group) @@ -161,6 +165,7 @@ class NixlBaseConnectorWorker: plan: TPMapping, src_blocks_data: np.ndarray, num_fa_descs: int, + block_size_ratio: int = 1, ) -> Iterator[list[tuple[int, int, int]]]: """Build split handle data for P_TP > D_TP scenario. @@ -187,6 +192,11 @@ class NixlBaseConnectorWorker: # Per-FA-descriptor replicate flag, in _build_fa_local emission order. fa_desc_replicated = self._fa_desc_replicated(num_fa_descs) + + assert block_size_ratio == 1 or fa_num_splits == 1 or all(fa_desc_replicated), ( + "Head-sharded attention reads with P_TP > D_TP and heterogeneous " + "block sizes are not supported" + ) src_blocks_list = src_blocks_data.tolist() for p_idx, p_rank in enumerate(plan.all_source_ranks): @@ -442,9 +452,12 @@ class NixlBaseConnectorWorker: # nixl_prepped_dlist_handle. self.src_xfer_handles_by_block_size: dict[int, int] = {} + # Local descriptor arrays per remote block size (block_size_ratio>1), + # kept for building per-tp-ratio splits at the same granularity. + self.src_blocks_data_by_block_size: dict[int, np.ndarray] = {} # Populated dynamically during handshake based on remote configuration. - # Keep track of regions at different tp_ratio values. tp_ratio->handles - self.src_xfer_handles_by_tp_ratio: dict[int, list[int]] = {} + # Per-source split handles, keyed by (tp_ratio, remote_block_size). + self.src_xfer_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) @@ -1258,11 +1271,7 @@ class NixlBaseConnectorWorker: agent_metadata_bytes=encoder.encode(agent_metadata), ) - def _build_mamba_local( - self, - base_addresses: list[int], - block_size_ratio: int, - ) -> np.ndarray: + def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray: """Build desc regions (conv sub-projections + ssm) per layer for local mamba blocks with DS conv layout, as an Nx3 uint64 array. @@ -1289,16 +1298,17 @@ class NixlBaseConnectorWorker: | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 | +-------------------+ +--------------------+ |1st_split-2nd_split| |1st_split-2nd_split | + + Mamba state blocks are indivisible (not token-extent data), so the + descriptors always use the local page geometry regardless of any + attention block-size ratio; their desc ids are likewise never + ratio-expanded (see _compute_desc_ids). """ - assert block_size_ratio == 1, ( - "Mamba 3-read transfer with block_size_ratio != 1 is not tested. " - f"Got block_size_ratio={block_size_ratio}." - ) assert base_addresses, "Local KV cache base addresses must not be empty." assert self._conv_decomp is not None conv_offsets = self._conv_decomp.local_conv_offsets conv_size, ssm_size = self._mamba_ssm_size - num_blocks = self._logical_num_blocks * block_size_ratio + num_blocks = self._logical_num_blocks physical_per_logical = self._physical_blocks_per_logical_kv_block device_id = self.device_id block_arange = np.arange(num_blocks, dtype=np.uint64) @@ -1307,9 +1317,7 @@ class NixlBaseConnectorWorker: for i, base_addr in enumerate(base_addresses): # Jump one page_size, but ssm page_size may be bigger when kernel # locks block size to a specific value (physical_per_logical scale). - page_stride = ( - self.block_len_per_layer[i] // block_size_ratio * physical_per_logical - ) + page_stride = self.block_len_per_layer[i] * physical_per_logical blk_addrs = base_addr + block_arange * page_stride for off, sz in conv_offsets: parts.append(self._stack_descs(blk_addrs + off, sz, device_id)) @@ -1459,7 +1467,7 @@ class NixlBaseConnectorWorker: self.device_id, ) if self._has_mamba: - assert self.num_descs == len(blocks_data) + assert self.num_descs * block_size_ratio == len(blocks_data) # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split # is unnecessary — a single conv desc per block suffices. Consider # adding a fast path that falls back to the standard 2-region @@ -1467,7 +1475,7 @@ class NixlBaseConnectorWorker: # remote has been seen. Currently we always register 4 regions # because local descs are created before knowing the remote TP. logger.debug("Registering local Mamba descriptors (4 regions/layer)") - mamba = self._build_mamba_local(local_base_addresses, block_size_ratio) + 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) @@ -1607,29 +1615,44 @@ class NixlBaseConnectorWorker: plan = self.tp_mappings[engine_id] + ### (Optional) Register a local handler at the remote engine's block + ### granularity (remote/prefill blocks smaller than local). + remote_block_size = nixl_agent_meta.block_size + src_blocks_data = self.src_blocks_data + if block_size_ratio > 1: + if remote_block_size not in self.src_xfer_handles_by_block_size: + handle, blocks_data = self.register_local_xfer_handler( + remote_block_size + ) + self.src_xfer_handles_by_block_size[remote_block_size] = handle + self.src_blocks_data_by_block_size[remote_block_size] = blocks_data + src_blocks_data = self.src_blocks_data_by_block_size[remote_block_size] + ### (Optional) Register local agent memory regions. MLA is not split. + split_key = (tp_ratio, remote_block_size) if ( tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1) - and tp_ratio not in self.src_xfer_handles_by_tp_ratio + and split_key not in self.src_xfer_handles_by_tp_ratio ): # Remote tp_size > local tp_size: read from multiple remote ranks. # Logically "split" own regions into per-source chunks. Hybrid # MLA+SSM also needs this path: MLA is replicated and read once, # while the SSM state is sharded across every remote TP rank. - # We only do this once per remote tp_size (replica-friendly). - self.src_xfer_handles_by_tp_ratio[tp_ratio] = [] + # We only do this once per remote (tp_size, block_size). + self.src_xfer_handles_by_tp_ratio[split_key] = [] for handle_data in self._build_local_splits_from_plan( plan, - self.src_blocks_data, - self.num_descs, + src_blocks_data, + self.num_descs * block_size_ratio, + block_size_ratio, ): descs = self.nixl_wrapper.get_xfer_descs( handle_data, self.nixl_memory_type ) handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs) - self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle) + self.src_xfer_handles_by_tp_ratio[split_key].append(handle) ### Register remote agent memory regions # With homogeneous TP, D pulls the whole kv cache from corresponding rank. With @@ -1665,13 +1688,6 @@ class NixlBaseConnectorWorker: self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs) ) - if block_size_ratio > 1: - # when prefill with smaller block_size, we need to init a - # new handler with same block_len to match - self.src_xfer_handles_by_block_size[nixl_agent_meta.block_size] = ( - self.register_local_xfer_handler(nixl_agent_meta.block_size)[0] - ) - return remote_agent_name def _validate_remote_agent_handshake( @@ -1716,9 +1732,13 @@ class NixlBaseConnectorWorker: "Disable prefix caching with --no-enable-prefix-caching." ) - if self._is_hma_required: - assert block_size_ratio == 1, ( - "HMA does not support different remote block size yet" + if block_size_ratio != 1: + # Heterogeneous block sizes transfer at remote-block granularity; + # the untransferred tail of the last local attention block is + # zeroed in the receive post-process, and mamba state pages + # transfer 1:1 (never sub-split). + assert not self.use_host_buffer, ( + "Heterogeneous block sizes are not supported with host buffer" ) kv_cache_layout = ( self.kv_cache_layout @@ -1875,27 +1895,57 @@ class NixlBaseConnectorWorker: "d2h", ) + @cached_property + def _attention_kv_caches(self) -> list[torch.Tensor]: + """Device KV caches of attention layers (mamba states excluded), + as consumed by the receive post-process.""" + assert self.device_kv_caches, ( + "_attention_kv_caches accessed before register_kv_caches" + ) + mamba_layers = { + name + for g, group in enumerate(self.kv_cache_config.kv_cache_groups) + if _is_ssm_spec(self._group_spec_types[g]) + for name in group.layer_names + } + kv_caches = self.device_kv_caches + return [cache for name, cache in kv_caches.items() if name not in mamba_layers] + def post_process_device_kv_on_receive( self, block_size_ratio: int, - block_ids_list: list[list[int]], + block_ids_list: list[tuple[list[int], int]], + convert: bool = True, ): """ Post process device kv cache after receiving from remote. - 3 types of post processing supported: + 3 types of conversion supported (``convert``): * kv_cache_postprocess_layout => convert from HND to NHD * kv_cache_postprocess_blksize => convert from small block size to large block size * kv_cache_postprocess_blksize_and_layout => convert from small block size to large block size and convert from HND to NHD + The transfer only covers ``covered_sub_blocks`` remote-sized + sub-blocks of each request's local attention blocks; the rest was + clipped, either by remote-block pairing (block-size ratio) or by the + hetero-ppl front trim in ``_apply_prefix_caching``. Those blocks were + excluded from the scheduler's alloc-time KV zeroing (which would race + the RDMA write), so everything past the covered range is zeroed here. + Stale bytes would otherwise surface as garbage or NaNs once decode + grows into the untransferred tail. """ if len(self.device_kv_caches) == 0: return assert block_size_ratio >= 1, "Only nP < nD supported currently." assert self.transfer_topo is not None - if self.enable_permute_local_kv and block_size_ratio > 1: + if not convert: + logger.debug( + "Post-processing device kv cache on receive by zeroing " + "untransferred blocks." + ) + elif self.enable_permute_local_kv and block_size_ratio > 1: logger.debug( "Post-processing device kv cache on receive by converting " "block_size with %sx bigger and permuting layout from HND" @@ -1914,18 +1964,45 @@ class NixlBaseConnectorWorker: block_size_ratio, ) - for block_ids in block_ids_list: - indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) + attn_caches = self._attention_kv_caches + device = attn_caches[0].device + for block_ids, covered_sub_blocks in block_ids_list: + # Blocks the transfer didn't write: the token tail of the last + # partially covered block, then everything beyond it. + covered_blocks, sub_blocks_in_last = divmod( + covered_sub_blocks, block_size_ratio + ) + first_stale = covered_blocks + (1 if sub_blocks_in_last else 0) + has_stale = first_stale < len(block_ids) + indices = None + if convert or has_stale: + indices = async_tensor_h2d(block_ids, device, torch.long) - for cache in self.device_kv_caches.values(): - if self.enable_permute_local_kv and block_size_ratio > 1: - kv_postprocess_blksize_and_layout_on_receive( - cache, indices, block_size_ratio - ) - elif self.enable_permute_local_kv: - kv_postprocess_layout_on_receive(cache, indices) - else: - kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio) + if convert: + for cache in attn_caches: + if self.enable_permute_local_kv and block_size_ratio > 1: + kv_postprocess_blksize_and_layout_on_receive( + cache, indices, block_size_ratio + ) + elif self.enable_permute_local_kv: + kv_postprocess_layout_on_receive(cache, indices) + else: + kv_postprocess_blksize_on_receive( + cache, indices, block_size_ratio + ) + + if sub_blocks_in_last: + last_block_id = block_ids[covered_blocks] + for cache in attn_caches: + # Both post-processed layouts leave tokens on dim 1. + sub_block_tokens = cache.shape[1] // block_size_ratio + zero_from = sub_blocks_in_last * sub_block_tokens + cache[last_block_id, zero_from:].zero_() + if has_stale: + assert indices is not None + stale_ids = indices[first_stale:] + for cache in attn_caches: + cache.index_fill_(0, stale_ids, 0) def post_process_device_kv_on_receive_heterogeneous_attn( self, block_ids: list[int] @@ -1995,18 +2072,33 @@ class NixlBaseConnectorWorker: if self.use_host_buffer: self.sync_recved_kv_to_device(req_id, meta) - # post processing for heteroblocksize + # Post processing for heteroblocksize/layout, and for blocks the + # transfer clipped. The latter happens either at remote-block + # granularity (block_size_ratio > 1) or at kernel-block + # granularity, when equal kernel pages meet differing logical + # block sizes and _apply_prefix_caching front-trims to the + # minimum count (hybrid heterogeneous TP). remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id) block_size_ratio = self.transfer_topo.block_size_ratio( remote_info.remote_block_size ) - if not self.use_mla and ( - block_size_ratio > 1 or self.enable_permute_local_kv - ): - assert not self._is_hma_required - block_ids_for_blocksize_post_process[block_size_ratio].append( - meta.local_physical_block_ids[0] - ) + hetero_ppl = ( + remote_info.remote_physical_blocks_per_logical + != self._physical_blocks_per_logical_kv_block + ) + if block_size_ratio > 1 or self.enable_permute_local_kv or hetero_ppl: + for g, local_group in enumerate(meta.local_physical_block_ids): + if not local_group or _is_ssm_spec(self._group_spec_types[g]): + continue + # Number of remote-sized sub-blocks the transfer covered; + # everything past this was clipped and must be zeroed. + covered_sub_blocks = min( + len(local_group) * block_size_ratio, + len(meta.remote.block_ids[g]), + ) + block_ids_for_blocksize_post_process[block_size_ratio].append( + (local_group, covered_sub_blocks) + ) # post processing for heterogeneous attention if self.enable_heterogeneous_attn_post_process: block_ids_for_heterogeneous_attn_post_process.append( @@ -2016,7 +2108,14 @@ class NixlBaseConnectorWorker: block_size_ratio, block_ids_list, ) in block_ids_for_blocksize_post_process.items(): - self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list) + # MLA never needs the block-size/layout conversion, but its + # clipped blocks still need zeroing. + convert = not self.use_mla and ( + block_size_ratio > 1 or self.enable_permute_local_kv + ) + self.post_process_device_kv_on_receive( + block_size_ratio, block_ids_list, convert + ) for block_ids in block_ids_for_heterogeneous_attn_post_process: self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids) @@ -2206,6 +2305,45 @@ class NixlBaseConnectorWorker: return mapped_2d.flatten().astype(np.int64) + def _map_block_ids_for_block_size_ratio( + self, + local_block_ids: BlockIds, + remote_block_ids: BlockIds, + block_size_ratio: int, + ) -> tuple[BlockIds, BlockIds]: + """Map attention-group block ids to remote-block granularity. + + Each local attention block is split into ``block_size_ratio`` + sub-blocks paired 1:1 with remote blocks. Sub-blocks beyond the + remote list — the untransferred tail of the last local block — are + clipped here and zeroed in the receive post-process. Mamba state + blocks are indivisible and transfer 1:1, unexpanded. + + ex: remote (prefill) block ids with block_size 4: + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + Local (decode) block ids with block_size 16: [1, 2, 3] expand to + [4, 5, ..., 15], then clip to the first 10 to pair 1:1 with remote. + """ + mapped_local: list[list[int]] = [] + mapped_remote: list[list[int]] = [] + for i, remote_group in enumerate(remote_block_ids): + local_group = local_block_ids[i] if local_block_ids else [] + if _is_ssm_spec(self._group_spec_types[i]): + mapped_local.append(list(local_group)) + mapped_remote.append(list(remote_group)) + continue + mapped = self.get_mapped_blocks( + np.asarray(local_group), block_size_ratio + ).tolist() + if len(mapped) > len(remote_group): + mapped = mapped[: len(remote_group)] + mapped_local.append(mapped) + mapped_remote.append(list(remote_group)) + if not any(mapped_local): + # Full prefix cache hit is indicated with an empty list. + return [], mapped_remote + return mapped_local, mapped_remote + def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> BlockIds: """ Convert block ids to kernel physical block ids. @@ -2300,13 +2438,18 @@ class NixlBaseConnectorWorker: remote_block_ids[i] = remote_group[-num_local_blocks:] else: # TODO Handle prefix caching with different block_sizes - max_padding = max( - self._physical_blocks_per_logical_kv_block, - remote_physical_per_logical, + # Allocation rounding legitimately leaves up to + # ppl - 1 trailing dead kernel blocks per side (plus one + # extra local block for the recomputed final token), so + # the counts may differ by up to the sum of the two + # ratios; anything larger indicates mismatched lists. + max_padding = ( + self._physical_blocks_per_logical_kv_block + + remote_physical_per_logical ) - assert abs(num_local_blocks - num_remote_blocks) < max_padding, ( + assert abs(num_local_blocks - num_remote_blocks) <= max_padding, ( f"Group {i}: |{num_local_blocks} - " - f"{num_remote_blocks}| >= {max_padding}" + f"{num_remote_blocks}| > {max_padding}" ) num_blocks = min(num_local_blocks, num_remote_blocks) local_block_ids[i] = local_block_ids[i][:num_blocks] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 40d6851769c..0c66596c4b7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -5,8 +5,6 @@ import time from typing import TYPE_CHECKING -import numpy as np - from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( NixlBaseConnectorWorker, ) @@ -178,10 +176,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): ) # Get side handles. if tp_ratio < 0 and (not self.use_mla or len(read_specs) > 1): - assert remote_block_size == self.block_size # Remote tp_size > local tp_size: we must perform multiple # reads. Get the memory chunk onto which we will write to. - local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i] + split_key = (tp_ratio, remote_block_size) + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i] else: # Single read from remote, we write to the whole memory region. # Also handle remote block size different from local block size. @@ -235,30 +233,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): remote_info.remote_block_size ) if block_size_ratio > 1: - # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups. - assert not self._is_hma_required - local_block_ids0 = local_block_ids[0] if local_block_ids else [] - remote_block_ids0 = remote_block_ids[0] - local_block_ids_mapped = self.get_mapped_blocks( - np.asarray(local_block_ids0), block_size_ratio - ).tolist() - if len(local_block_ids_mapped) > len(remote_block_ids0): - # NOTE: - # get_mapped_blocks will always expand block_ids for n times. - # ex: - # prefill block_ids with block_size as 4: - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - # Local decode block_ids with block_size as 16: [1, 2, 3] - # expanded decode block_ids with get_mapped_blocks from [1, 2, 3] to - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] - # Then we clip local to align with prefill - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - local_block_ids_mapped = local_block_ids_mapped[ - : len(remote_block_ids0) - ] - local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else [] - remote_block_ids = [remote_block_ids0] + local_block_ids, remote_block_ids = ( + self._map_block_ids_for_block_size_ratio( + local_block_ids, remote_block_ids, block_size_ratio + ) + ) # 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). # after we detect the txn is complete (which means we cannot make the diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 8cb710e7e58..6bc48963c58 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -39,7 +39,6 @@ from concurrent.futures import Future from typing import TYPE_CHECKING, Any import msgspec -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 ( @@ -553,8 +552,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): req_id, ) if tp_ratio < 0 and not self.use_mla: - assert remote_block_size == self.block_size - local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i] + split_key = (tp_ratio, remote_block_size) + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i] else: local_xfer_side_handle = self.src_xfer_handles_by_block_size[ remote_block_size @@ -606,18 +605,11 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): remote_info.remote_block_size ) if block_size_ratio > 1: - assert not self._is_hma_required - local_block_ids0 = local_block_ids[0] if local_block_ids else [] - remote_block_ids0 = remote_block_ids[0] - local_block_ids_mapped = self.get_mapped_blocks( - np.asarray(local_block_ids0), block_size_ratio - ).tolist() - if len(local_block_ids_mapped) > len(remote_block_ids0): - local_block_ids_mapped = local_block_ids_mapped[ - : len(remote_block_ids0) - ] - local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else [] - remote_block_ids = [remote_block_ids0] + local_block_ids, remote_block_ids = ( + self._map_block_ids_for_block_size_ratio( + local_block_ids, remote_block_ids, block_size_ratio + ) + ) notif_id = f"{remote_request_id}:{self.world_size}".encode()