From df8e63f4edabc06e51c10fe479fc9b95c303542f Mon Sep 17 00:00:00 2001 From: zhanqiuhu <49648934+ZhanqiuHu@users.noreply.github.com> Date: Wed, 6 May 2026 09:16:25 -0400 Subject: [PATCH] nixl refactor: new transfer design (#40731) Signed-off-by: ZhanqiuHu Signed-off-by: NickLucche Co-authored-by: NickLucche --- .../kv_connector/unit/test_nixl_connector.py | 5 + .../unit/test_nixl_connector_hma.py | 176 ++-- tests/v1/kv_connector/unit/test_tp_mapping.py | 141 ++++ .../kv_transfer/kv_connector/utils.py | 381 +-------- .../kv_connector/v1/nixl/metadata.py | 4 +- .../kv_connector/v1/nixl/tp_mapping.py | 142 ++++ .../kv_connector/v1/nixl/worker.py | 768 ++++++++---------- .../v1/ssm_conv_transfer_utils.py | 13 +- 8 files changed, 753 insertions(+), 877 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_tp_mapping.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index fb4b641e13..3803e4fd38 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -527,6 +527,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): block_size=self.block_size, ssm_sizes=(0, 0), attn_backend_name=self.backend_name, + physical_blocks_per_logical_kv_block=1, ), remote_tp_rank=remote_tp_rank, remote_tp_size=remote_tp_size, @@ -726,6 +727,7 @@ class TestNixlHandshake: worker.num_blocks = 1 worker.dst_num_blocks[worker.engine_id] = worker.num_blocks worker.src_blocks_data = [(0, worker.block_len_per_layer[0], worker.tp_rank)] + worker.num_descs = len(worker.src_blocks_data) def check_handshake(remote_tp_size: int): tp_ratio = remote_tp_size // local_tp_size @@ -978,6 +980,7 @@ class TestNixlHandshake: block_size=worker.block_size, ssm_sizes=(0, 0), attn_backend_name=worker.backend_name, + physical_blocks_per_logical_kv_block=1, ) with pytest.raises(RuntimeError): @@ -1035,6 +1038,7 @@ class TestNixlHandshake: block_size=worker.block_size, ssm_sizes=(0, 0), attn_backend_name=worker.backend_name, + physical_blocks_per_logical_kv_block=1, ) # We don't check layout for homogeneous TP and MLA for now, as the @@ -2354,6 +2358,7 @@ def test_compatibility_hash_validation( block_size=prefill_block_size, ssm_sizes=(0, 0), attn_backend_name=decode_worker.backend_name, + physical_blocks_per_logical_kv_block=1, ) handshake_payload = NixlHandshakePayload( compatibility_hash=remote_hash, 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 3f5a9b9cc0..edc591e060 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -93,75 +93,82 @@ def test_logical_to_kernel_block_ids_with_hma(): @pytest.mark.cpu_test @pytest.mark.parametrize( - "has_mamba,swa_enabled,mamba_enabled,remote_ratio," - "remote_block_ids,expected_remote_block_ids", + "group_spec_types,expansion_stride,remote_block_ids,expected_remote_block_ids", [ - # Non-mamba (FA+SWA): both groups expanded via _logical_to_kernel_block_ids. - # Regression for https://github.com/vllm-project/vllm/pull/39724 - ( - False, - True, - False, - 1, + pytest.param( + ("FullAttentionSpec", "SlidingWindowSpec"), + 2, ([0, 1, 2], [3, 4]), [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]], + id="dense_fa_swa", ), - # Mamba (FA+Mamba): FA expanded via _logical_to_remote_kernel_block_ids, - # Mamba passed through unchanged. - # remote_ratio=261 (Nemotron 30B TP=1) != local_ratio=2 so that using - # the wrong conversion method produces different FA results. - ( - True, - False, - True, + pytest.param( + ("FullAttentionSpec", "MambaSpec"), 261, ([0, 1, 2], [10, 11]), [[0, 1, 261, 262, 522, 523], [10, 11]], + id="mamba_fa_ssm", ), ], - ids=["non_mamba_fa_swa", "mamba_fa_ssm"], ) def test_read_blocks_for_req_expands_remote_ids( - has_mamba, - swa_enabled, - mamba_enabled, - remote_ratio, + group_spec_types, + expansion_stride, remote_block_ids, expected_remote_block_ids, ): """_read_blocks_for_req must expand remote logical block IDs to kernel block IDs when kernel block size != logical block size. - Non-mamba path uses _logical_to_kernel_block_ids (all groups expanded). - Mamba path uses _logical_to_remote_kernel_block_ids (FA expanded, Mamba - passed through). + The hot path always calls _logical_to_remote_kernel_block_ids with + remote_info.remote_physical_blocks_per_logical (model-agnostic). """ from unittest.mock import MagicMock from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( NixlConnectorMetadata, ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( + TPMapping, + ) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + MambaSpec, + SlidingWindowSpec, + ) + + spec_name_to_type = { + "FullAttentionSpec": FullAttentionSpec, + "SlidingWindowSpec": SlidingWindowSpec, + "MambaSpec": MambaSpec, + } + resolved_types = tuple(spec_name_to_type[n] for n in group_spec_types) worker = object.__new__(NixlConnectorWorker) - worker._has_mamba = has_mamba worker._physical_blocks_per_logical_kv_block = 2 + + has_mamba = any(t is MambaSpec for t in resolved_types) + has_swa = any(t is SlidingWindowSpec for t in resolved_types) worker.kv_cache_config = make_kv_cache_config( - block_size=16, swa_enabled=swa_enabled, mamba_enabled=mamba_enabled + block_size=16, swa_enabled=has_swa, mamba_enabled=has_mamba ) remote_engine_id = "remote-engine" - if has_mamba: - worker._physical_blocks_per_logical = {remote_engine_id: remote_ratio} - # Mock transfer_topo: empty remote ranks skips the transfer machinery - # entirely, isolating the block-ID expansion logic. worker.transfer_topo = MagicMock() - worker.transfer_topo.target_remote_ranks.return_value = [] - worker.transfer_topo.get_engine_info.return_value = MagicMock(remote_tp_size=1) worker.transfer_topo.tp_ratio.return_value = 1 + remote_info = MagicMock() + remote_info.remote_physical_blocks_per_logical = expansion_stride + worker.transfer_topo.get_engine_info.return_value = remote_info + worker.use_mla = False + + mock_plan = MagicMock(spec=TPMapping) + mock_plan.all_source_ranks = () + mock_plan.source_ranks_per_group = () + worker.tp_mappings = {remote_engine_id: mock_plan} metadata = NixlConnectorMetadata() metadata.add_new_req_to_recv( @@ -306,75 +313,82 @@ def test_nixl_metadata_hma_block_ids_structure(): assert list(req_meta.remote.block_ids[1]) == [18, 19, 20, 21] -@pytest.mark.cpu_test -def test_get_block_descs_ids_hybrid_ssm(): - """Test _get_block_descs_ids uses per-group strides for hybrid FA+SSM - when ratio=1 (no kernel block size mismatch).""" +def _make_mock_worker_for_desc_ids( + num_regions: int, + has_mamba: bool, + group_spec_types: tuple, + block_len_per_layer: list[int] | None = None, +): + """Build a mock NixlConnectorWorker with attrs needed by _compute_desc_ids.""" + from unittest.mock import MagicMock + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) - worker = object.__new__(NixlConnectorWorker) + worker = MagicMock(spec=NixlConnectorWorker) + worker.num_regions = num_regions + worker._has_mamba = has_mamba + worker._group_spec_types = group_spec_types + worker.block_len_per_layer = block_len_per_layer or [100] + worker._compute_desc_ids = NixlConnectorWorker._compute_desc_ids.__get__( + worker, NixlConnectorWorker + ) + return worker - num_blocks = 100 - engine_id = "test-engine" - worker.num_regions = 2 - worker.dst_num_blocks = {engine_id: num_blocks} - worker._has_mamba = True - worker._is_mamba_group = [False, True] - worker._physical_blocks_per_logical_kv_block = 1 - worker._physical_blocks_per_logical = {engine_id: 1} - worker.block_len_per_layer = [100] - # num_descs = num_regions * num_blocks (no blocks_first doubling) - worker.num_descs = 2 * num_blocks + +@pytest.mark.cpu_test +def test_get_block_descs_ids_hybrid_ssm(): + """Test _compute_desc_ids uses per-group strides for hybrid + FA+SSM when ratio=1 (no kernel block size mismatch).""" + 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], + ) fa_blocks = [3, 5] ssm_blocks = [1, 2] - result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks)) + result = worker._compute_desc_ids( + block_ids=(fa_blocks, ssm_blocks), + dst_num_blocks=100, + block_size_ratio=None, + physical_blocks_per_logical=1, + ) - # FA group: stride=num_blocks=100, offset=0 - # region0: [3, 5], region1: [103, 105] - # SSM group: stride=logical_blocks=100 (=num_blocks/ratio=100/1), - # offset=num_fa_descs=200, 4 regions per Mamba layer (x, B, C, ssm) - # region0: [201, 202], region1: [301, 302], - # region2: [401, 402], region3: [501, 502] expected = [3, 5, 103, 105, 201, 202, 301, 302, 401, 402, 501, 502] assert list(result) == expected, f"Expected {expected}, got {list(result)}" @pytest.mark.cpu_test def test_get_block_descs_ids_kernel_block_mismatch(): - """Test _get_block_descs_ids uses different strides for FA (kernel blocks) - vs SSM (logical blocks) when ratio > 1.""" - from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( - NixlConnectorWorker, - ) - - worker = object.__new__(NixlConnectorWorker) + """Test _compute_desc_ids uses different strides for FA + (kernel blocks) vs SSM (logical blocks) when ratio > 1.""" + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec ratio = 4 logical_blocks = 100 num_blocks = logical_blocks * ratio # 400 kernel blocks - engine_id = "test-engine" - worker.num_regions = 2 - worker.dst_num_blocks = {engine_id: num_blocks} - worker._has_mamba = True - worker._is_mamba_group = [False, True] - worker._physical_blocks_per_logical_kv_block = ratio - worker._physical_blocks_per_logical = {engine_id: ratio} - worker.block_len_per_layer = [100] - worker.num_descs = 2 * num_blocks # 800 - fa_blocks = [3, 7] # kernel-level block IDs - ssm_blocks = [1, 2] # logical block IDs - result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks)) + worker = _make_mock_worker_for_desc_ids( + num_regions=2, + has_mamba=True, + group_spec_types=(FullAttentionSpec, MambaSpec), + block_len_per_layer=[100], + ) + + fa_blocks = [3, 7] + ssm_blocks = [1, 2] + result = worker._compute_desc_ids( + block_ids=(fa_blocks, ssm_blocks), + dst_num_blocks=num_blocks, + block_size_ratio=None, + physical_blocks_per_logical=ratio, + ) - # FA group: stride=num_blocks=400, offset=0 - # region0: [3, 7], region1: [403, 407] - # SSM group: stride=logical_blocks=400//4=100, offset=num_fa_descs=800, - # 4 regions per Mamba layer (x, B, C, ssm) - # region0: [801, 802], region1: [901, 902], - # region2: [1001, 1002], region3: [1101, 1102] expected = [3, 7, 403, 407, 801, 802, 901, 902, 1001, 1002, 1101, 1102] assert list(result) == expected, f"Expected {expected}, got {list(result)}" diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py new file mode 100644 index 0000000000..e57244a31f --- /dev/null +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for TP mapping and transfer plan utilities. + +These tests verify that TP mapping produces correct outputs +(source ranks, split handles, desc IDs). +No GPU or NIXL required. +""" + +from __future__ import annotations + +import pytest + +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( + TPMapping, + compute_tp_mapping, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, +) +from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + +# ====================================================================== +# Test fixtures / helpers +# ====================================================================== + + +def _compute_mapping( + tp_rank: int = 0, + tp_size: int = 1, + remote_tp_size: int = 1, + is_mla: bool = False, + num_kv_heads: int = 8, + group_spec_types: tuple[type, ...] = (FullAttentionSpec,), +) -> TPMapping: + return compute_tp_mapping( + tp_rank=tp_rank, + tp_size=tp_size, + remote_tp_size=remote_tp_size, + is_mla=is_mla, + total_num_kv_heads=num_kv_heads, + group_spec_types=group_spec_types, + ) + + +# ====================================================================== +# TP mapping structure tests +# ====================================================================== + + +class TestTPMappingStructure: + def test_source_ranks_homogeneous(self): + m = _compute_mapping(tp_size=2, tp_rank=1, remote_tp_size=2) + assert m.all_source_ranks == (1,) + + def test_source_ranks_d_gt_p(self): + m = _compute_mapping(tp_size=4, tp_rank=2, remote_tp_size=2) + assert m.all_source_ranks == (1,) + + def test_source_ranks_p_gt_d(self): + m = _compute_mapping(tp_size=1, tp_rank=0, remote_tp_size=2) + assert m.all_source_ranks == (0, 1) + + +# ====================================================================== +# Split handle tests +# ====================================================================== + + +def _make_mock_worker_for_splits(group_spec_types): + """Build a mock NixlConnectorWorker with _group_spec_types for split tests.""" + worker = object.__new__(NixlConnectorWorker) + worker._group_spec_types = group_spec_types + return worker + + +class TestBuildSrcSplitHandles: + @pytest.mark.parametrize("remote_tp_size", [2, 4]) + def test_build_src_split_handles(self, remote_tp_size): + tp_rank = 0 + tp_size = 1 + + plan = _compute_mapping( + tp_rank=tp_rank, + tp_size=tp_size, + remote_tp_size=remote_tp_size, + ) + + worker = _make_mock_worker_for_splits((FullAttentionSpec,)) + src_blocks_data = [(0x2000 + i * 1024, 1024, 0) for i in range(8)] + num_descs = len(src_blocks_data) + splits = list( + worker._build_local_splits_from_plan( + plan, + src_blocks_data, + num_descs, + ) + ) + + assert len(splits) == remote_tp_size + for handle in splits: + assert len(handle) == len(src_blocks_data) + for _, length, _ in handle: + assert length == 1024 // remote_tp_size + + +class TestMambaPlanSplitHandles: + """Verify split handles for Mamba with FA/SSM distinction.""" + + def test_fa_and_ssm_different_split_factors(self): + """Section 0 split by num_attn_reads, section 1 by abs_tp.""" + fa_readers = (0,) + ssm_readers = (0, 1) + plan = TPMapping( + source_ranks_per_group=(fa_readers, ssm_readers), + 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 descs + 1 SSM desc + src_blocks_data = [ + (1000, 200, 0), # FA desc 0 + (2000, 200, 0), # FA desc 1 + (3000, 400, 0), # SSM desc 0 + ] + + splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2)) + + assert len(splits) == 2 # 2 source ranks + + # Rank 0 (FA source, p_idx=0): + # FA: chunk=200//1=200, slot=0 → (1000, 200, 0), (2000, 200, 0) + # SSM: chunk=400//2=200, idx=0 → (3000, 200, 0) + assert splits[0] == [(1000, 200, 0), (2000, 200, 0), (3000, 200, 0)] + + # Rank 1 (not FA source, p_idx=1): + # 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)] diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index 63b56eddfa..b85416ab30 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -386,38 +386,6 @@ class EngineTransferInfo: """Physical blocks per logical block.""" -@dataclass(frozen=True) -class MambaEngineTransferInfo(EngineTransferInfo): - """Extends ``EngineTransferInfo`` with Mamba-hybrid transfer geometry. - - For hybrid SSM+Attention models, FA and Mamba layers may require - different numbers of reads from different remote ranks. This - dataclass captures that per-engine transfer plan. - """ - - remote_fa_source_ranks: tuple[int, ...] - """Remote ranks carrying unique FA heads for this local rank.""" - - remote_all_source_ranks: tuple[int, ...] - """All remote ranks this local rank reads from (FA + Mamba).""" - - remote_num_fa_reads: int - """Number of distinct remote ranks needed for FA data.""" - - remote_num_mamba_reads: int - """Number of distinct remote ranks needed for Mamba data.""" - - remote_fa_descriptor_bytes: int - """Byte size of one FA K (or V) descriptor entry.""" - - is_remote_replicated: bool - """Whether the remote engine has replicated KV heads - (remote_tp_size > total_num_kv_heads).""" - - remote_physical_heads: int - """Physical KV heads stored per remote rank.""" - - # ---- Transfer topology ---- @@ -439,8 +407,6 @@ class TransferTopology: self.local_physical_heads = max(1, self.total_num_kv_heads // self.tp_size) self._engines: dict[EngineId, EngineTransferInfo] = {} - self._fa_source_sets: dict[EngineId, frozenset[int]] = {} - self._fa_source_indices: dict[EngineId, dict[int, int]] = {} # Figure out whether the first dimension of the cache is K/V # or num_blocks. @@ -487,24 +453,12 @@ class TransferTopology: def register_remote_engine( self, remote_engine_id: EngineId, - remote_tp_size: int, - remote_block_size: int, - remote_block_len: int, - remote_physical_blocks_per_logical: int, - *, - local_block_len: int = 0, + info: EngineTransferInfo, ) -> EngineTransferInfo: """Register a remote engine, unifying worker dicts state. - Only remote engines should be registered here — the local engine's - identity (tp_size, block_size, etc.) is set via ``__init__`` params. - - For Mamba models, also computes the Mamba transfer plan and - builds the FA source lookup caches. - - Args: - local_block_len: Local representative block_len (bytes). - Required for Mamba models to compute ``fa_descriptor_bytes``. + The caller (worker) is responsible for computing the info via + the transfer policy. This method only stores and deduplicates. """ assert remote_engine_id != self.engine_id, ( f"Cannot register local engine {self.engine_id} as remote. " @@ -512,29 +466,6 @@ class TransferTopology: ) if remote_engine_id in self._engines: return self._engines[remote_engine_id] - info: EngineTransferInfo - if self.is_mamba: - info = self._build_mamba_info( - remote_tp_size=remote_tp_size, - remote_block_size=remote_block_size, - remote_block_len=remote_block_len, - remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical), - local_block_len=local_block_len, - ) - assert isinstance(info, MambaEngineTransferInfo) - self._fa_source_sets[remote_engine_id] = frozenset( - info.remote_fa_source_ranks - ) - self._fa_source_indices[remote_engine_id] = { - r: i for i, r in enumerate(info.remote_fa_source_ranks) - } - else: - info = EngineTransferInfo( - remote_tp_size=remote_tp_size, - remote_block_len=remote_block_len, - remote_block_size=remote_block_size, - remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical), - ) self._engines[remote_engine_id] = info return info @@ -622,14 +553,8 @@ class TransferTopology: """Get the remote TP rank(s) that the current local TP rank will read from. When remote tp_size > local tp_size, reads from multiple remote ranks. - - For Mamba models, returns the precomputed ``all_source_ranks`` - (FA + Mamba union). """ info = self._engines[remote_engine_id] - if isinstance(info, MambaEngineTransferInfo): - return list(info.remote_all_source_ranks) - tp_ratio = self.tp_ratio(info.remote_tp_size) if tp_ratio > 0: return [self.tp_rank // tp_ratio] @@ -662,309 +587,15 @@ class TransferTopology: # Regular case: backends like FA register K/V in separate regions return cache if self.split_k_and_v else [cache] - # ============================================================ - # Mamba-specific methods - # ============================================================ - - def should_skip_fa(self, remote_engine_id: EngineId, remote_rank: int) -> bool: - """Whether to skip FA groups for this remote rank (mamba-only).""" - return remote_rank not in self._fa_source_sets[remote_engine_id] - - def fa_head_slot(self, remote_engine_id: EngineId, remote_rank: int) -> int: - """Index into local FA block for this remote rank's head data. - - For remote ranks in ``fa_source_ranks``, returns 0, 1, …, reads-1. - For ranks NOT in ``fa_source_ranks`` (replicated duplicates), - returns the slot of the matching source rank with the same head. - """ - fa_index = self._fa_source_indices[remote_engine_id] - if remote_rank in fa_index: - return fa_index[remote_rank] - mamba_info = self._engines[remote_engine_id] - assert isinstance(mamba_info, MambaEngineTransferInfo) - K = self.total_num_kv_heads - remote_tp = mamba_info.remote_tp_size - r_head = self._physical_head_range(remote_tp, K, remote_rank) - for target in mamba_info.remote_fa_source_ranks: - t_head = self._physical_head_range(remote_tp, K, target) - if self._range_overlap(r_head, t_head): - return fa_index[target] - return 0 - - def fa_rank_offset( - self, remote_engine_id: EngineId, remote_kv_block_len: int - ) -> int: - """Byte offset into remote FA block for this local rank. - - When local TP is replicated (local_tp > K), multiple local ranks - share a head. Computes offset *relative to the target remote - rank's first head* so it works regardless of how many heads the - remote has. Returns 0 when local does not index into remote. - """ - mamba_info = self._engines[remote_engine_id] - assert isinstance(mamba_info, MambaEngineTransferInfo) - tp_ratio = self.tp_ratio(mamba_info.remote_tp_size) - if self.is_mla or tp_ratio <= 0: - return 0 - K = self.total_num_kv_heads - is_local_replicated = self.tp_size > K - if is_local_replicated: - local_head = self.tp_rank * K // self.tp_size - p_rank = mamba_info.remote_fa_source_ranks[0] - p_start = p_rank * K // mamba_info.remote_tp_size - return (local_head - p_start) * remote_kv_block_len - return self.tp_rank % tp_ratio * remote_kv_block_len - - def needs_split_handles(self, remote_engine_id: EngineId) -> bool: - """Whether per-remote-rank split handles are needed. - - True when FA and mamba have different read counts, requiring - different splitting factors in the local handle. - """ - mamba_info = self._engines[remote_engine_id] - assert isinstance(mamba_info, MambaEngineTransferInfo) - tp_ratio = self.tp_ratio(mamba_info.remote_tp_size) - return ( - tp_ratio < 0 - and not self.is_mla - and len(mamba_info.remote_all_source_ranks) > 1 - ) - - def compute_split_handle_data( - self, - remote_engine_id: EngineId, - src_blocks_data: list[tuple[int, int, int]], - num_fa_descs: int, - abs_tp: int, - ) -> list[list[tuple[int, int, int]]]: - """Per-remote-rank (addr, len, dev) triples for Mamba-HMA split - handles. - - FA descriptors (indices < num_fa_descs) are sliced by - ``remote_num_fa_reads``; mamba descriptors are sliced uniformly - by ``abs_tp``. - """ - mamba_info = self._engines[remote_engine_id] - assert isinstance(mamba_info, MambaEngineTransferInfo) - all_handle_data: list[list[tuple[int, int, int]]] = [] - for p_idx, p_rank in enumerate(mamba_info.remote_all_source_ranks): - handle_data: list[tuple[int, int, int]] = [] - skip_fa = self.should_skip_fa(remote_engine_id, p_rank) - fa_slot = self.fa_head_slot(remote_engine_id, p_rank) if not skip_fa else 0 - for j, (addr, local_len, dev) in enumerate(src_blocks_data): - if j < num_fa_descs: - assert mamba_info.remote_num_fa_reads >= 1 - fa_chunk = local_len // mamba_info.remote_num_fa_reads - handle_data.append((addr + fa_slot * fa_chunk, fa_chunk, dev)) - else: - mamba_chunk = local_len // abs_tp - handle_data.append((addr + p_idx * mamba_chunk, mamba_chunk, dev)) - all_handle_data.append(handle_data) - return all_handle_data - - def filter_block_ids_for_rank( - self, - remote_engine_id: EngineId, - remote_rank: int, - local_ids: BlockIds, - remote_ids: BlockIds, - is_mamba_group: list[bool], - ) -> tuple[BlockIds, BlockIds]: - """Zero out FA groups for remote ranks outside ``fa_source_ranks``. - - Returns (filtered_local_ids, filtered_remote_ids). When the - remote rank carries FA data for this local rank, returns the - inputs unchanged. - """ - if not self.should_skip_fa(remote_engine_id, remote_rank): - return local_ids, remote_ids - num_groups = len(local_ids) - filtered_local: list[list[int]] = [ - [] if not is_mamba_group[g] else local_ids[g] for g in range(num_groups) - ] - filtered_remote: list[list[int]] = [ - [] if not is_mamba_group[g] else remote_ids[g] for g in range(num_groups) - ] - return filtered_local, filtered_remote - def describe(self, remote_engine_id: EngineId) -> str: """One-line summary of transfer config for logging.""" info = self._engines[remote_engine_id] - base = ( + return ( + f"TransferTopology(" f"tp_ratio={self.tp_ratio(info.remote_tp_size)}, " f"K={self.total_num_kv_heads}, " f"local_tp={self.tp_size}, " f"remote_tp={info.remote_tp_size}, " f"local_rank={self.tp_rank}, " - f"remote_block_len={info.remote_block_len}" - ) - if isinstance(info, MambaEngineTransferInfo): - return ( - f"TransferTopology.mamba({base}, " - f"fa_reads={info.remote_num_fa_reads}, " - f"mamba_reads={info.remote_num_mamba_reads}, " - f"fa_sources={list(info.remote_fa_source_ranks)}, " - f"all_sources={list(info.remote_all_source_ranks)}, " - f"fa_desc_bytes={info.remote_fa_descriptor_bytes})" - ) - return f"TransferTopology({base})" - - # ============================================================ - # Private helpers - # ============================================================ - # Mamba-HMA hetero-TP transfer config: - # With hetero-TP (P_TP > D_TP), FA KV cache may be replicated across - # P ranks (when P_TP > num_kv_heads), but Mamba conv/SSM state is - # almost always uniquely sharded per P rank. So the number of P - # ranks D must read from can differ between FA and Mamba, and they - # must be handled separately. - - @staticmethod - def _physical_head_range(tp_size: int, num_heads: int, rank: int) -> range: - """Physical KV head range stored in a rank's KV cache tensor. - - When ``tp_size <= num_heads``: sharded, K/TP contiguous heads per rank. - When ``tp_size > num_heads``: 1 physical head per rank. Heads are - distributed **contiguously** (matching vLLM's GQA weight partitioning): - consecutive ranks share a head before moving to the next one. - """ - if tp_size <= num_heads: - assert num_heads % tp_size == 0 - per_rank = num_heads // tp_size - return range(rank * per_rank, (rank + 1) * per_rank) - else: - h = rank * num_heads // tp_size - return range(h, h + 1) - - @staticmethod - def _range_overlap(a: range, b: range) -> range: - start = max(a.start, b.start) - stop = min(a.stop, b.stop) - return range(start, max(start, stop)) - - # ============================================================ - # Private: build Mamba transfer info - # ============================================================ - - def _build_mamba_info( - self, - remote_tp_size: int, - remote_block_size: int, - remote_block_len: int, - remote_physical_blocks_per_logical: int, - local_block_len: int, - ) -> MambaEngineTransferInfo: - """Compute Mamba transfer plan.""" - K = self.total_num_kv_heads - local_tp = self.tp_size - local_rank = self.tp_rank - - is_remote_replicated = remote_tp_size > K - remote_physical_heads = max(1, K // remote_tp_size) - - if local_tp >= remote_tp_size: - assert local_tp % remote_tp_size == 0 - tp_ratio = local_tp // remote_tp_size - else: - assert remote_tp_size % local_tp == 0 - tp_ratio = -(remote_tp_size // local_tp) - - abs_tp = -tp_ratio if tp_ratio < 0 else 1 - - mamba_range: range | None = None - if tp_ratio < 0: - mamba_range = range(local_rank * abs_tp, (local_rank + 1) * abs_tp) - - # ---- FA read targets ---- - if self.is_mla or tp_ratio >= 0: - num_fa_reads = 1 - fa_source_ranks: list[int] = ( - [0] - if self.is_mla - else [local_rank // tp_ratio if tp_ratio > 0 else local_rank] - ) - else: - local_needs = self._physical_head_range(local_tp, K, local_rank) - search_range = ( - mamba_range if mamba_range is not None else range(remote_tp_size) - ) - seen: set[tuple[int, int]] = set() - fa_source_ranks = [] - for p in search_range: - p_has = self._physical_head_range(remote_tp_size, K, p) - ov = self._range_overlap(local_needs, p_has) - if len(ov) > 0: - key = (ov.start, ov.stop) - if key not in seen: - seen.add(key) - fa_source_ranks.append(p) - if not fa_source_ranks: - for p in range(remote_tp_size): - p_has = self._physical_head_range(remote_tp_size, K, p) - ov = self._range_overlap(local_needs, p_has) - if len(ov) > 0: - key = (ov.start, ov.stop) - if key not in seen: - seen.add(key) - fa_source_ranks.append(p) - num_fa_reads = len(fa_source_ranks) - - # ---- All source ranks (mamba + FA) ---- - if mamba_range is not None and abs_tp > num_fa_reads: - num_mamba_reads = abs_tp - all_source_ranks = list(mamba_range) - else: - num_mamba_reads = num_fa_reads - all_source_ranks = list(fa_source_ranks) - - # ---- FA descriptor bytes ---- - effective_block_len = min(local_block_len, remote_block_len) - if self.is_kv_layout_blocks_first: - fa_descriptor_bytes = effective_block_len // 2 - else: - fa_descriptor_bytes = effective_block_len - - # ---- Validation ---- - is_local_replicated = local_tp > K - if is_local_replicated and is_remote_replicated and tp_ratio > 0: - logger.info( - "Both-replicated hetero-TP: local_tp=%d > remote_tp=%d > K=%d.", - local_tp, - remote_tp_size, - K, - ) - tt_set = set(all_source_ranks) - for t in fa_source_ranks: - if t not in tt_set: - logger.error( - "FA source rank %d NOT in all_source_ranks %s.", - t, - all_source_ranks, - ) - if self.is_kv_layout_blocks_first and tp_ratio < 0 and num_fa_reads > 0: - local_k_half = local_block_len // 2 - remote_k_half = remote_block_len // 2 - expected = local_k_half // num_fa_reads - if expected != remote_k_half: - logger.warning( - "FA size mismatch: local_k_half=%d / reads=%d = %d, " - "but remote_k_half=%d.", - local_k_half, - num_fa_reads, - expected, - remote_k_half, - ) - - return MambaEngineTransferInfo( - remote_tp_size=remote_tp_size, - remote_block_len=remote_block_len, - remote_block_size=remote_block_size, - remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical), - remote_fa_source_ranks=tuple(fa_source_ranks), - remote_all_source_ranks=tuple(all_source_ranks), - remote_num_fa_reads=num_fa_reads, - remote_num_mamba_reads=num_mamba_reads, - remote_fa_descriptor_bytes=fa_descriptor_bytes, - is_remote_replicated=is_remote_replicated, - remote_physical_heads=remote_physical_heads, + f"remote_block_len={info.remote_block_len})" ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index 71ebbf1174..c56e373ba9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -32,8 +32,9 @@ GET_META_MSG = b"get_meta_msg" # Version History: # 1: Initial version with compatibility checking # 2: Add remote_request_id to kv_transfer_params +# 3: Add physical_blocks_per_logical_kv_block to NixlAgentMetadata # -NIXL_CONNECTOR_VERSION: int = 2 +NIXL_CONNECTOR_VERSION: int = 3 @dataclass @@ -48,6 +49,7 @@ class NixlAgentMetadata: block_size: int ssm_sizes: tuple[int, int] attn_backend_name: str + physical_blocks_per_logical_kv_block: int @dataclass diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py new file mode 100644 index 0000000000..7115b8bed5 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TP mapping computation for NIXL KV cache transfers.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from vllm.distributed.kv_transfer.kv_connector.utils import ( + BlockIds, +) +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, MambaSpec + +# ====================================================================== +# Data structures +# ====================================================================== + + +@dataclass(frozen=True) +class ReadSpec: + """Specification for a single remote block read operation.""" + + remote_rank: int + local_block_ids: BlockIds + remote_block_ids: BlockIds + + +def _is_attention_spec(spec_type: type[KVCacheSpec]) -> bool: + return issubclass(spec_type, AttentionSpec) + + +def _is_ssm_spec(spec_type: type[KVCacheSpec]) -> bool: + return issubclass(spec_type, MambaSpec) + + +@dataclass(frozen=True) +class TPMapping: + """Complete local-to-remote TP mapping for one remote engine. + + Generated once per remote engine during handshake. + """ + + # Remote TP ranks that this local rank reads from, per group. + # Position = local piece index. + source_ranks_per_group: tuple[tuple[int, ...], ...] + + # Superset of all source ranks (union of all groups). + all_source_ranks: tuple[int, ...] + + # Maps each source rank to its FA head slot index. + rank_to_attention_slot: dict[int, int] + + # FA head offset factor for hetero-TP (D_TP > P_TP). + rank_offset_factor: int + + +# ====================================================================== +# TP mapping computation +# ====================================================================== + + +def compute_tp_mapping( + tp_rank: int, + tp_size: int, + remote_tp_size: int, + is_mla: bool, + total_num_kv_heads: int, + group_spec_types: tuple[type[KVCacheSpec], ...], +) -> TPMapping: + """Build the complete local-to-remote TP mapping. + + Computes source ranks, head slot assignments, and the rank offset + factor in a single pass. + """ + # --- Attention source ranks --- + if is_mla: + # All heads replicated across all ranks. + attn_ranks = [0] + elif tp_size >= remote_tp_size: + # D (local TP) > P (remote TP): multiple local ranks read different chunks from + # *one* remote rank, corresponding to different kv heads. + attn_ranks = [tp_rank * remote_tp_size // tp_size] + else: + # P (remote TP) > D (local TP): one local rank + # reads from multiple remote ranks. + # GQA dedup: when K < remote_tp_size, several remote ranks + # hold the same KV head. np.unique keeps only the first + # rank per unique head so we don't issue redundant reads. + abs_tp = remote_tp_size // tp_size + start = tp_rank * abs_tp + heads = np.arange(start, start + abs_tp) * total_num_kv_heads // remote_tp_size + _, unique_idx = np.unique(heads, return_index=True) + attn_ranks = (start + np.sort(unique_idx)).tolist() + + # --- SSM source ranks --- + has_ssm = any(_is_ssm_spec(t) for t in group_spec_types) + if has_ssm: + if tp_size < remote_tp_size: + abs_tp = remote_tp_size // tp_size + ssm_ranks = list(range(tp_rank * abs_tp, (tp_rank + 1) * abs_tp)) + else: + ssm_ranks = list(attn_ranks) + else: + ssm_ranks = [] + + all_ranks = sorted(set(attn_ranks) | set(ssm_ranks)) + + # --- Per-group ordered source ranks --- + source_ranks_per_group = tuple( + tuple(ssm_ranks) if _is_ssm_spec(t) else tuple(attn_ranks) + for t in group_spec_types + ) + + # --- Attention head slots --- + head_to_slot: dict[int, int] = {} + for i, r in enumerate(attn_ranks): + head_to_slot[r * total_num_kv_heads // remote_tp_size] = i + rank_to_attention_slot = { + r: head_to_slot.get(r * total_num_kv_heads // remote_tp_size, 0) + for r in all_ranks + } + + # --- Rank offset factor --- + if is_mla or tp_size <= remote_tp_size: + # We don't index into remote for reading, no offset needed. + rank_offset_factor = 0 + elif tp_size > total_num_kv_heads: + local_head = tp_rank * total_num_kv_heads // tp_size + p_start = attn_ranks[0] * total_num_kv_heads // remote_tp_size + rank_offset_factor = local_head - p_start + else: + # D TP > P TP: we index into remote to read different heads depending on rank. + rank_offset_factor = tp_rank % (tp_size // remote_tp_size) + + return TPMapping( + source_ranks_per_group=source_ranks_per_group, + all_source_ranks=tuple(all_ranks), + rank_to_attention_slot=rank_to_attention_slot, + rank_offset_factor=rank_offset_factor, + ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py index 607bf4b988..caa5f432c5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py @@ -9,6 +9,7 @@ import threading import time import uuid from collections import defaultdict +from collections.abc import Iterator from concurrent.futures import Future, ThreadPoolExecutor from typing import TYPE_CHECKING, Any, cast @@ -21,7 +22,7 @@ from vllm import envs from vllm.distributed.kv_transfer.kv_connector.utils import ( BlockIds, EngineId, - MambaEngineTransferInfo, + EngineTransferInfo, TransferTopology, get_current_attn_backends, kv_postprocess_blksize_and_layout_on_receive, @@ -43,13 +44,19 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import ( NixlKVConnectorStats, ) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( + ReadSpec, + TPMapping, + _is_attention_spec, + _is_ssm_spec, + compute_tp_mapping, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( _NIXL_SUPPORTED_DEVICE, zmq_ctx, ) from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( MambaConvSplitInfo, - compute_physical_blocks_per_logical, derive_mamba_conv_split, ) from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config @@ -58,7 +65,6 @@ from vllm.distributed.parallel_state import ( get_tensor_model_parallel_world_size, ) from vllm.logger import init_logger -from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first from vllm.platforms import current_platform from vllm.utils.network_utils import make_zmq_path from vllm.v1.attention.backends.utils import get_kv_cache_layout @@ -80,6 +86,111 @@ logger = init_logger(__name__) class NixlConnectorWorker: """Implementation of Worker side methods""" + def _compute_desc_ids( + self, + block_ids: BlockIds, + dst_num_blocks: int, + block_size_ratio: float | None, + physical_blocks_per_logical: int, + ) -> np.ndarray: + """Compute NIXL descriptor IDs for given block IDs.""" + num_fa_regions = self.num_regions + num_ssm_regions = len(self.block_len_per_layer) * 4 if self._has_mamba else 0 + + num_blocks = dst_num_blocks + if block_size_ratio is not None: + num_blocks = int(num_blocks * block_size_ratio) + ratio = physical_blocks_per_logical + logical_blocks = num_blocks // ratio + + num_fa_descs = num_fa_regions * num_blocks + + # All-attention fast path: single vectorized broadcast. + if num_ssm_regions == 0: + block_arr = np.concatenate(block_ids)[None, :] + region_ids = np.arange(num_fa_regions)[:, None] + return (region_ids * num_blocks + block_arr).flatten() + + # NOTE (NickLucche) With HMA, every kv group has the same number + # of layers and layers from different groups share the same kv + # tensor. Therefore we 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). + all_descs: list[np.ndarray] = [] + for i, group in enumerate(block_ids): + group_arr = np.asarray(group) + if _is_attention_spec(self._group_spec_types[i]): + fa_region_ids = np.arange(num_fa_regions)[:, None] + all_descs.append( + (fa_region_ids * num_blocks + group_arr[None, :]).flatten() + ) + elif _is_ssm_spec(self._group_spec_types[i]): + # NOTE (NickLucche) SSM and Attention block regions can + # be exchanged arbitrarily by manager. Therefore, descs + # are laid out as: + # [descs_fa (all regions) | descs_ssm (all regions)]. + # num_fa_descs offset must be computed per-engine since + # P and D can have different num_blocks (and thus + # different FA desc counts). + ssm_region_ids = np.arange(num_ssm_regions)[:, None] + all_descs.append( + ( + ssm_region_ids * logical_blocks + + group_arr[None, :] + + num_fa_descs + ).flatten() + ) + else: + raise ValueError( + f"Unknown spec type {self._group_spec_types[i]} at index {i}" + ) + + return np.concatenate(all_descs) + + def _build_local_splits_from_plan( + self, + plan: TPMapping, + src_blocks_data: list[tuple[int, int, int]], + num_fa_descs: int, + ) -> Iterator[list[tuple[int, int, int]]]: + """Build split handle data for P_TP > D_TP scenario. + + num_fa_descs is the boundary between FA and SSM descriptors. + Split counts are derived from source_ranks_per_group lengths. + FA uses rank_to_attention_slot for the slot offset; + SSM uses the rank's positional index. + """ + fa_idx = next( + i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t) + ) + fa_num_splits = len(plan.source_ranks_per_group[fa_idx]) + + has_ssm_descs = num_fa_descs < len(src_blocks_data) + ssm_idx = next( + (i for i, t in enumerate(self._group_spec_types) if _is_ssm_spec(t)), + None, + ) + ssm_num_splits = ( + len(plan.source_ranks_per_group[ssm_idx]) + if has_ssm_descs and ssm_idx is not None + else 0 + ) + + for p_idx, p_rank in enumerate(plan.all_source_ranks): + fa_slot = plan.rank_to_attention_slot.get(p_rank, 0) + + handle: list[tuple[int, int, int]] = [] + for j, (addr, local_len, dev) in enumerate(src_blocks_data): + if j < num_fa_descs: + chunk = local_len // fa_num_splits + handle.append((addr + fa_slot * chunk, chunk, dev)) + else: + chunk = local_len // ssm_num_splits + handle.append((addr + p_idx * chunk, chunk, dev)) + yield handle + def __init__( self, vllm_config: "VllmConfig", @@ -119,44 +230,37 @@ class NixlConnectorWorker: } self.hma_group_size = len(kv_cache_config.kv_cache_tensors) - # ---- Mamba model state (derived from model config) ---- - self._is_mamba_group = [ - isinstance(group.kv_cache_spec, MambaSpec) - for group in kv_cache_config.kv_cache_groups - ] + # ---- Model state (derived from model config) ---- mamba_ssm_size = (0, 0) - self._has_mamba = any(self._is_mamba_group) + # Conv state sub-projection decomposition (None when no Mamba). + # The 3-read transfer requires DS (dim, state_len) conv layout so + # that x/B/C sub-projections are contiguous in memory. + self._conv_decomp: MambaConvSplitInfo | None = None + self._has_mamba = any( + isinstance(g.kv_cache_spec, MambaSpec) + for g in kv_cache_config.kv_cache_groups + ) if self._has_mamba: assert self._is_hma_required + from vllm.model_executor.layers.mamba.mamba_utils import ( + is_conv_state_dim_first, + ) + + assert is_conv_state_dim_first(), ( + "3-read Mamba conv transfer requires DS conv state layout. " + "Set VLLM_SSM_CONV_STATE_LAYOUT=DS" + ) mamba_spec = next( spec for spec in self._layer_specs.values() if isinstance(spec, MambaSpec) ) - conv_nbytes, ssm_nbytes = ( - torch.tensor([], dtype=mamba_spec.dtypes[0]).element_size(), # type: ignore[misc] - torch.tensor([], dtype=mamba_spec.dtypes[1]).element_size(), # type: ignore[misc] - ) - conv_shape, ssm_shape = ( - torch.Size(mamba_spec.shapes[0]), - torch.Size(mamba_spec.shapes[1]), - ) - mamba_ssm_size = ( - conv_shape.numel() * conv_nbytes, - ssm_shape.numel() * ssm_nbytes, + self._conv_decomp = derive_mamba_conv_split( + mamba_spec, + vllm_config.parallel_config.tensor_parallel_size, ) + mamba_ssm_size = self._conv_decomp.ssm_sizes self._mamba_ssm_size = mamba_ssm_size - # Conv state sub-projection decomposition (None when no Mamba). - # The 3-read transfer requires DS (dim, state_len) conv layout so - # that x/B/C sub-projections are contiguous in memory. - self._conv_decomp: MambaConvSplitInfo | None = None - if self._has_mamba: - assert is_conv_state_dim_first(), ( - "3-read Mamba conv transfer requires DS conv state layout. " - "Set VLLM_SSM_CONV_STATE_LAYOUT=DS" - ) - local_tp = vllm_config.parallel_config.tensor_parallel_size - self._conv_decomp = derive_mamba_conv_split(mamba_spec, local_tp) # Agent. non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"] @@ -268,14 +372,6 @@ class NixlConnectorWorker: self.dst_num_blocks: dict[EngineId, int] = {} self._registered_descs: list[Any] = [] - # ---- Mamba-HMA per-engine state (only used when self._has_mamba) ---- - # NOTE (ZhanqiuHu): _physical_blocks_per_logical MUST be per-engine. - # physical_blocks_per_logical = ceil((conv_bytes + ssm_bytes) / block_len) - # where conv/ssm bytes are per-TP-rank (dimension-sharded). With - # heterogeneous TP the per-rank sizes differ, so the ratio differs: - # e.g. Nemotron 30B: P(TP=4) → 131, D(TP=1) → 261. - self._physical_blocks_per_logical: dict[EngineId, int] = {} - # In progress transfers. # [req_id -> list[handle]] self._recving_metadata: dict[ReqId, ReqMeta] = {} @@ -330,6 +426,13 @@ class NixlConnectorWorker: self._physical_blocks_per_logical_kv_block = 1 self._sync_block_size_with_kernel() + self._group_spec_types = tuple( + type(g.kv_cache_spec) for g in self.kv_cache_config.kv_cache_groups + ) + + # Per-engine TP mappings. Generated during handshake. + self.tp_mappings: dict[EngineId, TPMapping] = {} + self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config( "enforce_handshake_compat", True ) @@ -812,9 +915,6 @@ class NixlConnectorWorker: self.dst_num_blocks[self.engine_id] = self.num_blocks if self._has_mamba: - self._physical_blocks_per_logical[self.engine_id] = ( - self._physical_blocks_per_logical_kv_block - ) logger.info( "Hybrid SSM registration: num_blocks=%s, " "logical_num_blocks=%s, ratio=%s, num_regions=%s, " @@ -847,6 +947,9 @@ class NixlConnectorWorker: block_size=self.block_size, ssm_sizes=self._mamba_ssm_size, attn_backend_name=self.backend_name, + physical_blocks_per_logical_kv_block=( + self._physical_blocks_per_logical_kv_block + ), ) # Wrap metadata in payload with hash for defensive decoding assert self.compat_hash is not None @@ -875,6 +978,8 @@ class NixlConnectorWorker: result: list[tuple[int, int, int]] = [] 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 ) @@ -894,67 +999,11 @@ class NixlConnectorWorker: ) return result - def _build_fa_remote_for_mamba( - self, - nixl_agent_meta: NixlAgentMetadata, - block_size_ratio: int, - transfer_topo: TransferTopology, - remote_engine_id: EngineId, - ) -> list[tuple[int, int, int]]: - """Build remote FA descriptors for mamba models. - - Uses TransferTopology for GQA-aware FA divisor and head-based rank - offset instead of the standard uniform tp_ratio split. - """ - 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}." - ) - # TODO (ZhanqiuHu): unify with register_remote_blocks when Mamba-HMA - # hetero-TP logic stabilizes. - mamba_info = transfer_topo.get_engine_info(remote_engine_id) - assert isinstance(mamba_info, MambaEngineTransferInfo) - tp_ratio = transfer_topo.tp_ratio(mamba_info.remote_tp_size) - result: list[tuple[int, int, int]] = [] - for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): - local_block_len = self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=True, mamba_view=False - ) - remote_kv_block_len = local_block_len // block_size_ratio - if block_size_ratio > 1: - local_block_len = remote_kv_block_len - - if tp_ratio < 0 and not self.use_mla: - local_block_len = local_block_len // mamba_info.remote_num_fa_reads - - rank_offset = transfer_topo.fa_rank_offset( - remote_engine_id, remote_kv_block_len - ) - - num_blocks = nixl_agent_meta.num_blocks - page_size = nixl_agent_meta.block_lens[i] - for block_id in range(num_blocks): - block_offset = block_id * page_size - addr = base_addr + block_offset + rank_offset - result.append((addr, local_block_len, nixl_agent_meta.device_id)) - - if transfer_topo.is_kv_layout_blocks_first: - second_split = self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=False, mamba_view=False - ) - if tp_ratio < 0 and not self.use_mla: - second_split = second_split // mamba_info.remote_num_fa_reads - for block_id in range(num_blocks): - block_offset = block_id * page_size - addr = base_addr + block_offset + rank_offset - v_addr = addr + nixl_agent_meta.block_lens[i] // 2 - result.append((v_addr, second_split, nixl_agent_meta.device_id)) - return result - def _build_mamba_remote( self, nixl_agent_meta: NixlAgentMetadata, tp_ratio: int, + transfer_info: EngineTransferInfo, ) -> list[tuple[int, int, int]]: """Build 4 remote desc regions (x, B, C, ssm) per layer for the 3-read transfer. For hetero-TP, each D rank reads only its @@ -974,17 +1023,15 @@ class NixlConnectorWorker: ssm_read_size = self._mamba_ssm_size[1] else: # NOTE (ZhanqiuHu): tp_ratio < 0 means P_TP > D_TP, so P pages - # are smaller than D's. self._conv_decomp has D-sized dimensions, - # but we need P-sized offsets. Scale down by |tp_ratio|. + # are smaller than D's. self._conv_decomp has D-sized dimensions, + # but we need P-sized offsets. Scale down by |tp_ratio|. abs_ratio = -tp_ratio xb_p = self._conv_decomp.x_bytes // abs_ratio bb_p = self._conv_decomp.b_bytes // abs_ratio conv_offsets = [(0, xb_p), (xb_p, bb_p), (xb_p + bb_p, bb_p)] ssm_read_size = nixl_agent_meta.ssm_sizes[1] - remote_physical_per_logical = self._physical_blocks_per_logical[ - nixl_agent_meta.engine_id - ] + remote_physical_per_logical = transfer_info.remote_physical_blocks_per_logical num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical device_id = nixl_agent_meta.device_id @@ -1007,6 +1054,91 @@ class NixlConnectorWorker: result.append((ssm_addr, ssm_read_size, device_id)) return result + def _build_fa_local( + self, + base_addresses: list[int], + block_size_ratio: int, + ) -> list[tuple[int, int, int]]: + """Build local FA descriptors for all layers.""" + assert self.transfer_topo is not None + num_blocks = self.num_blocks * block_size_ratio + result: list[tuple[int, int, int]] = [] + for i, base_addr in enumerate(base_addresses): + kv_block_len = ( + self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=True, mamba_view=False + ) + // block_size_ratio + ) + page_stride = self.block_len_per_layer[i] // block_size_ratio + for block_id in range(num_blocks): + block_offset = block_id * page_stride + addr = base_addr + block_offset + result.append((addr, kv_block_len, self.device_id)) + + if self.transfer_topo.is_kv_layout_blocks_first: + # Separate and interleave K/V regions to maintain the same + # descs ordering. This is needed for selecting contiguous heads + # when split across TP ranks. + second_split = self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=False, mamba_view=False + ) + for block_id in range(num_blocks): + block_offset = block_id * page_stride + addr = base_addr + block_offset + v_addr = addr + kv_block_len + result.append((v_addr, second_split, self.device_id)) + return result + + def _build_fa_remote( + self, + plan: TPMapping, + nixl_agent_meta: NixlAgentMetadata, + block_size_ratio: int, + ) -> list[tuple[int, int, int]]: + """Build remote FA descriptors for all layers.""" + assert self.transfer_topo is not None + fa_group_idx = next( + i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t) + ) + num_attn_reads = len(plan.source_ranks_per_group[fa_group_idx]) + num_blocks = nixl_agent_meta.num_blocks + result: list[tuple[int, int, int]] = [] + for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): + # Read our whole local region size from remote.. + local_block_len = self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=True, mamba_view=False + ) + remote_kv_block_len = local_block_len // block_size_ratio + if block_size_ratio > 1: + # ..using remote kv_block_len as transfer unit + local_block_len = remote_kv_block_len + + local_block_len = local_block_len // num_attn_reads + rank_offset = plan.rank_offset_factor * remote_kv_block_len + + page_size = nixl_agent_meta.block_lens[i] + for block_id in range(num_blocks): + block_offset = block_id * page_size + # For each block, grab the kv heads chunk belonging to current local + # tp rank of size local_block_len. + addr = base_addr + block_offset + rank_offset + result.append((addr, local_block_len, nixl_agent_meta.device_id)) + + if self.transfer_topo.is_kv_layout_blocks_first: + # With FlashInfer index V separately to allow head splitting. + second_split = self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=False, mamba_view=False + ) + second_split = second_split // num_attn_reads + for block_id in range(num_blocks): + block_offset = block_id * page_size + addr = base_addr + block_offset + rank_offset + # Hop over the first split of remote page, K, to read V. + v_addr = addr + nixl_agent_meta.block_lens[i] // 2 + result.append((v_addr, second_split, nixl_agent_meta.device_id)) + return result + def register_local_xfer_handler( self, block_size: int, @@ -1023,68 +1155,23 @@ class NixlConnectorWorker: data copy correctness. """ assert self.transfer_topo is not None - transfer_topo = self.transfer_topo - block_size_ratio = self.block_size // block_size - blocks_data: list[tuple[int, int, int]] = [] local_base_addresses = self.kv_caches_base_addr[self.engine_id][self.tp_rank] - def register_blocks(blocks_data: list[tuple[int, int, int]], mamba: bool): - for i, base_addr in enumerate(local_base_addresses): - # The new block_len is using prefill block_len; - # and num_blocks is multiple with N - kv_block_len = ( - self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=True, mamba_view=mamba - ) - // block_size_ratio - ) - # Jump one page_size, but ssm page_size may be bigger when kernel - # locks block size to a specific value. - block_len_per_layer = ( - self.block_len_per_layer[i] - // block_size_ratio - * (1 if not mamba else self._physical_blocks_per_logical_kv_block) - ) - num_blocks = self._logical_num_blocks if mamba else self.num_blocks - num_blocks = num_blocks * block_size_ratio - for block_id in range(num_blocks): - block_offset = block_id * block_len_per_layer - addr = base_addr + block_offset - # (addr, len, device id) - blocks_data.append((addr, kv_block_len, self.device_id)) - - if transfer_topo.is_kv_layout_blocks_first: - second_split = self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=False, mamba_view=mamba - ) - # Separate and interleave K/V regions to maintain the same - # descs ordering. This is needed for selecting contiguous heads - # when split across TP ranks. - for block_id in range(num_blocks): - block_offset = block_id * block_len_per_layer - addr = base_addr + block_offset - # Register addresses for V cache (K registered first). - v_addr = addr + kv_block_len - blocks_data.append((v_addr, second_split, self.device_id)) - logger.debug( - "Created %s blocks for src engine %s and rank %s on device id %s", - len(blocks_data), - self.engine_id, - self.tp_rank, - self.device_id, - ) - - # NOTE (ZhanqiuHu): mamba=True path in register_blocks is not used - # right now — we use _build_mamba_local instead for the 3-read - # approach. However, we might still need this as a fallback for homogeneous TP. - register_blocks(blocks_data, mamba=False) + blocks_data = self._build_fa_local(local_base_addresses, block_size_ratio) + logger.debug( + "Created %s blocks for src engine %s and rank %s on device id %s", + len(blocks_data), + self.engine_id, + self.tp_rank, + self.device_id, + ) if self._has_mamba: assert self.num_descs == len(blocks_data) - # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-read split is - # unnecessary — a single conv desc per block suffices. Consider + # 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 - # registration (register_blocks mamba=True) when no hetero-TP + # registration (_build_fa_local mamba=True) when no hetero-TP # 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)") @@ -1160,26 +1247,26 @@ class NixlConnectorWorker: assert self.transfer_topo is not None transfer_topo = self.transfer_topo physical_blocks_per_logical = ( - compute_physical_blocks_per_logical( - nixl_agent_meta.ssm_sizes, - nixl_agent_meta.block_lens[0], - ) - if self._has_mamba - else 1 + nixl_agent_meta.physical_blocks_per_logical_kv_block ) - transfer_topo.register_remote_engine( - remote_engine_id=engine_id, + transfer_info = EngineTransferInfo( remote_tp_size=remote_tp_size, remote_block_size=nixl_agent_meta.block_size, remote_block_len=nixl_agent_meta.block_lens[0], remote_physical_blocks_per_logical=physical_blocks_per_logical, - local_block_len=self.block_len_per_layer[0], ) - if self._has_mamba and engine_id not in self._physical_blocks_per_logical: - self._physical_blocks_per_logical[engine_id] = physical_blocks_per_logical - + transfer_topo.register_remote_engine(engine_id, transfer_info) logger.info("Transfer plan: %s", transfer_topo.describe(engine_id)) + self.tp_mappings[engine_id] = compute_tp_mapping( + transfer_topo.tp_rank, + transfer_topo.tp_size, + transfer_info.remote_tp_size, + transfer_topo.is_mla, + transfer_topo.total_num_kv_heads, + self._group_spec_types, + ) + remote_agent_name = self.nixl_wrapper.add_remote_agent( nixl_agent_meta.agent_metadata ) @@ -1206,11 +1293,6 @@ class NixlConnectorWorker: # this is the ratio between the two sizes. tp_ratio = transfer_topo.tp_ratio(remote_tp_size) - # Handle tp_size>num_kv_heads: replicate KV cache. - indexes_into_remote = ( - not transfer_topo.replicates_kv_cache(engine_id) and tp_ratio > 0 - ) - logger.debug( "Registering remote agent (%s, rank %s) memory regions with tp_ratio %s", engine_id, @@ -1218,6 +1300,8 @@ class NixlConnectorWorker: tp_ratio, ) + plan = self.tp_mappings[engine_id] + ### (Optional) Register local agent memory regions. MLA is not split. if ( tp_ratio < 0 @@ -1227,149 +1311,51 @@ class NixlConnectorWorker: # Remote tp_size > local tp_size: read from multiple remote ranks. # Logically "split" own regions into |tp_ratio| chunks. Mind that # we only do this once per remote tp_size (replica-friendly). - abs_tp = -tp_ratio self.src_xfer_handles_by_tp_ratio[tp_ratio] = [] - if self._has_mamba: - if transfer_topo.needs_split_handles(engine_id): - # Mamba-HMA: FA and Mamba use different split factors. - for handle_data in transfer_topo.compute_split_handle_data( - engine_id, self.src_blocks_data, self.num_descs, abs_tp - ): - 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) - - logger.info( - "Mamba-HMA split handles: %s, num_descs=%s", - transfer_topo.describe(engine_id), - self.num_descs, - ) - else: - # Original path: uniform divide by abs_tp (non-Mamba-HMA). - for i in range(abs_tp): - blocks_data = [] - for memory_region in self.src_blocks_data: - addr, local_block_len, own_tp_rank = memory_region - remote_block_len = local_block_len // abs_tp - addr = addr + i * remote_block_len - blocks_data.append((addr, remote_block_len, own_tp_rank)) - descs = self.nixl_wrapper.get_xfer_descs( - blocks_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) + for handle_data in self._build_local_splits_from_plan( + plan, + self.src_blocks_data, + self.num_descs, + ): + 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) ### Register remote agent memory regions - blocks_data = [] - # With homogeneous TP, D pulls the whole kv cache from corresponding - # rank. With heterogeneous TP, prepare the descriptors by splitting the - # P KV cache along kv_head dim, of D worker's kv_head size (D>P). + # With homogeneous TP, D pulls the whole kv cache from corresponding rank. With + # heterogeneous TP, prepare the descriptors by splitting the P KV cache along + # kv_head dim, of D worker's kv_head size (D>P). # Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..]. # Register all remote blocks, but only the corresponding kv heads. - def register_remote_blocks( - blocks_data: list[tuple[int, int, int]], mamba: bool - ): - for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): - # Read our whole local region size from remote. - local_block_len = self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=True, mamba_view=mamba - ) - remote_kv_block_len = local_block_len // block_size_ratio - if block_size_ratio > 1: - # using remote kv_block_len as transfer unit - local_block_len = remote_kv_block_len - - if tp_ratio < 0 and not self.use_mla: - # Remote tp is bigger: read a chunk of local region from remote - local_block_len = local_block_len // (-tp_ratio) - rank_offset = ( - self.tp_rank % tp_ratio * remote_kv_block_len - if indexes_into_remote - else 0 - ) - - # Assume same num_blocks for mamba and fa - num_blocks = ( - nixl_agent_meta.num_blocks - if not mamba - else nixl_agent_meta.num_blocks - // self._physical_blocks_per_logical_kv_block - ) - page_size = nixl_agent_meta.block_lens[i] * ( - 1 if not mamba else self._physical_blocks_per_logical_kv_block - ) - for block_id in range(num_blocks): - block_offset = block_id * page_size - # For each block, grab the heads chunk belonging to rank_i - # of size remote_nheads // tp_ratio, which correspond to - # self.block_len == remote_block_len//tp_ratio bytes. - addr = base_addr + block_offset + rank_offset - # (addr, len, device id) - blocks_data.append( - (addr, local_block_len, nixl_agent_meta.device_id) - ) - - if transfer_topo.is_kv_layout_blocks_first: - # With FlashInfer index V separately to allow head splitting. - second_split = self.get_backend_aware_kv_block_len( - layer_idx=i, first_split=False, mamba_view=mamba - ) - # Apply the same scaling as local_block_len above for when we read - # a chunk of local V from `tp_ratio` separate remote workers. - if tp_ratio < 0 and not self.use_mla: - second_split = second_split // (-tp_ratio) - for block_id in range(num_blocks): - block_offset = block_id * page_size - addr = base_addr + block_offset + rank_offset - # Hop over the first split of remote page: either K or Conv. - if mamba: - v_addr = addr + nixl_agent_meta.ssm_sizes[0] - else: - v_addr = addr + nixl_agent_meta.block_lens[i] // 2 - blocks_data.append( - (v_addr, second_split, nixl_agent_meta.device_id) - ) - - logger.debug( - "Created %s blocks for dst engine %s" - " with remote rank %s and local rank %s", - len(blocks_data), - engine_id, - remote_tp_rank, - self.tp_rank, - ) - + blocks_data = self._build_fa_remote( + plan, + nixl_agent_meta, + block_size_ratio, + ) + logger.debug( + "Created %s blocks for dst engine %s with remote rank %s and local rank %s", + len(blocks_data), + engine_id, + remote_tp_rank, + self.tp_rank, + ) if self._has_mamba: - # Mamba-HMA: separate FA registration with GQA-aware sizing, - # plus mamba 3-read registration for the Mamba "view" of the - # same KV cache tensors. logger.debug( "Registering remote Mamba blocks for engine %s rank %s", engine_id, remote_tp_rank, ) - blocks_data.extend( - self._build_fa_remote_for_mamba( - nixl_agent_meta, - block_size_ratio, - transfer_topo, - engine_id, - ) - ) blocks_data.extend( self._build_mamba_remote( nixl_agent_meta, tp_ratio, + transfer_info, ) ) - else: - register_remote_blocks(blocks_data, mamba=False) # Register with NIXL. descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) @@ -1897,33 +1883,49 @@ class NixlConnectorWorker: def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): assert meta.remote is not None and self.transfer_topo is not None engine_id = meta.remote.engine_id - remote_ranks = self.transfer_topo.target_remote_ranks(engine_id) + plan = self.tp_mappings[engine_id] remote_info = self.transfer_topo.get_engine_info(engine_id) tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size) - if self._has_mamba: - # Expand remote logical → kernel block IDs. - meta.remote.block_ids = self._logical_to_remote_kernel_block_ids( - meta.remote.block_ids, - self._physical_blocks_per_logical[meta.remote.engine_id], + meta.remote.block_ids = self._logical_to_remote_kernel_block_ids( + meta.remote.block_ids, + remote_info.remote_physical_blocks_per_logical, + ) + remote_block_ids = meta.remote.block_ids + local_block_ids = meta.local_physical_block_ids + num_groups = len(local_block_ids) + read_specs = [ + ReadSpec( + remote_rank=rank, + local_block_ids=[ + list(local_block_ids[g]) + if rank in plan.source_ranks_per_group[g] + else [] + for g in range(num_groups) + ], + remote_block_ids=[ + list(remote_block_ids[g]) + if rank in plan.source_ranks_per_group[g] + else [] + for g in range(num_groups) + ], ) - else: - meta.remote.block_ids = self._logical_to_kernel_block_ids( - meta.remote.block_ids - ) - # D may have to perform multiple reads from different remote ranks. - for i, remote_rank in enumerate(remote_ranks): - if self.use_mla and tp_ratio < 0 and i > 0: - # MLA opt: when P TP > D TP, only a single read is executed for - # the first remote rank (cache is duplicated).. - break + for rank in plan.all_source_ranks + ] + # D may have to perform multiple reads from different remote ranks. + # MLA opt: when P TP > D TP, only a single read is executed for + # the first remote rank (cache is duplicated). + if self.use_mla and tp_ratio < 0: + read_specs = read_specs[:1] + + for i, spec in enumerate(read_specs): remote_block_size = remote_info.remote_block_size logger.debug( "Remote agent %s available, calling _read_blocks" " on remote rank %s with remote block size %s for req %s", meta.remote.engine_id, - remote_rank, + spec.remote_rank, remote_block_size, req_id, ) @@ -1942,49 +1944,34 @@ class NixlConnectorWorker: # Destination handle: remote_engine_id -> remote_rank -> handle. remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][ - remote_rank + spec.remote_rank ] - local_ids: BlockIds = meta.local_physical_block_ids - remote_ids: BlockIds = meta.remote.block_ids - if self._has_mamba: - # Mamba-HMA: zero out FA groups for P ranks outside fa_read_targets. - local_ids, remote_ids = self.transfer_topo.filter_block_ids_for_rank( - engine_id, - remote_rank, - local_ids, - remote_ids, - self._is_mamba_group, - ) - 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_block_ids=local_ids, - remote_block_ids=remote_ids, - remote_rank=remote_rank, local_xfer_side_handle=local_xfer_side_handle, remote_xfer_side_handle=remote_xfer_side_handle, ) - if self.use_mla and tp_ratio < 0: - # ..but we still need to notify the other remote ranks that we - # have the blocks we need so they can update the request state. - notif_id = f"{meta.remote.request_id}:{self.world_size}".encode() - remote_agents = self._remote_agents[meta.remote.engine_id] - for rank_to_notify, agent in remote_agents.items(): - if rank_to_notify != remote_rank: - self.nixl_wrapper.send_notif(agent, notif_msg=notif_id) + if self.use_mla and tp_ratio < 0 and read_specs: + # ..but we still need to notify the other remote ranks that we + # have the blocks we need so they can update the request state. + notif_id = f"{req_id}:{self.world_size}".encode() + remote_agents = self._remote_agents[meta.remote.engine_id] + read_ranks = {s.remote_rank for s in read_specs} + for rank_to_notify, agent in remote_agents.items(): + if rank_to_notify not in read_ranks: + self.nixl_wrapper.send_notif(agent, notif_msg=notif_id) def _read_blocks( self, - local_block_ids: BlockIds, - remote_block_ids: BlockIds, + read_spec: ReadSpec, dst_engine_id: str, request_id: str, remote_request_id: str, - remote_rank: int, local_xfer_side_handle: int, remote_xfer_side_handle: int, ): @@ -1993,6 +1980,10 @@ class NixlConnectorWorker: a single remote worker. """ assert self.transfer_topo is not None + remote_rank = read_spec.remote_rank + local_block_ids = read_spec.local_block_ids + remote_block_ids = read_spec.remote_block_ids + remote_info = self.transfer_topo.get_engine_info(dst_engine_id) block_size_ratio = self.transfer_topo.block_size_ratio( remote_info.remote_block_size @@ -2061,16 +2052,14 @@ class NixlConnectorWorker: == len(local_block_ids) == len(self.kv_cache_config.kv_cache_groups) ) + # Partial prefix cache hit: just read uncomputed blocks. + # Skip mamba groups — their blocks represent full state (conv+ssm), + # not per-token data, so trimming would corrupt the transfer. remote_block_ids = list(remote_block_ids) for i, remote_group in enumerate(remote_block_ids): - num_remote_blocks = len(remote_group) num_local_blocks = len(local_block_ids[i]) - if not self._is_mamba_group[i]: - assert num_local_blocks <= num_remote_blocks - # Partial prefix cache hit: just read uncomputed blocks. - # Skip mamba groups — their blocks represent full state (conv+ssm), - # not per-token data, so trimming would corrupt the transfer. - if num_local_blocks < num_remote_blocks and not self._is_mamba_group[i]: + assert num_local_blocks <= len(remote_group) + if num_local_blocks < len(remote_group): remote_block_ids[i] = remote_group[-num_local_blocks:] # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from @@ -2078,14 +2067,17 @@ class NixlConnectorWorker: # workers will issue xfers to parts of the P worker remote kv caches. # Get descs ids. - remote_block_descs_ids = self._get_block_descs_ids( - dst_engine_id, - remote_block_ids, + remote_block_descs_ids = self._compute_desc_ids( + block_ids=remote_block_ids, + 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, ) - local_block_descs_ids = self._get_block_descs_ids( - self.engine_id, - local_block_ids, + 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, ) assert len(local_block_descs_ids) == len(remote_block_descs_ids) @@ -2147,63 +2139,6 @@ class NixlConnectorWorker: return mapped_2d.flatten().astype(np.int64) - def _get_block_descs_ids( - self, - engine_id: str, - block_ids: BlockIds, - block_size_ratio: float | None = None, - ) -> np.ndarray: - """ - Get the descs ids for a set of block ids. - When HMA is enabled number of descriptors across kv cache groups might differ. - A single flattened array is returned for all groups anyway. - """ - region_ids = np.arange(self.num_regions) - - # 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. - num_blocks = self.dst_num_blocks[engine_id] - if block_size_ratio is not None: - num_blocks = int(num_blocks * block_size_ratio) - - # 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). - region_ids = region_ids[:, None] - if not self._has_mamba: - block_ids = np.concatenate(block_ids)[None, :] - descs_ids = region_ids * num_blocks + block_ids - return descs_ids.flatten() - else: - # NOTE (NickLucche) SSM and Attention blocks regions can be exchanged - # arbitrarily by manager. Therefore, descs are duplicated for SSM and - # Attention like so: - # desc_handle->[descs_fa (all regions) | descs_ssm (all regions)]. - # This is like having two "low-level views" of the same storage. - # `num_fa_descs` offset must be computed per-engine since P and D can - # have different num_blocks (and thus different FA descs counts). - physical_per_logical = self._physical_blocks_per_logical[engine_id] - logical_blocks = num_blocks // physical_per_logical - num_fa_descs = self.num_regions * num_blocks - # 3-read mamba: 4 regions per unique cache tensor (x, B, C, ssm). - mamba_region_ids = np.arange(len(self.block_len_per_layer) * 4)[:, None] - all_descs = [] - for i, group in enumerate(block_ids): - group_arr = np.asarray(group)[None, :] - if self._is_mamba_group[i]: - all_descs.append( - ( - mamba_region_ids * logical_blocks + group_arr + num_fa_descs - ).flatten() - ) - else: - all_descs.append((region_ids * num_blocks + group_arr).flatten()) - return np.concatenate(all_descs) - def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds: """ Convert logical block ids to kernel physical block ids. @@ -2295,10 +2230,7 @@ class NixlConnectorWorker: """ assert self.transfer_topo is not None if self.transfer_topo.is_kv_layout_blocks_first: - # For indexing only half (either just the K or V part). if mamba_view: - # NOTE (NickLucche) Mamba Opt: this is already skipping the padding so - # we're only transferring the minimum required bytes. block_len = self._mamba_ssm_size[not first_split] else: block_len = self.block_len_per_layer[layer_idx] // 2 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py index 309426814c..00b8e2bb72 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py @@ -31,6 +31,7 @@ class MambaConvSplitInfo: x_local: int # intermediate_size / TP (columns for x) b_local: int # groups_ss / TP (columns for B; C is same size) conv_dtype_size: int # bytes per element (e.g. 2 for float16) + ssm_sizes: tuple[int, int] # (conv_state_bytes, ssm_state_bytes) @property def conv_dim_local(self) -> int: @@ -99,8 +100,8 @@ def derive_mamba_conv_split( local_tp: this engine's tensor-parallel size. Returns: - MambaConvSplitInfo with per-rank x_local, b_local, conv_rows, and - conv_dtype_size. + MambaConvSplitInfo with per-rank x_local, b_local, conv_rows, + conv_dtype_size, and ssm_sizes (conv_state_bytes, ssm_state_bytes). """ if mamba_spec.mamba_type != "mamba2": raise NotImplementedError( @@ -142,12 +143,20 @@ def derive_mamba_conv_split( dtype=mamba_spec.dtypes[0], # type: ignore[misc] ).element_size() + ssm_dtype_size = torch.tensor( + [], + dtype=mamba_spec.dtypes[1], # type: ignore[misc] + ).element_size() + conv_state_bytes = torch.Size(mamba_spec.shapes[0]).numel() * conv_dtype_size + ssm_state_bytes = torch.Size(mamba_spec.shapes[1]).numel() * ssm_dtype_size + # Divide by TP to get per-rank column counts. return MambaConvSplitInfo( conv_rows=conv_rows, x_local=intermediate_size // local_tp, b_local=groups_ss // local_tp, conv_dtype_size=conv_dtype_size, + ssm_sizes=(conv_state_bytes, ssm_state_bytes), )