@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 80608ba5af

Fork 0

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 14, 05:30 UTC

[NIXL] Add support for MLA caches with different latent dim (#25902)

Signed-off-by: NickLucche <nlucches@redhat.com> Signed-off-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Chen Zhang <zhangch99@outlook.com>

Nicolò Lucchesi committed Sep 30, 2025 at 14:18 UTC 80608ba5afe6559b3b06f56905bb5cb8c691e7c2

2 files changed +66 -42

tests/v1/kv_connector/unit/test_nixl_connector.py

+7 -6

@@ -255,8 +255,9 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

255 time.sleep(self._hand_shake_latency)

256 # These should've been done in register_kv_caches(), called by

257 # gpu_model_runner. Here we just hardcode some dummy values.

258 - self.slot_size_bytes = 4096

259 - self.block_len = self.slot_size_bytes * self.block_size

258 + slot_size_bytes = 4096

259 + self.slot_size_per_layer = [slot_size_bytes]

260 + self.block_len_per_layer = [slot_size_bytes * self.block_size]

261 self.num_blocks = 1

262 self.dst_num_blocks[self.engine_id] = self.num_blocks

263

@@ -268,7 +269,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

269 agent_metadata=FakeNixlWrapper.AGENT_METADATA,

270 kv_caches_base_addr=[0],

271 num_blocks=1,

271 - block_len=self.block_len,

272 + block_lens=self.block_len_per_layer,

273 attn_backend_name=self.backend_name,

274 # `self.kv_cache_layout` is only forced to HND when vllm engine

275 # is started. We mock HND here.

@@ -485,8 +486,8 @@ class TestNixlHandshake:

486 worker = connector.connector_worker

487

488 # Minimal local registration params used by add_remote_agent

488 - worker.slot_size_bytes = 4096

489 - worker.block_len = worker.slot_size_bytes * worker.block_size

489 + worker.slot_size_per_layer = [4096]

490 + worker.block_len_per_layer = [4096 * worker.block_size]

491 worker.num_blocks = 1

492 worker.dst_num_blocks[worker.engine_id] = worker.num_blocks

493

@@ -498,7 +499,7 @@ class TestNixlHandshake:

499 agent_metadata=FakeNixlWrapper.AGENT_METADATA,

500 kv_caches_base_addr=[0],

501 num_blocks=1,

501 - block_len=worker.block_len,

502 + block_lens=worker.block_len_per_layer,

503 attn_backend_name=worker.backend_name,

504 kv_cache_layout=mismatched_layout,

505 )

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+59 -36

@@ -87,7 +87,7 @@ class NixlAgentMetadata(

87 agent_metadata: bytes

88 kv_caches_base_addr: list[int]

89 num_blocks: int

90 - block_len: int

90 + block_lens: list[int]

91 attn_backend_name: str

92 kv_cache_layout: str

@@ -772,6 +772,9 @@ class NixlConnectorWorker:

772 split_k_and_v = not (self.use_mla or self._use_pallas

773 or self._use_flashinfer)

774 tensor_size_bytes = None

775 + # Enable different block lengths for different layers when MLA is used.

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

777 + self.slot_size_per_layer = list[int]() # HD bytes in kv terms

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

779 cache_list = cache_or_caches if split_k_and_v else [

780 cache_or_caches

@@ -789,10 +792,25 @@ class NixlConnectorWorker:

792 tensor_size_bytes = curr_tensor_size_bytes

793 self.num_blocks = cache.shape[0]

794

792 - assert tensor_size_bytes == curr_tensor_size_bytes, \

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

795 + assert cache.shape[0] == self.num_blocks, \

796 + "All kv cache tensors must have the same number of blocks"

797 +

798 + self.block_len_per_layer.append(curr_tensor_size_bytes //

799 + self.num_blocks)

800 + self.slot_size_per_layer.append(self.block_len_per_layer[-1] //

801 + self.block_size)

802 +

803 + if not self.use_mla:

804 + # Different kv cache shape is not supported by HeteroTP

805 + assert tensor_size_bytes == curr_tensor_size_bytes, \

806 + "All kv cache tensors must have the same size"

807 caches_data.append(

795 - (base_addr, tensor_size_bytes, self.tp_rank, ""))

808 + (base_addr, curr_tensor_size_bytes, self.tp_rank, ""))

809 +

810 + logger.debug("Different block lengths collected: %s",

811 + set(self.block_len_per_layer))

812 + assert len(self.block_len_per_layer) == len(seen_base_addresses)

813 + assert self.num_blocks != 0

814

815 self.kv_caches_base_addr[self.engine_id] = seen_base_addresses

816 self.num_regions = len(caches_data)

@@ -805,16 +823,12 @@ class NixlConnectorWorker:

823 logger.debug("Done registering descs")

824 self._registered_descs.append(descs)

825

808 - assert tensor_size_bytes is not None

809 - assert self.num_blocks != 0

810 - assert tensor_size_bytes % self.num_blocks == 0

811 - self.block_len = tensor_size_bytes // self.num_blocks

812 - self.slot_size_bytes = self.block_len // self.block_size

826 self.device_kv_caches = kv_caches

827 self.dst_num_blocks[self.engine_id] = self.num_blocks

828 if self._use_flashinfer:

816 - assert self.slot_size_bytes % 2 == 0

817 - self.slot_size_bytes /= 2

829 + for i in range(len(self.slot_size_per_layer)):

830 + assert self.slot_size_per_layer[i] % 2 == 0

831 + self.slot_size_per_layer[i] //= 2

832

833 # NOTE (NickLucche) When FlashInfer is used, memory is registered

834 # with joint KV for each block. This minimizes the overhead in

@@ -824,17 +838,17 @@ class NixlConnectorWorker:

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

839 self.num_regions *= 2

840

827 - kv_block_len = self.get_backend_aware_kv_block_len()

841 # Register local/src descr for NIXL xfer.

842 blocks_data = []

830 - for base_addr in seen_base_addresses:

843 + for i, base_addr in enumerate(seen_base_addresses):

844 + kv_block_len = self.get_backend_aware_kv_block_len(layer_idx=i)

845 # NOTE With heter-TP, more blocks are prepared than what are

846 # needed as self.num_blocks >= nixl_agent_meta.num_blocks. We

847 # could create fewer, but then _get_block_descs_ids needs to

848 # select agent_meta.num_blocks instead of self.num_blocks for

849 # local descr, and that makes handling regular flow less clean.

850 for block_id in range(self.num_blocks):

837 - block_offset = block_id * self.block_len

851 + block_offset = block_id * self.block_len_per_layer[i]

852 addr = base_addr + block_offset

853 # (addr, len, device id)

854 blocks_data.append((addr, kv_block_len, self.tp_rank))

@@ -844,7 +858,7 @@ class NixlConnectorWorker:

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

859 # when split across TP ranks.

860 for block_id in range(self.num_blocks):

847 - block_offset = block_id * self.block_len

861 + block_offset = block_id * self.block_len_per_layer[i]

862 addr = base_addr + block_offset

863 # Register addresses for V cache (K registered first).

864 v_addr = addr + kv_block_len

@@ -884,7 +898,7 @@ class NixlConnectorWorker:

898 agent_metadata=self.nixl_wrapper.get_agent_metadata(),

899 kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id],

900 num_blocks=self.num_blocks,

887 - block_len=self.block_len,

901 + block_lens=self.block_len_per_layer,

902 attn_backend_name=self.backend_name,

903 kv_cache_layout=self.kv_cache_layout)

904 ready_event = threading.Event()

@@ -909,7 +923,7 @@ class NixlConnectorWorker:

923 The latter, assuming D.world_size > P.world_size, requires that two or

924 more local TP worker share the xfer from a single TP worker.

925

912 - Here's an example:

926 + Here's an example (non-MLA case):

927

928 rank_offset p_remote_tp_rank

929 (kv split no)

@@ -965,14 +979,20 @@ class NixlConnectorWorker:

979 total_num_kv_heads = self.model_config.get_total_num_kv_heads()

980 is_kv_replicated = self._tp_size[engine_id] // total_num_kv_heads >= 1

981

982 + remote_block_len = nixl_agent_meta.block_lens[0]

983 if self.use_mla or is_kv_replicated:

969 - # With MLA the only difference is in the number of blocks.

970 - remote_block_size = nixl_agent_meta.block_len // (

971 - self.slot_size_bytes)

972 - assert self.block_len == nixl_agent_meta.block_len

984 + # With replicated KV cache, only the number of blocks can differ.

985 + assert self.block_len_per_layer == nixl_agent_meta.block_lens, \

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

987 + remote_block_size = remote_block_len // (

988 + self.slot_size_per_layer[0])

989 else:

974 - remote_block_size = nixl_agent_meta.block_len // (

975 - self.slot_size_bytes * tp_ratio)

990 + # When MLA is not used, this is a list of the same block length

991 + for block_len in nixl_agent_meta.block_lens:

992 + assert block_len == remote_block_len, \

993 + "All remote layers must have the same block size"

994 + remote_block_size = remote_block_len // (

995 + self.slot_size_per_layer[0] * tp_ratio)

996 if self._use_flashinfer:

997 # With flashinfer, KV are sent in the same message.

998 remote_block_size //= 2

@@ -983,14 +1003,14 @@ class NixlConnectorWorker:

1003 raise ValueError(

1004 "Heterogeneous TP is not supported on XPU")

1005

986 - assert nixl_agent_meta.block_len == self.block_len * tp_ratio, (

1006 + assert remote_block_len == self.block_len_per_layer[0] * tp_ratio, (

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

1008 "local_kv_heads*tp_ratio, block_size, head_dim] and same dtype."

1009 )

1010

1011 assert self.block_size == remote_block_size, (

992 - "Remote P worker with different block size is not supported "

993 - f"{self.block_size=} {remote_block_size=}")

1012 + "Remote P worker with different page/block size is not supported "

1013 + f"{self.block_size=}, {remote_block_size=}")

1014

1015 # Create dst descs and xfer side handles. TP workers have same #blocks.

1016 if engine_id in self.dst_num_blocks:

@@ -1005,13 +1025,16 @@ class NixlConnectorWorker:

1025 # Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..].

1026 self.kv_caches_base_addr[

1027 engine_id] = nixl_agent_meta.kv_caches_base_addr

1008 - kv_block_len = self.get_backend_aware_kv_block_len()

1009 - rank_offset = self.tp_rank % tp_ratio * kv_block_len \

1010 - if not (self.use_mla or is_kv_replicated) else 0

1028 +

1029 + assert len(nixl_agent_meta.kv_caches_base_addr) == len(

1030 + self.block_len_per_layer)

1031 # Register all remote blocks, but only the corresponding kv heads.

1012 - for base_addr in nixl_agent_meta.kv_caches_base_addr:

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

1033 + kv_block_len = self.get_backend_aware_kv_block_len(layer_idx=i)

1034 + rank_offset = self.tp_rank % tp_ratio * kv_block_len \

1035 + if not (self.use_mla or is_kv_replicated) else 0

1036 for block_id in range(nixl_agent_meta.num_blocks):

1014 - block_offset = block_id * nixl_agent_meta.block_len

1037 + block_offset = block_id * nixl_agent_meta.block_lens[i]

1038 # For each block, grab the heads chunk belonging to rank_i

1039 # of size remote_nheads // tp_ratio, which correspond to

1040 # self.block_len == remote_block_len//tp_ratio bytes.

@@ -1022,9 +1045,9 @@ class NixlConnectorWorker:

1045 if self._use_flashinfer:

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

1047 for block_id in range(nixl_agent_meta.num_blocks):

1025 - block_offset = block_id * nixl_agent_meta.block_len

1048 + block_offset = block_id * nixl_agent_meta.block_lens[i]

1049 addr = base_addr + block_offset + rank_offset

1027 - v_addr = addr + nixl_agent_meta.block_len // 2

1050 + v_addr = addr + nixl_agent_meta.block_lens[i] // 2

1051 blocks_data.append((v_addr, kv_block_len, remote_tp_rank))

1052

1053 logger.debug(

@@ -1351,7 +1374,7 @@ class NixlConnectorWorker:

1374 descs_ids = region_ids * num_blocks + block_ids

1375 return descs_ids.flatten()

1376

1354 - def get_backend_aware_kv_block_len(self):

1377 + def get_backend_aware_kv_block_len(self, layer_idx: int):

1378 """

1379 Get the block length for one K/V element (K and V have the same size).

1380

@@ -1362,9 +1385,9 @@ class NixlConnectorWorker:

1385 """

1386 if self._use_flashinfer:

1387 # For indexing only half (either just the K or V part).

1365 - block_len = self.block_len // 2

1388 + block_len = self.block_len_per_layer[layer_idx] // 2

1389 else:

1367 - block_len = self.block_len

1390 + block_len = self.block_len_per_layer[layer_idx]

1391 return block_len

1392

1393 def get_kv_connector_stats(self) -> Optional[KVConnectorStats]: