sglang 0.5.3rc0__py3-none-any.whl → 0.5.4__py3-none-any.whl

sglang/srt/disaggregation/mooncake/conn.py

@@ -1,33 +1,27 @@
 from __future__ import annotations
 
-import asyncio
 import concurrent.futures
 import ctypes
 import dataclasses
 import logging
 import os
-import queue
-import socket
 import struct
 import threading
 import time
 from collections import defaultdict
-from functools import cache
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Set, Tuple
 
 import numpy as np
 import numpy.typing as npt
 import requests
 import zmq
-from aiohttp import web
-
-from sglang.srt.disaggregation.base.conn import (
-    BaseKVBootstrapServer,
-    BaseKVManager,
-    BaseKVReceiver,
-    BaseKVSender,
-    KVArgs,
-    KVPoll,
+
+from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll
+from sglang.srt.disaggregation.common.conn import (
+    CommonKVBootstrapServer,
+    CommonKVManager,
+    CommonKVReceiver,
+    CommonKVSender,
 )
 from sglang.srt.disaggregation.common.utils import (
     FastQueue,
@@ -35,23 +29,12 @@ from sglang.srt.disaggregation.common.utils import (
 )
 from sglang.srt.disaggregation.mooncake.transfer_engine import MooncakeTransferEngine
 from sglang.srt.disaggregation.utils import DisaggregationMode
-from sglang.srt.distributed import get_pp_group
-from sglang.srt.layers.dp_attention import (
-    get_attention_dp_rank,
-    get_attention_dp_size,
-    get_attention_tp_rank,
-    get_attention_tp_size,
-)
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
     format_tcp_address,
     get_bool_env_var,
-    get_free_port,
     get_int_env_var,
-    get_ip,
-    get_local_ip_auto,
     is_valid_ipv6_address,
-    maybe_wrap_ipv6_address,
 )
 
 logger = logging.getLogger(__name__)
@@ -75,6 +58,7 @@ class TransferKVChunk:
     index_slice: slice
     is_last: bool
     prefill_aux_index: Optional[int]
+    state_indices: Optional[List[int]]
 
 
 # decode
@@ -86,6 +70,7 @@ class TransferInfo:
     mooncake_session_id: str
     dst_kv_indices: npt.NDArray[np.int32]
     dst_aux_index: int
+    dst_state_indices: List[int]
     required_dst_info_num: int
     is_dummy: bool
 
@@ -95,9 +80,14 @@ class TransferInfo:
             is_dummy = True
             dst_kv_indices = np.array([], dtype=np.int32)
             dst_aux_index = None
+            dst_state_indices = []
         else:
             dst_kv_indices = np.frombuffer(msg[4], dtype=np.int32)
             dst_aux_index = int(msg[5].decode("ascii"))
+            if msg[6] == b"":
+                dst_state_indices = []
+            else:
+                dst_state_indices = list(np.frombuffer(msg[6], dtype=np.int32))
             is_dummy = False
         return cls(
             room=int(msg[0].decode("ascii")),
@@ -106,7 +96,8 @@ class TransferInfo:
             mooncake_session_id=msg[3].decode("ascii"),
             dst_kv_indices=dst_kv_indices,
             dst_aux_index=dst_aux_index,
-            required_dst_info_num=int(msg[6].decode("ascii")),
+            dst_state_indices=dst_state_indices,
+            required_dst_info_num=int(msg[7].decode("ascii")),
             is_dummy=is_dummy,
         )
 
@@ -120,6 +111,7 @@ class KVArgsRegisterInfo:
     mooncake_session_id: str
     dst_kv_ptrs: list[int]
     dst_aux_ptrs: list[int]
+    dst_state_data_ptrs: list[int]
     dst_tp_rank: int
     dst_attn_tp_size: int
     dst_kv_item_len: int
@@ -133,9 +125,10 @@ class KVArgsRegisterInfo:
             mooncake_session_id=msg[3].decode("ascii"),
             dst_kv_ptrs=list(struct.unpack(f"{len(msg[4])//8}Q", msg[4])),
             dst_aux_ptrs=list(struct.unpack(f"{len(msg[5])//8}Q", msg[5])),
-            dst_tp_rank=int(msg[6].decode("ascii")),
-            dst_attn_tp_size=int(msg[7].decode("ascii")),
-            dst_kv_item_len=int(msg[8].decode("ascii")),
+            dst_state_data_ptrs=list(struct.unpack(f"{len(msg[6])//8}Q", msg[6])),
+            dst_tp_rank=int(msg[7].decode("ascii")),
+            dst_attn_tp_size=int(msg[8].decode("ascii")),
+            dst_kv_item_len=int(msg[9].decode("ascii")),
         )
 
 
@@ -159,7 +152,7 @@ class AuxDataCodec:
         return
 
 
-class MooncakeKVManager(BaseKVManager):
+class MooncakeKVManager(CommonKVManager):
     AUX_DATA_HEADER = b"AUX_DATA"
 
     def __init__(
@@ -169,48 +162,19 @@ class MooncakeKVManager(BaseKVManager):
         server_args: ServerArgs,
         is_mla_backend: Optional[bool] = False,
     ):
-        self.kv_args = args
-        self.local_ip = get_local_ip_auto()
-        self.is_mla_backend = is_mla_backend
-        self.disaggregation_mode = disaggregation_mode
+        super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
         self.init_engine()
-        # for p/d multi node infer
-        self.bootstrap_host = server_args.host
-        self.bootstrap_port = server_args.disaggregation_bootstrap_port
-        self.dist_init_addr = server_args.dist_init_addr
-        self.attn_tp_size = get_attention_tp_size()
-        self.attn_tp_rank = get_attention_tp_rank()
-        self.attn_dp_size = get_attention_dp_size()
-        self.attn_dp_rank = get_attention_dp_rank()
-        self.system_dp_size = (
-            1 if server_args.enable_dp_attention else server_args.dp_size
-        )
-        self.system_dp_rank = (
-            self.kv_args.system_dp_rank if self.kv_args.system_dp_rank else 0
-        )
-        self.pp_size = server_args.pp_size
-        self.pp_rank = self.kv_args.pp_rank
-        self.request_status: Dict[int, KVPoll] = {}
-        self.rank_port = None
-        self.server_socket = zmq.Context().socket(zmq.PULL)
-        if is_valid_ipv6_address(self.local_ip):
-            self.server_socket.setsockopt(zmq.IPV6, 1)
-
         self.register_buffer_to_engine()
         if self.disaggregation_mode == DisaggregationMode.PREFILL:
-            self.transfer_infos: Dict[int, Dict[str, TransferInfo]] = {}
-            self.decode_kv_args_table: Dict[str, KVArgsRegisterInfo] = {}
             self.start_prefill_thread()
-            self._register_to_bootstrap()
             self.session_failures = defaultdict(int)
             self.failed_sessions = set()
             self.session_lock = threading.Lock()
-            self.pp_group = get_pp_group()
             # Determine the number of threads to use for kv sender
             cpu_count = os.cpu_count()
             transfer_thread_pool_size = get_int_env_var(
                 "SGLANG_DISAGGREGATION_THREAD_POOL_SIZE",
-                min(max(4, int(0.75 * cpu_count) // 8), 12),
+                min(max(4, int(0.5 * cpu_count) // 8), 12),
             )
             transfer_queue_size = get_int_env_var("SGLANG_DISAGGREGATION_QUEUE_SIZE", 4)
             self.transfer_queues: List[FastQueue] = [
@@ -245,8 +209,6 @@ class MooncakeKVManager(BaseKVManager):
             self.session_pool = defaultdict(requests.Session)
             self.session_pool_lock = threading.Lock()
             self.addr_to_rooms_tracker = defaultdict(set)
-            self.connection_lock = threading.Lock()
-            self.required_prefill_response_num_table: Dict[int, int] = {}
             self.prefill_response_tracker: Dict[int, Set[int]] = defaultdict(set)
             # Heartbeat interval should be at least 2 seconds
             self.heartbeat_interval = max(
@@ -257,20 +219,12 @@ class MooncakeKVManager(BaseKVManager):
                 get_int_env_var("SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE", 2), 1
             )
             self.start_decode_thread()
-            self.connection_pool: Dict[str, Dict[str, Union[str, int]]] = {}
-            self.prefill_attn_tp_size_table: Dict[str, int] = {}
-            self.prefill_dp_size_table: Dict[str, int] = {}
-            self.prefill_pp_size_table: Dict[str, int] = {}
             # If a timeout happens on the decode side, it means decode instances
             # fail to receive the KV Cache transfer done signal after bootstrapping.
             # These timeout requests should be aborted to release the tree cache.
             self.waiting_timeout = get_int_env_var(
                 "SGLANG_DISAGGREGATION_WAITING_TIMEOUT", 300
             )
-        else:
-            raise ValueError(
-                f"Unsupported DisaggregationMode: {self.disaggregation_mode}"
-            )
 
         self.failure_records: Dict[int, str] = {}
         self.failure_lock = threading.Lock()
@@ -295,13 +249,11 @@ class MooncakeKVManager(BaseKVManager):
                 self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens
             )
 
-    @cache
-    def _connect(self, endpoint: str, is_ipv6: bool = False):
-        socket = zmq.Context().socket(zmq.PUSH)
-        if is_ipv6:
-            socket.setsockopt(zmq.IPV6, 1)
-        socket.connect(endpoint)
-        return socket
+        # Batch register state/extra pool data buffers
+        if self.kv_args.state_data_ptrs and self.kv_args.state_data_lens:
+            self.engine.batch_register(
+                self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
+            )
 
     def _transfer_data(self, mooncake_session_id, transfer_blocks):
         if not transfer_blocks:
@@ -312,62 +264,60 @@ class MooncakeKVManager(BaseKVManager):
             mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
         )
 
-    def send_kvcache(
+    def _send_kvcache_generic(
         self,
         mooncake_session_id: str,
-        prefill_kv_indices: npt.NDArray[np.int32],
-        dst_kv_ptrs: list[int],
-        dst_kv_indices: npt.NDArray[np.int32],
+        src_data_ptrs: list[int],
+        dst_data_ptrs: list[int],
+        item_lens: list[int],
+        prefill_data_indices: npt.NDArray[np.int32],
+        dst_data_indices: npt.NDArray[np.int32],
         executor: concurrent.futures.ThreadPoolExecutor,
-    ):
-        # Group by indices
+    ) -> int:
+        """
+        Generic KV cache transfer supporting both MHA and MLA architectures.
+        This method is used by both send_kvcache (full pool) and maybe_send_extra.
+        """
+        # Group by indices for optimization
         prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
-            prefill_kv_indices, dst_kv_indices
+            prefill_data_indices, dst_data_indices
         )
 
         layers_params = None
 
         # pp is not supported on the decode side yet
-        start_layer = self.kv_args.prefill_start_layer
-        end_layer = start_layer + len(self.kv_args.kv_data_ptrs)
         if self.is_mla_backend:
-            src_kv_ptrs = self.kv_args.kv_data_ptrs
-            layers_per_pp_stage = len(src_kv_ptrs)
-            dst_kv_ptrs = dst_kv_ptrs[start_layer:end_layer]
-            kv_item_len = self.kv_args.kv_item_lens[0]
+            src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = (
+                self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
+            )
+            kv_item_len = item_lens[0]
             layers_params = [
                 (
                     src_kv_ptrs[layer_id],
                     dst_kv_ptrs[layer_id],
                     kv_item_len,
                 )
-                for layer_id in range(layers_per_pp_stage)
+                for layer_id in range(layers_current_pp_stage)
             ]
         else:
-            num_kv_layers = len(self.kv_args.kv_data_ptrs) // 2
-            dst_num_total_layers = num_kv_layers * self.pp_size
-            src_k_ptrs = self.kv_args.kv_data_ptrs[:num_kv_layers]
-            src_v_ptrs = self.kv_args.kv_data_ptrs[num_kv_layers:]
-            layers_per_pp_stage = len(src_k_ptrs)
-            dst_k_ptrs = dst_kv_ptrs[start_layer:end_layer]
-            dst_v_ptrs = dst_kv_ptrs[
-                dst_num_total_layers + start_layer : dst_num_total_layers + end_layer
-            ]
-            kv_item_len = self.kv_args.kv_item_lens[0]
+            src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = (
+                self.get_mha_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
+            )
+            kv_item_len = item_lens[0]
             layers_params = [
                 (
                     src_k_ptrs[layer_id],
                     dst_k_ptrs[layer_id],
                     kv_item_len,
                 )
-                for layer_id in range(layers_per_pp_stage)
+                for layer_id in range(layers_current_pp_stage)
             ] + [
                 (
                     src_v_ptrs[layer_id],
                     dst_v_ptrs[layer_id],
                     kv_item_len,
                 )
-                for layer_id in range(layers_per_pp_stage)
+                for layer_id in range(layers_current_pp_stage)
             ]
         assert layers_params is not None
 
@@ -417,6 +367,24 @@ class MooncakeKVManager(BaseKVManager):
 
         return 0
 
+    def send_kvcache(
+        self,
+        mooncake_session_id: str,
+        prefill_kv_indices: npt.NDArray[np.int32],
+        dst_kv_ptrs: list[int],
+        dst_kv_indices: npt.NDArray[np.int32],
+        executor: concurrent.futures.ThreadPoolExecutor,
+    ):
+        return self._send_kvcache_generic(
+            mooncake_session_id=mooncake_session_id,
+            src_data_ptrs=self.kv_args.kv_data_ptrs,
+            dst_data_ptrs=dst_kv_ptrs,
+            item_lens=self.kv_args.kv_item_lens,
+            prefill_data_indices=prefill_kv_indices,
+            dst_data_indices=dst_kv_indices,
+            executor=executor,
+        )
+
     def send_kvcache_slice(
         self,
         mooncake_session_id: str,
@@ -465,18 +433,9 @@ class MooncakeKVManager(BaseKVManager):
             num_heads_to_send = dst_heads_per_rank
             dst_head_start_offset = 0
 
-        # pp is not supported on the decode side yet
-        num_kv_layers = len(self.kv_args.kv_data_ptrs) // 2
-        dst_num_total_layers = num_kv_layers * self.pp_size
-        src_k_ptrs = self.kv_args.kv_data_ptrs[:num_kv_layers]
-        src_v_ptrs = self.kv_args.kv_data_ptrs[num_kv_layers:]
-        layers_per_pp_stage = len(src_k_ptrs)
-        start_layer = self.pp_rank * layers_per_pp_stage
-        end_layer = start_layer + layers_per_pp_stage
-        dst_k_ptrs = dst_kv_ptrs[start_layer:end_layer]
-        dst_v_ptrs = dst_kv_ptrs[
-            dst_num_total_layers + start_layer : dst_num_total_layers + end_layer
-        ]
+        src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = (
+            self.get_mha_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs)
+        )
 
         # Calculate precise byte offset and length for the sub-slice within the token
         src_head_slice_offset = src_head_start_offset * bytes_per_head_slice_to_send
@@ -502,7 +461,7 @@ class MooncakeKVManager(BaseKVManager):
                 dst_head_slice_offset,
                 heads_bytes_per_token_to_send,
             )
-            for layer_id in range(layers_per_pp_stage)
+            for layer_id in range(layers_current_pp_stage)
         ] + [
             (
                 src_v_ptrs[layer_id],
@@ -513,7 +472,7 @@ class MooncakeKVManager(BaseKVManager):
                 dst_head_slice_offset,
                 heads_bytes_per_token_to_send,
             )
-            for layer_id in range(layers_per_pp_stage)
+            for layer_id in range(layers_current_pp_stage)
         ]
 
         def process_layer_tp_aware(layer_params):
@@ -654,6 +613,79 @@ class MooncakeKVManager(BaseKVManager):
             ]
         )
 
+    def _handle_aux_data(self, msg: List[bytes]):
+        """Handle AUX_DATA messages received by the decode thread."""
+        room = int(msg[1].decode("ascii"))
+        buffer_index = int(msg[2].decode("ascii"))
+        aux_index = int(msg[3].decode("ascii"))
+        data_length = struct.unpack(">I", msg[4])[0]
+        data = msg[5]
+
+        if len(data) != data_length:
+            logger.error(f"AUX_DATA length mismatch for bootstrap_room {room}")
+            return
+
+        AuxDataCodec.deserialize_data_to_buffer(
+            self.kv_args, buffer_index, aux_index, data
+        )
+
+        logger.debug(
+            f"Received AUX_DATA for bootstrap_room {room} with length:{len(data)}"
+        )
+
+    def maybe_send_extra(
+        self,
+        req: TransferInfo,
+        prefill_state_indices: list[int],
+        dst_state_data_ptrs: list[int],
+        executor: concurrent.futures.ThreadPoolExecutor,
+    ):
+        """Send state or extra pool data with type-specific handling."""
+        state_type = getattr(self.kv_args, "state_type", "none")
+
+        if state_type == "mamba":
+            return self._send_mamba_state(
+                req,
+                prefill_state_indices,
+                dst_state_data_ptrs,
+            )
+        elif state_type in ["swa", "nsa"]:
+            # Reuse _send_kvcache_generic interface to send extra pool data
+            prefill_state_indices = np.array(prefill_state_indices, dtype=np.int32)
+            dst_state_indices = np.array(req.dst_state_indices, dtype=np.int32)
+            return self._send_kvcache_generic(
+                mooncake_session_id=req.mooncake_session_id,
+                src_data_ptrs=self.kv_args.state_data_ptrs,
+                dst_data_ptrs=dst_state_data_ptrs,
+                item_lens=self.kv_args.state_item_lens,
+                prefill_data_indices=prefill_state_indices,
+                dst_data_indices=dst_state_indices,
+                executor=executor,
+            )
+        else:
+            return 0
+
+    def _send_mamba_state(
+        self,
+        req: TransferInfo,
+        prefill_mamba_index: list[int],
+        dst_state_data_ptrs: list[int],
+    ):
+        """Transfer Mamba states."""
+        assert len(prefill_mamba_index) == 1, "Mamba should have single state index"
+
+        transfer_blocks = []
+        prefill_state_data_ptrs = self.kv_args.state_data_ptrs
+        prefill_state_item_lens = self.kv_args.state_item_lens
+
+        for i, dst_state_ptr in enumerate(dst_state_data_ptrs):
+            length = prefill_state_item_lens[i]
+            src_addr = prefill_state_data_ptrs[i] + length * int(prefill_mamba_index[0])
+            dst_addr = dst_state_ptr + length * int(req.dst_state_indices[0])
+            transfer_blocks.append((src_addr, dst_addr, length))
+
+        return self._transfer_data(req.mooncake_session_id, transfer_blocks)
+
     def sync_status_to_decode_endpoint(
         self, remote: str, dst_port: int, room: int, status: int, prefill_rank: int
     ):
@@ -763,6 +795,22 @@ class MooncakeKVManager(BaseKVManager):
                             break
 
                         if kv_chunk.is_last:
+                            if kv_chunk.state_indices is not None:
+                                if not self.is_mla_backend and (
+                                    self.attn_tp_size
+                                    != target_rank_registration_info.dst_attn_tp_size
+                                ):
+                                    raise RuntimeError(
+                                        f"PD Disaggregation does NOT support PD different TP sizes for non-MLA hybrid models yet."
+                                    )
+
+                                self.maybe_send_extra(
+                                    req,
+                                    kv_chunk.state_indices,
+                                    target_rank_registration_info.dst_state_data_ptrs,
+                                    executor,
+                                )
+
                             if self.pp_group.is_last_rank:
                                 # Only the last chunk we need to send the aux data
                                 ret = self.send_aux(
@@ -802,11 +850,7 @@ class MooncakeKVManager(BaseKVManager):
                     f"Transfer thread failed because of {e}. Prefill instance with bootstrap_port={self.bootstrap_port} is dead."
                 )
 
-    def _bind_server_socket(self):
-        self.server_socket.bind(format_tcp_address(self.local_ip, self.rank_port))
-
     def start_prefill_thread(self):
-        self.rank_port = get_free_port()
         self._bind_server_socket()
 
         def bootstrap_thread():
@@ -830,7 +874,7 @@ class MooncakeKVManager(BaseKVManager):
                     )
                     continue
                 else:
-                    required_dst_info_num = int(waiting_req_bytes[6].decode("ascii"))
+                    required_dst_info_num = int(waiting_req_bytes[7].decode("ascii"))
                     room = int(room)
                     if room not in self.transfer_infos:
                         self.transfer_infos[room] = {}
@@ -844,28 +888,7 @@ class MooncakeKVManager(BaseKVManager):
 
         threading.Thread(target=bootstrap_thread).start()
 
-    def _handle_aux_data(self, msg: List[bytes]):
-        """Handle AUX_DATA messages received by the decode thread."""
-        room = int(msg[1].decode("ascii"))
-        buffer_index = int(msg[2].decode("ascii"))
-        aux_index = int(msg[3].decode("ascii"))
-        data_length = struct.unpack(">I", msg[4])[0]
-        data = msg[5]
-
-        if len(data) != data_length:
-            logger.error(f"AUX_DATA length mismatch for bootstrap_room {room}")
-            return
-
-        AuxDataCodec.deserialize_data_to_buffer(
-            self.kv_args, buffer_index, aux_index, data
-        )
-
-        logger.debug(
-            f"Received AUX_DATA for bootstrap_room {room} with length:{len(data)}"
-        )
-
     def start_decode_thread(self):
-        self.rank_port = get_free_port()
         self._bind_server_socket()
 
         def decode_thread():
@@ -962,6 +985,7 @@ class MooncakeKVManager(BaseKVManager):
         index_slice: slice,
         is_last: bool,
         aux_index: Optional[int] = None,
+        state_indices: Optional[List[int]] = None,
     ):
         assert self.disaggregation_mode == DisaggregationMode.PREFILL
         assert not is_last or (is_last and aux_index is not None)
@@ -995,6 +1019,7 @@ class MooncakeKVManager(BaseKVManager):
                 index_slice=index_slice,
                 is_last=is_last,
                 prefill_aux_index=aux_index,
+                state_indices=state_indices,
             )
         )
 
@@ -1020,51 +1045,6 @@ class MooncakeKVManager(BaseKVManager):
     def get_session_id(self):
         return self.engine.get_session_id()
 
-    def _register_to_bootstrap(self):
-        """Register KVSender to bootstrap server via HTTP POST."""
-        if self.dist_init_addr:
-            # multi node case: bootstrap server's host is dist_init_addr
-            if self.dist_init_addr.startswith("["):  # [ipv6]:port or [ipv6]
-                if self.dist_init_addr.endswith("]"):
-                    host = self.dist_init_addr
-                else:
-                    host, _ = self.dist_init_addr.rsplit(":", 1)
-            else:
-                host = socket.gethostbyname(self.dist_init_addr.rsplit(":", 1)[0])
-        else:
-            # single node case: bootstrap server's host is same as http server's host
-            host = self.bootstrap_host
-            host = maybe_wrap_ipv6_address(host)
-
-        bootstrap_server_url = f"{host}:{self.bootstrap_port}"
-        url = f"http://{bootstrap_server_url}/route"
-        payload = {
-            "role": "Prefill",
-            "attn_tp_size": self.attn_tp_size,
-            "attn_tp_rank": self.attn_tp_rank,
-            "attn_dp_size": self.attn_dp_size,
-            "attn_dp_rank": self.attn_dp_rank,
-            "pp_size": self.pp_size,
-            "pp_rank": self.pp_rank,
-            "system_dp_size": self.system_dp_size,
-            "system_dp_rank": self.system_dp_rank,
-            "rank_ip": self.local_ip,
-            "rank_port": self.rank_port,
-        }
-
-        try:
-            response = requests.put(url, json=payload, timeout=5)
-            if response.status_code == 200:
-                logger.debug("Prefill successfully registered to bootstrap server.")
-            else:
-                logger.error(
-                    f"Prefill instance failed to connect to bootstrap server: {response.status_code}, {response.text}"
-                )
-        except Exception as e:
-            logger.error(
-                f"Prefill instance failed to register to bootstrap server: {e}"
-            )
-
     def _handle_node_failure(self, failed_bootstrap_addr):
         with self.connection_lock:
             keys_to_remove = [
@@ -1103,7 +1083,7 @@ class MooncakeKVManager(BaseKVManager):
         )
 
 
-class MooncakeKVSender(BaseKVSender):
+class MooncakeKVSender(CommonKVSender):
 
     def __init__(
         self,
@@ -1113,23 +1093,14 @@ class MooncakeKVSender(BaseKVSender):
         dest_tp_ranks: List[int],
         pp_rank: int,
     ):
-        self.kv_mgr = mgr
-        self.bootstrap_room = bootstrap_room
-        self.kv_mgr.update_status(bootstrap_room, KVPoll.Bootstrapping)
-        self.aux_index = None
-        self.bootstrap_server_url = bootstrap_addr
+        super().__init__(mgr, bootstrap_addr, bootstrap_room, dest_tp_ranks, pp_rank)
         self.conclude_state = None
         self.init_time = time.time()
-        # inner state
-        self.curr_idx = 0
-
-    def init(self, num_kv_indices: int, aux_index: Optional[int] = None):
-        self.num_kv_indices = num_kv_indices
-        self.aux_index = aux_index
 
     def send(
         self,
         kv_indices: npt.NDArray[np.int32],
+        state_indices: Optional[List[int]] = None,
     ):
         index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices))
         self.curr_idx += len(kv_indices)
@@ -1149,6 +1120,7 @@ class MooncakeKVSender(BaseKVSender):
                 index_slice,
                 True,
                 aux_index=self.aux_index,
+                state_indices=state_indices,
             )
 
     def poll(self) -> KVPoll:
@@ -1203,7 +1175,7 @@ class MooncakeKVSender(BaseKVSender):
         self.conclude_state = KVPoll.Failed
 
 
-class MooncakeKVReceiver(BaseKVReceiver):
+class MooncakeKVReceiver(CommonKVReceiver):
     _ctx = zmq.Context()
     _socket_cache = {}
     _socket_locks = {}
@@ -1216,166 +1188,11 @@ class MooncakeKVReceiver(BaseKVReceiver):
         bootstrap_room: Optional[int] = None,
         prefill_dp_rank: Optional[int] = None,
     ):
-        self.bootstrap_room = bootstrap_room
-        self.bootstrap_addr = bootstrap_addr
-        self.kv_mgr = mgr
-        self.session_id = self.kv_mgr.get_session_id()
-        self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping)
+        self.session_id = mgr.get_session_id()
         self.conclude_state = None
         self.init_time = None
+        super().__init__(mgr, bootstrap_addr, bootstrap_room, prefill_dp_rank)
 
-        if self.bootstrap_addr not in self.kv_mgr.prefill_dp_size_table:
-            (
-                self.prefill_attn_tp_size,
-                self.prefill_dp_size,
-                self.prefill_pp_size,
-            ) = self._get_prefill_parallel_info_from_server()
-            if (
-                self.prefill_attn_tp_size is None
-                or self.prefill_dp_size is None
-                or self.prefill_pp_size is None
-            ):
-                self.kv_mgr.record_failure(
-                    self.bootstrap_room,
-                    f"Could not fetch prefill parallel info from bootstrap_addr: {self.bootstrap_addr}",
-                )
-                self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
-                return
-            else:
-                logger.debug(
-                    f"Fetch prefill parallel info from [{self.bootstrap_addr}]: DP size:{self.prefill_dp_size}, TP size:{self.prefill_attn_tp_size} PP size:{self.prefill_pp_size}"
-                )
-                self.kv_mgr.prefill_attn_tp_size_table[self.bootstrap_addr] = (
-                    self.prefill_attn_tp_size
-                )
-                self.kv_mgr.prefill_dp_size_table[self.bootstrap_addr] = (
-                    self.prefill_dp_size
-                )
-                self.kv_mgr.prefill_pp_size_table[self.bootstrap_addr] = (
-                    self.prefill_pp_size
-                )
-        else:
-            self.prefill_attn_tp_size = self.kv_mgr.prefill_attn_tp_size_table[
-                self.bootstrap_addr
-            ]
-            self.prefill_dp_size = self.kv_mgr.prefill_dp_size_table[
-                self.bootstrap_addr
-            ]
-            self.prefill_pp_size = self.kv_mgr.prefill_pp_size_table[
-                self.bootstrap_addr
-            ]
-
-        # Currently, we don't allow prefill instance and decode instance to
-        # have different TP sizes per DP rank, except for models using MLA.
-        if self.kv_mgr.attn_tp_size == self.prefill_attn_tp_size:
-            self.target_tp_rank = (
-                self.kv_mgr.kv_args.engine_rank % self.kv_mgr.attn_tp_size
-            )
-            self.required_dst_info_num = 1
-            self.required_prefill_response_num = 1 * (
-                self.prefill_pp_size // self.kv_mgr.pp_size
-            )
-            self.target_tp_ranks = [self.target_tp_rank]
-        elif self.kv_mgr.attn_tp_size > self.prefill_attn_tp_size:
-            if not self.kv_mgr.is_mla_backend:
-                logger.warning_once(
-                    "Performance is NOT guaranteed when using different TP sizes for non-MLA models. "
-                )
-            self.target_tp_rank = (
-                self.kv_mgr.kv_args.engine_rank % self.kv_mgr.attn_tp_size
-            ) // (self.kv_mgr.attn_tp_size // self.prefill_attn_tp_size)
-            self.required_dst_info_num = (
-                self.kv_mgr.attn_tp_size // self.prefill_attn_tp_size
-            )
-            self.required_prefill_response_num = 1 * (
-                self.prefill_pp_size // self.kv_mgr.pp_size
-            )
-            self.target_tp_ranks = [self.target_tp_rank]
-        else:
-            if not self.kv_mgr.is_mla_backend:
-                logger.warning_once(
-                    "Performance is NOT guaranteed when using different TP sizes for non-MLA models. "
-                )
-            # For non-MLA models, one decode rank needs to retrieve KVCache from multiple prefill ranks for non MLA models;
-            self.target_tp_ranks = [
-                rank
-                for rank in range(
-                    (self.kv_mgr.kv_args.engine_rank % self.kv_mgr.attn_tp_size)
-                    * (self.prefill_attn_tp_size // self.kv_mgr.attn_tp_size),
-                    (self.kv_mgr.kv_args.engine_rank % self.kv_mgr.attn_tp_size + 1)
-                    * (self.prefill_attn_tp_size // self.kv_mgr.attn_tp_size),
-                )
-            ]
-
-            # For MLA models, we can retrieve KVCache from only one prefill rank, but we still need to maintain
-            # multiple connections in the connection pool and have to send dummy requests to other prefill ranks,
-            # or the KVPoll will never be set correctly
-            self.target_tp_rank = self.target_tp_ranks[0]
-            self.required_dst_info_num = 1
-            if self.kv_mgr.is_mla_backend:
-                self.required_prefill_response_num = (
-                    self.prefill_pp_size // self.kv_mgr.pp_size
-                )
-            else:
-                self.required_prefill_response_num = (
-                    self.prefill_attn_tp_size // self.kv_mgr.attn_tp_size
-                ) * (self.prefill_pp_size // self.kv_mgr.pp_size)
-
-        if prefill_dp_rank is not None:
-            logger.debug(f"Targeting DP rank: {prefill_dp_rank}")
-            self.prefill_dp_rank = prefill_dp_rank
-        else:
-            self.prefill_dp_rank = bootstrap_room % self.prefill_dp_size
-
-        # FIXME: alias here: target_dp_group -> prefill_dp_rank
-        self.target_dp_group = self.prefill_dp_rank
-
-        self.kv_mgr.required_prefill_response_num_table[self.bootstrap_room] = (
-            self.required_prefill_response_num
-        )
-        # NOTE: key distinguished by bootstrap_addr, target_dp_group, and target_tp_rank
-        bootstrap_key = (
-            f"{self.bootstrap_addr}_{self.target_dp_group}_{self.target_tp_rank}"
-        )
-
-        if bootstrap_key not in self.kv_mgr.connection_pool:
-            bootstrap_infos = []
-            for target_tp_rank in self.target_tp_ranks:
-                for target_pp_rank in range(self.prefill_pp_size):
-                    bootstrap_info = self._get_bootstrap_info_from_server(
-                        target_tp_rank, self.target_dp_group, target_pp_rank
-                    )
-                    if bootstrap_info is not None:
-                        if self.kv_mgr.is_mla_backend:
-                            # For MLA: target_tp_rank is the selected real rank, others are dummy ranks
-                            bootstrap_info["is_dummy"] = not bool(
-                                target_tp_rank == self.target_tp_rank
-                                or self.target_tp_rank is None
-                            )
-                        else:
-                            # For non-MLA: all target_tp_ranks are selected real ranks
-                            bootstrap_info["is_dummy"] = False
-                        logger.debug(
-                            f"Fetched bootstrap info: {bootstrap_info} for DP {self.target_dp_group} TP {target_tp_rank} PP {target_pp_rank}"
-                        )
-                        bootstrap_infos.append(bootstrap_info)
-                    else:
-                        self.kv_mgr.record_failure(
-                            self.bootstrap_room,
-                            f"Could not fetch bootstrap info for engine rank: {self.kv_mgr.kv_args.engine_rank} and target_dp_group: {self.target_dp_group} and target_pp_rank {target_pp_rank}",
-                        )
-                        self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
-                        return
-
-            self.bootstrap_infos = bootstrap_infos
-            self.kv_mgr.connection_pool[bootstrap_key] = self.bootstrap_infos
-
-            # Register kv_args only once to prefill KVManager according to the info fetched from the bootstrap server
-            self._register_kv_args()
-        else:
-            self.bootstrap_infos = self.kv_mgr.connection_pool[bootstrap_key]
-
-        assert len(self.bootstrap_infos) > 0
         self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room)
         self.kv_mgr.update_status(self.bootstrap_room, KVPoll.WaitingForInput)
 
@@ -1398,29 +1215,6 @@ class MooncakeKVReceiver(BaseKVReceiver):
             logger.error(f"Error fetching prefill info from bootstrap: {e}")
             return None
 
-    def _get_prefill_parallel_info_from_server(
-        self,
-    ) -> Tuple[Optional[int], Optional[int], Optional[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}&target_pp_rank={-1}"
-            response = requests.get(url)
-            if response.status_code == 200:
-                prefill_parallel_info = response.json()
-                return (
-                    int(prefill_parallel_info["prefill_attn_tp_size"]),
-                    int(prefill_parallel_info["prefill_dp_size"]),
-                    int(prefill_parallel_info["prefill_pp_size"]),
-                )
-            else:
-                logger.error(
-                    f"Failed to get prefill parallel info: {response.status_code}, {response.text}"
-                )
-                return None, None, None
-        except Exception as e:
-            logger.error(f"Error fetching prefill parallel info from bootstrap: {e}")
-            return None, None, None
-
     def _register_kv_args(self):
         for bootstrap_info in self.bootstrap_infos:
             packed_kv_data_ptrs = b"".join(
@@ -1429,6 +1223,9 @@ class MooncakeKVReceiver(BaseKVReceiver):
             packed_aux_data_ptrs = b"".join(
                 struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.aux_data_ptrs
             )
+            packed_state_data_ptrs = b"".join(
+                struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.state_data_ptrs
+            )
             # Note(shangming): No need to add pp rank here since pp is not supported on the decode side yet
             tp_rank = self.kv_mgr.kv_args.engine_rank
             kv_item_len = self.kv_mgr.kv_args.kv_item_lens[0]
@@ -1446,35 +1243,27 @@ class MooncakeKVReceiver(BaseKVReceiver):
                         self.session_id.encode("ascii"),
                         packed_kv_data_ptrs,
                         packed_aux_data_ptrs,
+                        packed_state_data_ptrs,
                         dst_tp_rank,
                         dst_attn_tp_size,
                         dst_kv_item_len,
                     ]
                 )
 
-    @classmethod
-    def _connect(cls, endpoint: str, is_ipv6: bool = False):
-        with cls._global_lock:
-            if endpoint not in cls._socket_cache:
-                sock = cls._ctx.socket(zmq.PUSH)
-                if is_ipv6:
-                    sock.setsockopt(zmq.IPV6, 1)
-                sock.connect(endpoint)
-                cls._socket_cache[endpoint] = sock
-                cls._socket_locks[endpoint] = threading.Lock()
-            return cls._socket_cache[endpoint], cls._socket_locks[endpoint]
-
-    @classmethod
-    def _connect_to_bootstrap_server(cls, bootstrap_info: dict):
-        ip_address = bootstrap_info["rank_ip"]
-        port = bootstrap_info["rank_port"]
-        is_ipv6_address = is_valid_ipv6_address(ip_address)
-        sock, lock = cls._connect(
-            format_tcp_address(ip_address, port), is_ipv6=is_ipv6_address
-        )
-        return sock, lock
+    def init(
+        self,
+        kv_indices: npt.NDArray[np.int32],
+        aux_index: Optional[int] = None,
+        state_indices: Optional[List[int]] = None,
+    ):
+        if self.bootstrap_infos is None:
+            self.kv_mgr.record_failure(
+                self.bootstrap_room,
+                f"Could not fetch prefill parallel info from bootstrap_addr: {self.bootstrap_addr}",
+            )
+            self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
+            return
 
-    def init(self, kv_indices: npt.NDArray[np.int32], aux_index: Optional[int] = None):
         for bootstrap_info in self.bootstrap_infos:
             sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
             is_dummy = bootstrap_info["is_dummy"]
@@ -1488,6 +1277,14 @@ class MooncakeKVReceiver(BaseKVReceiver):
                         self.session_id.encode("ascii"),
                         kv_indices.tobytes() if not is_dummy else b"",
                         str(aux_index).encode("ascii") if not is_dummy else b"",
+                        (
+                            np.array(
+                                state_indices,
+                                dtype=np.int32,
+                            ).tobytes()
+                            if not is_dummy and state_indices is not None
+                            else b""
+                        ),
                         str(self.required_dst_info_num).encode("ascii"),
                     ]
                 )
@@ -1551,154 +1348,5 @@ class MooncakeKVReceiver(BaseKVReceiver):
         self.conclude_state = KVPoll.Failed
 
 
-class MooncakeKVBootstrapServer(BaseKVBootstrapServer):
-    def __init__(self, host: str, port: int):
-        self.host = host
-        self.port = port
-        self.app = web.Application()
-        self.store = dict()
-        self.lock = asyncio.Lock()
-        self._setup_routes()
-        self.pp_size = None
-        self.attn_tp_size = None
-        self.dp_size = None
-        self.prefill_port_table: Dict[
-            int, Dict[int, Dict[int, Dict[str, Union[str, int]]]]
-        ] = {}
-
-        # Start bootstrap server
-        self.thread = threading.Thread(target=self._run_server, daemon=True)
-        self.run()
-
-    def run(self):
-        self.thread.start()
-
-    def _setup_routes(self):
-        self.app.router.add_route("*", "/route", self._handle_route)
-        self.app.router.add_get("/health", self._handle_health_check)
-
-    async def _handle_health_check(self, request):
-        return web.Response(text="OK", status=200)
-
-    async def _handle_route(self, request: web.Request):
-        method = request.method
-        if method == "PUT":
-            return await self._handle_route_put(request)
-        elif method == "GET":
-            return await self._handle_route_get(request)
-        else:
-            return web.Response(
-                text="Method not allowed", status=405, content_type="application/json"
-            )
-
-    async def _handle_route_put(self, request: web.Request):
-        data = await request.json()
-        role = data["role"]
-        attn_tp_size = data["attn_tp_size"]
-        attn_tp_rank = data["attn_tp_rank"]
-        attn_dp_size = data["attn_dp_size"]
-        attn_dp_rank = data["attn_dp_rank"]
-        pp_size = data["pp_size"]
-        pp_rank = data["pp_rank"]
-        system_dp_size = data["system_dp_size"]
-        system_dp_rank = data["system_dp_rank"]
-        rank_ip = data["rank_ip"]
-        rank_port = int(data["rank_port"])
-
-        if self.attn_tp_size is None:
-            self.attn_tp_size = attn_tp_size
-
-        if self.dp_size is None:
-            self.dp_size = attn_dp_size if system_dp_size == 1 else system_dp_size
-
-        if self.pp_size is None:
-            self.pp_size = pp_size
-
-        if role == "Prefill":
-            if system_dp_size == 1:
-                dp_group = attn_dp_rank
-            else:
-                dp_group = system_dp_rank
-
-            # Add lock to make sure thread-safe
-            async with self.lock:
-                if dp_group not in self.prefill_port_table:
-                    self.prefill_port_table[dp_group] = {}
-                if attn_tp_rank not in self.prefill_port_table[dp_group]:
-                    self.prefill_port_table[dp_group][attn_tp_rank] = {}
-
-            self.prefill_port_table[dp_group][attn_tp_rank][pp_rank] = {
-                "rank_ip": rank_ip,
-                "rank_port": rank_port,
-            }
-            logger.debug(
-                f"Register prefill bootstrap: DP{dp_group} TP{attn_tp_rank} PP{pp_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")
-        target_dp_group = request.query.get("target_dp_group")
-        target_pp_rank = request.query.get("target_pp_rank")
-        if not engine_rank or not target_dp_group or not target_pp_rank:
-            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
-            and int(target_pp_rank) == -1
-        ):
-            prefill_parallel_info = {
-                "prefill_attn_tp_size": self.attn_tp_size,
-                "prefill_dp_size": self.dp_size,
-                "prefill_pp_size": self.pp_size,
-            }
-            return web.json_response(prefill_parallel_info, status=200)
-
-        # Find corresponding prefill info
-        async with self.lock:
-            bootstrap_info = self.prefill_port_table[int(target_dp_group)][
-                int(engine_rank)
-            ][int(target_pp_rank)]
-
-        if bootstrap_info is not None:
-            return web.json_response(bootstrap_info, status=200)
-        else:
-            return web.Response(text="Bootstrap info not Found", status=404)
-
-    def _run_server(self):
-        try:
-            # Event Loop
-            self._loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(self._loop)
-
-            access_log = None
-            if logging.getLogger(__name__).getEffectiveLevel() <= logging.DEBUG:
-                access_log = self.app.logger
-
-            self._runner = web.AppRunner(self.app, access_log=access_log)
-            self._loop.run_until_complete(self._runner.setup())
-
-            site = web.TCPSite(self._runner, host=self.host, port=self.port)
-            self._loop.run_until_complete(site.start())
-            self._loop.run_forever()
-        except Exception as e:
-            logger.error(f"Server error: {str(e)}")
-        finally:
-            # Cleanup
-            self._loop.run_until_complete(self._runner.cleanup())
-            self._loop.close()
-
-    def close(self):
-        """Shutdown"""
-        if self._loop is not None and self._loop.is_running():
-            self._loop.call_soon_threadsafe(self._loop.stop)
-            logger.info("Stopping server loop...")
-
-        if self.thread.is_alive():
-            self.thread.join(timeout=2)
-            logger.info("Server thread stopped")
-
-    def poll(self) -> KVPoll: ...
+class MooncakeKVBootstrapServer(CommonKVBootstrapServer):
+    pass