lightseekorg--tokenspeed/python/tokenspeed/runtime/pd/mooncake/receiver.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import struct
import threading
import time
from dataclasses import dataclass

import numpy as np
import numpy.typing as npt
import requests
import zmq

from tokenspeed.runtime.pd.base.status import TransferPoll
from tokenspeed.runtime.pd.mooncake.entities import KVTransferError
from tokenspeed.runtime.pd.transfer_plan import (
    BufferKind,
    BufferLayout,
    ParallelLayout,
    PDTransferPlanner,
    RankTransferPlan,
    encode_transfer_fragments,
)
from tokenspeed.runtime.pd.utils import (
    PageTransferMetadata,
)
from tokenspeed.runtime.utils import (
    get_colorful_logger,
)
from tokenspeed.runtime.utils.network import get_local_ip_by_remote

logger = get_colorful_logger(__name__)

from tokenspeed.runtime.pd.mooncake.decode import (
    MooncakeKVManagerDecode,
    PrefillParallelInfo,
)


def _get_prefill_parallel_info_from_server(
    bootstrap_addr,
) -> PrefillParallelInfo | None:
    """Fetch the prefill parallel info from the bootstrap server."""
    try:
        url = f"http://{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 PrefillParallelInfo(
                tp_size=int(prefill_parallel_info["prefill_tp_size"]),
                dp_size=int(prefill_parallel_info["prefill_dp_size"]),
                enable_mla_l1_5_cache=bool(
                    prefill_parallel_info["enable_mla_l1_5_cache"]
                ),
                kv_item_lens=tuple(
                    int(x) for x in prefill_parallel_info.get("kv_item_lens", [])
                ),
                kv_unit_lens=tuple(
                    int(x) for x in prefill_parallel_info.get("kv_unit_lens", [])
                ),
                state_item_lens=tuple(
                    int(x) for x in prefill_parallel_info.get("state_item_lens", [])
                ),
                state_unit_lens=tuple(
                    int(x) for x in prefill_parallel_info.get("state_unit_lens", [])
                ),
            )
        else:
            logger.error(
                "Failed to get prefill parallel info: %s, %s",
                response.status_code,
                response.text,
            )
            return None
    except Exception as exc:
        logger.error("Error fetching prefill parallel info from bootstrap: %s", exc)
        return None


def _get_bootstrap_info_from_server(bootstrap_addr, engine_rank, target_dp_group):
    """Fetch the bootstrap info from the bootstrap server."""
    try:
        url = f"http://{bootstrap_addr}/route?engine_rank={engine_rank}&target_dp_group={target_dp_group}"
        response = requests.get(url, timeout=5)
        if response.status_code == 200:
            bootstrap_info = response.json()
            return bootstrap_info
        else:
            logger.error(
                "Failed to get prefill server info: %s, %s",
                response.status_code,
                response.text,
            )
            return None
    except Exception as exc:
        logger.error("Error fetching prefill info from bootstrap: %s", exc)
        return None


@dataclass(frozen=True)
class ReceiverRoutePlan:
    target_tp_rank: int | None
    target_tp_ranks: tuple[int, ...]
    required_prefill_response_num: int
    default_required_dst_info_num: int
    transfer_plan: RankTransferPlan | None = None
    supports_remote_spec_candidates: bool = True

    def required_dst_info_num_for_tp_rank(self, tp_rank: int) -> int:
        if self.transfer_plan is None:
            return self.default_required_dst_info_num
        return self.transfer_plan.required_dst_info_num_for_prefill_rank(tp_rank)

    def fragments_for_tp_rank(self, tp_rank: int):
        if self.transfer_plan is None:
            return ()
        return self.transfer_plan.fragments_by_prefill_rank.get(tp_rank, ())


def _buffer_kind_for_layer_offset(
    kv_args, layer_index: int, offset_index: int
) -> BufferKind:
    is_draft = layer_index >= getattr(kv_args, "target_layer_num", len(kv_args.offsets))
    if is_draft:
        return BufferKind.DRAFT_K if offset_index == 0 else BufferKind.DRAFT_V
    return BufferKind.TARGET_K if offset_index == 0 else BufferKind.TARGET_V


def _unit_lens_or_default(unit_lens, item_lens):
    if unit_lens:
        return tuple(int(x) for x in unit_lens)
    return tuple(1 for _ in item_lens)


def _build_buffer_layout_pair(
    *,
    buffer_index: int,
    buffer_kind: BufferKind,
    sharded_axis: str,
    prefill_item_len: int,
    decode_item_len: int,
    prefill_unit_len: int,
    decode_unit_len: int,
    prefill_tp_size: int,
    decode_tp_size: int,
):
    """Build compatible logical layouts for one prefill/decode buffer pair.

    Besides normal TP sharding and fully replicated buffers, this handles GQA
    KV caches where prefill TP is larger than the number of distinct KV heads.
    In that case multiple prefill TP ranks carry the same KV head, so the
    transfer plan uses one representative rank from each replica group.
    """

    if prefill_unit_len != decode_unit_len:
        raise ValueError(
            f"prefill/decode unit sizes differ for {buffer_kind.value}: "
            f"prefill={prefill_unit_len}, decode={decode_unit_len}"
        )
    if prefill_item_len % prefill_unit_len != 0:
        raise ValueError(
            f"prefill item length is not unit-aligned for {buffer_kind.value}: "
            f"item={prefill_item_len}, unit={prefill_unit_len}"
        )
    if decode_item_len % decode_unit_len != 0:
        raise ValueError(
            f"decode item length is not unit-aligned for {buffer_kind.value}: "
            f"item={decode_item_len}, unit={decode_unit_len}"
        )

    prefill_local_units = prefill_item_len // prefill_unit_len
    decode_local_units = decode_item_len // decode_unit_len
    prefill_global_units = prefill_local_units * prefill_tp_size
    decode_global_units = decode_local_units * decode_tp_size
    prefill_tp_replica_group_size = 1
    if prefill_global_units == decode_global_units:
        logical_axis = sharded_axis
        logical_size = decode_global_units
    elif prefill_item_len == decode_item_len:
        logical_axis = "replicated"
        logical_size = decode_local_units
    elif (
        sharded_axis == "kv_head"
        and decode_global_units % prefill_local_units == 0
        and decode_global_units // prefill_local_units <= prefill_tp_size
        and prefill_tp_size % (decode_global_units // prefill_local_units) == 0
    ):
        logical_axis = sharded_axis
        logical_size = decode_global_units
        prefill_distinct_tp_size = decode_global_units // prefill_local_units
        prefill_tp_replica_group_size = prefill_tp_size // prefill_distinct_tp_size
    else:
        raise ValueError(
            f"unsupported heterogeneous TP buffer layout for {buffer_kind.value}: "
            f"prefill_item={prefill_item_len}, decode_item={decode_item_len}, "
            f"prefill_tp={prefill_tp_size}, decode_tp={decode_tp_size}, "
            f"unit={decode_unit_len}"
        )

    return (
        BufferLayout(
            buffer_index=buffer_index,
            buffer_kind=buffer_kind,
            logical_axis=logical_axis,
            logical_size=logical_size,
            page_size=1,
            bytes_per_logical_unit=decode_unit_len,
            item_stride_bytes=prefill_item_len,
            tp_replica_group_size=prefill_tp_replica_group_size,
        ),
        BufferLayout(
            buffer_index=buffer_index,
            buffer_kind=buffer_kind,
            logical_axis=logical_axis,
            logical_size=logical_size,
            page_size=1,
            bytes_per_logical_unit=decode_unit_len,
            item_stride_bytes=decode_item_len,
        ),
    )


def _build_kv_buffer_layouts(
    kv_args, prefill_parallel_info: PrefillParallelInfo, decode_tp_size: int
):
    prefill_tp_size = prefill_parallel_info.prefill_tp_size_per_dp_rank
    prefill_kv_item_lens = tuple(prefill_parallel_info.kv_item_lens)
    if not prefill_kv_item_lens:
        prefill_kv_item_lens = tuple(
            int(x) * decode_tp_size // prefill_tp_size for x in kv_args.kv_item_lens
        )
    prefill_kv_unit_lens = _unit_lens_or_default(
        prefill_parallel_info.kv_unit_lens, prefill_kv_item_lens
    )
    decode_kv_unit_lens = _unit_lens_or_default(
        getattr(kv_args, "kv_unit_lens", []), kv_args.kv_item_lens
    )

    prefill_buffers = []
    decode_buffers = []
    for layer_index, ptr_offsets in enumerate(kv_args.offsets):
        for offset_index, ptr_offset in enumerate(ptr_offsets):
            buffer_kind = _buffer_kind_for_layer_offset(
                kv_args, layer_index, offset_index
            )
            prefill_buffer, decode_buffer = _build_buffer_layout_pair(
                buffer_index=ptr_offset,
                buffer_kind=buffer_kind,
                sharded_axis="kv_head",
                prefill_item_len=int(prefill_kv_item_lens[ptr_offset]),
                decode_item_len=int(kv_args.kv_item_lens[ptr_offset]),
                prefill_unit_len=int(prefill_kv_unit_lens[ptr_offset]),
                decode_unit_len=int(decode_kv_unit_lens[ptr_offset]),
                prefill_tp_size=prefill_tp_size,
                decode_tp_size=decode_tp_size,
            )
            prefill_buffers.append(prefill_buffer)
            decode_buffers.append(decode_buffer)

    decode_state_item_lens = tuple(getattr(kv_args, "state_item_lens", []) or [])
    prefill_state_item_lens = tuple(prefill_parallel_info.state_item_lens)
    if not prefill_state_item_lens:
        prefill_state_item_lens = tuple(
            int(x) * decode_tp_size // prefill_tp_size for x in decode_state_item_lens
        )
    prefill_state_unit_lens = _unit_lens_or_default(
        prefill_parallel_info.state_unit_lens, prefill_state_item_lens
    )
    decode_state_unit_lens = _unit_lens_or_default(
        getattr(kv_args, "state_unit_lens", []), decode_state_item_lens
    )

    for state_index, decode_item_len in enumerate(decode_state_item_lens):
        prefill_buffer, decode_buffer = _build_buffer_layout_pair(
            buffer_index=state_index,
            buffer_kind=BufferKind.MAMBA_STATE,
            sharded_axis="state_channel",
            prefill_item_len=int(prefill_state_item_lens[state_index]),
            decode_item_len=int(decode_item_len),
            prefill_unit_len=int(prefill_state_unit_lens[state_index]),
            decode_unit_len=int(decode_state_unit_lens[state_index]),
            prefill_tp_size=prefill_tp_size,
            decode_tp_size=decode_tp_size,
        )
        prefill_buffers.append(prefill_buffer)
        decode_buffers.append(decode_buffer)
    return tuple(prefill_buffers), tuple(decode_buffers)


def _build_non_mla_route_plan(kv_mgr, prefill_parallel_info: PrefillParallelInfo):
    prefill_tp_size = prefill_parallel_info.prefill_tp_size_per_dp_rank
    decode_tp_size = kv_mgr.world_size // kv_mgr.dp_size
    decode_tp_rank = kv_mgr.kv_args.engine_rank % decode_tp_size
    prefill_buffers, decode_buffers = _build_kv_buffer_layouts(
        kv_mgr.kv_args,
        prefill_parallel_info,
        decode_tp_size,
    )
    planner = PDTransferPlanner(
        prefill_layout=ParallelLayout(
            role="prefill",
            world_size=prefill_tp_size,
            dp_size=1,
        ),
        decode_layout=ParallelLayout(
            role="decode",
            world_size=decode_tp_size,
            dp_size=1,
        ),
        prefill_buffers=prefill_buffers,
        decode_buffers=decode_buffers,
    )
    transfer_plan = planner.plan_for_decode_rank(decode_tp_rank)
    target_tp_ranks = tuple(transfer_plan.target_prefill_ranks)
    target_tp_rank = (
        target_tp_ranks[0] if transfer_plan.plan_kind == "identity" else None
    )
    default_required_dst_info_num = (
        transfer_plan.required_dst_info_num_for_prefill_rank(target_tp_ranks[0])
        if target_tp_ranks
        else 1
    )
    return ReceiverRoutePlan(
        target_tp_rank=target_tp_rank,
        target_tp_ranks=target_tp_ranks,
        required_prefill_response_num=transfer_plan.required_prefill_response_num,
        default_required_dst_info_num=default_required_dst_info_num,
        transfer_plan=transfer_plan,
        supports_remote_spec_candidates=transfer_plan.plan_kind == "identity",
    )


def _legacy_mla_route_plan(
    *,
    target_tp_rank: int | None,
    target_tp_ranks,
    required_dst_info_num: int,
    required_prefill_response_num: int,
) -> ReceiverRoutePlan:
    return ReceiverRoutePlan(
        target_tp_rank=target_tp_rank,
        target_tp_ranks=tuple(target_tp_ranks),
        required_prefill_response_num=required_prefill_response_num,
        default_required_dst_info_num=required_dst_info_num,
    )


def _calc(kv_mgr, prefill_parallel_info: PrefillParallelInfo) -> ReceiverRoutePlan:
    prefill_tp_size_per_dp_rank = prefill_parallel_info.prefill_tp_size_per_dp_rank
    local_tp_size_per_dp_rank = kv_mgr.world_size // kv_mgr.dp_size

    if prefill_parallel_info.enable_mla_l1_5_cache:
        if not kv_mgr.is_mla_backend:
            raise RuntimeError(
                "PD with MLA L1.5 cache is not yet supported for non-MLA models"
            )
        return _legacy_mla_route_plan(
            target_tp_rank=None,
            target_tp_ranks=range(prefill_tp_size_per_dp_rank),
            required_dst_info_num=local_tp_size_per_dp_rank,
            required_prefill_response_num=prefill_tp_size_per_dp_rank,
        )

    if not kv_mgr.is_mla_backend:
        return _build_non_mla_route_plan(kv_mgr, prefill_parallel_info)

    if local_tp_size_per_dp_rank == prefill_tp_size_per_dp_rank:
        target_tp_rank = kv_mgr.kv_args.engine_rank % local_tp_size_per_dp_rank
        return _legacy_mla_route_plan(
            target_tp_rank=target_tp_rank,
            target_tp_ranks=(target_tp_rank,),
            required_dst_info_num=1,
            required_prefill_response_num=1,
        )

    if local_tp_size_per_dp_rank > prefill_tp_size_per_dp_rank:
        target_tp_rank = (kv_mgr.kv_args.engine_rank % local_tp_size_per_dp_rank) // (
            local_tp_size_per_dp_rank // prefill_tp_size_per_dp_rank
        )
        return _legacy_mla_route_plan(
            target_tp_rank=target_tp_rank,
            target_tp_ranks=(target_tp_rank,),
            required_dst_info_num=local_tp_size_per_dp_rank
            // prefill_tp_size_per_dp_rank,
            required_prefill_response_num=1,
        )

    target_tp_ranks = tuple(
        range(
            (kv_mgr.kv_args.engine_rank % local_tp_size_per_dp_rank)
            * (prefill_tp_size_per_dp_rank // local_tp_size_per_dp_rank),
            (kv_mgr.kv_args.engine_rank % local_tp_size_per_dp_rank + 1)
            * (prefill_tp_size_per_dp_rank // local_tp_size_per_dp_rank),
        )
    )
    return _legacy_mla_route_plan(
        target_tp_rank=target_tp_ranks[0],
        target_tp_ranks=target_tp_ranks,
        required_dst_info_num=1,
        required_prefill_response_num=1,
    )


class MooncakeKVReceiver:
    _ctx = zmq.Context()
    _socket_cache = {}
    _socket_locks = {}
    _global_lock = threading.Lock()

    def __init__(
        self, mgr: MooncakeKVManagerDecode, bootstrap_addr: str, bootstrap_room: int
    ):
        self.kv_mgr = mgr
        self.bootstrap_addr = bootstrap_addr
        self.bootstrap_room = bootstrap_room

        self.session_id = self.kv_mgr.get_session_id()
        self.conclude_state = None
        self.init_time = None
        self.prefill_enable_mla_l1_5_cache = None
        self.dst_enable_mla_l1_5_cache = False

        self.kv_mgr.update_status(self.bootstrap_room, TransferPoll.Bootstrapping)
        logger.info(
            "[MooncakeKVReceiver.__init__] bootstrap_addr=%s bootstrap_room=%s session_id=%s",
            bootstrap_addr,
            bootstrap_room,
            self.session_id,
        )

        prefill_parallel_info = self._get_prefill_parallel_info()
        if prefill_parallel_info 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, TransferPoll.Failed)

        route_plan = _calc(self.kv_mgr, prefill_parallel_info)
        self.route_plan = route_plan
        self.supports_remote_spec_candidates = (
            route_plan.supports_remote_spec_candidates
        )
        self.required_dst_info_num = route_plan.default_required_dst_info_num
        self.kv_mgr.required_prefill_response_num_table[self.bootstrap_room] = (
            route_plan.required_prefill_response_num
        )
        target_dp_group = self.bootstrap_room % prefill_parallel_info.dp_size
        target_tp_key = ",".join(str(rank) for rank in route_plan.target_tp_ranks)
        bootstrap_key = f"{self.bootstrap_addr}_{target_dp_group}_{target_tp_key}"
        if bootstrap_key not in self.kv_mgr.connection_pool:
            bootstrap_infos = self._get_bootstrap_infos(target_dp_group, route_plan)
            if bootstrap_infos is None:
                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: {target_dp_group}",
                )
                self.kv_mgr.update_status(self.bootstrap_room, TransferPoll.Failed)
            else:
                if not bootstrap_infos:
                    raise RuntimeError("Could not fetch bootstrap info.")
                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]

        self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room)
        self.kv_mgr.update_status(self.bootstrap_room, TransferPoll.Bootstrapped)
        logger.info(
            "[MooncakeKVReceiver.__init__] done, status set to Bootstrapped. "
            "bootstrap_room=%s bootstrap_addr=%s session_id=%s",
            self.bootstrap_room,
            self.bootstrap_addr,
            self.session_id,
        )

    def _get_prefill_parallel_info(self):
        prefill_parallel_info = self.kv_mgr.prefill_parallel_info.get(
            self.bootstrap_addr
        )

        if prefill_parallel_info is not None:
            return prefill_parallel_info
        else:
            prefill_parallel_info = _get_prefill_parallel_info_from_server(
                self.bootstrap_addr
            )

            if prefill_parallel_info is None:
                return None
            else:
                logger.debug(
                    "Fetch prefill parallel info from [%s]: DP size:%s, TP size:%s",
                    self.bootstrap_addr,
                    prefill_parallel_info.dp_size,
                    prefill_parallel_info.tp_size,
                )
                self.kv_mgr.prefill_parallel_info[self.bootstrap_addr] = (
                    prefill_parallel_info
                )
                return prefill_parallel_info

    def _get_bootstrap_infos(self, target_dp_group, route_plan: ReceiverRoutePlan):
        bootstrap_infos = []
        for _target_tp_rank in route_plan.target_tp_ranks:
            bootstrap_info = _get_bootstrap_info_from_server(
                self.bootstrap_addr,
                _target_tp_rank,
                target_dp_group,
            )
            if bootstrap_info is not None:
                #  only support MLA for now: select one prefill rank as real rank
                bootstrap_info["is_dummy"] = not bool(
                    _target_tp_rank == route_plan.target_tp_rank
                    or route_plan.target_tp_rank is None
                )
                bootstrap_info["required_dst_info_num"] = (
                    route_plan.required_dst_info_num_for_tp_rank(_target_tp_rank)
                )
                bootstrap_info["transfer_fragments"] = route_plan.fragments_for_tp_rank(
                    _target_tp_rank
                )
                logger.debug(
                    "Fetched bootstrap info: %s for DP %s TP %s",
                    bootstrap_info,
                    target_dp_group,
                    _target_tp_rank,
                )
                bootstrap_infos.append(bootstrap_info)
            else:
                return None
        return bootstrap_infos

    def _register_kv_args(self):
        for bootstrap_info in self.bootstrap_infos:
            self.prefill_server_url = (
                f"{bootstrap_info['rank_ip']}:{bootstrap_info['rank_port']}"
            )
            logger.info(
                "[MooncakeKVReceiver._register_kv_args] sending kv_args to prefill=%s bootstrap_room=%s session_id=%s",
                self.prefill_server_url,
                self.bootstrap_room,
                self.session_id,
            )
            packed_kv_data_ptrs = b"".join(
                struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.kv_data_ptrs
            )
            packed_state_data_ptrs = b"".join(
                struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.state_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,
                        b"",  # aux_data_ptrs removed; kept as empty frame for protocol compat
                        packed_state_data_ptrs,
                        # Include decode_prefix_len for kv_args registration
                        str(getattr(self, "decode_prefix_len", 0)).encode("ascii"),
                    ]
                )

    @classmethod
    def _connect(cls, endpoint: str):
        with cls._global_lock:
            if endpoint not in cls._socket_cache:
                sock = cls._ctx.socket(zmq.PUSH)
                sock.connect(endpoint)
                cls._socket_cache[endpoint] = sock
                cls._socket_locks[endpoint] = threading.Lock()
            return cls._socket_cache[endpoint], cls._socket_locks[endpoint]

    def prefill(
        self,
        kv_indices: npt.NDArray[np.int64],
        aux_index: int | None = None,
        decode_prefix_len: int | None = 0,
        mla_l1_5_args: PageTransferMetadata | None = None,
        mamba_indices: npt.NDArray[np.int64] | None = None,
    ):
        logger.info(
            "[MooncakeKVReceiver.init] bootstrap_room=%s kv_indices_len=%d aux_index=%s decode_prefix_len=%s",
            self.bootstrap_room,
            len(kv_indices),
            aux_index,
            decode_prefix_len,
        )
        # Store decode_prefix_len to be sent back to prefill
        self.decode_prefix_len = decode_prefix_len
        dst_page_transfer_mask = None
        dst_page_local_indices = None
        if mla_l1_5_args is not None:
            dst_page_transfer_mask = mla_l1_5_args.page_transfer_mask
            dst_page_local_indices = mla_l1_5_args.page_local_indices

        for bootstrap_info in self.bootstrap_infos:
            self.prefill_server_url = (
                f"{bootstrap_info['rank_ip']}:{bootstrap_info['rank_port']}"
            )
            is_dummy = bootstrap_info["is_dummy"]

            logger.info(
                "[MooncakeKVReceiver.init] sending pre-alloc multipart to prefill=%s bootstrap_room=%s is_dummy=%s",
                self.prefill_server_url,
                self.bootstrap_room,
                bootstrap_info["is_dummy"],
            )
            sock, lock = self._connect("tcp://" + self.prefill_server_url)
            with lock:
                message_parts = [
                    str(self.bootstrap_room).encode("ascii"),
                    get_local_ip_by_remote().encode("ascii"),
                    str(self.kv_mgr.rank_port).encode("ascii"),
                    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"",
                    str(
                        bootstrap_info.get(
                            "required_dst_info_num", self.required_dst_info_num
                        )
                    ).encode("ascii"),
                    # Send decode_prefix_len as additional message part
                    (
                        str(self.decode_prefix_len).encode("ascii")
                        if not is_dummy
                        else b""
                    ),
                    (
                        str(int(self.dst_enable_mla_l1_5_cache)).encode("ascii")
                        if not is_dummy
                        else b""
                    ),
                    (
                        dst_page_transfer_mask.tobytes()
                        if (not is_dummy and dst_page_transfer_mask is not None)
                        else b""
                    ),
                    (
                        dst_page_local_indices.tobytes()
                        if (not is_dummy and dst_page_local_indices is not None)
                        else b""
                    ),
                    (
                        mamba_indices.tobytes()
                        if (not is_dummy and mamba_indices is not None)
                        else b""
                    ),
                ]
                transfer_fragments = bootstrap_info.get("transfer_fragments", ())
                if not is_dummy and transfer_fragments:
                    message_parts.extend(encode_transfer_fragments(transfer_fragments))
                sock.send_multipart(message_parts)
            self.init_time = time.time()

    def poll(self) -> TransferPoll:
        if self.conclude_state is None:
            status = self.kv_mgr.check_status(self.bootstrap_room)
            if status in (TransferPoll.Success, TransferPoll.Failed):
                self.conclude_state = status
            elif status == TransferPoll.WaitingForInput:
                if self.init_time is not None:
                    now = time.time()
                    elapsed = now - self.init_time
                    if elapsed >= self.kv_mgr.waiting_timeout:
                        logger.warning_once(
                            "Some requests fail to receive KV Cache transfer done signal after bootstrapping. "
                            "If a greater mean TTFT is acceptable, you can 'export TOKENSPEED_DISAGGREGATION_WAITING_TIMEOUT=600' (10 minutes) to relax the timeout condition. "
                        )
                        self.kv_mgr.record_failure(
                            self.bootstrap_room,
                            f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s in TransferPoll.WaitingForInput",
                        )
                        self.conclude_state = TransferPoll.Failed
                        return TransferPoll.Failed
            elif status == TransferPoll.Transferring:
                logger.warning(
                    "Req(room=%s) in Transferring, which is unexpected",
                    self.bootstrap_room,
                )

            return status
        else:
            return self.conclude_state

    def clear(self) -> None:
        if self.bootstrap_room in self.kv_mgr.request_status:
            self.kv_mgr.request_status.pop(self.bootstrap_room)

        if self.bootstrap_room in self.kv_mgr.required_prefill_response_num_table:
            self.kv_mgr.required_prefill_response_num_table.pop(self.bootstrap_room)

        if self.bootstrap_room in self.kv_mgr.prefill_response_tracker:
            self.kv_mgr.prefill_response_tracker.pop(self.bootstrap_room)

    def failure_exception(self):
        # Explicitly set the status to failure since this request has failed in another rank
        if self.conclude_state is None:
            self.conclude_state = TransferPoll.Failed

        self.clear()

        with self.kv_mgr.failure_lock:
            failure_reason = self.kv_mgr.failure_records.pop(
                self.bootstrap_room, "Failed due to an unknown reason from another rank"
            )
        raise KVTransferError(self.bootstrap_room, failure_reason, self.bootstrap_addr)