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, 01:15 UTC
[Bugfix][PD] Fix DSv4 Disaggregated (#41957)
Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: ZhanqiuHu <zhu@redhat.com>
Nicolò Lucchesi committed
May 9, 2026 at 18:48 UTC
171d59ae8d1b8d6718ec3cd6558334ce8bbc0eb4
5 files changed
+49
-35
tests/v1/kv_connector/unit/test_tp_mapping.py
+7
-2
@@ -9,6 +9,8 @@ No GPU or NIXL required.
9
10
from __future__ import annotations
11
12
+from types import SimpleNamespace
13
+
14
import pytest
15
16
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
@@ -33,12 +35,15 @@ def _compute_mapping(
35
num_kv_heads: int = 8,
36
group_spec_types: tuple[type, ...] = (FullAttentionSpec,),
37
) -> TPMapping:
36
- return compute_tp_mapping(
38
+ transfer_topology = SimpleNamespace(
39
tp_rank=tp_rank,
40
tp_size=tp_size,
39
- remote_tp_size=remote_tp_size,
41
is_mla=is_mla,
42
total_num_kv_heads=num_kv_heads,
43
+ )
44
+ return compute_tp_mapping(
45
+ transfer_topology=transfer_topology,
46
+ remote_tp_size=remote_tp_size,
47
group_spec_types=group_spec_types,
48
)
49
vllm/distributed/kv_transfer/kv_connector/utils.py
+1
-1
@@ -593,7 +593,7 @@ class TransferTopology:
593
return (
594
f"TransferTopology("
595
f"tp_ratio={self.tp_ratio(info.remote_tp_size)}, "
596
- f"K={self.total_num_kv_heads}, "
596
+ f"num_kv_heads={self.total_num_kv_heads if not self.is_mla else 1}, "
597
f"local_tp={self.tp_size}, "
598
f"remote_tp={info.remote_tp_size}, "
599
f"local_rank={self.tp_rank}, "
vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py
+9
-9
@@ -10,6 +10,7 @@ import numpy as np
10
11
from vllm.distributed.kv_transfer.kv_connector.utils import (
12
BlockIds,
13
+ TransferTopology,
14
)
15
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, MambaSpec
16
@@ -62,11 +63,8 @@ class TPMapping:
63
64
65
def compute_tp_mapping(
65
- tp_rank: int,
66
- tp_size: int,
66
+ transfer_topology: TransferTopology,
67
remote_tp_size: int,
68
- is_mla: bool,
69
- total_num_kv_heads: int,
68
group_spec_types: tuple[type[KVCacheSpec], ...],
69
) -> TPMapping:
70
"""Build the complete local-to-remote TP mapping.
@@ -74,13 +72,15 @@ def compute_tp_mapping(
72
Computes source ranks, head slot assignments, and the rank offset
73
factor in a single pass.
74
"""
75
+ tp_rank = transfer_topology.tp_rank
76
+ tp_size = transfer_topology.tp_size
77
+ total_num_kv_heads = transfer_topology.total_num_kv_heads
78
# --- 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:
79
+ if transfer_topology.is_mla or tp_size >= remote_tp_size:
80
# D (local TP) > P (remote TP): multiple local ranks read different chunks from
81
# *one* remote rank, corresponding to different kv heads.
82
+ # For MLA, we only need one remote since cache is duplicated. When P TP=k*TP k,
83
+ # this will spread mla ranks to read from remote k*tp_rank.
84
attn_ranks = [tp_rank * remote_tp_size // tp_size]
85
else:
86
# P (remote TP) > D (local TP): one local rank
@@ -123,7 +123,7 @@ def compute_tp_mapping(
123
}
124
125
# --- Rank offset factor ---
126
- if is_mla or tp_size <= remote_tp_size:
126
+ if transfer_topology.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:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py
+9
@@ -10,6 +10,7 @@ import zmq
10
11
from vllm.platforms import current_platform
12
from vllm.utils.network_utils import make_zmq_socket
13
+from vllm.v1.kv_cache_interface import KVCacheSpec, UniformTypeKVCacheSpecs
14
15
# Supported platforms and types of kv transfer buffer.
16
# {device: tuple of supported kv buffer types}
@@ -46,3 +47,11 @@ def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
47
finally:
48
if ctx is not None:
49
ctx.destroy(linger=0)
50
+
51
+
52
+def get_representative_spec_type(spec: KVCacheSpec) -> type[KVCacheSpec]:
53
+ if isinstance(spec, UniformTypeKVCacheSpecs):
54
+ # All inner specs are the same type; pick any.
55
+ inner = next(iter(spec.kv_cache_specs.values()))
56
+ return type(inner)
57
+ return type(spec)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+23
-23
@@ -53,6 +53,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
53
)
54
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
55
_NIXL_SUPPORTED_DEVICE,
56
+ get_representative_spec_type,
57
zmq_ctx,
58
)
59
from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
@@ -100,24 +101,24 @@ class NixlConnectorWorker:
101
num_blocks = dst_num_blocks
102
if block_size_ratio is not None:
103
num_blocks = int(num_blocks * block_size_ratio)
103
- ratio = physical_blocks_per_logical
104
- logical_blocks = num_blocks // ratio
105
-
104
num_fa_descs = num_fa_regions * num_blocks
105
106
# All-attention fast path: single vectorized broadcast.
107
if num_ssm_regions == 0:
108
+ # NOTE (NickLucche) With HMA, every kv group has the same number of layers
109
+ # and layers from different groups share the same kv tensor.
110
+ # eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be
111
+ # read across all regions, same for [3], but group0-group1 blocks will
112
+ # always differ (different areas). Therefore we can just flatten the
113
+ # block_ids and compute the descs ids for all groups at once.
114
block_arr = np.concatenate(block_ids)[None, :]
115
region_ids = np.arange(num_fa_regions)[:, None]
116
return (region_ids * num_blocks + block_arr).flatten()
117
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).
118
+ # Compute desc ids per group using the right stride: FA descs have
119
+ # num_blocks entries per region (kernel granularity), SSM descs have
120
+ # logical_blocks entries per region (no kernel splitting).
121
+ logical_blocks = num_blocks // physical_blocks_per_logical
122
all_descs: list[np.ndarray] = []
123
for i, group in enumerate(block_ids):
124
group_arr = np.asarray(group)
@@ -426,8 +427,10 @@ class NixlConnectorWorker:
427
self._physical_blocks_per_logical_kv_block = 1
428
self._sync_block_size_with_kernel()
429
430
+ # Unwrap UniformTypeKVCacheSpecs to get the representative spec type
431
self._group_spec_types = tuple(
430
- type(g.kv_cache_spec) for g in self.kv_cache_config.kv_cache_groups
432
+ get_representative_spec_type(g.kv_cache_spec)
433
+ for g in self.kv_cache_config.kv_cache_groups
434
)
435
436
# Per-engine TP mappings. Generated during handshake.
@@ -1259,12 +1262,9 @@ class NixlConnectorWorker:
1262
logger.info("Transfer plan: %s", transfer_topo.describe(engine_id))
1263
1264
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,
1265
+ transfer_topology=transfer_topo,
1266
+ remote_tp_size=remote_tp_size,
1267
+ group_spec_types=self._group_spec_types,
1268
)
1269
1270
remote_agent_name = self.nixl_wrapper.add_remote_agent(
@@ -1391,7 +1391,8 @@ class NixlConnectorWorker:
1391
)
1392
# num_kv_heads > tp_size with P_TP > D_TP not supported for non-mamba.
1393
# Mamba models can have replicated FA KV with tp_ratio < 0.
1394
- if not self._has_mamba:
1394
+ # MLA models do not need to handle kv replication.
1395
+ if not self.use_mla and not self._has_mamba:
1396
assert not (
1397
tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id)
1398
)
@@ -1915,9 +1916,9 @@ class NixlConnectorWorker:
1916
1917
# D may have to perform multiple reads from different remote ranks.
1918
# MLA opt: when P TP > D TP, only a single read is executed for
1918
- # the first remote rank (cache is duplicated).
1919
+ # the first remote rank (cache is duplicated)..
1920
if self.use_mla and tp_ratio < 0:
1920
- read_specs = read_specs[:1]
1921
+ assert len(read_specs) == 1
1922
1923
for i, spec in enumerate(read_specs):
1924
remote_block_size = remote_info.remote_block_size
@@ -1959,11 +1960,10 @@ class NixlConnectorWorker:
1960
if self.use_mla and tp_ratio < 0 and read_specs:
1961
# ..but we still need to notify the other remote ranks that we
1962
# have the blocks we need so they can update the request state.
1962
- notif_id = f"{req_id}:{self.world_size}".encode()
1963
+ notif_id = f"{meta.remote.request_id}:{self.world_size}".encode()
1964
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:
1966
+ if rank_to_notify != read_specs[0].remote_rank:
1967
self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)
1968
1969
def _read_blocks(