A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 06:15 UTC
nixl refactor: new transfer design (#40731)
Signed-off-by: ZhanqiuHu <zhu@redhat.com> Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: NickLucche <nlucches@redhat.com>
zhanqiuhu committed
May 6, 2026 at 09:16 UTC
df8e63f4edabc06e51c10fe479fc9b95c303542f
8 files changed
+755
-879
tests/v1/kv_connector/unit/test_nixl_connector.py
+5
@@ -527,6 +527,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
527
block_size=self.block_size,
528
ssm_sizes=(0, 0),
529
attn_backend_name=self.backend_name,
530
+ physical_blocks_per_logical_kv_block=1,
531
),
532
remote_tp_rank=remote_tp_rank,
533
remote_tp_size=remote_tp_size,
@@ -726,6 +727,7 @@ class TestNixlHandshake:
727
worker.num_blocks = 1
728
worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
729
worker.src_blocks_data = [(0, worker.block_len_per_layer[0], worker.tp_rank)]
730
+ worker.num_descs = len(worker.src_blocks_data)
731
732
def check_handshake(remote_tp_size: int):
733
tp_ratio = remote_tp_size // local_tp_size
@@ -978,6 +980,7 @@ class TestNixlHandshake:
980
block_size=worker.block_size,
981
ssm_sizes=(0, 0),
982
attn_backend_name=worker.backend_name,
983
+ physical_blocks_per_logical_kv_block=1,
984
)
985
986
with pytest.raises(RuntimeError):
@@ -1035,6 +1038,7 @@ class TestNixlHandshake:
1038
block_size=worker.block_size,
1039
ssm_sizes=(0, 0),
1040
attn_backend_name=worker.backend_name,
1041
+ physical_blocks_per_logical_kv_block=1,
1042
)
1043
1044
# We don't check layout for homogeneous TP and MLA for now, as the
@@ -2354,6 +2358,7 @@ def test_compatibility_hash_validation(
2358
block_size=prefill_block_size,
2359
ssm_sizes=(0, 0),
2360
attn_backend_name=decode_worker.backend_name,
2361
+ physical_blocks_per_logical_kv_block=1,
2362
)
2363
handshake_payload = NixlHandshakePayload(
2364
compatibility_hash=remote_hash,
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+98
-84
@@ -93,75 +93,82 @@ def test_logical_to_kernel_block_ids_with_hma():
93
94
@pytest.mark.cpu_test
95
@pytest.mark.parametrize(
96
- "has_mamba,swa_enabled,mamba_enabled,remote_ratio,"
97
- "remote_block_ids,expected_remote_block_ids",
96
+ "group_spec_types,expansion_stride,remote_block_ids,expected_remote_block_ids",
97
[
99
- # Non-mamba (FA+SWA): both groups expanded via _logical_to_kernel_block_ids.
100
- # Regression for https://github.com/vllm-project/vllm/pull/39724
101
- (
102
- False,
103
- True,
104
- False,
105
- 1,
98
+ pytest.param(
99
+ ("FullAttentionSpec", "SlidingWindowSpec"),
100
+ 2,
101
([0, 1, 2], [3, 4]),
102
[[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]],
103
+ id="dense_fa_swa",
104
),
109
- # Mamba (FA+Mamba): FA expanded via _logical_to_remote_kernel_block_ids,
110
- # Mamba passed through unchanged.
111
- # remote_ratio=261 (Nemotron 30B TP=1) != local_ratio=2 so that using
112
- # the wrong conversion method produces different FA results.
113
- (
114
- True,
115
- False,
116
- True,
105
+ pytest.param(
106
+ ("FullAttentionSpec", "MambaSpec"),
107
261,
108
([0, 1, 2], [10, 11]),
109
[[0, 1, 261, 262, 522, 523], [10, 11]],
110
+ id="mamba_fa_ssm",
111
),
112
],
122
- ids=["non_mamba_fa_swa", "mamba_fa_ssm"],
113
)
114
def test_read_blocks_for_req_expands_remote_ids(
125
- has_mamba,
126
- swa_enabled,
127
- mamba_enabled,
128
- remote_ratio,
115
+ group_spec_types,
116
+ expansion_stride,
117
remote_block_ids,
118
expected_remote_block_ids,
119
):
120
"""_read_blocks_for_req must expand remote logical block IDs to kernel
121
block IDs when kernel block size != logical block size.
122
135
- Non-mamba path uses _logical_to_kernel_block_ids (all groups expanded).
136
- Mamba path uses _logical_to_remote_kernel_block_ids (FA expanded, Mamba
137
- passed through).
123
+ The hot path always calls _logical_to_remote_kernel_block_ids with
124
+ remote_info.remote_physical_blocks_per_logical (model-agnostic).
125
"""
126
from unittest.mock import MagicMock
127
128
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
129
NixlConnectorMetadata,
130
)
131
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
132
+ TPMapping,
133
+ )
134
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
135
NixlConnectorWorker,
136
)
137
+ from vllm.v1.kv_cache_interface import (
138
+ FullAttentionSpec,
139
+ MambaSpec,
140
+ SlidingWindowSpec,
141
+ )
142
+
143
+ spec_name_to_type = {
144
+ "FullAttentionSpec": FullAttentionSpec,
145
+ "SlidingWindowSpec": SlidingWindowSpec,
146
+ "MambaSpec": MambaSpec,
147
+ }
148
+ resolved_types = tuple(spec_name_to_type[n] for n in group_spec_types)
149
150
worker = object.__new__(NixlConnectorWorker)
149
- worker._has_mamba = has_mamba
151
worker._physical_blocks_per_logical_kv_block = 2
152
+
153
+ has_mamba = any(t is MambaSpec for t in resolved_types)
154
+ has_swa = any(t is SlidingWindowSpec for t in resolved_types)
155
worker.kv_cache_config = make_kv_cache_config(
152
- block_size=16, swa_enabled=swa_enabled, mamba_enabled=mamba_enabled
156
+ block_size=16, swa_enabled=has_swa, mamba_enabled=has_mamba
157
)
158
159
remote_engine_id = "remote-engine"
156
- if has_mamba:
157
- worker._physical_blocks_per_logical = {remote_engine_id: remote_ratio}
160
159
- # Mock transfer_topo: empty remote ranks skips the transfer machinery
160
- # entirely, isolating the block-ID expansion logic.
161
worker.transfer_topo = MagicMock()
162
- worker.transfer_topo.target_remote_ranks.return_value = []
163
- worker.transfer_topo.get_engine_info.return_value = MagicMock(remote_tp_size=1)
162
worker.transfer_topo.tp_ratio.return_value = 1
163
+ remote_info = MagicMock()
164
+ remote_info.remote_physical_blocks_per_logical = expansion_stride
165
+ worker.transfer_topo.get_engine_info.return_value = remote_info
166
+ worker.use_mla = False
167
+
168
+ mock_plan = MagicMock(spec=TPMapping)
169
+ mock_plan.all_source_ranks = ()
170
+ mock_plan.source_ranks_per_group = ()
171
+ worker.tp_mappings = {remote_engine_id: mock_plan}
172
173
metadata = NixlConnectorMetadata()
174
metadata.add_new_req_to_recv(
@@ -306,75 +313,82 @@ def test_nixl_metadata_hma_block_ids_structure():
313
assert list(req_meta.remote.block_ids[1]) == [18, 19, 20, 21]
314
315
309
-@pytest.mark.cpu_test
310
-def test_get_block_descs_ids_hybrid_ssm():
311
- """Test _get_block_descs_ids uses per-group strides for hybrid FA+SSM
312
- when ratio=1 (no kernel block size mismatch)."""
316
+def _make_mock_worker_for_desc_ids(
317
+ num_regions: int,
318
+ has_mamba: bool,
319
+ group_spec_types: tuple,
320
+ block_len_per_layer: list[int] | None = None,
321
+):
322
+ """Build a mock NixlConnectorWorker with attrs needed by _compute_desc_ids."""
323
+ from unittest.mock import MagicMock
324
+
325
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
326
NixlConnectorWorker,
327
)
328
317
- worker = object.__new__(NixlConnectorWorker)
329
+ worker = MagicMock(spec=NixlConnectorWorker)
330
+ worker.num_regions = num_regions
331
+ worker._has_mamba = has_mamba
332
+ worker._group_spec_types = group_spec_types
333
+ worker.block_len_per_layer = block_len_per_layer or [100]
334
+ worker._compute_desc_ids = NixlConnectorWorker._compute_desc_ids.__get__(
335
+ worker, NixlConnectorWorker
336
+ )
337
+ return worker
338
319
- num_blocks = 100
320
- engine_id = "test-engine"
321
- worker.num_regions = 2
322
- worker.dst_num_blocks = {engine_id: num_blocks}
323
- worker._has_mamba = True
324
- worker._is_mamba_group = [False, True]
325
- worker._physical_blocks_per_logical_kv_block = 1
326
- worker._physical_blocks_per_logical = {engine_id: 1}
327
- worker.block_len_per_layer = [100]
328
- # num_descs = num_regions * num_blocks (no blocks_first doubling)
329
- worker.num_descs = 2 * num_blocks
339
+
340
+@pytest.mark.cpu_test
341
+def test_get_block_descs_ids_hybrid_ssm():
342
+ """Test _compute_desc_ids uses per-group strides for hybrid
343
+ FA+SSM when ratio=1 (no kernel block size mismatch)."""
344
+ from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
345
+
346
+ worker = _make_mock_worker_for_desc_ids(
347
+ num_regions=2,
348
+ has_mamba=True,
349
+ group_spec_types=(FullAttentionSpec, MambaSpec),
350
+ block_len_per_layer=[100],
351
+ )
352
353
fa_blocks = [3, 5]
354
ssm_blocks = [1, 2]
333
- result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks))
334
-
335
- # FA group: stride=num_blocks=100, offset=0
336
- # region0: [3, 5], region1: [103, 105]
337
- # SSM group: stride=logical_blocks=100 (=num_blocks/ratio=100/1),
338
- # offset=num_fa_descs=200, 4 regions per Mamba layer (x, B, C, ssm)
339
- # region0: [201, 202], region1: [301, 302],
340
- # region2: [401, 402], region3: [501, 502]
355
+ result = worker._compute_desc_ids(
356
+ block_ids=(fa_blocks, ssm_blocks),
357
+ dst_num_blocks=100,
358
+ block_size_ratio=None,
359
+ physical_blocks_per_logical=1,
360
+ )
361
+
362
expected = [3, 5, 103, 105, 201, 202, 301, 302, 401, 402, 501, 502]
363
assert list(result) == expected, f"Expected {expected}, got {list(result)}"
364
365
366
@pytest.mark.cpu_test
367
def test_get_block_descs_ids_kernel_block_mismatch():
347
- """Test _get_block_descs_ids uses different strides for FA (kernel blocks)
348
- vs SSM (logical blocks) when ratio > 1."""
349
- from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
350
- NixlConnectorWorker,
351
- )
352
-
353
- worker = object.__new__(NixlConnectorWorker)
368
+ """Test _compute_desc_ids uses different strides for FA
369
+ (kernel blocks) vs SSM (logical blocks) when ratio > 1."""
370
+ from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
371
372
ratio = 4
373
logical_blocks = 100
374
num_blocks = logical_blocks * ratio # 400 kernel blocks
358
- engine_id = "test-engine"
359
- worker.num_regions = 2
360
- worker.dst_num_blocks = {engine_id: num_blocks}
361
- worker._has_mamba = True
362
- worker._is_mamba_group = [False, True]
363
- worker._physical_blocks_per_logical_kv_block = ratio
364
- worker._physical_blocks_per_logical = {engine_id: ratio}
365
- worker.block_len_per_layer = [100]
366
- worker.num_descs = 2 * num_blocks # 800
367
-
368
- fa_blocks = [3, 7] # kernel-level block IDs
369
- ssm_blocks = [1, 2] # logical block IDs
370
- result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks))
371
-
372
- # FA group: stride=num_blocks=400, offset=0
373
- # region0: [3, 7], region1: [403, 407]
374
- # SSM group: stride=logical_blocks=400//4=100, offset=num_fa_descs=800,
375
- # 4 regions per Mamba layer (x, B, C, ssm)
376
- # region0: [801, 802], region1: [901, 902],
377
- # region2: [1001, 1002], region3: [1101, 1102]
375
+
376
+ worker = _make_mock_worker_for_desc_ids(
377
+ num_regions=2,
378
+ has_mamba=True,
379
+ group_spec_types=(FullAttentionSpec, MambaSpec),
380
+ block_len_per_layer=[100],
381
+ )
382
+
383
+ fa_blocks = [3, 7]
384
+ ssm_blocks = [1, 2]
385
+ result = worker._compute_desc_ids(
386
+ block_ids=(fa_blocks, ssm_blocks),
387
+ dst_num_blocks=num_blocks,
388
+ block_size_ratio=None,
389
+ physical_blocks_per_logical=ratio,
390
+ )
391
+
392
expected = [3, 7, 403, 407, 801, 802, 901, 902, 1001, 1002, 1101, 1102]
393
assert list(result) == expected, f"Expected {expected}, got {list(result)}"
394
tests/v1/kv_connector/unit/test_tp_mapping.py
new
+141
@@ -0,0 +1,141 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for TP mapping and transfer plan utilities.
4
+
5
+These tests verify that TP mapping produces correct outputs
6
+(source ranks, split handles, desc IDs).
7
+No GPU or NIXL required.
8
+"""
9
+
10
+from __future__ import annotations
11
+
12
+import pytest
13
+
14
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
15
+ TPMapping,
16
+ compute_tp_mapping,
17
+)
18
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
19
+ NixlConnectorWorker,
20
+)
21
+from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
22
+
23
+# ======================================================================
24
+# Test fixtures / helpers
25
+# ======================================================================
26
+
27
+
28
+def _compute_mapping(
29
+ tp_rank: int = 0,
30
+ tp_size: int = 1,
31
+ remote_tp_size: int = 1,
32
+ is_mla: bool = False,
33
+ num_kv_heads: int = 8,
34
+ group_spec_types: tuple[type, ...] = (FullAttentionSpec,),
35
+) -> TPMapping:
36
+ return compute_tp_mapping(
37
+ tp_rank=tp_rank,
38
+ tp_size=tp_size,
39
+ remote_tp_size=remote_tp_size,
40
+ is_mla=is_mla,
41
+ total_num_kv_heads=num_kv_heads,
42
+ group_spec_types=group_spec_types,
43
+ )
44
+
45
+
46
+# ======================================================================
47
+# TP mapping structure tests
48
+# ======================================================================
49
+
50
+
51
+class TestTPMappingStructure:
52
+ def test_source_ranks_homogeneous(self):
53
+ m = _compute_mapping(tp_size=2, tp_rank=1, remote_tp_size=2)
54
+ assert m.all_source_ranks == (1,)
55
+
56
+ def test_source_ranks_d_gt_p(self):
57
+ m = _compute_mapping(tp_size=4, tp_rank=2, remote_tp_size=2)
58
+ assert m.all_source_ranks == (1,)
59
+
60
+ def test_source_ranks_p_gt_d(self):
61
+ m = _compute_mapping(tp_size=1, tp_rank=0, remote_tp_size=2)
62
+ assert m.all_source_ranks == (0, 1)
63
+
64
+
65
+# ======================================================================
66
+# Split handle tests
67
+# ======================================================================
68
+
69
+
70
+def _make_mock_worker_for_splits(group_spec_types):
71
+ """Build a mock NixlConnectorWorker with _group_spec_types for split tests."""
72
+ worker = object.__new__(NixlConnectorWorker)
73
+ worker._group_spec_types = group_spec_types
74
+ return worker
75
+
76
+
77
+class TestBuildSrcSplitHandles:
78
+ @pytest.mark.parametrize("remote_tp_size", [2, 4])
79
+ def test_build_src_split_handles(self, remote_tp_size):
80
+ tp_rank = 0
81
+ tp_size = 1
82
+
83
+ plan = _compute_mapping(
84
+ tp_rank=tp_rank,
85
+ tp_size=tp_size,
86
+ remote_tp_size=remote_tp_size,
87
+ )
88
+
89
+ worker = _make_mock_worker_for_splits((FullAttentionSpec,))
90
+ src_blocks_data = [(0x2000 + i * 1024, 1024, 0) for i in range(8)]
91
+ num_descs = len(src_blocks_data)
92
+ splits = list(
93
+ worker._build_local_splits_from_plan(
94
+ plan,
95
+ src_blocks_data,
96
+ num_descs,
97
+ )
98
+ )
99
+
100
+ assert len(splits) == remote_tp_size
101
+ for handle in splits:
102
+ assert len(handle) == len(src_blocks_data)
103
+ for _, length, _ in handle:
104
+ assert length == 1024 // remote_tp_size
105
+
106
+
107
+class TestMambaPlanSplitHandles:
108
+ """Verify split handles for Mamba with FA/SSM distinction."""
109
+
110
+ def test_fa_and_ssm_different_split_factors(self):
111
+ """Section 0 split by num_attn_reads, section 1 by abs_tp."""
112
+ fa_readers = (0,)
113
+ ssm_readers = (0, 1)
114
+ plan = TPMapping(
115
+ source_ranks_per_group=(fa_readers, ssm_readers),
116
+ all_source_ranks=(0, 1),
117
+ rank_to_attention_slot={0: 0, 1: 0},
118
+ rank_offset_factor=0,
119
+ )
120
+
121
+ worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec))
122
+ # 2 FA descs + 1 SSM desc
123
+ src_blocks_data = [
124
+ (1000, 200, 0), # FA desc 0
125
+ (2000, 200, 0), # FA desc 1
126
+ (3000, 400, 0), # SSM desc 0
127
+ ]
128
+
129
+ splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2))
130
+
131
+ assert len(splits) == 2 # 2 source ranks
132
+
133
+ # Rank 0 (FA source, p_idx=0):
134
+ # FA: chunk=200//1=200, slot=0 → (1000, 200, 0), (2000, 200, 0)
135
+ # SSM: chunk=400//2=200, idx=0 → (3000, 200, 0)
136
+ assert splits[0] == [(1000, 200, 0), (2000, 200, 0), (3000, 200, 0)]
137
+
138
+ # Rank 1 (not FA source, p_idx=1):
139
+ # FA: chunk=200//1=200, slot=0 (skip_fa) → (1000, 200, 0), (2000, 200, 0)
140
+ # SSM: chunk=400//2=200, idx=1 → (3200, 200, 0)
141
+ assert splits[1] == [(1000, 200, 0), (2000, 200, 0), (3200, 200, 0)]
vllm/distributed/kv_transfer/kv_connector/utils.py
+6
-375
@@ -386,38 +386,6 @@ class EngineTransferInfo:
386
"""Physical blocks per logical block."""
387
388
389
-@dataclass(frozen=True)
390
-class MambaEngineTransferInfo(EngineTransferInfo):
391
- """Extends ``EngineTransferInfo`` with Mamba-hybrid transfer geometry.
392
-
393
- For hybrid SSM+Attention models, FA and Mamba layers may require
394
- different numbers of reads from different remote ranks. This
395
- dataclass captures that per-engine transfer plan.
396
- """
397
-
398
- remote_fa_source_ranks: tuple[int, ...]
399
- """Remote ranks carrying unique FA heads for this local rank."""
400
-
401
- remote_all_source_ranks: tuple[int, ...]
402
- """All remote ranks this local rank reads from (FA + Mamba)."""
403
-
404
- remote_num_fa_reads: int
405
- """Number of distinct remote ranks needed for FA data."""
406
-
407
- remote_num_mamba_reads: int
408
- """Number of distinct remote ranks needed for Mamba data."""
409
-
410
- remote_fa_descriptor_bytes: int
411
- """Byte size of one FA K (or V) descriptor entry."""
412
-
413
- is_remote_replicated: bool
414
- """Whether the remote engine has replicated KV heads
415
- (remote_tp_size > total_num_kv_heads)."""
416
-
417
- remote_physical_heads: int
418
- """Physical KV heads stored per remote rank."""
419
-
420
-
389
# ---- Transfer topology ----
390
391
@@ -439,8 +407,6 @@ class TransferTopology:
407
self.local_physical_heads = max(1, self.total_num_kv_heads // self.tp_size)
408
409
self._engines: dict[EngineId, EngineTransferInfo] = {}
442
- self._fa_source_sets: dict[EngineId, frozenset[int]] = {}
443
- self._fa_source_indices: dict[EngineId, dict[int, int]] = {}
410
411
# Figure out whether the first dimension of the cache is K/V
412
# or num_blocks.
@@ -487,24 +453,12 @@ class TransferTopology:
453
def register_remote_engine(
454
self,
455
remote_engine_id: EngineId,
490
- remote_tp_size: int,
491
- remote_block_size: int,
492
- remote_block_len: int,
493
- remote_physical_blocks_per_logical: int,
494
- *,
495
- local_block_len: int = 0,
456
+ info: EngineTransferInfo,
457
) -> EngineTransferInfo:
458
"""Register a remote engine, unifying worker dicts state.
459
499
- Only remote engines should be registered here — the local engine's
500
- identity (tp_size, block_size, etc.) is set via ``__init__`` params.
501
-
502
- For Mamba models, also computes the Mamba transfer plan and
503
- builds the FA source lookup caches.
504
-
505
- Args:
506
- local_block_len: Local representative block_len (bytes).
507
- Required for Mamba models to compute ``fa_descriptor_bytes``.
460
+ The caller (worker) is responsible for computing the info via
461
+ the transfer policy. This method only stores and deduplicates.
462
"""
463
assert remote_engine_id != self.engine_id, (
464
f"Cannot register local engine {self.engine_id} as remote. "
@@ -512,29 +466,6 @@ class TransferTopology:
466
)
467
if remote_engine_id in self._engines:
468
return self._engines[remote_engine_id]
515
- info: EngineTransferInfo
516
- if self.is_mamba:
517
- info = self._build_mamba_info(
518
- remote_tp_size=remote_tp_size,
519
- remote_block_size=remote_block_size,
520
- remote_block_len=remote_block_len,
521
- remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical),
522
- local_block_len=local_block_len,
523
- )
524
- assert isinstance(info, MambaEngineTransferInfo)
525
- self._fa_source_sets[remote_engine_id] = frozenset(
526
- info.remote_fa_source_ranks
527
- )
528
- self._fa_source_indices[remote_engine_id] = {
529
- r: i for i, r in enumerate(info.remote_fa_source_ranks)
530
- }
531
- else:
532
- info = EngineTransferInfo(
533
- remote_tp_size=remote_tp_size,
534
- remote_block_len=remote_block_len,
535
- remote_block_size=remote_block_size,
536
- remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical),
537
- )
469
self._engines[remote_engine_id] = info
470
return info
471
@@ -622,14 +553,8 @@ class TransferTopology:
553
"""Get the remote TP rank(s) that the current local TP rank will
554
read from. When remote tp_size > local tp_size, reads from
555
multiple remote ranks.
625
-
626
- For Mamba models, returns the precomputed ``all_source_ranks``
627
- (FA + Mamba union).
556
"""
557
info = self._engines[remote_engine_id]
630
- if isinstance(info, MambaEngineTransferInfo):
631
- return list(info.remote_all_source_ranks)
632
-
558
tp_ratio = self.tp_ratio(info.remote_tp_size)
559
if tp_ratio > 0:
560
return [self.tp_rank // tp_ratio]
@@ -662,309 +587,15 @@ class TransferTopology:
587
# Regular case: backends like FA register K/V in separate regions
588
return cache if self.split_k_and_v else [cache]
589
665
- # ============================================================
666
- # Mamba-specific methods
667
- # ============================================================
668
-
669
- def should_skip_fa(self, remote_engine_id: EngineId, remote_rank: int) -> bool:
670
- """Whether to skip FA groups for this remote rank (mamba-only)."""
671
- return remote_rank not in self._fa_source_sets[remote_engine_id]
672
-
673
- def fa_head_slot(self, remote_engine_id: EngineId, remote_rank: int) -> int:
674
- """Index into local FA block for this remote rank's head data.
675
-
676
- For remote ranks in ``fa_source_ranks``, returns 0, 1, …, reads-1.
677
- For ranks NOT in ``fa_source_ranks`` (replicated duplicates),
678
- returns the slot of the matching source rank with the same head.
679
- """
680
- fa_index = self._fa_source_indices[remote_engine_id]
681
- if remote_rank in fa_index:
682
- return fa_index[remote_rank]
683
- mamba_info = self._engines[remote_engine_id]
684
- assert isinstance(mamba_info, MambaEngineTransferInfo)
685
- K = self.total_num_kv_heads
686
- remote_tp = mamba_info.remote_tp_size
687
- r_head = self._physical_head_range(remote_tp, K, remote_rank)
688
- for target in mamba_info.remote_fa_source_ranks:
689
- t_head = self._physical_head_range(remote_tp, K, target)
690
- if self._range_overlap(r_head, t_head):
691
- return fa_index[target]
692
- return 0
693
-
694
- def fa_rank_offset(
695
- self, remote_engine_id: EngineId, remote_kv_block_len: int
696
- ) -> int:
697
- """Byte offset into remote FA block for this local rank.
698
-
699
- When local TP is replicated (local_tp > K), multiple local ranks
700
- share a head. Computes offset *relative to the target remote
701
- rank's first head* so it works regardless of how many heads the
702
- remote has. Returns 0 when local does not index into remote.
703
- """
704
- mamba_info = self._engines[remote_engine_id]
705
- assert isinstance(mamba_info, MambaEngineTransferInfo)
706
- tp_ratio = self.tp_ratio(mamba_info.remote_tp_size)
707
- if self.is_mla or tp_ratio <= 0:
708
- return 0
709
- K = self.total_num_kv_heads
710
- is_local_replicated = self.tp_size > K
711
- if is_local_replicated:
712
- local_head = self.tp_rank * K // self.tp_size
713
- p_rank = mamba_info.remote_fa_source_ranks[0]
714
- p_start = p_rank * K // mamba_info.remote_tp_size
715
- return (local_head - p_start) * remote_kv_block_len
716
- return self.tp_rank % tp_ratio * remote_kv_block_len
717
-
718
- def needs_split_handles(self, remote_engine_id: EngineId) -> bool:
719
- """Whether per-remote-rank split handles are needed.
720
-
721
- True when FA and mamba have different read counts, requiring
722
- different splitting factors in the local handle.
723
- """
724
- mamba_info = self._engines[remote_engine_id]
725
- assert isinstance(mamba_info, MambaEngineTransferInfo)
726
- tp_ratio = self.tp_ratio(mamba_info.remote_tp_size)
727
- return (
728
- tp_ratio < 0
729
- and not self.is_mla
730
- and len(mamba_info.remote_all_source_ranks) > 1
731
- )
732
-
733
- def compute_split_handle_data(
734
- self,
735
- remote_engine_id: EngineId,
736
- src_blocks_data: list[tuple[int, int, int]],
737
- num_fa_descs: int,
738
- abs_tp: int,
739
- ) -> list[list[tuple[int, int, int]]]:
740
- """Per-remote-rank (addr, len, dev) triples for Mamba-HMA split
741
- handles.
742
-
743
- FA descriptors (indices < num_fa_descs) are sliced by
744
- ``remote_num_fa_reads``; mamba descriptors are sliced uniformly
745
- by ``abs_tp``.
746
- """
747
- mamba_info = self._engines[remote_engine_id]
748
- assert isinstance(mamba_info, MambaEngineTransferInfo)
749
- all_handle_data: list[list[tuple[int, int, int]]] = []
750
- for p_idx, p_rank in enumerate(mamba_info.remote_all_source_ranks):
751
- handle_data: list[tuple[int, int, int]] = []
752
- skip_fa = self.should_skip_fa(remote_engine_id, p_rank)
753
- fa_slot = self.fa_head_slot(remote_engine_id, p_rank) if not skip_fa else 0
754
- for j, (addr, local_len, dev) in enumerate(src_blocks_data):
755
- if j < num_fa_descs:
756
- assert mamba_info.remote_num_fa_reads >= 1
757
- fa_chunk = local_len // mamba_info.remote_num_fa_reads
758
- handle_data.append((addr + fa_slot * fa_chunk, fa_chunk, dev))
759
- else:
760
- mamba_chunk = local_len // abs_tp
761
- handle_data.append((addr + p_idx * mamba_chunk, mamba_chunk, dev))
762
- all_handle_data.append(handle_data)
763
- return all_handle_data
764
-
765
- def filter_block_ids_for_rank(
766
- self,
767
- remote_engine_id: EngineId,
768
- remote_rank: int,
769
- local_ids: BlockIds,
770
- remote_ids: BlockIds,
771
- is_mamba_group: list[bool],
772
- ) -> tuple[BlockIds, BlockIds]:
773
- """Zero out FA groups for remote ranks outside ``fa_source_ranks``.
774
-
775
- Returns (filtered_local_ids, filtered_remote_ids). When the
776
- remote rank carries FA data for this local rank, returns the
777
- inputs unchanged.
778
- """
779
- if not self.should_skip_fa(remote_engine_id, remote_rank):
780
- return local_ids, remote_ids
781
- num_groups = len(local_ids)
782
- filtered_local: list[list[int]] = [
783
- [] if not is_mamba_group[g] else local_ids[g] for g in range(num_groups)
784
- ]
785
- filtered_remote: list[list[int]] = [
786
- [] if not is_mamba_group[g] else remote_ids[g] for g in range(num_groups)
787
- ]
788
- return filtered_local, filtered_remote
789
-
590
def describe(self, remote_engine_id: EngineId) -> str:
591
"""One-line summary of transfer config for logging."""
592
info = self._engines[remote_engine_id]
793
- base = (
593
+ return (
594
+ f"TransferTopology("
595
f"tp_ratio={self.tp_ratio(info.remote_tp_size)}, "
596
f"K={self.total_num_kv_heads}, "
597
f"local_tp={self.tp_size}, "
598
f"remote_tp={info.remote_tp_size}, "
599
f"local_rank={self.tp_rank}, "
799
- f"remote_block_len={info.remote_block_len}"
800
- )
801
- if isinstance(info, MambaEngineTransferInfo):
802
- return (
803
- f"TransferTopology.mamba({base}, "
804
- f"fa_reads={info.remote_num_fa_reads}, "
805
- f"mamba_reads={info.remote_num_mamba_reads}, "
806
- f"fa_sources={list(info.remote_fa_source_ranks)}, "
807
- f"all_sources={list(info.remote_all_source_ranks)}, "
808
- f"fa_desc_bytes={info.remote_fa_descriptor_bytes})"
809
- )
810
- return f"TransferTopology({base})"
811
-
812
- # ============================================================
813
- # Private helpers
814
- # ============================================================
815
- # Mamba-HMA hetero-TP transfer config:
816
- # With hetero-TP (P_TP > D_TP), FA KV cache may be replicated across
817
- # P ranks (when P_TP > num_kv_heads), but Mamba conv/SSM state is
818
- # almost always uniquely sharded per P rank. So the number of P
819
- # ranks D must read from can differ between FA and Mamba, and they
820
- # must be handled separately.
821
-
822
- @staticmethod
823
- def _physical_head_range(tp_size: int, num_heads: int, rank: int) -> range:
824
- """Physical KV head range stored in a rank's KV cache tensor.
825
-
826
- When ``tp_size <= num_heads``: sharded, K/TP contiguous heads per rank.
827
- When ``tp_size > num_heads``: 1 physical head per rank. Heads are
828
- distributed **contiguously** (matching vLLM's GQA weight partitioning):
829
- consecutive ranks share a head before moving to the next one.
830
- """
831
- if tp_size <= num_heads:
832
- assert num_heads % tp_size == 0
833
- per_rank = num_heads // tp_size
834
- return range(rank * per_rank, (rank + 1) * per_rank)
835
- else:
836
- h = rank * num_heads // tp_size
837
- return range(h, h + 1)
838
-
839
- @staticmethod
840
- def _range_overlap(a: range, b: range) -> range:
841
- start = max(a.start, b.start)
842
- stop = min(a.stop, b.stop)
843
- return range(start, max(start, stop))
844
-
845
- # ============================================================
846
- # Private: build Mamba transfer info
847
- # ============================================================
848
-
849
- def _build_mamba_info(
850
- self,
851
- remote_tp_size: int,
852
- remote_block_size: int,
853
- remote_block_len: int,
854
- remote_physical_blocks_per_logical: int,
855
- local_block_len: int,
856
- ) -> MambaEngineTransferInfo:
857
- """Compute Mamba transfer plan."""
858
- K = self.total_num_kv_heads
859
- local_tp = self.tp_size
860
- local_rank = self.tp_rank
861
-
862
- is_remote_replicated = remote_tp_size > K
863
- remote_physical_heads = max(1, K // remote_tp_size)
864
-
865
- if local_tp >= remote_tp_size:
866
- assert local_tp % remote_tp_size == 0
867
- tp_ratio = local_tp // remote_tp_size
868
- else:
869
- assert remote_tp_size % local_tp == 0
870
- tp_ratio = -(remote_tp_size // local_tp)
871
-
872
- abs_tp = -tp_ratio if tp_ratio < 0 else 1
873
-
874
- mamba_range: range | None = None
875
- if tp_ratio < 0:
876
- mamba_range = range(local_rank * abs_tp, (local_rank + 1) * abs_tp)
877
-
878
- # ---- FA read targets ----
879
- if self.is_mla or tp_ratio >= 0:
880
- num_fa_reads = 1
881
- fa_source_ranks: list[int] = (
882
- [0]
883
- if self.is_mla
884
- else [local_rank // tp_ratio if tp_ratio > 0 else local_rank]
885
- )
886
- else:
887
- local_needs = self._physical_head_range(local_tp, K, local_rank)
888
- search_range = (
889
- mamba_range if mamba_range is not None else range(remote_tp_size)
890
- )
891
- seen: set[tuple[int, int]] = set()
892
- fa_source_ranks = []
893
- for p in search_range:
894
- p_has = self._physical_head_range(remote_tp_size, K, p)
895
- ov = self._range_overlap(local_needs, p_has)
896
- if len(ov) > 0:
897
- key = (ov.start, ov.stop)
898
- if key not in seen:
899
- seen.add(key)
900
- fa_source_ranks.append(p)
901
- if not fa_source_ranks:
902
- for p in range(remote_tp_size):
903
- p_has = self._physical_head_range(remote_tp_size, K, p)
904
- ov = self._range_overlap(local_needs, p_has)
905
- if len(ov) > 0:
906
- key = (ov.start, ov.stop)
907
- if key not in seen:
908
- seen.add(key)
909
- fa_source_ranks.append(p)
910
- num_fa_reads = len(fa_source_ranks)
911
-
912
- # ---- All source ranks (mamba + FA) ----
913
- if mamba_range is not None and abs_tp > num_fa_reads:
914
- num_mamba_reads = abs_tp
915
- all_source_ranks = list(mamba_range)
916
- else:
917
- num_mamba_reads = num_fa_reads
918
- all_source_ranks = list(fa_source_ranks)
919
-
920
- # ---- FA descriptor bytes ----
921
- effective_block_len = min(local_block_len, remote_block_len)
922
- if self.is_kv_layout_blocks_first:
923
- fa_descriptor_bytes = effective_block_len // 2
924
- else:
925
- fa_descriptor_bytes = effective_block_len
926
-
927
- # ---- Validation ----
928
- is_local_replicated = local_tp > K
929
- if is_local_replicated and is_remote_replicated and tp_ratio > 0:
930
- logger.info(
931
- "Both-replicated hetero-TP: local_tp=%d > remote_tp=%d > K=%d.",
932
- local_tp,
933
- remote_tp_size,
934
- K,
935
- )
936
- tt_set = set(all_source_ranks)
937
- for t in fa_source_ranks:
938
- if t not in tt_set:
939
- logger.error(
940
- "FA source rank %d NOT in all_source_ranks %s.",
941
- t,
942
- all_source_ranks,
943
- )
944
- if self.is_kv_layout_blocks_first and tp_ratio < 0 and num_fa_reads > 0:
945
- local_k_half = local_block_len // 2
946
- remote_k_half = remote_block_len // 2
947
- expected = local_k_half // num_fa_reads
948
- if expected != remote_k_half:
949
- logger.warning(
950
- "FA size mismatch: local_k_half=%d / reads=%d = %d, "
951
- "but remote_k_half=%d.",
952
- local_k_half,
953
- num_fa_reads,
954
- expected,
955
- remote_k_half,
956
- )
957
-
958
- return MambaEngineTransferInfo(
959
- remote_tp_size=remote_tp_size,
960
- remote_block_len=remote_block_len,
961
- remote_block_size=remote_block_size,
962
- remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical),
963
- remote_fa_source_ranks=tuple(fa_source_ranks),
964
- remote_all_source_ranks=tuple(all_source_ranks),
965
- remote_num_fa_reads=num_fa_reads,
966
- remote_num_mamba_reads=num_mamba_reads,
967
- remote_fa_descriptor_bytes=fa_descriptor_bytes,
968
- is_remote_replicated=is_remote_replicated,
969
- remote_physical_heads=remote_physical_heads,
600
+ f"remote_block_len={info.remote_block_len})"
601
)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+3
-1
@@ -32,8 +32,9 @@ GET_META_MSG = b"get_meta_msg"
32
# Version History:
33
# 1: Initial version with compatibility checking
34
# 2: Add remote_request_id to kv_transfer_params
35
+# 3: Add physical_blocks_per_logical_kv_block to NixlAgentMetadata
36
#
36
-NIXL_CONNECTOR_VERSION: int = 2
37
+NIXL_CONNECTOR_VERSION: int = 3
38
39
40
@dataclass
@@ -48,6 +49,7 @@ class NixlAgentMetadata:
49
block_size: int
50
ssm_sizes: tuple[int, int]
51
attn_backend_name: str
52
+ physical_blocks_per_logical_kv_block: int
53
54
55
@dataclass
vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py
new
+142
@@ -0,0 +1,142 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""TP mapping computation for NIXL KV cache transfers."""
4
+
5
+from __future__ import annotations
6
+
7
+from dataclasses import dataclass
8
+
9
+import numpy as np
10
+
11
+from vllm.distributed.kv_transfer.kv_connector.utils import (
12
+ BlockIds,
13
+)
14
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, MambaSpec
15
+
16
+# ======================================================================
17
+# Data structures
18
+# ======================================================================
19
+
20
+
21
+@dataclass(frozen=True)
22
+class ReadSpec:
23
+ """Specification for a single remote block read operation."""
24
+
25
+ remote_rank: int
26
+ local_block_ids: BlockIds
27
+ remote_block_ids: BlockIds
28
+
29
+
30
+def _is_attention_spec(spec_type: type[KVCacheSpec]) -> bool:
31
+ return issubclass(spec_type, AttentionSpec)
32
+
33
+
34
+def _is_ssm_spec(spec_type: type[KVCacheSpec]) -> bool:
35
+ return issubclass(spec_type, MambaSpec)
36
+
37
+
38
+@dataclass(frozen=True)
39
+class TPMapping:
40
+ """Complete local-to-remote TP mapping for one remote engine.
41
+
42
+ Generated once per remote engine during handshake.
43
+ """
44
+
45
+ # Remote TP ranks that this local rank reads from, per group.
46
+ # Position = local piece index.
47
+ source_ranks_per_group: tuple[tuple[int, ...], ...]
48
+
49
+ # Superset of all source ranks (union of all groups).
50
+ all_source_ranks: tuple[int, ...]
51
+
52
+ # Maps each source rank to its FA head slot index.
53
+ rank_to_attention_slot: dict[int, int]
54
+
55
+ # FA head offset factor for hetero-TP (D_TP > P_TP).
56
+ rank_offset_factor: int
57
+
58
+
59
+# ======================================================================
60
+# TP mapping computation
61
+# ======================================================================
62
+
63
+
64
+def compute_tp_mapping(
65
+ tp_rank: int,
66
+ tp_size: int,
67
+ remote_tp_size: int,
68
+ is_mla: bool,
69
+ total_num_kv_heads: int,
70
+ group_spec_types: tuple[type[KVCacheSpec], ...],
71
+) -> TPMapping:
72
+ """Build the complete local-to-remote TP mapping.
73
+
74
+ Computes source ranks, head slot assignments, and the rank offset
75
+ factor in a single pass.
76
+ """
77
+ # --- Attention source ranks ---
78
+ if is_mla:
79
+ # All heads replicated across all ranks.
80
+ attn_ranks = [0]
81
+ elif tp_size >= remote_tp_size:
82
+ # D (local TP) > P (remote TP): multiple local ranks read different chunks from
83
+ # *one* remote rank, corresponding to different kv heads.
84
+ attn_ranks = [tp_rank * remote_tp_size // tp_size]
85
+ else:
86
+ # P (remote TP) > D (local TP): one local rank
87
+ # reads from multiple remote ranks.
88
+ # GQA dedup: when K < remote_tp_size, several remote ranks
89
+ # hold the same KV head. np.unique keeps only the first
90
+ # rank per unique head so we don't issue redundant reads.
91
+ abs_tp = remote_tp_size // tp_size
92
+ start = tp_rank * abs_tp
93
+ heads = np.arange(start, start + abs_tp) * total_num_kv_heads // remote_tp_size
94
+ _, unique_idx = np.unique(heads, return_index=True)
95
+ attn_ranks = (start + np.sort(unique_idx)).tolist()
96
+
97
+ # --- SSM source ranks ---
98
+ has_ssm = any(_is_ssm_spec(t) for t in group_spec_types)
99
+ if has_ssm:
100
+ if tp_size < remote_tp_size:
101
+ abs_tp = remote_tp_size // tp_size
102
+ ssm_ranks = list(range(tp_rank * abs_tp, (tp_rank + 1) * abs_tp))
103
+ else:
104
+ ssm_ranks = list(attn_ranks)
105
+ else:
106
+ ssm_ranks = []
107
+
108
+ all_ranks = sorted(set(attn_ranks) | set(ssm_ranks))
109
+
110
+ # --- Per-group ordered source ranks ---
111
+ source_ranks_per_group = tuple(
112
+ tuple(ssm_ranks) if _is_ssm_spec(t) else tuple(attn_ranks)
113
+ for t in group_spec_types
114
+ )
115
+
116
+ # --- Attention head slots ---
117
+ head_to_slot: dict[int, int] = {}
118
+ for i, r in enumerate(attn_ranks):
119
+ head_to_slot[r * total_num_kv_heads // remote_tp_size] = i
120
+ rank_to_attention_slot = {
121
+ r: head_to_slot.get(r * total_num_kv_heads // remote_tp_size, 0)
122
+ for r in all_ranks
123
+ }
124
+
125
+ # --- Rank offset factor ---
126
+ if is_mla or tp_size <= remote_tp_size:
127
+ # We don't index into remote for reading, no offset needed.
128
+ rank_offset_factor = 0
129
+ elif tp_size > total_num_kv_heads:
130
+ local_head = tp_rank * total_num_kv_heads // tp_size
131
+ p_start = attn_ranks[0] * total_num_kv_heads // remote_tp_size
132
+ rank_offset_factor = local_head - p_start
133
+ else:
134
+ # D TP > P TP: we index into remote to read different heads depending on rank.
135
+ rank_offset_factor = tp_rank % (tp_size // remote_tp_size)
136
+
137
+ return TPMapping(
138
+ source_ranks_per_group=source_ranks_per_group,
139
+ all_source_ranks=tuple(all_ranks),
140
+ rank_to_attention_slot=rank_to_attention_slot,
141
+ rank_offset_factor=rank_offset_factor,
142
+ )
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+349
-417
@@ -9,6 +9,7 @@ import threading
9
import time
10
import uuid
11
from collections import defaultdict
12
+from collections.abc import Iterator
13
from concurrent.futures import Future, ThreadPoolExecutor
14
from typing import TYPE_CHECKING, Any, cast
15
@@ -21,7 +22,7 @@ from vllm import envs
22
from vllm.distributed.kv_transfer.kv_connector.utils import (
23
BlockIds,
24
EngineId,
24
- MambaEngineTransferInfo,
25
+ EngineTransferInfo,
26
TransferTopology,
27
get_current_attn_backends,
28
kv_postprocess_blksize_and_layout_on_receive,
@@ -43,13 +44,19 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
44
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
45
NixlKVConnectorStats,
46
)
47
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
48
+ ReadSpec,
49
+ TPMapping,
50
+ _is_attention_spec,
51
+ _is_ssm_spec,
52
+ compute_tp_mapping,
53
+)
54
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
55
_NIXL_SUPPORTED_DEVICE,
56
zmq_ctx,
57
)
58
from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
59
MambaConvSplitInfo,
52
- compute_physical_blocks_per_logical,
60
derive_mamba_conv_split,
61
)
62
from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
@@ -58,7 +65,6 @@ from vllm.distributed.parallel_state import (
65
get_tensor_model_parallel_world_size,
66
)
67
from vllm.logger import init_logger
61
-from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
68
from vllm.platforms import current_platform
69
from vllm.utils.network_utils import make_zmq_path
70
from vllm.v1.attention.backends.utils import get_kv_cache_layout
@@ -80,6 +86,111 @@ logger = init_logger(__name__)
86
class NixlConnectorWorker:
87
"""Implementation of Worker side methods"""
88
89
+ def _compute_desc_ids(
90
+ self,
91
+ block_ids: BlockIds,
92
+ dst_num_blocks: int,
93
+ block_size_ratio: float | None,
94
+ physical_blocks_per_logical: int,
95
+ ) -> np.ndarray:
96
+ """Compute NIXL descriptor IDs for given block IDs."""
97
+ num_fa_regions = self.num_regions
98
+ num_ssm_regions = len(self.block_len_per_layer) * 4 if self._has_mamba else 0
99
+
100
+ num_blocks = dst_num_blocks
101
+ if block_size_ratio is not None:
102
+ num_blocks = int(num_blocks * block_size_ratio)
103
+ ratio = physical_blocks_per_logical
104
+ logical_blocks = num_blocks // ratio
105
+
106
+ num_fa_descs = num_fa_regions * num_blocks
107
+
108
+ # All-attention fast path: single vectorized broadcast.
109
+ if num_ssm_regions == 0:
110
+ block_arr = np.concatenate(block_ids)[None, :]
111
+ region_ids = np.arange(num_fa_regions)[:, None]
112
+ return (region_ids * num_blocks + block_arr).flatten()
113
+
114
+ # NOTE (NickLucche) With HMA, every kv group has the same number
115
+ # of layers and layers from different groups share the same kv
116
+ # tensor. Therefore we compute desc IDs per group using the
117
+ # right stride:
118
+ # FA descs have num_blocks entries per region (kernel granularity),
119
+ # SSM descs have logical_blocks entries per region (no kernel
120
+ # splitting).
121
+ all_descs: list[np.ndarray] = []
122
+ for i, group in enumerate(block_ids):
123
+ group_arr = np.asarray(group)
124
+ if _is_attention_spec(self._group_spec_types[i]):
125
+ fa_region_ids = np.arange(num_fa_regions)[:, None]
126
+ all_descs.append(
127
+ (fa_region_ids * num_blocks + group_arr[None, :]).flatten()
128
+ )
129
+ elif _is_ssm_spec(self._group_spec_types[i]):
130
+ # NOTE (NickLucche) SSM and Attention block regions can
131
+ # be exchanged arbitrarily by manager. Therefore, descs
132
+ # are laid out as:
133
+ # [descs_fa (all regions) | descs_ssm (all regions)].
134
+ # num_fa_descs offset must be computed per-engine since
135
+ # P and D can have different num_blocks (and thus
136
+ # different FA desc counts).
137
+ ssm_region_ids = np.arange(num_ssm_regions)[:, None]
138
+ all_descs.append(
139
+ (
140
+ ssm_region_ids * logical_blocks
141
+ + group_arr[None, :]
142
+ + num_fa_descs
143
+ ).flatten()
144
+ )
145
+ else:
146
+ raise ValueError(
147
+ f"Unknown spec type {self._group_spec_types[i]} at index {i}"
148
+ )
149
+
150
+ return np.concatenate(all_descs)
151
+
152
+ def _build_local_splits_from_plan(
153
+ self,
154
+ plan: TPMapping,
155
+ src_blocks_data: list[tuple[int, int, int]],
156
+ num_fa_descs: int,
157
+ ) -> Iterator[list[tuple[int, int, int]]]:
158
+ """Build split handle data for P_TP > D_TP scenario.
159
+
160
+ num_fa_descs is the boundary between FA and SSM descriptors.
161
+ Split counts are derived from source_ranks_per_group lengths.
162
+ FA uses rank_to_attention_slot for the slot offset;
163
+ SSM uses the rank's positional index.
164
+ """
165
+ fa_idx = next(
166
+ i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t)
167
+ )
168
+ fa_num_splits = len(plan.source_ranks_per_group[fa_idx])
169
+
170
+ has_ssm_descs = num_fa_descs < len(src_blocks_data)
171
+ ssm_idx = next(
172
+ (i for i, t in enumerate(self._group_spec_types) if _is_ssm_spec(t)),
173
+ None,
174
+ )
175
+ ssm_num_splits = (
176
+ len(plan.source_ranks_per_group[ssm_idx])
177
+ if has_ssm_descs and ssm_idx is not None
178
+ else 0
179
+ )
180
+
181
+ for p_idx, p_rank in enumerate(plan.all_source_ranks):
182
+ fa_slot = plan.rank_to_attention_slot.get(p_rank, 0)
183
+
184
+ handle: list[tuple[int, int, int]] = []
185
+ for j, (addr, local_len, dev) in enumerate(src_blocks_data):
186
+ if j < num_fa_descs:
187
+ chunk = local_len // fa_num_splits
188
+ handle.append((addr + fa_slot * chunk, chunk, dev))
189
+ else:
190
+ chunk = local_len // ssm_num_splits
191
+ handle.append((addr + p_idx * chunk, chunk, dev))
192
+ yield handle
193
+
194
def __init__(
195
self,
196
vllm_config: "VllmConfig",
@@ -119,44 +230,37 @@ class NixlConnectorWorker:
230
}
231
self.hma_group_size = len(kv_cache_config.kv_cache_tensors)
232
122
- # ---- Mamba model state (derived from model config) ----
123
- self._is_mamba_group = [
124
- isinstance(group.kv_cache_spec, MambaSpec)
125
- for group in kv_cache_config.kv_cache_groups
126
- ]
233
+ # ---- Model state (derived from model config) ----
234
mamba_ssm_size = (0, 0)
128
- self._has_mamba = any(self._is_mamba_group)
129
- if self._has_mamba:
130
- assert self._is_hma_required
131
- mamba_spec = next(
132
- spec
133
- for spec in self._layer_specs.values()
134
- if isinstance(spec, MambaSpec)
135
- )
136
- conv_nbytes, ssm_nbytes = (
137
- torch.tensor([], dtype=mamba_spec.dtypes[0]).element_size(), # type: ignore[misc]
138
- torch.tensor([], dtype=mamba_spec.dtypes[1]).element_size(), # type: ignore[misc]
139
- )
140
- conv_shape, ssm_shape = (
141
- torch.Size(mamba_spec.shapes[0]),
142
- torch.Size(mamba_spec.shapes[1]),
143
- )
144
- mamba_ssm_size = (
145
- conv_shape.numel() * conv_nbytes,
146
- ssm_shape.numel() * ssm_nbytes,
147
- )
148
- self._mamba_ssm_size = mamba_ssm_size
235
# Conv state sub-projection decomposition (None when no Mamba).
236
# The 3-read transfer requires DS (dim, state_len) conv layout so
237
# that x/B/C sub-projections are contiguous in memory.
238
self._conv_decomp: MambaConvSplitInfo | None = None
239
+ self._has_mamba = any(
240
+ isinstance(g.kv_cache_spec, MambaSpec)
241
+ for g in kv_cache_config.kv_cache_groups
242
+ )
243
if self._has_mamba:
244
+ assert self._is_hma_required
245
+ from vllm.model_executor.layers.mamba.mamba_utils import (
246
+ is_conv_state_dim_first,
247
+ )
248
+
249
assert is_conv_state_dim_first(), (
250
"3-read Mamba conv transfer requires DS conv state layout. "
251
"Set VLLM_SSM_CONV_STATE_LAYOUT=DS"
252
)
158
- local_tp = vllm_config.parallel_config.tensor_parallel_size
159
- self._conv_decomp = derive_mamba_conv_split(mamba_spec, local_tp)
253
+ mamba_spec = next(
254
+ spec
255
+ for spec in self._layer_specs.values()
256
+ if isinstance(spec, MambaSpec)
257
+ )
258
+ self._conv_decomp = derive_mamba_conv_split(
259
+ mamba_spec,
260
+ vllm_config.parallel_config.tensor_parallel_size,
261
+ )
262
+ mamba_ssm_size = self._conv_decomp.ssm_sizes
263
+ self._mamba_ssm_size = mamba_ssm_size
264
265
# Agent.
266
non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"]
@@ -268,14 +372,6 @@ class NixlConnectorWorker:
372
self.dst_num_blocks: dict[EngineId, int] = {}
373
self._registered_descs: list[Any] = []
374
271
- # ---- Mamba-HMA per-engine state (only used when self._has_mamba) ----
272
- # NOTE (ZhanqiuHu): _physical_blocks_per_logical MUST be per-engine.
273
- # physical_blocks_per_logical = ceil((conv_bytes + ssm_bytes) / block_len)
274
- # where conv/ssm bytes are per-TP-rank (dimension-sharded). With
275
- # heterogeneous TP the per-rank sizes differ, so the ratio differs:
276
- # e.g. Nemotron 30B: P(TP=4) → 131, D(TP=1) → 261.
277
- self._physical_blocks_per_logical: dict[EngineId, int] = {}
278
-
375
# In progress transfers.
376
# [req_id -> list[handle]]
377
self._recving_metadata: dict[ReqId, ReqMeta] = {}
@@ -330,6 +426,13 @@ class NixlConnectorWorker:
426
self._physical_blocks_per_logical_kv_block = 1
427
self._sync_block_size_with_kernel()
428
429
+ self._group_spec_types = tuple(
430
+ type(g.kv_cache_spec) for g in self.kv_cache_config.kv_cache_groups
431
+ )
432
+
433
+ # Per-engine TP mappings. Generated during handshake.
434
+ self.tp_mappings: dict[EngineId, TPMapping] = {}
435
+
436
self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
437
"enforce_handshake_compat", True
438
)
@@ -812,9 +915,6 @@ class NixlConnectorWorker:
915
self.dst_num_blocks[self.engine_id] = self.num_blocks
916
917
if self._has_mamba:
815
- self._physical_blocks_per_logical[self.engine_id] = (
816
- self._physical_blocks_per_logical_kv_block
817
- )
918
logger.info(
919
"Hybrid SSM registration: num_blocks=%s, "
920
"logical_num_blocks=%s, ratio=%s, num_regions=%s, "
@@ -847,6 +947,9 @@ class NixlConnectorWorker:
947
block_size=self.block_size,
948
ssm_sizes=self._mamba_ssm_size,
949
attn_backend_name=self.backend_name,
950
+ physical_blocks_per_logical_kv_block=(
951
+ self._physical_blocks_per_logical_kv_block
952
+ ),
953
)
954
# Wrap metadata in payload with hash for defensive decoding
955
assert self.compat_hash is not None
@@ -875,6 +978,8 @@ class NixlConnectorWorker:
978
979
result: list[tuple[int, int, int]] = []
980
for i, base_addr in enumerate(base_addresses):
981
+ # Jump one page_size, but ssm page_size may be bigger when kernel
982
+ # locks block size to a specific value (physical_per_logical scale).
983
page_stride = (
984
self.block_len_per_layer[i] // block_size_ratio * physical_per_logical
985
)
@@ -894,67 +999,11 @@ class NixlConnectorWorker:
999
)
1000
return result
1001
897
- def _build_fa_remote_for_mamba(
898
- self,
899
- nixl_agent_meta: NixlAgentMetadata,
900
- block_size_ratio: int,
901
- transfer_topo: TransferTopology,
902
- remote_engine_id: EngineId,
903
- ) -> list[tuple[int, int, int]]:
904
- """Build remote FA descriptors for mamba models.
905
-
906
- Uses TransferTopology for GQA-aware FA divisor and head-based rank
907
- offset instead of the standard uniform tp_ratio split.
908
- """
909
- assert block_size_ratio == 1, (
910
- "Mamba 3-read transfer with block_size_ratio != 1 is not tested. "
911
- f"Got block_size_ratio={block_size_ratio}."
912
- )
913
- # TODO (ZhanqiuHu): unify with register_remote_blocks when Mamba-HMA
914
- # hetero-TP logic stabilizes.
915
- mamba_info = transfer_topo.get_engine_info(remote_engine_id)
916
- assert isinstance(mamba_info, MambaEngineTransferInfo)
917
- tp_ratio = transfer_topo.tp_ratio(mamba_info.remote_tp_size)
918
- result: list[tuple[int, int, int]] = []
919
- for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
920
- local_block_len = self.get_backend_aware_kv_block_len(
921
- layer_idx=i, first_split=True, mamba_view=False
922
- )
923
- remote_kv_block_len = local_block_len // block_size_ratio
924
- if block_size_ratio > 1:
925
- local_block_len = remote_kv_block_len
926
-
927
- if tp_ratio < 0 and not self.use_mla:
928
- local_block_len = local_block_len // mamba_info.remote_num_fa_reads
929
-
930
- rank_offset = transfer_topo.fa_rank_offset(
931
- remote_engine_id, remote_kv_block_len
932
- )
933
-
934
- num_blocks = nixl_agent_meta.num_blocks
935
- page_size = nixl_agent_meta.block_lens[i]
936
- for block_id in range(num_blocks):
937
- block_offset = block_id * page_size
938
- addr = base_addr + block_offset + rank_offset
939
- result.append((addr, local_block_len, nixl_agent_meta.device_id))
940
-
941
- if transfer_topo.is_kv_layout_blocks_first:
942
- second_split = self.get_backend_aware_kv_block_len(
943
- layer_idx=i, first_split=False, mamba_view=False
944
- )
945
- if tp_ratio < 0 and not self.use_mla:
946
- second_split = second_split // mamba_info.remote_num_fa_reads
947
- for block_id in range(num_blocks):
948
- block_offset = block_id * page_size
949
- addr = base_addr + block_offset + rank_offset
950
- v_addr = addr + nixl_agent_meta.block_lens[i] // 2
951
- result.append((v_addr, second_split, nixl_agent_meta.device_id))
952
- return result
953
-
1002
def _build_mamba_remote(
1003
self,
1004
nixl_agent_meta: NixlAgentMetadata,
1005
tp_ratio: int,
1006
+ transfer_info: EngineTransferInfo,
1007
) -> list[tuple[int, int, int]]:
1008
"""Build 4 remote desc regions (x, B, C, ssm) per layer for
1009
the 3-read transfer. For hetero-TP, each D rank reads only its
@@ -974,17 +1023,15 @@ class NixlConnectorWorker:
1023
ssm_read_size = self._mamba_ssm_size[1]
1024
else:
1025
# NOTE (ZhanqiuHu): tp_ratio < 0 means P_TP > D_TP, so P pages
977
- # are smaller than D's. self._conv_decomp has D-sized dimensions,
978
- # but we need P-sized offsets. Scale down by |tp_ratio|.
1026
+ # are smaller than D's. self._conv_decomp has D-sized dimensions,
1027
+ # but we need P-sized offsets. Scale down by |tp_ratio|.
1028
abs_ratio = -tp_ratio
1029
xb_p = self._conv_decomp.x_bytes // abs_ratio
1030
bb_p = self._conv_decomp.b_bytes // abs_ratio
1031
conv_offsets = [(0, xb_p), (xb_p, bb_p), (xb_p + bb_p, bb_p)]
1032
ssm_read_size = nixl_agent_meta.ssm_sizes[1]
1033
985
- remote_physical_per_logical = self._physical_blocks_per_logical[
986
- nixl_agent_meta.engine_id
987
- ]
1034
+ remote_physical_per_logical = transfer_info.remote_physical_blocks_per_logical
1035
num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical
1036
device_id = nixl_agent_meta.device_id
1037
@@ -1007,6 +1054,91 @@ class NixlConnectorWorker:
1054
result.append((ssm_addr, ssm_read_size, device_id))
1055
return result
1056
1057
+ def _build_fa_local(
1058
+ self,
1059
+ base_addresses: list[int],
1060
+ block_size_ratio: int,
1061
+ ) -> list[tuple[int, int, int]]:
1062
+ """Build local FA descriptors for all layers."""
1063
+ assert self.transfer_topo is not None
1064
+ num_blocks = self.num_blocks * block_size_ratio
1065
+ result: list[tuple[int, int, int]] = []
1066
+ for i, base_addr in enumerate(base_addresses):
1067
+ kv_block_len = (
1068
+ self.get_backend_aware_kv_block_len(
1069
+ layer_idx=i, first_split=True, mamba_view=False
1070
+ )
1071
+ // block_size_ratio
1072
+ )
1073
+ page_stride = self.block_len_per_layer[i] // block_size_ratio
1074
+ for block_id in range(num_blocks):
1075
+ block_offset = block_id * page_stride
1076
+ addr = base_addr + block_offset
1077
+ result.append((addr, kv_block_len, self.device_id))
1078
+
1079
+ if self.transfer_topo.is_kv_layout_blocks_first:
1080
+ # Separate and interleave K/V regions to maintain the same
1081
+ # descs ordering. This is needed for selecting contiguous heads
1082
+ # when split across TP ranks.
1083
+ second_split = self.get_backend_aware_kv_block_len(
1084
+ layer_idx=i, first_split=False, mamba_view=False
1085
+ )
1086
+ for block_id in range(num_blocks):
1087
+ block_offset = block_id * page_stride
1088
+ addr = base_addr + block_offset
1089
+ v_addr = addr + kv_block_len
1090
+ result.append((v_addr, second_split, self.device_id))
1091
+ return result
1092
+
1093
+ def _build_fa_remote(
1094
+ self,
1095
+ plan: TPMapping,
1096
+ nixl_agent_meta: NixlAgentMetadata,
1097
+ block_size_ratio: int,
1098
+ ) -> list[tuple[int, int, int]]:
1099
+ """Build remote FA descriptors for all layers."""
1100
+ assert self.transfer_topo is not None
1101
+ fa_group_idx = next(
1102
+ i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t)
1103
+ )
1104
+ num_attn_reads = len(plan.source_ranks_per_group[fa_group_idx])
1105
+ num_blocks = nixl_agent_meta.num_blocks
1106
+ result: list[tuple[int, int, int]] = []
1107
+ for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
1108
+ # Read our whole local region size from remote..
1109
+ local_block_len = self.get_backend_aware_kv_block_len(
1110
+ layer_idx=i, first_split=True, mamba_view=False
1111
+ )
1112
+ remote_kv_block_len = local_block_len // block_size_ratio
1113
+ if block_size_ratio > 1:
1114
+ # ..using remote kv_block_len as transfer unit
1115
+ local_block_len = remote_kv_block_len
1116
+
1117
+ local_block_len = local_block_len // num_attn_reads
1118
+ rank_offset = plan.rank_offset_factor * remote_kv_block_len
1119
+
1120
+ page_size = nixl_agent_meta.block_lens[i]
1121
+ for block_id in range(num_blocks):
1122
+ block_offset = block_id * page_size
1123
+ # For each block, grab the kv heads chunk belonging to current local
1124
+ # tp rank of size local_block_len.
1125
+ addr = base_addr + block_offset + rank_offset
1126
+ result.append((addr, local_block_len, nixl_agent_meta.device_id))
1127
+
1128
+ if self.transfer_topo.is_kv_layout_blocks_first:
1129
+ # With FlashInfer index V separately to allow head splitting.
1130
+ second_split = self.get_backend_aware_kv_block_len(
1131
+ layer_idx=i, first_split=False, mamba_view=False
1132
+ )
1133
+ second_split = second_split // num_attn_reads
1134
+ for block_id in range(num_blocks):
1135
+ block_offset = block_id * page_size
1136
+ addr = base_addr + block_offset + rank_offset
1137
+ # Hop over the first split of remote page, K, to read V.
1138
+ v_addr = addr + nixl_agent_meta.block_lens[i] // 2
1139
+ result.append((v_addr, second_split, nixl_agent_meta.device_id))
1140
+ return result
1141
+
1142
def register_local_xfer_handler(
1143
self,
1144
block_size: int,
@@ -1023,68 +1155,23 @@ class NixlConnectorWorker:
1155
data copy correctness.
1156
"""
1157
assert self.transfer_topo is not None
1026
- transfer_topo = self.transfer_topo
1027
-
1158
block_size_ratio = self.block_size // block_size
1029
- blocks_data: list[tuple[int, int, int]] = []
1159
local_base_addresses = self.kv_caches_base_addr[self.engine_id][self.tp_rank]
1160
1032
- def register_blocks(blocks_data: list[tuple[int, int, int]], mamba: bool):
1033
- for i, base_addr in enumerate(local_base_addresses):
1034
- # The new block_len is using prefill block_len;
1035
- # and num_blocks is multiple with N
1036
- kv_block_len = (
1037
- self.get_backend_aware_kv_block_len(
1038
- layer_idx=i, first_split=True, mamba_view=mamba
1039
- )
1040
- // block_size_ratio
1041
- )
1042
- # Jump one page_size, but ssm page_size may be bigger when kernel
1043
- # locks block size to a specific value.
1044
- block_len_per_layer = (
1045
- self.block_len_per_layer[i]
1046
- // block_size_ratio
1047
- * (1 if not mamba else self._physical_blocks_per_logical_kv_block)
1048
- )
1049
- num_blocks = self._logical_num_blocks if mamba else self.num_blocks
1050
- num_blocks = num_blocks * block_size_ratio
1051
- for block_id in range(num_blocks):
1052
- block_offset = block_id * block_len_per_layer
1053
- addr = base_addr + block_offset
1054
- # (addr, len, device id)
1055
- blocks_data.append((addr, kv_block_len, self.device_id))
1056
-
1057
- if transfer_topo.is_kv_layout_blocks_first:
1058
- second_split = self.get_backend_aware_kv_block_len(
1059
- layer_idx=i, first_split=False, mamba_view=mamba
1060
- )
1061
- # Separate and interleave K/V regions to maintain the same
1062
- # descs ordering. This is needed for selecting contiguous heads
1063
- # when split across TP ranks.
1064
- for block_id in range(num_blocks):
1065
- block_offset = block_id * block_len_per_layer
1066
- addr = base_addr + block_offset
1067
- # Register addresses for V cache (K registered first).
1068
- v_addr = addr + kv_block_len
1069
- blocks_data.append((v_addr, second_split, self.device_id))
1070
- logger.debug(
1071
- "Created %s blocks for src engine %s and rank %s on device id %s",
1072
- len(blocks_data),
1073
- self.engine_id,
1074
- self.tp_rank,
1075
- self.device_id,
1076
- )
1077
-
1078
- # NOTE (ZhanqiuHu): mamba=True path in register_blocks is not used
1079
- # right now — we use _build_mamba_local instead for the 3-read
1080
- # approach. However, we might still need this as a fallback for homogeneous TP.
1081
- register_blocks(blocks_data, mamba=False)
1161
+ blocks_data = self._build_fa_local(local_base_addresses, block_size_ratio)
1162
+ logger.debug(
1163
+ "Created %s blocks for src engine %s and rank %s on device id %s",
1164
+ len(blocks_data),
1165
+ self.engine_id,
1166
+ self.tp_rank,
1167
+ self.device_id,
1168
+ )
1169
if self._has_mamba:
1170
assert self.num_descs == len(blocks_data)
1084
- # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-read split is
1085
- # unnecessary — a single conv desc per block suffices. Consider
1171
+ # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split
1172
+ # is unnecessary — a single conv desc per block suffices. Consider
1173
# adding a fast path that falls back to the standard 2-region
1087
- # registration (register_blocks mamba=True) when no hetero-TP
1174
+ # registration (_build_fa_local mamba=True) when no hetero-TP
1175
# remote has been seen. Currently we always register 4 regions
1176
# because local descs are created before knowing the remote TP.
1177
logger.debug("Registering local Mamba descriptors (4 regions/layer)")
@@ -1160,26 +1247,26 @@ class NixlConnectorWorker:
1247
assert self.transfer_topo is not None
1248
transfer_topo = self.transfer_topo
1249
physical_blocks_per_logical = (
1163
- compute_physical_blocks_per_logical(
1164
- nixl_agent_meta.ssm_sizes,
1165
- nixl_agent_meta.block_lens[0],
1166
- )
1167
- if self._has_mamba
1168
- else 1
1250
+ nixl_agent_meta.physical_blocks_per_logical_kv_block
1251
)
1170
- transfer_topo.register_remote_engine(
1171
- remote_engine_id=engine_id,
1252
+ transfer_info = EngineTransferInfo(
1253
remote_tp_size=remote_tp_size,
1254
remote_block_size=nixl_agent_meta.block_size,
1255
remote_block_len=nixl_agent_meta.block_lens[0],
1256
remote_physical_blocks_per_logical=physical_blocks_per_logical,
1176
- local_block_len=self.block_len_per_layer[0],
1257
)
1178
- if self._has_mamba and engine_id not in self._physical_blocks_per_logical:
1179
- self._physical_blocks_per_logical[engine_id] = physical_blocks_per_logical
1180
-
1258
+ transfer_topo.register_remote_engine(engine_id, transfer_info)
1259
logger.info("Transfer plan: %s", transfer_topo.describe(engine_id))
1260
1261
+ self.tp_mappings[engine_id] = compute_tp_mapping(
1262
+ transfer_topo.tp_rank,
1263
+ transfer_topo.tp_size,
1264
+ transfer_info.remote_tp_size,
1265
+ transfer_topo.is_mla,
1266
+ transfer_topo.total_num_kv_heads,
1267
+ self._group_spec_types,
1268
+ )
1269
+
1270
remote_agent_name = self.nixl_wrapper.add_remote_agent(
1271
nixl_agent_meta.agent_metadata
1272
)
@@ -1206,11 +1293,6 @@ class NixlConnectorWorker:
1293
# this is the ratio between the two sizes.
1294
tp_ratio = transfer_topo.tp_ratio(remote_tp_size)
1295
1209
- # Handle tp_size>num_kv_heads: replicate KV cache.
1210
- indexes_into_remote = (
1211
- not transfer_topo.replicates_kv_cache(engine_id) and tp_ratio > 0
1212
- )
1213
-
1296
logger.debug(
1297
"Registering remote agent (%s, rank %s) memory regions with tp_ratio %s",
1298
engine_id,
@@ -1218,6 +1300,8 @@ class NixlConnectorWorker:
1300
tp_ratio,
1301
)
1302
1303
+ plan = self.tp_mappings[engine_id]
1304
+
1305
### (Optional) Register local agent memory regions. MLA is not split.
1306
if (
1307
tp_ratio < 0
@@ -1227,149 +1311,51 @@ class NixlConnectorWorker:
1311
# Remote tp_size > local tp_size: read from multiple remote ranks.
1312
# Logically "split" own regions into |tp_ratio| chunks. Mind that
1313
# we only do this once per remote tp_size (replica-friendly).
1230
- abs_tp = -tp_ratio
1314
self.src_xfer_handles_by_tp_ratio[tp_ratio] = []
1315
1233
- if self._has_mamba:
1234
- if transfer_topo.needs_split_handles(engine_id):
1235
- # Mamba-HMA: FA and Mamba use different split factors.
1236
- for handle_data in transfer_topo.compute_split_handle_data(
1237
- engine_id, self.src_blocks_data, self.num_descs, abs_tp
1238
- ):
1239
- descs = self.nixl_wrapper.get_xfer_descs(
1240
- handle_data, self.nixl_memory_type
1241
- )
1242
- handle = self.nixl_wrapper.prep_xfer_dlist(
1243
- "NIXL_INIT_AGENT", descs
1244
- )
1245
- self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
1246
-
1247
- logger.info(
1248
- "Mamba-HMA split handles: %s, num_descs=%s",
1249
- transfer_topo.describe(engine_id),
1250
- self.num_descs,
1251
- )
1252
- else:
1253
- # Original path: uniform divide by abs_tp (non-Mamba-HMA).
1254
- for i in range(abs_tp):
1255
- blocks_data = []
1256
- for memory_region in self.src_blocks_data:
1257
- addr, local_block_len, own_tp_rank = memory_region
1258
- remote_block_len = local_block_len // abs_tp
1259
- addr = addr + i * remote_block_len
1260
- blocks_data.append((addr, remote_block_len, own_tp_rank))
1261
- descs = self.nixl_wrapper.get_xfer_descs(
1262
- blocks_data, self.nixl_memory_type
1263
- )
1264
- handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)
1265
- self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
1316
+ for handle_data in self._build_local_splits_from_plan(
1317
+ plan,
1318
+ self.src_blocks_data,
1319
+ self.num_descs,
1320
+ ):
1321
+ descs = self.nixl_wrapper.get_xfer_descs(
1322
+ handle_data, self.nixl_memory_type
1323
+ )
1324
+ handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)
1325
+ self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
1326
1327
### Register remote agent memory regions
1268
- blocks_data = []
1269
- # With homogeneous TP, D pulls the whole kv cache from corresponding
1270
- # rank. With heterogeneous TP, prepare the descriptors by splitting the
1271
- # P KV cache along kv_head dim, of D worker's kv_head size (D>P).
1328
+ # With homogeneous TP, D pulls the whole kv cache from corresponding rank. With
1329
+ # heterogeneous TP, prepare the descriptors by splitting the P KV cache along
1330
+ # kv_head dim, of D worker's kv_head size (D>P).
1331
# Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..].
1332
1333
# Register all remote blocks, but only the corresponding kv heads.
1275
- def register_remote_blocks(
1276
- blocks_data: list[tuple[int, int, int]], mamba: bool
1277
- ):
1278
- for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
1279
- # Read our whole local region size from remote.
1280
- local_block_len = self.get_backend_aware_kv_block_len(
1281
- layer_idx=i, first_split=True, mamba_view=mamba
1282
- )
1283
- remote_kv_block_len = local_block_len // block_size_ratio
1284
- if block_size_ratio > 1:
1285
- # using remote kv_block_len as transfer unit
1286
- local_block_len = remote_kv_block_len
1287
-
1288
- if tp_ratio < 0 and not self.use_mla:
1289
- # Remote tp is bigger: read a chunk of local region from remote
1290
- local_block_len = local_block_len // (-tp_ratio)
1291
- rank_offset = (
1292
- self.tp_rank % tp_ratio * remote_kv_block_len
1293
- if indexes_into_remote
1294
- else 0
1295
- )
1296
-
1297
- # Assume same num_blocks for mamba and fa
1298
- num_blocks = (
1299
- nixl_agent_meta.num_blocks
1300
- if not mamba
1301
- else nixl_agent_meta.num_blocks
1302
- // self._physical_blocks_per_logical_kv_block
1303
- )
1304
- page_size = nixl_agent_meta.block_lens[i] * (
1305
- 1 if not mamba else self._physical_blocks_per_logical_kv_block
1306
- )
1307
- for block_id in range(num_blocks):
1308
- block_offset = block_id * page_size
1309
- # For each block, grab the heads chunk belonging to rank_i
1310
- # of size remote_nheads // tp_ratio, which correspond to
1311
- # self.block_len == remote_block_len//tp_ratio bytes.
1312
- addr = base_addr + block_offset + rank_offset
1313
- # (addr, len, device id)
1314
- blocks_data.append(
1315
- (addr, local_block_len, nixl_agent_meta.device_id)
1316
- )
1317
-
1318
- if transfer_topo.is_kv_layout_blocks_first:
1319
- # With FlashInfer index V separately to allow head splitting.
1320
- second_split = self.get_backend_aware_kv_block_len(
1321
- layer_idx=i, first_split=False, mamba_view=mamba
1322
- )
1323
- # Apply the same scaling as local_block_len above for when we read
1324
- # a chunk of local V from `tp_ratio` separate remote workers.
1325
- if tp_ratio < 0 and not self.use_mla:
1326
- second_split = second_split // (-tp_ratio)
1327
- for block_id in range(num_blocks):
1328
- block_offset = block_id * page_size
1329
- addr = base_addr + block_offset + rank_offset
1330
- # Hop over the first split of remote page: either K or Conv.
1331
- if mamba:
1332
- v_addr = addr + nixl_agent_meta.ssm_sizes[0]
1333
- else:
1334
- v_addr = addr + nixl_agent_meta.block_lens[i] // 2
1335
- blocks_data.append(
1336
- (v_addr, second_split, nixl_agent_meta.device_id)
1337
- )
1338
-
1339
- logger.debug(
1340
- "Created %s blocks for dst engine %s"
1341
- " with remote rank %s and local rank %s",
1342
- len(blocks_data),
1343
- engine_id,
1344
- remote_tp_rank,
1345
- self.tp_rank,
1346
- )
1347
-
1334
+ blocks_data = self._build_fa_remote(
1335
+ plan,
1336
+ nixl_agent_meta,
1337
+ block_size_ratio,
1338
+ )
1339
+ logger.debug(
1340
+ "Created %s blocks for dst engine %s with remote rank %s and local rank %s",
1341
+ len(blocks_data),
1342
+ engine_id,
1343
+ remote_tp_rank,
1344
+ self.tp_rank,
1345
+ )
1346
if self._has_mamba:
1349
- # Mamba-HMA: separate FA registration with GQA-aware sizing,
1350
- # plus mamba 3-read registration for the Mamba "view" of the
1351
- # same KV cache tensors.
1347
logger.debug(
1348
"Registering remote Mamba blocks for engine %s rank %s",
1349
engine_id,
1350
remote_tp_rank,
1351
)
1357
- blocks_data.extend(
1358
- self._build_fa_remote_for_mamba(
1359
- nixl_agent_meta,
1360
- block_size_ratio,
1361
- transfer_topo,
1362
- engine_id,
1363
- )
1364
- )
1352
blocks_data.extend(
1353
self._build_mamba_remote(
1354
nixl_agent_meta,
1355
tp_ratio,
1356
+ transfer_info,
1357
)
1358
)
1371
- else:
1372
- register_remote_blocks(blocks_data, mamba=False)
1359
1360
# Register with NIXL.
1361
descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
@@ -1897,33 +1883,49 @@ class NixlConnectorWorker:
1883
def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
1884
assert meta.remote is not None and self.transfer_topo is not None
1885
engine_id = meta.remote.engine_id
1900
- remote_ranks = self.transfer_topo.target_remote_ranks(engine_id)
1886
+ plan = self.tp_mappings[engine_id]
1887
remote_info = self.transfer_topo.get_engine_info(engine_id)
1888
tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)
1889
1904
- if self._has_mamba:
1905
- # Expand remote logical → kernel block IDs.
1906
- meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(
1907
- meta.remote.block_ids,
1908
- self._physical_blocks_per_logical[meta.remote.engine_id],
1909
- )
1910
- else:
1911
- meta.remote.block_ids = self._logical_to_kernel_block_ids(
1912
- meta.remote.block_ids
1890
+ meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(
1891
+ meta.remote.block_ids,
1892
+ remote_info.remote_physical_blocks_per_logical,
1893
+ )
1894
+ remote_block_ids = meta.remote.block_ids
1895
+ local_block_ids = meta.local_physical_block_ids
1896
+ num_groups = len(local_block_ids)
1897
+ read_specs = [
1898
+ ReadSpec(
1899
+ remote_rank=rank,
1900
+ local_block_ids=[
1901
+ list(local_block_ids[g])
1902
+ if rank in plan.source_ranks_per_group[g]
1903
+ else []
1904
+ for g in range(num_groups)
1905
+ ],
1906
+ remote_block_ids=[
1907
+ list(remote_block_ids[g])
1908
+ if rank in plan.source_ranks_per_group[g]
1909
+ else []
1910
+ for g in range(num_groups)
1911
+ ],
1912
)
1913
+ for rank in plan.all_source_ranks
1914
+ ]
1915
+
1916
# D may have to perform multiple reads from different remote ranks.
1915
- for i, remote_rank in enumerate(remote_ranks):
1916
- if self.use_mla and tp_ratio < 0 and i > 0:
1917
- # MLA opt: when P TP > D TP, only a single read is executed for
1918
- # the first remote rank (cache is duplicated)..
1919
- break
1917
+ # MLA opt: when P TP > D TP, only a single read is executed for
1918
+ # the first remote rank (cache is duplicated).
1919
+ if self.use_mla and tp_ratio < 0:
1920
+ read_specs = read_specs[:1]
1921
1922
+ for i, spec in enumerate(read_specs):
1923
remote_block_size = remote_info.remote_block_size
1924
logger.debug(
1925
"Remote agent %s available, calling _read_blocks"
1926
" on remote rank %s with remote block size %s for req %s",
1927
meta.remote.engine_id,
1926
- remote_rank,
1928
+ spec.remote_rank,
1929
remote_block_size,
1930
req_id,
1931
)
@@ -1942,49 +1944,34 @@ class NixlConnectorWorker:
1944
1945
# Destination handle: remote_engine_id -> remote_rank -> handle.
1946
remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][
1945
- remote_rank
1947
+ spec.remote_rank
1948
]
1949
1948
- local_ids: BlockIds = meta.local_physical_block_ids
1949
- remote_ids: BlockIds = meta.remote.block_ids
1950
- if self._has_mamba:
1951
- # Mamba-HMA: zero out FA groups for P ranks outside fa_read_targets.
1952
- local_ids, remote_ids = self.transfer_topo.filter_block_ids_for_rank(
1953
- engine_id,
1954
- remote_rank,
1955
- local_ids,
1956
- remote_ids,
1957
- self._is_mamba_group,
1958
- )
1959
-
1950
self._read_blocks(
1951
+ read_spec=spec,
1952
request_id=req_id,
1953
dst_engine_id=meta.remote.engine_id,
1954
remote_request_id=meta.remote.request_id,
1964
- local_block_ids=local_ids,
1965
- remote_block_ids=remote_ids,
1966
- remote_rank=remote_rank,
1955
local_xfer_side_handle=local_xfer_side_handle,
1956
remote_xfer_side_handle=remote_xfer_side_handle,
1957
)
1958
1971
- if self.use_mla and tp_ratio < 0:
1972
- # ..but we still need to notify the other remote ranks that we
1973
- # have the blocks we need so they can update the request state.
1974
- notif_id = f"{meta.remote.request_id}:{self.world_size}".encode()
1975
- remote_agents = self._remote_agents[meta.remote.engine_id]
1976
- for rank_to_notify, agent in remote_agents.items():
1977
- if rank_to_notify != remote_rank:
1978
- self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)
1959
+ if self.use_mla and tp_ratio < 0 and read_specs:
1960
+ # ..but we still need to notify the other remote ranks that we
1961
+ # have the blocks we need so they can update the request state.
1962
+ notif_id = f"{req_id}:{self.world_size}".encode()
1963
+ remote_agents = self._remote_agents[meta.remote.engine_id]
1964
+ read_ranks = {s.remote_rank for s in read_specs}
1965
+ for rank_to_notify, agent in remote_agents.items():
1966
+ if rank_to_notify not in read_ranks:
1967
+ self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)
1968
1969
def _read_blocks(
1970
self,
1982
- local_block_ids: BlockIds,
1983
- remote_block_ids: BlockIds,
1971
+ read_spec: ReadSpec,
1972
dst_engine_id: str,
1973
request_id: str,
1974
remote_request_id: str,
1987
- remote_rank: int,
1975
local_xfer_side_handle: int,
1976
remote_xfer_side_handle: int,
1977
):
@@ -1993,6 +1980,10 @@ class NixlConnectorWorker:
1980
a single remote worker.
1981
"""
1982
assert self.transfer_topo is not None
1983
+ remote_rank = read_spec.remote_rank
1984
+ local_block_ids = read_spec.local_block_ids
1985
+ remote_block_ids = read_spec.remote_block_ids
1986
+
1987
remote_info = self.transfer_topo.get_engine_info(dst_engine_id)
1988
block_size_ratio = self.transfer_topo.block_size_ratio(
1989
remote_info.remote_block_size
@@ -2061,16 +2052,14 @@ class NixlConnectorWorker:
2052
== len(local_block_ids)
2053
== len(self.kv_cache_config.kv_cache_groups)
2054
)
2055
+ # Partial prefix cache hit: just read uncomputed blocks.
2056
+ # Skip mamba groups — their blocks represent full state (conv+ssm),
2057
+ # not per-token data, so trimming would corrupt the transfer.
2058
remote_block_ids = list(remote_block_ids)
2059
for i, remote_group in enumerate(remote_block_ids):
2066
- num_remote_blocks = len(remote_group)
2060
num_local_blocks = len(local_block_ids[i])
2068
- if not self._is_mamba_group[i]:
2069
- assert num_local_blocks <= num_remote_blocks
2070
- # Partial prefix cache hit: just read uncomputed blocks.
2071
- # Skip mamba groups — their blocks represent full state (conv+ssm),
2072
- # not per-token data, so trimming would corrupt the transfer.
2073
- if num_local_blocks < num_remote_blocks and not self._is_mamba_group[i]:
2061
+ assert num_local_blocks <= len(remote_group)
2062
+ if num_local_blocks < len(remote_group):
2063
remote_block_ids[i] = remote_group[-num_local_blocks:]
2064
2065
# NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
@@ -2078,14 +2067,17 @@ class NixlConnectorWorker:
2067
# workers will issue xfers to parts of the P worker remote kv caches.
2068
2069
# Get descs ids.
2081
- remote_block_descs_ids = self._get_block_descs_ids(
2082
- dst_engine_id,
2083
- remote_block_ids,
2070
+ remote_block_descs_ids = self._compute_desc_ids(
2071
+ block_ids=remote_block_ids,
2072
+ dst_num_blocks=self.dst_num_blocks[dst_engine_id],
2073
+ block_size_ratio=None,
2074
+ physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
2075
)
2085
- local_block_descs_ids = self._get_block_descs_ids(
2086
- self.engine_id,
2087
- local_block_ids,
2076
+ local_block_descs_ids = self._compute_desc_ids(
2077
+ block_ids=local_block_ids,
2078
+ dst_num_blocks=self.dst_num_blocks[self.engine_id],
2079
block_size_ratio=block_size_ratio,
2080
+ physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
2081
)
2082
2083
assert len(local_block_descs_ids) == len(remote_block_descs_ids)
@@ -2147,63 +2139,6 @@ class NixlConnectorWorker:
2139
2140
return mapped_2d.flatten().astype(np.int64)
2141
2150
- def _get_block_descs_ids(
2151
- self,
2152
- engine_id: str,
2153
- block_ids: BlockIds,
2154
- block_size_ratio: float | None = None,
2155
- ) -> np.ndarray:
2156
- """
2157
- Get the descs ids for a set of block ids.
2158
- When HMA is enabled number of descriptors across kv cache groups might differ.
2159
- A single flattened array is returned for all groups anyway.
2160
- """
2161
- region_ids = np.arange(self.num_regions)
2162
-
2163
- # NOTE (NickLucche) With HMA, every kv group has the same number of layers and
2164
- # layers from different groups share the same kv tensor.
2165
- # eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be read across all regions,
2166
- # same for [3], but group0-group1 blocks will always differ (different areas).
2167
- # Therefore we can just flatten the block_ids and compute the descs ids for all
2168
- # groups at once.
2169
- num_blocks = self.dst_num_blocks[engine_id]
2170
- if block_size_ratio is not None:
2171
- num_blocks = int(num_blocks * block_size_ratio)
2172
-
2173
- # Compute desc ids per group using the right stride: FA descs have
2174
- # num_blocks entries per region (kernel granularity), SSM descs have
2175
- # logical_blocks entries per region (no kernel splitting).
2176
- region_ids = region_ids[:, None]
2177
- if not self._has_mamba:
2178
- block_ids = np.concatenate(block_ids)[None, :]
2179
- descs_ids = region_ids * num_blocks + block_ids
2180
- return descs_ids.flatten()
2181
- else:
2182
- # NOTE (NickLucche) SSM and Attention blocks regions can be exchanged
2183
- # arbitrarily by manager. Therefore, descs are duplicated for SSM and
2184
- # Attention like so:
2185
- # desc_handle->[descs_fa (all regions) | descs_ssm (all regions)].
2186
- # This is like having two "low-level views" of the same storage.
2187
- # `num_fa_descs` offset must be computed per-engine since P and D can
2188
- # have different num_blocks (and thus different FA descs counts).
2189
- physical_per_logical = self._physical_blocks_per_logical[engine_id]
2190
- logical_blocks = num_blocks // physical_per_logical
2191
- num_fa_descs = self.num_regions * num_blocks
2192
- # 3-read mamba: 4 regions per unique cache tensor (x, B, C, ssm).
2193
- mamba_region_ids = np.arange(len(self.block_len_per_layer) * 4)[:, None]
2194
- all_descs = []
2195
- for i, group in enumerate(block_ids):
2196
- group_arr = np.asarray(group)[None, :]
2197
- if self._is_mamba_group[i]:
2198
- all_descs.append(
2199
- (
2200
- mamba_region_ids * logical_blocks + group_arr + num_fa_descs
2201
- ).flatten()
2202
- )
2203
- else:
2204
- all_descs.append((region_ids * num_blocks + group_arr).flatten())
2205
- return np.concatenate(all_descs)
2206
-
2142
def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds:
2143
"""
2144
Convert logical block ids to kernel physical block ids.
@@ -2295,10 +2230,7 @@ class NixlConnectorWorker:
2230
"""
2231
assert self.transfer_topo is not None
2232
if self.transfer_topo.is_kv_layout_blocks_first:
2298
- # For indexing only half (either just the K or V part).
2233
if mamba_view:
2300
- # NOTE (NickLucche) Mamba Opt: this is already skipping the padding so
2301
- # we're only transferring the minimum required bytes.
2234
block_len = self._mamba_ssm_size[not first_split]
2235
else:
2236
block_len = self.block_len_per_layer[layer_idx] // 2
vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py
+11
-2
@@ -31,6 +31,7 @@ class MambaConvSplitInfo:
31
x_local: int # intermediate_size / TP (columns for x)
32
b_local: int # groups_ss / TP (columns for B; C is same size)
33
conv_dtype_size: int # bytes per element (e.g. 2 for float16)
34
+ ssm_sizes: tuple[int, int] # (conv_state_bytes, ssm_state_bytes)
35
36
@property
37
def conv_dim_local(self) -> int:
@@ -99,8 +100,8 @@ def derive_mamba_conv_split(
100
local_tp: this engine's tensor-parallel size.
101
102
Returns:
102
- MambaConvSplitInfo with per-rank x_local, b_local, conv_rows, and
103
- conv_dtype_size.
103
+ MambaConvSplitInfo with per-rank x_local, b_local, conv_rows,
104
+ conv_dtype_size, and ssm_sizes (conv_state_bytes, ssm_state_bytes).
105
"""
106
if mamba_spec.mamba_type != "mamba2":
107
raise NotImplementedError(
@@ -142,12 +143,20 @@ def derive_mamba_conv_split(
143
dtype=mamba_spec.dtypes[0], # type: ignore[misc]
144
).element_size()
145
146
+ ssm_dtype_size = torch.tensor(
147
+ [],
148
+ dtype=mamba_spec.dtypes[1], # type: ignore[misc]
149
+ ).element_size()
150
+ conv_state_bytes = torch.Size(mamba_spec.shapes[0]).numel() * conv_dtype_size
151
+ ssm_state_bytes = torch.Size(mamba_spec.shapes[1]).numel() * ssm_dtype_size
152
+
153
# Divide by TP to get per-rank column counts.
154
return MambaConvSplitInfo(
155
conv_rows=conv_rows,
156
x_local=intermediate_size // local_tp,
157
b_local=groups_ss // local_tp,
158
conv_dtype_size=conv_dtype_size,
159
+ ssm_sizes=(conv_state_bytes, ssm_state_bytes),
160
)
161
162