sglang 0.4.5.post3__py3-none-any.whl → 0.4.6__py3-none-any.whl

sglang/srt/disaggregation/decode.py

@@ -137,7 +137,7 @@ class DecodePreallocQueue:
         kv_receiver_class = get_kv_class(self.transfer_backend, KVClassType.RECEIVER)
         kv_receiver = kv_receiver_class(
             mgr=self.kv_manager,
-            bootstrap_addr=f"{req.bootstrap_host}:{self.bootstrap_port}",
+            bootstrap_addr=f"{req.bootstrap_host}:{req.bootstrap_port}",
             bootstrap_room=req.bootstrap_room,
         )
         self.queue.append(DecodeRequest(req=req, kv_receiver=kv_receiver))
@@ -307,7 +307,7 @@ class DecodeTransferQueue:
     def extend(self, req_conns) -> None:
         self.queue.extend(req_conns)
 
-    def pop_transferred(self) -> List[Req]:
+    def pop_transferred(self) -> List[DecodeRequest]:
         if not self.queue:
             return []
 
@@ -330,7 +330,7 @@ class DecodeTransferQueue:
                 assert len(decode_req.req.output_ids) == 0
                 assert decode_req.req.transferred_output_id is None
                 decode_req.req.transferred_output_id = output_id
-                transferred_reqs.append(decode_req.req)
+                transferred_reqs.append(decode_req)
                 indices_to_remove.add(i)
             elif poll in [
                 KVPoll.Bootstrapping,
@@ -444,8 +444,17 @@ class ScheduleBatchDisaggregationDecodeMixin:
 
 class SchedulerDisaggregationDecodeMixin:
 
+    def _prepare_idle_batch_and_run(self, batch, delay_process=False):
+        batch, _ = self.prepare_dp_attn_batch(batch)
+        result = None
+        if batch:
+            result = self.run_batch(batch)
+            if not delay_process:
+                self.process_batch_result(batch, result)
+        return batch, result
+
     @torch.no_grad()
-    def event_loop_normal_disagg_decode(self):
+    def event_loop_normal_disagg_decode(self: Scheduler):
         """A normal scheduler loop for decode worker in disaggregation mode."""
 
         while True:
@@ -456,14 +465,25 @@ class SchedulerDisaggregationDecodeMixin:
             batch = self.get_next_disagg_decode_batch_to_run()
             self.cur_batch = batch
 
+            prepare_dp_attn_flag = (
+                self.server_args.enable_dp_attention
+                or self.server_args.enable_sp_layernorm
+            )
+
             if batch:
                 # Generate fake extend output.
                 if batch.forward_mode.is_extend():
                     # Note: Logprobs should be handled on the prefill engine.
                     self.stream_output(batch.reqs, False)
+                    if prepare_dp_attn_flag:
+                        self._prepare_idle_batch_and_run(None)
                 else:
+                    if prepare_dp_attn_flag:
+                        self.prepare_dp_attn_batch(batch)
                     result = self.run_batch(batch)
                     self.process_batch_result(batch, result)
+            elif prepare_dp_attn_flag:
+                batch, _ = self._prepare_idle_batch_and_run(None)
 
             if batch is None and (
                 len(self.disagg_decode_transfer_queue.queue)
@@ -477,10 +497,10 @@ class SchedulerDisaggregationDecodeMixin:
             self.last_batch = batch
 
     @torch.no_grad()
-    def event_loop_overlap_disagg_decode(self):
+    def event_loop_overlap_disagg_decode(self: Scheduler):
         result_queue = deque()
         self.last_batch: Optional[ScheduleBatch] = None
-        self.last_batch_is_extend = False  # last batch is modifed in-place, so we need another variable to track if it's extend
+        self.last_batch_in_queue = False  # last batch is modifed in-place, so we need another variable to track if it's extend
 
         while True:
             recv_reqs = self.recv_requests()
@@ -489,20 +509,41 @@ class SchedulerDisaggregationDecodeMixin:
             self.process_decode_queue()
             batch = self.get_next_disagg_decode_batch_to_run()
             self.cur_batch = batch
-            last_batch_is_extend = False
+            last_batch_in_queue = False
+
+            prepare_dp_attn_flag = (
+                self.server_args.enable_dp_attention
+                or self.server_args.enable_sp_layernorm
+            )
 
             if batch:
                 # Generate fake extend output.
                 if batch.forward_mode.is_extend():
                     # Note: Logprobs should be handled on the prefill engine.
                     self.stream_output(batch.reqs, False)
-                    last_batch_is_extend = True
+                    if prepare_dp_attn_flag:
+                        batch_, result = self._prepare_idle_batch_and_run(
+                            None, delay_process=True
+                        )
+                        if batch_:
+                            result_queue.append((batch_.copy(), result))
+                            last_batch_in_queue = True
                 else:
+                    if prepare_dp_attn_flag:
+                        self.prepare_dp_attn_batch(batch)
                     result = self.run_batch(batch)
                     result_queue.append((batch.copy(), result))
+                    last_batch_in_queue = True
+            elif prepare_dp_attn_flag:
+                batch, result = self._prepare_idle_batch_and_run(
+                    None, delay_process=True
+                )
+                if batch:
+                    result_queue.append((batch.copy(), result))
+                    last_batch_in_queue = True
 
             # Process the results of the previous batch but skip if the last batch is extend
-            if self.last_batch and not self.last_batch_is_extend:
+            if self.last_batch and self.last_batch_in_queue:
                 tmp_batch, tmp_result = result_queue.popleft()
                 self.process_batch_result(tmp_batch, tmp_result)
 
@@ -516,7 +557,7 @@ class SchedulerDisaggregationDecodeMixin:
                 self.new_token_ratio = self.init_new_token_ratio
 
             self.last_batch = batch
-            self.last_batch_is_extend = last_batch_is_extend
+            self.last_batch_in_queue = last_batch_in_queue
 
     def get_next_disagg_decode_batch_to_run(
         self: Scheduler,
@@ -600,8 +641,15 @@ class SchedulerDisaggregationDecodeMixin:
 
     def process_decode_queue(self: Scheduler):
         req_conns = self.disagg_decode_prealloc_queue.pop_preallocated()
+
+        def _num_pre_alloc(req):
+            return len(req.req.origin_input_ids) + max(len(req.req.output_ids) - 1, 0)
+
+        self.num_tokens_pre_allocated += sum(_num_pre_alloc(req) for req in req_conns)
         self.disagg_decode_transfer_queue.extend(req_conns)
         alloc_reqs = (
             self.disagg_decode_transfer_queue.pop_transferred()
         )  # the requests which kv has arrived
-        self.waiting_queue.extend(alloc_reqs)
+        self.num_tokens_pre_allocated -= sum(_num_pre_alloc(req) for req in alloc_reqs)
+
+        self.waiting_queue.extend([req.req for req in alloc_reqs])

sglang/srt/disaggregation/mini_lb.py

@@ -6,6 +6,7 @@ import asyncio
 import random
 import urllib
 from itertools import chain
+from typing import List
 
 import aiohttp
 import orjson
@@ -14,13 +15,22 @@ from fastapi import FastAPI, HTTPException
 from fastapi.responses import ORJSONResponse, Response, StreamingResponse
 
 
+class PrefillConfig:
+    def __init__(self, url: str, bootstrap_port: int):
+        self.url = url
+        self.bootstrap_port = bootstrap_port
+
+
 class MiniLoadBalancer:
-    def __init__(self, prefill_servers, decode_servers):
-        self.prefill_servers = prefill_servers
+    def __init__(self, prefill_configs: List[PrefillConfig], decode_servers: List[str]):
+        self.prefill_configs = prefill_configs
+        self.prefill_servers = [p.url for p in prefill_configs]
         self.decode_servers = decode_servers
 
     def select_pair(self):
-        return random.choice(self.prefill_servers), random.choice(self.decode_servers)
+        prefill_config = random.choice(self.prefill_configs)
+        decode_server = random.choice(self.decode_servers)
+        return prefill_config.url, prefill_config.bootstrap_port, decode_server
 
     async def generate(
         self, modified_request, prefill_server, decode_server, endpoint
@@ -160,7 +170,7 @@ async def get_model_info():
 
 @app.post("/generate")
 async def handle_generate_request(request_data: dict):
-    prefill_server, decode_server = load_balancer.select_pair()
+    prefill_server, bootstrap_port, decode_server = load_balancer.select_pair()
 
     # Parse and transform prefill_server for bootstrap data
     parsed_url = urllib.parse.urlparse(prefill_server)
@@ -172,6 +182,7 @@ async def handle_generate_request(request_data: dict):
         modified_request.update(
             {
                 "bootstrap_host": [hostname] * batch_size,
+                "bootstrap_port": [bootstrap_port] * batch_size,
                 "bootstrap_room": [
                     _generate_bootstrap_room() for _ in range(batch_size)
                 ],
@@ -181,6 +192,7 @@ async def handle_generate_request(request_data: dict):
         modified_request.update(
             {
                 "bootstrap_host": hostname,
+                "bootstrap_port": bootstrap_port,
                 "bootstrap_room": _generate_bootstrap_room(),
             }
         )
@@ -197,7 +209,7 @@ async def handle_generate_request(request_data: dict):
 
 @app.post("/v1/chat/completions")
 async def handle_completion_request(request_data: dict):
-    prefill_server, decode_server = load_balancer.select_pair()
+    prefill_server, bootstrap_port, decode_server = load_balancer.select_pair()
 
     # Parse and transform prefill_server for bootstrap data
     parsed_url = urllib.parse.urlparse(prefill_server)
@@ -206,6 +218,7 @@ async def handle_completion_request(request_data: dict):
     modified_request.update(
         {
             "bootstrap_host": hostname,
+            "bootstrap_port": bootstrap_port,
             "bootstrap_room": random.randint(0, 2**63 - 1),
         }
     )
@@ -255,9 +268,9 @@ async def get_models():
             raise HTTPException(status_code=500, detail=str(e))
 
 
-def run(prefill_addrs, decode_addrs, host, port):
+def run(prefill_configs, decode_addrs, host, port):
     global load_balancer
-    load_balancer = MiniLoadBalancer(prefill_addrs, decode_addrs)
+    load_balancer = MiniLoadBalancer(prefill_configs, decode_addrs)
     uvicorn.run(app, host=host, port=port)
 
 
@@ -268,6 +281,11 @@ if __name__ == "__main__":
     parser.add_argument(
         "--prefill", required=True, help="Comma-separated URLs for prefill servers"
     )
+    parser.add_argument(
+        "--prefill-bootstrap-ports",
+        help="Comma-separated bootstrap ports for prefill servers",
+        default="8998",
+    )
     parser.add_argument(
         "--decode", required=True, help="Comma-separated URLs for decode servers"
     )
@@ -278,4 +296,23 @@ if __name__ == "__main__":
         "--port", type=int, default=8000, help="Port to bind the server (default: 8000)"
     )
     args = parser.parse_args()
-    run(args.prefill.split(","), args.decode.split(","), args.host, args.port)
+
+    prefill_urls = args.prefill.split(",")
+    bootstrap_ports = [int(p) for p in args.prefill_bootstrap_ports.split(",")]
+
+    if len(bootstrap_ports) == 1:
+        bootstrap_ports = bootstrap_ports * len(prefill_urls)
+    else:
+        if len(bootstrap_ports) != len(prefill_urls):
+            raise ValueError(
+                "Number of prefill URLs must match number of bootstrap ports"
+            )
+            exit(1)
+
+    prefill_configs = []
+    for url, port in zip(prefill_urls, bootstrap_ports):
+        prefill_configs.append(PrefillConfig(url, port))
+
+    decode_addrs = args.decode.split(",")
+
+    run(prefill_configs, decode_addrs, args.host, args.port)

sglang/srt/disaggregation/mooncake/conn.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
 import dataclasses
 import logging
+import os
 import queue
 import socket
 import struct
@@ -73,9 +75,7 @@ class TransferInfo:
     endpoint: str
     dst_port: int
     mooncake_session_id: str
-    dst_kv_ptrs: list[int]
     dst_kv_indices: npt.NDArray[np.int64]
-    dst_aux_ptrs: list[int]
     dst_aux_index: int
 
     @classmethod
@@ -85,10 +85,29 @@ class TransferInfo:
             endpoint=msg[1].decode("ascii"),
             dst_port=int(msg[2].decode("ascii")),
             mooncake_session_id=msg[3].decode("ascii"),
+            dst_kv_indices=np.frombuffer(msg[4], dtype=np.int64),
+            dst_aux_index=int(msg[5].decode("ascii")),
+        )
+
+
+@dataclasses.dataclass
+class KVArgsRegisterInfo:
+    room: str
+    endpoint: str
+    dst_port: int
+    mooncake_session_id: str
+    dst_kv_ptrs: list[int]
+    dst_aux_ptrs: list[int]
+
+    @classmethod
+    def from_zmq(cls, msg: List[bytes]):
+        return cls(
+            room=str(msg[0].decode("ascii")),
+            endpoint=msg[1].decode("ascii"),
+            dst_port=int(msg[2].decode("ascii")),
+            mooncake_session_id=msg[3].decode("ascii"),
             dst_kv_ptrs=list(struct.unpack(f"{len(msg[4])//8}Q", msg[4])),
-            dst_kv_indices=np.frombuffer(msg[5], dtype=np.int64),
-            dst_aux_ptrs=list(struct.unpack(f"{len(msg[6])//8}Q", msg[6])),
-            dst_aux_index=int(msg[7].decode("ascii")),
+            dst_aux_ptrs=list(struct.unpack(f"{len(msg[5])//8}Q", msg[5])),
         )
 
 
@@ -109,6 +128,13 @@ class MooncakeKVManager(BaseKVManager):
         # for p/d multi node infer
         self.bootstrap_port = server_args.disaggregation_bootstrap_port
         self.dist_init_addr = server_args.dist_init_addr
+        self.tp_size = server_args.tp_size
+        self.dp_size = server_args.dp_size
+        self.enable_dp_attention = server_args.enable_dp_attention
+        if not server_args.enable_dp_attention and server_args.dp_size != 1:
+            raise ValueError(
+                "If dp_attention is not enabled, dp size must be 1 in disaggregation mode."
+            )
         self.request_status: Dict[int, KVPoll] = {}
         self.rank_port = None
         self.server_socket = zmq.Context().socket(zmq.PULL)
@@ -116,11 +142,19 @@ class MooncakeKVManager(BaseKVManager):
         if self.disaggregation_mode == DisaggregationMode.PREFILL:
             self.transfer_queue = queue.Queue()
             self.transfer_infos: Dict[int, TransferInfo] = {}
+            self.decode_kv_args_table: Dict[str, KVArgsRegisterInfo] = {}
             self.start_prefill_thread()
             self._register_to_bootstrap()
+
+            # Determine the number of threads to use for kv sender
+            cpu_count = os.cpu_count()
+            self.executor = concurrent.futures.ThreadPoolExecutor(
+                min(cpu_count // 4, 16)
+            )
         elif self.disaggregation_mode == DisaggregationMode.DECODE:
             self.start_decode_thread()
             self.connection_pool: Dict[str, Dict[str, Union[str, int]]] = {}
+            self.prefill_dp_size_table: Dict[str, int] = {}
         else:
             raise ValueError(
                 f"Unsupported DisaggregationMode: {self.disaggregation_mode}"
@@ -150,28 +184,53 @@ class MooncakeKVManager(BaseKVManager):
         dst_kv_ptrs: list[int],
         dst_kv_indices: npt.NDArray[np.int64],
     ):
-        # group by indices
+        # Group by indices
         prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
             prefill_kv_indices, dst_kv_indices
         )
 
         num_layers = len(self.kv_args.kv_data_ptrs)
-        for layer_id in range(num_layers):
-            src_ptr = self.kv_args.kv_data_ptrs[layer_id]
-            dst_ptr = dst_kv_ptrs[layer_id]
-            item_len = self.kv_args.kv_item_lens[layer_id]
+        layers_params = [
+            (
+                self.kv_args.kv_data_ptrs[layer_id],
+                dst_kv_ptrs[layer_id],
+                self.kv_args.kv_item_lens[layer_id],
+            )
+            for layer_id in range(num_layers)
+        ]
 
+        # Worker function for processing a single layer
+        def process_layer(src_ptr: int, dst_ptr: int, item_len: int) -> int:
             for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks):
                 src_addr = src_ptr + int(prefill_index[0]) * item_len
                 dst_addr = dst_ptr + int(decode_index[0]) * item_len
                 length = item_len * len(prefill_index)
 
-                # TODO: make async later
                 status = self.engine.transfer_sync(
                     mooncake_session_id, src_addr, dst_addr, length
                 )
                 if status != 0:
                     return status
+            return 0
+
+        futures = [
+            self.executor.submit(
+                process_layer,
+                src_ptr,
+                dst_ptr,
+                item_len,
+            )
+            for (src_ptr, dst_ptr, item_len) in layers_params
+        ]
+
+        for future in concurrent.futures.as_completed(futures):
+            status = future.result()
+            if status != 0:
+                # Immediate shutdown on first error (existing tasks will finish)
+                executor.shutdown(wait=False)
+                for f in futures:
+                    f.cancel()
+                return status
 
         return 0
 
@@ -215,6 +274,13 @@ class MooncakeKVManager(BaseKVManager):
                 waiting_req_bytes = self.server_socket.recv_multipart()
                 room = waiting_req_bytes[0].decode("ascii")
                 if room == "None":
+                    mooncake_session_id = waiting_req_bytes[3].decode("ascii")
+                    self.decode_kv_args_table[mooncake_session_id] = (
+                        KVArgsRegisterInfo.from_zmq(waiting_req_bytes)
+                    )
+                    logger.debug(
+                        f"Register KVArgs from {mooncake_session_id} successfully"
+                    )
                     continue
                 room = int(room)
                 self.transfer_infos[room] = TransferInfo.from_zmq(waiting_req_bytes)
@@ -236,7 +302,7 @@ class MooncakeKVManager(BaseKVManager):
                     ret = self.send_kvcache(
                         req.mooncake_session_id,
                         kv_chunk.prefill_kv_indices,
-                        req.dst_kv_ptrs,
+                        self.decode_kv_args_table[req.mooncake_session_id].dst_kv_ptrs,
                         chunked_dst_kv_indice,
                     )
                     if ret != 0:
@@ -251,7 +317,9 @@ class MooncakeKVManager(BaseKVManager):
                         ret = self.send_aux(
                             req.mooncake_session_id,
                             kv_chunk.prefill_aux_index,
-                            req.dst_aux_ptrs,
+                            self.decode_kv_args_table[
+                                req.mooncake_session_id
+                            ].dst_aux_ptrs,
                             req.dst_aux_index,
                         )
                         self.request_status[req.room] = (
@@ -331,6 +399,8 @@ class MooncakeKVManager(BaseKVManager):
         url = f"http://{bootstrap_server_url}/route"
         payload = {
             "role": "Prefill",
+            "tp_size": self.tp_size,
+            "dp_size": self.dp_size,
             "rank_ip": get_local_ip_by_remote(),
             "rank_port": self.rank_port,
             "engine_rank": self.kv_args.engine_rank,
@@ -408,12 +478,41 @@ class MooncakeKVReceiver(BaseKVReceiver):
         self.session_id = self.kv_mgr.get_session_id()
         self.kv_mgr.update_status(bootstrap_room, KVPoll.Bootstrapping)
 
+        if not self.kv_mgr.enable_dp_attention:
+            # We assume dp_attention should be activated simultaneously for
+            # both prefill role and decode role. If the decode instance does
+            # not enable dp_attention, then dp_attention is not enabled on the
+            # prefill instance as well. Therefore, we should skip questioning
+            # the prefill dp size to reduce bootstrap overhead.
+            self.prefill_dp_size = 1
+        elif self.bootstrap_addr not in self.kv_mgr.prefill_dp_size_table:
+            self.prefill_dp_size, tp_size_per_dp_rank = (
+                self._get_prefill_dp_size_from_server()
+            )
+            # Currently, we don't allow prefill instance and decode instance to
+            # have different TP sizes per DP rank.
+            assert tp_size_per_dp_rank == self.kv_mgr.tp_size // self.kv_mgr.dp_size
+            if self.prefill_dp_size is None:
+                logger.error(
+                    f"Could not fetch prefill dp_size for bootstrap_addr: {self.bootstrap_addr}"
+                )
+            else:
+                self.kv_mgr.prefill_dp_size_table[self.bootstrap_addr] = (
+                    self.prefill_dp_size
+                )
+        else:
+            self.prefill_dp_size = self.kv_mgr.prefill_dp_size_table[
+                self.bootstrap_addr
+            ]
+
         # NOTE: key distinguished by bootstrap_addr and engine_rank
+        self.target_dp_group = bootstrap_room % self.prefill_dp_size
         bootstrap_key = f"{self.bootstrap_addr}_{self.kv_mgr.kv_args.engine_rank}"
 
         if bootstrap_key not in self.kv_mgr.connection_pool:
             self.bootstrap_info = self._get_bootstrap_info_from_server(
-                self.kv_mgr.kv_args.engine_rank
+                self.kv_mgr.kv_args.engine_rank,
+                self.target_dp_group,
             )
             if self.bootstrap_info is None:
                 logger.error(
@@ -421,16 +520,18 @@ class MooncakeKVReceiver(BaseKVReceiver):
                 )
             else:
                 self.kv_mgr.connection_pool[bootstrap_key] = self.bootstrap_info
+                # Register kv_args only once to prefill KVManager according to the info fetched from the bootstrap server
+                self._register_kv_args()
         else:
             self.bootstrap_info = self.kv_mgr.connection_pool[bootstrap_key]
 
         assert self.bootstrap_info is not None
         self.kv_mgr.update_status(bootstrap_room, KVPoll.WaitingForInput)
 
-    def _get_bootstrap_info_from_server(self, engine_rank):
+    def _get_bootstrap_info_from_server(self, engine_rank, target_dp_group):
         """Fetch the bootstrap info from the bootstrap server."""
         try:
-            url = f"http://{self.bootstrap_addr}/route?engine_rank={engine_rank}"
+            url = f"http://{self.bootstrap_addr}/route?engine_rank={engine_rank}&target_dp_group={target_dp_group}"
             response = requests.get(url)
             if response.status_code == 200:
                 bootstrap_info = response.json()
@@ -444,6 +545,49 @@ class MooncakeKVReceiver(BaseKVReceiver):
             logger.error(f"Error fetching prefill info from bootstrap: {e}")
             return None
 
+    def _get_prefill_dp_size_from_server(self) -> int:
+        """Fetch the prefill parallel info from the bootstrap server."""
+        try:
+            url = f"http://{self.bootstrap_addr}/route?engine_rank={-1}&target_dp_group={-1}"
+            response = requests.get(url)
+            if response.status_code == 200:
+                prefill_parallel_info = response.json()
+                return int(prefill_parallel_info["prefill_dp_size"]), int(
+                    prefill_parallel_info["tp_size_per_dp_rank"]
+                )
+            else:
+                logger.error(
+                    f"Failed to get prefill parallel info: {response.status_code}, {response.text}"
+                )
+                return None
+        except Exception as e:
+            logger.error(f"Error fetching prefill parallel info from bootstrap: {e}")
+            return None
+
+    def _register_kv_args(self):
+        self.prefill_server_url = (
+            f"{self.bootstrap_info['rank_ip']}:{self.bootstrap_info['rank_port']}"
+        )
+
+        packed_kv_data_ptrs = b"".join(
+            struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.kv_data_ptrs
+        )
+        packed_aux_data_ptrs = b"".join(
+            struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.aux_data_ptrs
+        )
+        sock, lock = self._connect("tcp://" + self.prefill_server_url)
+        with lock:
+            sock.send_multipart(
+                [
+                    "None".encode("ascii"),
+                    get_local_ip_by_remote().encode("ascii"),
+                    str(self.kv_mgr.rank_port).encode("ascii"),
+                    self.session_id.encode("ascii"),
+                    packed_kv_data_ptrs,
+                    packed_aux_data_ptrs,
+                ]
+            )
+
     @classmethod
     def _connect(cls, endpoint: str):
         with cls._global_lock:
@@ -462,12 +606,6 @@ class MooncakeKVReceiver(BaseKVReceiver):
             f"Fetched bootstrap info: {self.bootstrap_info} for engine rank: {self.kv_mgr.kv_args.engine_rank}"
         )
 
-        packed_kv_data_ptrs = b"".join(
-            struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.kv_data_ptrs
-        )
-        packed_aux_data_ptrs = b"".join(
-            struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.aux_data_ptrs
-        )
         sock, lock = self._connect("tcp://" + self.prefill_server_url)
         with lock:
             sock.send_multipart(
@@ -476,9 +614,7 @@ class MooncakeKVReceiver(BaseKVReceiver):
                     get_local_ip_by_remote().encode("ascii"),
                     str(self.kv_mgr.rank_port).encode("ascii"),
                     self.session_id.encode("ascii"),
-                    packed_kv_data_ptrs,
                     kv_indices.tobytes(),
-                    packed_aux_data_ptrs,
                     str(aux_index).encode("ascii"),
                 ]
             )
@@ -497,7 +633,9 @@ class MooncakeKVBootstrapServer(BaseKVBootstrapServer):
         self.store = dict()
         self.lock = asyncio.Lock()
         self._setup_routes()
-        self.prefill_port_table: Dict[int, Dict[str, Union[str, int]]] = {}
+        self.dp_size = None
+        self.tp_size_per_dp_rank = None
+        self.prefill_port_table: Dict[int, Dict[int, Dict[str, Union[str, int]]]] = {}
 
         # Start bootstrap server
         self.thread = threading.Thread(target=self._run_server, daemon=True)
@@ -523,35 +661,64 @@ class MooncakeKVBootstrapServer(BaseKVBootstrapServer):
     async def _handle_route_put(self, request: web.Request):
         data = await request.json()
         role = data["role"]
+        tp_size = data["tp_size"]
+        dp_size = data["dp_size"]
         rank_ip = data["rank_ip"]
         rank_port = int(data["rank_port"])
         engine_rank = int(data["engine_rank"])
 
+        if self.dp_size is None:
+            self.dp_size = dp_size
+
+        tp_size_per_dp_rank = tp_size // dp_size
+        if self.tp_size_per_dp_rank == None:
+            self.tp_size_per_dp_rank = tp_size_per_dp_rank
+
         # Add lock to make sure thread-safe
         if role == "Prefill":
-            self.prefill_port_table[engine_rank] = {
+            dp_group = engine_rank // tp_size_per_dp_rank
+            tp_rank_in_dp_group = engine_rank % tp_size_per_dp_rank
+
+            async with self.lock:
+                if dp_group not in self.prefill_port_table:
+                    self.prefill_port_table[dp_group] = {}
+
+            self.prefill_port_table[dp_group][tp_rank_in_dp_group] = {
                 "rank_ip": rank_ip,
                 "rank_port": rank_port,
             }
             logger.debug(
-                f"Registered Prefill boostrap: {engine_rank} with rank_ip: {rank_ip} and rank_port: {rank_port}"
+                f"Register Prefill bootstrap: {engine_rank} with rank_ip: {rank_ip} and rank_port: {rank_port}"
             )
 
         return web.Response(text="OK", status=200)
 
     async def _handle_route_get(self, request: web.Request):
         engine_rank = request.query.get("engine_rank")
-        if not engine_rank:
-            return web.Response(text="Missing rank", status=400)
+        target_dp_group = request.query.get("target_dp_group")
+        if not engine_rank or not target_dp_group:
+            return web.Response(text="Missing inputs for bootstrap server.", status=400)
+
+        # Currently we use engine_rank == -1 and target_dp_group == -1 to sync dp size
+        if int(engine_rank) == -1 and int(target_dp_group) == -1:
+            prefill_parallel_info = {
+                "prefill_dp_size": self.dp_size,
+                "tp_size_per_dp_rank": self.tp_size_per_dp_rank,
+            }
+            return web.json_response(prefill_parallel_info, status=200)
 
         # Find corresponding prefill info
+        tp_rank_in_dp_group = int(engine_rank) % self.tp_size_per_dp_rank
+
         async with self.lock:
-            bootstrap_info = self.prefill_port_table.get(int(engine_rank))
+            bootstrap_info = self.prefill_port_table[int(target_dp_group)][
+                tp_rank_in_dp_group
+            ]
 
         if bootstrap_info is not None:
             return web.json_response(bootstrap_info, status=200)
         else:
-            return web.Response(text="Not Found", status=404)
+            return web.Response(text="Bootstrap info not Found", status=404)
 
     def _run_server(self):
         try: