b745e8b5d3 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import asyncio

4 +import logging

5 import threading

6 import time

7 from collections import defaultdict

@@ -66,6 +67,179 @@ TransferId = str # KV transfer coordination ID (shared by P/D)

67 logger = init_logger(__name__)

68

69

70 +@dataclass(frozen=True)

71 +class TransferRegion:

72 + base_addr: int

73 + block_len: int

74 + kv_block_len: int

75 +

76 +

77 +def _get_tp_ratio(local_tp_size: int, remote_tp_size: int) -> int:

78 + """Return the TP ratio used by heterogeneous TP transfer planning.

79 +

80 + Positive values mean one local rank maps into a larger remote KV region.

81 + Negative values mean one local rank must gather from multiple remote KV

82 + regions.

83 + """

84 + if local_tp_size >= remote_tp_size:

85 + assert local_tp_size % remote_tp_size == 0, (

86 + f"Local tensor parallel size {local_tp_size} is not divisible "

87 + f"by remote tensor parallel size {remote_tp_size}."

88 + )

89 + return local_tp_size // remote_tp_size

90 +

91 + assert remote_tp_size % local_tp_size == 0, (

92 + f"Remote tensor parallel size {remote_tp_size} is not divisible "

93 + f"by local tensor parallel size {local_tp_size}."

94 + )

95 + return -(remote_tp_size // local_tp_size)

96 +

97 +

98 +def _expand_transfer_regions(

99 + base_addrs: list[int],

100 + block_lens: list[int],

101 + is_kv_layout_blocks_first: bool,

102 +) -> list[TransferRegion]:

103 + """Expand registered KV tensors into the regions transferred by Mooncake."""

104 + assert len(base_addrs) == len(block_lens), (

105 + "Mooncake transfer regions require matching numbers of base addresses "

106 + f"and block lengths, got {len(base_addrs)} and {len(block_lens)}."

107 + )

108 + regions: list[TransferRegion] = []

109 + for base_addr, block_len in zip(base_addrs, block_lens):

110 + kv_block_len = block_len // 2 if is_kv_layout_blocks_first else block_len

111 + regions.append(

112 + TransferRegion(

113 + base_addr=base_addr,

114 + block_len=block_len,

115 + kv_block_len=kv_block_len,

116 + )

117 + )

118 + if is_kv_layout_blocks_first:

119 + regions.append(

120 + TransferRegion(

121 + base_addr=base_addr + kv_block_len,

122 + block_len=block_len,

123 + kv_block_len=kv_block_len,

124 + )

125 + )

126 + return regions

127 +

128 +

129 +def _compute_sender_transfer_plan(

130 + local_tp_rank: int,

131 + local_tp_size: int,

132 + remote_tp_rank: int,

133 + remote_tp_size: int,

134 + local_kv_block_len: int,

135 + remote_kv_block_len: int,

136 + producer_cache_replicated: bool,

137 +) -> tuple[bool, int, int, int]:

138 + """Plan one producer-rank to one consumer-rank copy for heterogeneous TP."""

139 + tp_ratio = _get_tp_ratio(local_tp_size, remote_tp_size)

140 +

141 + if tp_ratio == 1:

142 + return True, 0, 0, local_kv_block_len

143 +

144 + if tp_ratio > 0:

145 + if producer_cache_replicated:

146 + return local_tp_rank % tp_ratio == 0, 0, 0, local_kv_block_len

147 + return (

148 + True,

149 + 0,

150 + (local_tp_rank % tp_ratio) * local_kv_block_len,

151 + local_kv_block_len,

152 + )

153 +

154 + if producer_cache_replicated:

155 + return True, 0, 0, local_kv_block_len

156 +

157 + ratio_abs = -tp_ratio

158 + return (

159 + True,

160 + (remote_tp_rank % ratio_abs) * remote_kv_block_len,

161 + 0,

162 + remote_kv_block_len,

163 + )

164 +

165 +

166 +def _can_coalesce_block_transfers(

167 + local_region_block_len: int,

168 + remote_region_block_len: int,

169 + src_region_offset: int,

170 + dst_region_offset: int,

171 + transfer_len: int,

172 +) -> bool:

173 + """Whether a contiguous block group can be emitted as one larger copy."""

174 + return (

175 + src_region_offset == 0

176 + and dst_region_offset == 0

177 + and transfer_len == local_region_block_len

178 + and transfer_len == remote_region_block_len

179 + )

180 +

181 +

182 +def _validate_asymmetric_region_lengths(

183 + local_regions: list[TransferRegion],

184 + remote_regions: list[TransferRegion],

185 + local_tp_size: int,

186 + remote_tp_size: int,

187 + producer_cache_replicated: bool,

188 +) -> str | None:

189 + """Validate transfer-region metadata for a fixed producer/consumer pair.

190 +

191 + This checks registered KV regions, not per-request block counts. A region

192 + corresponds to one registered KV tensor, or one K/V half after expansion

193 + for layouts that store K and V together.

194 + """

195 + if len(local_regions) != len(remote_regions):

196 + return (

197 + "Mooncake asymmetric TP requires matching KV region counts between "

198 + "producer and consumer."

199 + )

200 +

201 + if producer_cache_replicated:

202 + return None

203 +

204 + tp_ratio = _get_tp_ratio(local_tp_size, remote_tp_size)

205 + for idx, (local_region, remote_region) in enumerate(

206 + zip(local_regions, remote_regions)

207 + ):

208 + if tp_ratio == 1:

209 + if local_region.kv_block_len != remote_region.kv_block_len:

210 + return (

211 + "Mooncake KV region length mismatch for homogeneous TP at "

212 + f"region {idx}: local={local_region.kv_block_len}, "

213 + f"remote={remote_region.kv_block_len}."

214 + )

215 + elif tp_ratio > 0:

216 + if remote_region.kv_block_len != local_region.kv_block_len * tp_ratio:

217 + return (

218 + "Mooncake destination KV region length does not match the "

219 + "producer TP ratio at region "

220 + f"{idx}: local={local_region.kv_block_len}, "

221 + f"remote={remote_region.kv_block_len}, tp_ratio={tp_ratio}."

222 + )

223 + else:

224 + ratio_abs = -tp_ratio

225 + if local_region.kv_block_len != remote_region.kv_block_len * ratio_abs:

226 + return (

227 + "Mooncake source KV region length does not match the "

228 + "consumer TP ratio at region "

229 + f"{idx}: local={local_region.kv_block_len}, "

230 + f"remote={remote_region.kv_block_len}, tp_ratio={tp_ratio}."

231 + )

232 +

233 + return None

234 +

235 +

236 +def _get_tensor_dense_flag(tensor: torch.Tensor) -> bool | None:

237 + is_dense = getattr(tensor, "is_non_overlapping_and_dense", None)

238 + if callable(is_dense):

239 + return bool(is_dense())

240 + return None

241 +

242 +

243 class MooncakeXferMetadata(

244 msgspec.Struct,

245 omit_defaults=True, # type: ignore[call-arg]

@@ -76,6 +250,7 @@ class MooncakeXferMetadata(

250 remote_tp_rank: int

251 req_blocks: dict[ReqId, tuple[TransferId, list[int]]]

252 kv_caches_base_addr: list[int]

253 + block_lens: list[int]

254

255

256 class MooncakeXferResponseStatus(IntEnum):

@@ -173,6 +348,24 @@ class MooncakeConnector(KVConnectorBase_V1):

348 self.connector_scheduler = None

349 self.connector_worker = MooncakeConnectorWorker(vllm_config, self.engine_id)

350

351 + @classmethod

352 + def get_required_kvcache_layout(cls, vllm_config: VllmConfig):

353 + if vllm_config.model_config is None:

354 + # This fallback mostly exists for unit tests that instantiate the

355 + # connector without a fully populated model config.

356 + logger.warning_once(

357 + "Unable to detect current VLLM config. "

358 + "Fallback to default kv cache layout."

359 + )

360 + return None

361 + if vllm_config.model_config.use_mla:

362 + return None

363 + logger.info_once(

364 + "MooncakeConnector setting KV cache layout to HND for "

365 + "heterogeneous TP-safe KV transfer."

366 + )

367 + return "HND"

368 +

369 ############################################################

370 # Scheduler Side Methods

371 ############################################################

@@ -487,6 +680,8 @@ class MooncakeConnectorWorker:

680 self.tp_rank = get_tensor_model_parallel_rank()

681 self.tp_size = get_tensor_model_parallel_world_size()

682 self.num_blocks = 0

683 + self.block_len_per_layer: list[int] = []

684 + self.seen_base_addresses: list[int] = []

685

686 assert (parallel_config := vllm_config.parallel_config)

687 dp_rank = parallel_config.data_parallel_index

@@ -685,9 +880,13 @@ class MooncakeConnectorWorker:

880 ):

881 pending_reqs: dict[ReqId, SendBlockMeta] = {}

882 remote_tp_ranks = self.kv_topo.get_target_remote_ranks(meta.remote_tp_size)

688 - if self.tp_rank not in remote_tp_ranks:

883 + if meta.remote_tp_rank not in remote_tp_ranks:

884 # This D worker does not pair with the P worker.

690 - msg = f"This P tp_rank {self.tp_rank} not in remote D target ranks {remote_tp_ranks}" # noqa: E501

885 + msg = (

886 + "This D tp_rank "

887 + f"{meta.remote_tp_rank} is not paired with P tp_rank "

888 + f"{self.tp_rank}; expected one of {remote_tp_ranks}."

889 + )

890 logger.error(msg)

891 response = MooncakeXferResponse(

892 status=MooncakeXferResponseStatus.ERROR,

@@ -695,6 +894,26 @@ class MooncakeConnectorWorker:

894 )

895 await sock.send_multipart((identity, self._encoder.encode(response)))

896 return

897 + local_regions = self._get_transfer_regions(

898 + self.kv_caches_base_addr, self.block_len_per_layer

899 + )

900 + remote_regions = self._get_transfer_regions(

901 + meta.kv_caches_base_addr, meta.block_lens

902 + )

903 + validation_err = _validate_asymmetric_region_lengths(

904 + local_regions=local_regions,

905 + remote_regions=remote_regions,

906 + local_tp_size=self.tp_size,

907 + remote_tp_size=meta.remote_tp_size,

908 + producer_cache_replicated=self._producer_cache_is_replicated(),

909 + )

910 + if validation_err is not None:

911 + response = MooncakeXferResponse(

912 + status=MooncakeXferResponseStatus.ERROR,

913 + err_msg=validation_err,

914 + )

915 + await sock.send_multipart((identity, self._encoder.encode(response)))

916 + return

917 for d_req_id, (transfer_id, _) in meta.req_blocks.items():

918 if transfer_id not in self.reqs_need_send:

919 # This req is not enqueued in P side yet, create it here.

@@ -763,17 +982,24 @@ class MooncakeConnectorWorker:

982 "Request %s expired before sending on P side.", d_req_id

983 )

984

766 - src_ptrs, dst_ptrs, lengths, err_reqs = await self._build_transfer_params(

767 - ready_reqs, meta

985 + (

986 + src_ptrs,

987 + dst_ptrs,

988 + lengths,

989 + err_reqs,

990 + err_msg,

991 + ) = await self._build_transfer_params(

992 + ready_reqs,

993 + meta,

994 + local_regions,

995 + remote_regions,

996 )

769 -

770 - if err_reqs:

771 - response = MooncakeXferResponse(

772 - status=response_status,

773 - err_reqs=err_reqs,

774 - err_msg="P num blocks less than D",

775 - )

776 - await sock.send_multipart((identity, self._encoder.encode(response)))

997 + err_req_set = set(err_reqs)

998 + ok_ready_reqs = [

999 + (d_req_id, send_meta)

1000 + for d_req_id, send_meta in ready_reqs

1001 + if d_req_id not in err_req_set

1002 + ]

1003

1004 if src_ptrs:

1005 remote_session = f"{meta.remote_hostname}:{meta.remote_port}"

@@ -787,58 +1013,61 @@ class MooncakeConnectorWorker:

1013 )

1014

1015 if ret_value != 0:

790 - err_reqs = []

791 - for d_req_id, send_meta in ready_reqs:

792 - send_meta.sending -= 1

793 - err_reqs.append(d_req_id)

794 - # Do best effort to transfer the remaining reqs.

795 - response = MooncakeXferResponse(

796 - status=response_status,

797 - err_reqs=err_reqs,

798 - err_msg=f"Mooncake transfer engine returned {ret_value}",

1016 + transfer_err_msg = f"Mooncake transfer engine returned {ret_value}"

1017 + err_msg = (

1018 + transfer_err_msg

1019 + if err_msg is None

1020 + else f"{err_msg}; {transfer_err_msg}"

1021 )

800 - await sock.send_multipart(

801 - (identity, self._encoder.encode(response))

802 - )

803 - continue

1022 + err_reqs = list(err_reqs)

1023 + for d_req_id, _ in ok_ready_reqs:

1024 + err_reqs.append(d_req_id)

1025 + err_req_set.add(d_req_id)

1026 + ok_ready_reqs = []

1027

1028 for d_req_id, send_meta in ready_reqs:

806 - # TODO: for heterogeneous TP (one P pairs to multiple D),

807 - # we need to check whether all headers are sent.

808 - # If not, we should set expire_time to normal and skip the below.

1029 send_meta.sending -= 1

1030 +

1031 + if d_req_id in err_req_set:

1032 + continue

1033 +

1034 send_meta.sent += 1

811 - if send_meta.sent == send_meta.need_send:

812 - del self.reqs_need_send[send_meta.transfer_id]

1035 + if (

1036 + send_meta.sent == send_meta.need_send

1037 + and self.reqs_need_send.pop(send_meta.transfer_id, None) is not None

1038 + ):

1039 self.finished_sending_reqs.add(send_meta.p_req_id)

1040

1041 response = MooncakeXferResponse(

1042 status=response_status,

817 - ok_reqs=[d_req_id for d_req_id, _ in ready_reqs],

1043 + ok_reqs=[d_req_id for d_req_id, _ in ok_ready_reqs] or None,

1044 + err_reqs=err_reqs or None,

1045 + err_msg=err_msg,

1046 )

1047 await sock.send_multipart((identity, self._encoder.encode(response)))

1048

1049 def resolve_need_send(self, send_meta: SendBlockMeta, remote_tp_ranks: list[int]):

1050 # Prepare for heterogeneous TP (one P pairs to multiple D)

1051 send_meta.need_send = len(remote_tp_ranks)

824 - if send_meta.need_send != 1:

825 - logger.error("Mooncake: Heterogeneous TP is not supported yet.")

826 - raise NotImplementedError(

827 - "Mooncake: Heterogeneous TP is not supported yet."

828 - )

1052 + logger.debug(

1053 + "Mooncake request %s will be served by %d consumer TP workers: %s",

1054 + send_meta.transfer_id,

1055 + send_meta.need_send,

1056 + remote_tp_ranks,

1057 + )

1058

1059 async def _build_transfer_params(

1060 self,

1061 ready_reqs: list[tuple[ReqId, SendBlockMeta]],

1062 agent_meta: MooncakeXferMetadata,

834 - ) -> tuple[list[int], list[int], list[int], list[ReqId]]:

1063 + local_regions: list[TransferRegion],

1064 + remote_regions: list[TransferRegion],

1065 + ) -> tuple[list[int], list[int], list[int], list[ReqId], str | None]:

1066 src_ptrs = []

1067 dst_ptrs = []

1068 lengths = []

1069 err_reqs: list[ReqId] = []

839 - local_base_addr = self.kv_caches_base_addr

840 - remote_base_addr = agent_meta.kv_caches_base_addr

841 - block_len = self.block_len

1070 + err_msg: str | None = None

1071 remote_session = f"{agent_meta.remote_hostname}:{agent_meta.remote_port}"

1072

1073 for d_req_id, send_meta in ready_reqs:

@@ -858,6 +1087,8 @@ class MooncakeConnectorWorker:

1087 num_remote_blocks,

1088 )

1089 err_reqs.append(d_req_id)

1090 + if err_msg is None:

1091 + err_msg = "P num blocks less than D"

1092 continue

1093 if num_local_blocks > num_remote_blocks:

1094 local_block_ids = local_block_ids[-num_remote_blocks:]

@@ -867,19 +1098,87 @@ class MooncakeConnectorWorker:

1098 local_block_ids, remote_block_ids

1099 )

1100

870 - for local_layer_addr, remote_layer_addr in zip(

871 - local_base_addr, remote_base_addr

872 - ):

1101 + for local_region, remote_region in zip(local_regions, remote_regions):

1102 + should_transfer, src_region_offset, dst_region_offset, transfer_len = (

1103 + self._get_sender_transfer_plan(

1104 + local_kv_block_len=local_region.kv_block_len,

1105 + remote_kv_block_len=remote_region.kv_block_len,

1106 + remote_tp_rank=agent_meta.remote_tp_rank,

1107 + remote_tp_size=agent_meta.remote_tp_size,

1108 + )

1109 + )

1110 + if not should_transfer:

1111 + # Replicated KV cache: only one producer rank in the TP group

1112 + # needs to send the actual bytes for this paired decoder rank.

1113 + # TODO: Account for replicated producer KV in

1114 + # get_target_remote_ranks() so we can avoid sending

1115 + # unnecessary ZMQ requests and remove this branch.

1116 + continue

1117 +

1118 + assert src_region_offset + transfer_len <= local_region.kv_block_len, (

1119 + "Computed source transfer region exceeds local KV block size."

1120 + )

1121 + assert dst_region_offset + transfer_len <= remote_region.kv_block_len, (

1122 + "Computed destination transfer region exceeds remote KV block size."

1123 + )

1124 + # Collapse one contiguous block group into a single larger

1125 + # transfer descriptor when the per-block copy is identical.

1126 + can_coalesce = _can_coalesce_block_transfers(

1127 + local_region_block_len=local_region.block_len,

1128 + remote_region_block_len=remote_region.block_len,

1129 + src_region_offset=src_region_offset,

1130 + dst_region_offset=dst_region_offset,

1131 + transfer_len=transfer_len,

1132 + )

1133 +

1134 for group_local_block_id, group_remote_block_id in zip(

1135 group_local_block_ids, group_remote_block_ids

1136 ):

876 - src_ptrs.append(

877 - local_layer_addr + group_local_block_id[0] * block_len

878 - )

879 - dst_ptrs.append(

880 - remote_layer_addr + group_remote_block_id[0] * block_len

1137 + if can_coalesce:

1138 + src_ptrs.append(

1139 + local_region.base_addr

1140 + + group_local_block_id[0] * local_region.block_len

1141 + + src_region_offset

1142 + )

1143 + dst_ptrs.append(

1144 + remote_region.base_addr

1145 + + group_remote_block_id[0] * remote_region.block_len

1146 + + dst_region_offset

1147 + )

1148 + lengths.append(transfer_len * len(group_local_block_id))

1149 + else:

1150 + for local_block_id, remote_block_id in zip(

1151 + group_local_block_id, group_remote_block_id

1152 + ):

1153 + src_ptrs.append(

1154 + local_region.base_addr

1155 + + local_block_id * local_region.block_len

1156 + + src_region_offset

1157 + )

1158 + dst_ptrs.append(

1159 + remote_region.base_addr

1160 + + remote_block_id * remote_region.block_len

1161 + + dst_region_offset

1162 + )

1163 + lengths.append(transfer_len)

1164 +

1165 + if local_region is local_regions[0]:

1166 + logger.debug(

1167 + "Mooncake transfer plan for request %s: local_tp=%d "

1168 + "remote_tp=%d remote_tp_rank=%d local_block_len=%d "

1169 + "remote_block_len=%d src_offset=%d dst_offset=%d "

1170 + "transfer_len=%d coalesce=%s",

1171 + d_req_id,

1172 + self.tp_size,

1173 + agent_meta.remote_tp_size,

1174 + agent_meta.remote_tp_rank,

1175 + local_region.block_len,

1176 + remote_region.block_len,

1177 + src_region_offset,

1178 + dst_region_offset,

1179 + transfer_len,

1180 + can_coalesce,

1181 )

882 - lengths.append(block_len * len(group_local_block_id))

1182

1183 logger.debug(

1184 "Sending kv_caches for request %s (%d blocks) to %s",

@@ -888,7 +1187,7 @@ class MooncakeConnectorWorker:

1187 remote_session,

1188 )

1189

891 - return src_ptrs, dst_ptrs, lengths, err_reqs

1190 + return src_ptrs, dst_ptrs, lengths, err_reqs, err_msg

1191

1192 def _send_blocks(

1193 self,

@@ -917,16 +1216,20 @@ class MooncakeConnectorWorker:

1216 kv_data_ptrs = []

1217 kv_data_lens = []

1218 seen_base_addresses = []

1219 + self.block_len_per_layer = []

1220

1221 split_k_and_v = self.kv_topo.split_k_and_v

1222 tensor_size_bytes = None

1223 for layer_name, cache_or_caches in kv_caches.items():

1224 + cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]

1225 logger.debug(

925 - "registering layer %s with shape %s", layer_name, cache_or_caches.shape

1226 + "registering layer %s with %d cache tensor(s)",

1227 + layer_name,

1228 + len(cache_list),

1229 )

927 - cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]

1230

1231 for cache in cache_list:

1232 + self._log_debug_cache_registration(layer_name, cache)

1233 base_addr = cache.data_ptr()

1234 if base_addr in seen_base_addresses:

1235 continue

@@ -937,16 +1240,24 @@ class MooncakeConnectorWorker:

1240 if tensor_size_bytes is None:

1241 tensor_size_bytes = curr_tensor_size_bytes

1242 self.num_blocks = cache.shape[0]

940 -

941 - assert tensor_size_bytes == curr_tensor_size_bytes, (

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

1243 + assert cache.shape[0] == self.num_blocks, (

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

1245 + )

1246 + assert curr_tensor_size_bytes % self.num_blocks == 0, (

1247 + "Mooncake expects each kv cache tensor size to be "

1248 + "divisible by the number of blocks."

1249 + )

1250 + self.block_len_per_layer.append(

1251 + curr_tensor_size_bytes // self.num_blocks

1252 )

1253 +

1254 kernel_block_size = cache.shape[-2 if self.use_mla else -3]

1255 assert self.block_size == kernel_block_size

1256 kv_data_ptrs.append(base_addr)

947 - kv_data_lens.append(tensor_size_bytes)

1257 + kv_data_lens.append(curr_tensor_size_bytes)

1258

1259 self.kv_caches_base_addr = seen_base_addresses

1260 + self.seen_base_addresses = seen_base_addresses

1261

1262 ret_value = self.engine.batch_register_memory(kv_data_ptrs, kv_data_lens)

1263 if ret_value != 0:

@@ -954,11 +1265,11 @@ class MooncakeConnectorWorker:

1265

1266 assert tensor_size_bytes is not None

1267 assert self.num_blocks != 0

957 - assert tensor_size_bytes % self.num_blocks == 0

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

1268 self.device_kv_caches = kv_caches

1269 logger.debug(

961 - "registered num_blocks=%d block_len=%d", self.num_blocks, self.block_len

1270 + "registered num_blocks=%d block_lens=%s",

1271 + self.num_blocks,

1272 + self.block_len_per_layer,

1273 )

1274

1275 # No need to launch server for D node.

@@ -1052,6 +1363,7 @@ class MooncakeConnectorWorker:

1363 for req_id, pull_meta in pull_metas.items()

1364 },

1365 kv_caches_base_addr=self.kv_caches_base_addr,

1366 + block_lens=self.block_len_per_layer,

1367 )

1368

1369 encoded_data = self._encoder.encode(metadata)

@@ -1152,11 +1464,11 @@ class MooncakeConnectorWorker:

1464 remote_engine_id

1465 )

1466 count = len(remote_tp_ranks)

1155 - if count != 1:

1156 - logger.error("Mooncake: Heterogeneous TP is not supported yet.")

1157 - raise NotImplementedError(

1158 - "Mooncake: Heterogeneous TP is not supported yet."

1159 - )

1467 + logger.debug(

1468 + "Receiving Mooncake KV for engine %s from producer TP ranks %s",

1469 + remote_engine_id,

1470 + remote_tp_ranks,

1471 + )

1472 for pull_meta in pull_metas.values():

1473 pull_meta.pull_tasks_count = count

1474 for remote_tp_rank in remote_tp_ranks:

@@ -1239,6 +1551,52 @@ class MooncakeConnectorWorker:

1551 self.record_send_reqs(metadata), self.sender_loop

1552 )

1553

1554 + def _producer_cache_is_replicated(self) -> bool:

1555 + return self.kv_topo.replicates_kv_cache(self.engine_id)

1556 +

1557 + def _get_transfer_regions(

1558 + self, base_addrs: list[int], block_lens: list[int]

1559 + ) -> list[TransferRegion]:

1560 + return _expand_transfer_regions(

1561 + base_addrs=base_addrs,

1562 + block_lens=block_lens,

1563 + is_kv_layout_blocks_first=self.kv_topo.is_kv_layout_blocks_first,

1564 + )

1565 +

1566 + def _get_sender_transfer_plan(

1567 + self,

1568 + local_kv_block_len: int,

1569 + remote_kv_block_len: int,

1570 + remote_tp_rank: int,

1571 + remote_tp_size: int,

1572 + ) -> tuple[bool, int, int, int]:

1573 + return _compute_sender_transfer_plan(

1574 + local_tp_rank=self.tp_rank,

1575 + local_tp_size=self.tp_size,

1576 + remote_tp_rank=remote_tp_rank,

1577 + remote_tp_size=remote_tp_size,

1578 + local_kv_block_len=local_kv_block_len,

1579 + remote_kv_block_len=remote_kv_block_len,

1580 + producer_cache_replicated=self._producer_cache_is_replicated(),

1581 + )

1582 +

1583 + def _log_debug_cache_registration(

1584 + self, layer_name: str, cache: torch.Tensor

1585 + ) -> None:

1586 + if not logger.isEnabledFor(logging.DEBUG):

1587 + return

1588 + logger.debug(

1589 + "Mooncake register view layer=%s shape=%s stride=%s "

1590 + "storage_offset=%d contiguous=%s dense=%s data_ptr=%d",

1591 + layer_name,

1592 + tuple(cache.shape),

1593 + tuple(cache.stride()),

1594 + cache.storage_offset(),

1595 + cache.is_contiguous(),

1596 + _get_tensor_dense_flag(cache),

1597 + cache.data_ptr(),

1598 + )

1599 +

1600

1601 def group_concurrent_contiguous(

1602 src_indices: list[int], dst_indices: list[int]

[KVTransfer][Mooncake] Add heterogeneous TP support for disaggregated P/D in MooncakeConnector (#36869)