@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6fbfdd1831

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[NIXL] Per-region KV transfer classification for mixed full-attn + MLA groups (#44583)

Dao007forever committed Jun 11, 2026 at 21:42 UTC 6fbfdd183145443274df49c09c46cd13ea27af5f

3 files changed +234 -68

tests/v1/kv_connector/unit/test_nixl_connector.py

+81

@@ -1063,6 +1063,87 @@ class TestNixlHandshake:

1063 # whole block is moved.

1064 worker.add_remote_agent(meta, remote_tp_size=1)

1065

1066 + @patch(

1067 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",

1068 + FakeNixlWrapper,

1069 + )

1070 + def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init):

1071 + """Mixed full-attn (SPLIT) + MLA (REPLICATE) single KV group under

1072 + heterogeneous TP must NOT raise (previously a NotImplementedError),

1073 + and the per-region gate must still reject a wrong block_len.

1074 + """

1075 + vllm_config = create_vllm_config()

1076 + with patch(

1077 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", # noqa: E501

1078 + return_value=2,

1079 + ):

1080 + connector = NixlConnector(

1081 + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)

1082 + )

1083 + connector.connector_worker = FakeNixlConnectorWorker(

1084 + vllm_config, connector.engine_id, hand_shake_latency=0

1085 + )

1086 + worker = connector.connector_worker

1087 +

1088 + # Region 0: full-attn (SPLIT). Region 1: MLA (REPLICATE).

1089 + fa_len = 4096 * worker.block_size

1090 + idx_len = 512 * worker.block_size

1091 + worker.slot_size_per_layer = [4096, 512]

1092 + worker.block_len_per_layer = [fa_len, idx_len]

1093 + worker._region_is_mla = [False, True]

1094 + worker.num_blocks = 1

1095 + worker.dst_num_blocks[worker.engine_id] = worker.num_blocks

1096 + worker.src_blocks_data = [

1097 + (0, fa_len, worker.tp_rank),

1098 + (0, idx_len, worker.tp_rank),

1099 + ]

1100 + worker.num_descs = len(worker.src_blocks_data)

1101 +

1102 + # D_TP=2, P_TP=1 -> tp_ratio=2. SPLIT region scales by tp_ratio;

1103 + # REPLICATE region is unchanged.

1104 + tp_ratio = 2

1105 + meta = NixlAgentMetadata(

1106 + engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

1107 + agent_metadata=FakeNixlWrapper.AGENT_METADATA,

1108 + kv_caches_base_addr=[0, 0],

1109 + device_id=0,

1110 + num_blocks=1,

1111 + block_lens=[fa_len * tp_ratio, idx_len],

1112 + kv_cache_layout=worker.kv_cache_layout,

1113 + block_size=worker.block_size,

1114 + ssm_sizes=(0, 0),

1115 + attn_backend_name=worker.backend_name,

1116 + physical_blocks_per_logical_kv_block=1,

1117 + )

1118 + worker.add_remote_agent(meta, remote_tp_size=1)

1119 + assert (

1120 + FakeNixlConnectorWorker.REMOTE_ENGINE_ID in worker.dst_xfer_side_handles

1121 + )

1122 + # Gate rejects an MLA region wrongly scaled by tp_ratio.

1123 + worker2 = FakeNixlConnectorWorker(

1124 + vllm_config, connector.engine_id, hand_shake_latency=0

1125 + )

1126 + worker2.block_len_per_layer = [fa_len, idx_len]

1127 + worker2._region_is_mla = [False, True]

1128 + worker2.num_blocks = 1

1129 + worker2.dst_num_blocks[worker2.engine_id] = worker2.num_blocks

1130 + bad_meta = NixlAgentMetadata(

1131 + engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

1132 + agent_metadata=FakeNixlWrapper.AGENT_METADATA,

1133 + kv_caches_base_addr=[0, 0],

1134 + device_id=0,

1135 + num_blocks=1,

1136 + # WRONG: MLA region scaled by tp_ratio (it should be replicated).

1137 + block_lens=[fa_len * tp_ratio, idx_len * tp_ratio],

1138 + kv_cache_layout=worker2.kv_cache_layout,

1139 + block_size=worker2.block_size,

1140 + ssm_sizes=(0, 0),

1141 + attn_backend_name=worker2.backend_name,

1142 + physical_blocks_per_logical_kv_block=1,

1143 + )

1144 + with pytest.raises(AssertionError):

1145 + worker2.add_remote_agent(bad_meta, remote_tp_size=1)

1146 +

1147

1148 # NOTE: resource cleanup in mp backend is a bit finicky, so the order in which

1149 # we put here is important. First run ray, it will clean up the resources, then

tests/v1/kv_connector/unit/test_tp_mapping.py

+11 -1

@@ -73,9 +73,19 @@ class TestTPMappingStructure:

75 def _make_mock_worker_for_splits(group_spec_types):

76 - """Build a mock NixlConnectorWorker with _group_spec_types for split tests."""

76 + """Build a mock NixlConnectorWorker with _group_spec_types for split tests.

77 +

78 + No per-region replicate flags are configured (``block_len_per_layer`` empty

79 + and ``num_regions == 0``), so ``_fa_desc_replicated`` takes its early-return

80 + path and treats every FA descriptor as SPLIT, matching the legacy behavior

81 + these tests assert.

82 + """

83 worker = object.__new__(NixlConnectorWorker)

84 worker._group_spec_types = group_spec_types

85 + worker.transfer_topo = SimpleNamespace(virtually_split_kv_in_blocks=False)

86 + worker.block_len_per_layer = []

87 + worker.num_regions = 0

88 + worker._region_is_mla = []

89 return worker

vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py

+142 -67

@@ -71,6 +71,7 @@ from vllm.v1.attention.backends.utils import get_kv_cache_layout

71 from vllm.v1.kv_cache_interface import (

72 FullAttentionSpec,

73 MambaSpec,

74 + MLAAttentionSpec,

75 UniformTypeKVCacheSpecs,

76 )

77 from vllm.v1.worker.block_table import BlockTable

@@ -178,19 +179,63 @@ class NixlConnectorWorker:

179 else 0

180 )

181

182 + # Per-FA-descriptor replicate flag, in _build_fa_local emission order.

183 + fa_desc_replicated = self._fa_desc_replicated(num_fa_descs)

184 +

185 for p_idx, p_rank in enumerate(plan.all_source_ranks):

186 fa_slot = plan.rank_to_attention_slot.get(p_rank, 0)

187

188 handle: list[tuple[int, int, int]] = []

189 for j, (addr, local_len, dev) in enumerate(src_blocks_data):

190 if j < num_fa_descs:

187 - chunk = local_len // fa_num_splits

188 - handle.append((addr + fa_slot * chunk, chunk, dev))

191 + if fa_desc_replicated[j]:

192 + # REPLICATE (MLA): whole block written on every rank.

193 + handle.append((addr, local_len, dev))

194 + else:

195 + # SPLIT (full-attn): this rank's head slice.

196 + chunk = local_len // fa_num_splits

197 + handle.append((addr + fa_slot * chunk, chunk, dev))

198 else:

199 chunk = local_len // ssm_num_splits

200 handle.append((addr + p_idx * chunk, chunk, dev))

201 yield handle

202

203 + def _fa_desc_replicated(self, num_fa_descs: int) -> list[bool]:

204 + """Per-FA-descriptor replicate flag, in _build_fa_local emission order

205 + (region-major; K then optional V per region). Length ``num_fa_descs``.

206 + """

207 + assert self.transfer_topo is not None

208 + n_regions = len(self.block_len_per_layer)

209 + # Unset only when the worker is built directly in unit tests; a real

210 + # model always registers regions (no-KV-cache crashes long before here).

211 + # Fall back to all-SPLIT to preserve the pre-per-region behavior.

212 + if n_regions == 0 or self.num_regions == 0:

213 + return [False] * num_fa_descs

214 + # Descriptors (blocks) per stream; all streams share the same count.

215 + nblk = num_fa_descs // self.num_regions

216 + virtually_split = self.transfer_topo.virtually_split_kv_in_blocks

217 + flags: list[bool] = []

218 + for i in range(n_regions):

219 + replicated = self._is_region_replicated(i)

220 + # REPLICATE (MLA) is key-only -> 1 stream; SPLIT emits K and V

221 + # (2 streams) under the virtually-split layout.

222 + num_streams = 1 if replicated or not virtually_split else 2

223 + flags.extend([replicated] * (num_streams * nblk))

224 + assert len(flags) == num_fa_descs, (

225 + f"FA desc flags {len(flags)} != num_fa_descs {num_fa_descs}"

226 + )

227 + return flags

228 +

229 + def _is_region_replicated(self, region_idx: int) -> bool:

230 + """Whether region ``region_idx`` is transferred REPLICATE vs SPLIT.

231 +

232 + REPLICATE (MLA): identical on every rank, whole block read from one

233 + rank at offset 0, key-only. SPLIT (full-attn): head-sharded across TP.

234 + Defaults to SPLIT when the per-region map is unset (e.g. tests that set

235 + block_len_per_layer without register_kv_caches).

236 + """

237 + return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx]

238 +

239 def __init__(

240 self,

241 vllm_config: "VllmConfig",

@@ -450,6 +495,15 @@ class NixlConnectorWorker:

495 for g in self.kv_cache_config.kv_cache_groups

496 )

497

498 + # Per-region MLA flag, 1:1 with block_len_per_layer. True -> REPLICATE

499 + # (MLA), False -> SPLIT (head-sharded full-attn). Mixed only for models

500 + # combining both (e.g. GQA main + MLA Eagle-3 draft).

501 + self._region_is_mla = list[bool]()

502 +

503 + # Enable different block lengths for different layers *only* when MLA is used.

504 + # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`.

505 + self.block_len_per_layer = list[int]()

506 +

507 # Per-engine TP mappings. Generated during handshake.

508 self.tp_mappings: dict[EngineId, TPMapping] = {}

509

@@ -849,9 +903,6 @@ class NixlConnectorWorker:

903 # to better exploit the memory layout (ie num_blocks is the first dim).

904 tensor_size_bytes = None

905

852 - # Enable different block lengths for different layers *only* when MLA is used.

853 - # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`.

854 - self.block_len_per_layer = list[int]()

906 for layer_name, cache_or_caches in xfer_buffers.items():

907 # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to

908 # that of FI, with block laid out as in `get_backend_aware_kv_block_len`.

@@ -895,8 +946,6 @@ class NixlConnectorWorker:

946 # `page_size` accounts for physical blocks, st KVCache is always

947 # [`num_blocks` * `page_size`]

948 curr_tensor_size_bytes = num_blocks * physical_page_size

898 - if tensor_size_bytes is None:

899 - tensor_size_bytes = curr_tensor_size_bytes

949

950 # TODO (NickLucche) we could eventually unify how we handle FA/FI regions,

951 # registering a single tensor for both K/V and splitting logically like FI.

@@ -920,6 +969,20 @@ class NixlConnectorWorker:

969 )

970 else:

971 self.block_len_per_layer.append(physical_page_size)

972 + is_mla_region = isinstance(layer_spec, MLAAttentionSpec)

973 + self._region_is_mla.append(is_mla_region)

974 +

975 + # HeteroTP cannot transfer differently-sized regions, so every

976 + # non-MLA region in a group must share one tensor size (this also

977 + # holds for Mamba-like models). The sole exception is the DeepSeek

978 + # MLA indexer, which sits in a UniformTypeKVCacheSpecs group at a

979 + # different size; MLA regions are therefore exempt.

980 + if not is_mla_region:

981 + if tensor_size_bytes is None:

982 + tensor_size_bytes = curr_tensor_size_bytes

983 + assert tensor_size_bytes == curr_tensor_size_bytes, (

984 + "All non-MLA kv cache tensors must have the same size"

985 + )

986

987 if cache.shape[0] != num_blocks:

988 raise AssertionError(

@@ -937,12 +1000,6 @@ class NixlConnectorWorker:

1000 f"{self.transfer_topo.is_kv_layout_blocks_first}"

1001 )

1002

940 - if not self.use_mla:

941 - # Different kv cache shape is not supported by HeteroTP.

942 - # This must also hold true for Mamba-like models.

943 - assert tensor_size_bytes == curr_tensor_size_bytes, (

944 - "All kv cache tensors must have the same size"

945 - )

1003 # Need to make sure the device ID is non-negative for NIXL,

1004 # Torch uses -1 to indicate CPU tensors.

1005 self.device_id = max(cache.get_device(), 0)

@@ -953,7 +1010,11 @@ class NixlConnectorWorker:

1010 logger.debug(

1011 "Different block lengths collected: %s", set(self.block_len_per_layer)

1012 )

956 - assert len(self.block_len_per_layer) == len(seen_base_addresses)

1013 + assert (

1014 + len(self.block_len_per_layer)

1015 + == len(seen_base_addresses)

1016 + == len(self._region_is_mla)

1017 + )

1018

1019 self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses

1020 self.num_regions = len(caches_data)

@@ -967,7 +1028,12 @@ class NixlConnectorWorker:

1028 # of 'virtual' regions here and halve `block_len` below.

1029 # Similarly for Mamba layers, we register SSM+Conv as a single region and

1030 # then duplicate it logically to be able to index SSM/Conv separately.

970 - self.num_regions *= 2

1031 + # Exception: key-only REPLICATE regions (MLA) have no V half, so

1032 + # they contribute a single desc stream and are not doubled.

1033 + self.num_regions = sum(

1034 + 1 if self._is_region_replicated(i) else 2

1035 + for i in range(len(self._region_is_mla))

1036 + )

1037

1038 # Total local FA descriptors (boundary between FA and mamba descs).

1039 self.num_descs = self.num_regions * self.num_blocks

@@ -1133,10 +1199,13 @@ class NixlConnectorWorker:

1199 addr = base_addr + block_offset

1200 result.append((addr, kv_block_len, self.device_id))

1201

1136 - if self.transfer_topo.virtually_split_kv_in_blocks:

1202 + if (

1203 + self.transfer_topo.virtually_split_kv_in_blocks

1204 + and not self._is_region_replicated(i)

1205 + ):

1206 # Separate and interleave K/V regions to maintain the same

1207 # descs ordering. This is needed for selecting contiguous heads

1139 - # when split across TP ranks.

1208 + # when split across TP ranks. (Skipped for key-only REPLICATE.)

1209 second_split = self.get_backend_aware_kv_block_len(

1210 layer_idx=i, first_split=False, mamba_view=False

1211 )

@@ -1158,10 +1227,13 @@ class NixlConnectorWorker:

1227 fa_group_idx = next(

1228 i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t)

1229 )

1161 - num_attn_reads = len(plan.source_ranks_per_group[fa_group_idx])

1230 + # SPLIT regions read their head slice from this many remote ranks at a

1231 + # per-rank offset; REPLICATE regions read the whole block once.

1232 + split_reads = len(plan.source_ranks_per_group[fa_group_idx])

1233 num_blocks = nixl_agent_meta.num_blocks

1234 result: list[tuple[int, int, int]] = []

1235 for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):

1236 + replicated = self._is_region_replicated(i)

1237 # Read our whole local region size from remote..

1238 local_block_len = self.get_backend_aware_kv_block_len(

1239 layer_idx=i, first_split=True, mamba_view=False

@@ -1171,8 +1243,13 @@ class NixlConnectorWorker:

1243 # ..using remote kv_block_len as transfer unit

1244 local_block_len = remote_kv_block_len

1245

1174 - local_block_len = local_block_len // num_attn_reads

1175 - rank_offset = plan.rank_offset_factor * remote_kv_block_len

1246 + # REPLICATE reads the whole block once at offset 0; SPLIT gathers

1247 + # its head slice from `split_reads` remote ranks at a per-rank offset.

1248 + num_reads = 1 if replicated else split_reads

1249 + rank_offset = (

1250 + 0 if replicated else plan.rank_offset_factor * remote_kv_block_len

1251 + )

1252 + local_block_len = local_block_len // num_reads

1253

1254 page_size = nixl_agent_meta.block_lens[i]

1255 for block_id in range(num_blocks):

@@ -1182,12 +1259,13 @@ class NixlConnectorWorker:

1259 addr = base_addr + block_offset + rank_offset

1260 result.append((addr, local_block_len, nixl_agent_meta.device_id))

1261

1185 - if self.transfer_topo.virtually_split_kv_in_blocks:

1262 + emits_v = self.transfer_topo.virtually_split_kv_in_blocks and not replicated

1263 + if emits_v:

1264 # With FlashInfer index V separately to allow head splitting.

1265 second_split = self.get_backend_aware_kv_block_len(

1266 layer_idx=i, first_split=False, mamba_view=False

1267 )

1190 - second_split = second_split // num_attn_reads

1268 + second_split = second_split // num_reads

1269 for block_id in range(num_blocks):

1270 block_offset = block_id * page_size

1271 addr = base_addr + block_offset + rank_offset

@@ -1527,49 +1605,43 @@ class NixlConnectorWorker:

1605 "Use HND layout on the prefill side."

1606 )

1607

1530 - # Block len can only vary across layers when using MLA.

1531 - remote_block_len = nixl_agent_meta.block_lens[0]

1532 - if self.use_mla or self.transfer_topo.is_kv_replicated(remote_engine_id):

1533 - # With replicated KV cache, only the number of blocks can differ.

1534 - # TODO (ZhanqiuHu): For mamba models, validate FA and mamba

1535 - # block_lens separately.

1536 - if not self._has_mamba:

1537 - for i in range(len(self.block_len_per_layer)):

1538 - assert (

1539 - self.block_len_per_layer[i] // block_size_ratio

1540 - == nixl_agent_meta.block_lens[i]

1541 - ), "KV cache sizes must match between P and D when replicated"

1542 - else:

1543 - # When MLA is not used, this is a list of the same block length

1544 - for block_len in nixl_agent_meta.block_lens:

1545 - assert block_len == remote_block_len, (

1546 - "All remote layers must have the same block size"

1547 - )

1548 -

1549 - # HMA hybrid models (mamba+attention) pad block_len to

1550 - # max(attn_page, mamba_page), so the linear tp_ratio scaling

1551 - # assumption only holds for pure-attention models.

1552 - if not self._has_mamba:

1553 - if tp_ratio > 0:

1554 - assert (

1555 - remote_block_len

1556 - == (self.block_len_per_layer[0] * tp_ratio) // block_size_ratio

1557 - ), (

1558 - "Remote P worker KV layer cache must be of shape [2, N,"

1559 - " local_kv_heads*tp_ratio, page_size, head_dim] and "

1560 - "same dtype."

1608 + # Per-region block_len validation enforcing the P/D invariant.

1609 + # REPLICATE regions (MLA, or a whole-model MLA / replicated-KV transfer)

1610 + # only allow the number of blocks to differ; SPLIT regions scale with

1611 + # tp_ratio. Mamba uses the ssm_sizes counterpart, so skip block_len here.

1612 + if not self._has_mamba:

1613 + assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), (

1614 + "Number of KV layers must match between prefill and decode"

1615 + )

1616 + model_replicated = self.use_mla or self.transfer_topo.is_kv_replicated(

1617 + remote_engine_id

1618 + )

1619 + for i, local_len in enumerate(self.block_len_per_layer):

1620 + replicated = model_replicated or self._is_region_replicated(i)

1621 + remote_len = nixl_agent_meta.block_lens[i]

1622 + if replicated:

1623 + # Whole block copied; only the number of blocks may differ.

1624 + assert local_len // block_size_ratio == remote_len, (

1625 + "KV cache sizes must match between P and D when "

1626 + f"replicated (region {i}: local={local_len}, "

1627 + f"remote={remote_len}, bsr={block_size_ratio})."

1628 + )

1629 + elif tp_ratio > 0:

1630 + # D_TP >= P_TP: remote holds tp_ratio x local heads.

1631 + assert remote_len == (local_len * tp_ratio) // block_size_ratio, (

1632 + f"SPLIT region {i}: remote P KV block_len {remote_len} "

1633 + f"must equal local {local_len} * tp_ratio {tp_ratio} "

1634 + f"// block_size_ratio {block_size_ratio}."

1635 )

1636 else:

1637 + # P_TP > D_TP: local holds |tp_ratio| x remote heads.

1638 assert block_size_ratio == 1, (

1564 - "Different local/remote block sizes are not supported"

1565 - " when P TP > D TP."

1639 + "Different local/remote block sizes are not supported "

1640 + "when P TP > D TP."

1641 )

1567 - assert remote_block_len == self.block_len_per_layer[0] // (

1568 - -tp_ratio

1569 - ), (

1570 - "Remote P worker KV layer cache must be of shape [2, N,"

1571 - " local_kv_heads/tp_ratio, page_size, head_dim] and "

1572 - "same dtype."

1642 + assert remote_len == local_len // (-tp_ratio), (

1643 + f"SPLIT region {i}: remote P KV block_len {remote_len} "

1644 + f"must equal local {local_len} // |tp_ratio| {-tp_ratio}."

1645 )

1646

1647 # TP workers that handhshake with same remote have same #blocks.

@@ -2450,13 +2522,16 @@ class NixlConnectorWorker:

2522 |1st_split-2nd_split| |1st_split-2nd_split |

2523 """

2524 assert self.transfer_topo is not None

2453 - if self.transfer_topo.virtually_split_kv_in_blocks:

2454 - if mamba_view:

2455 - block_len = self._mamba_ssm_size[not first_split]

2456 - else:

2457 - block_len = self.block_len_per_layer[layer_idx] // 2

2525 + virtually_split = self.transfer_topo.virtually_split_kv_in_blocks

2526 + if virtually_split and mamba_view:

2527 + block_len = self._mamba_ssm_size[not first_split]

2528 else:

2459 - block_len = self.block_len_per_layer[layer_idx]

2529 + # Per-descriptor block length: a SPLIT region (full-attn under the

2530 + # virtually-split layout) emits separate K and V and uses

2531 + # block_len//2; REPLICATE (MLA, key-only) and non-split layouts use

2532 + # the whole block.

2533 + half_block = virtually_split and not self._is_region_replicated(layer_idx)

2534 + block_len = self.block_len_per_layer[layer_idx] // (2 if half_block else 1)

2535 return block_len

2536

2537 def get_kv_connector_stats(self) -> KVConnectorStats | None: