lightseekorg--tokenspeed/python/tokenspeed/runtime/pd/mooncake/decode.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 threading
import time
from collections import defaultdict
from dataclasses import dataclass

import numpy as np
import requests

from tokenspeed.runtime.pd.base.status import TransferPoll
from tokenspeed.runtime.pd.mooncake.conn import MooncakeKVManagerBase
from tokenspeed.runtime.pd.mooncake.entities import KVArgs, KVManagerArgs
from tokenspeed.runtime.pd.utils import DisaggregationMode
from tokenspeed.runtime.utils import (
    get_colorful_logger,
)
from tokenspeed.runtime.utils.env import envs
from tokenspeed.runtime.utils.network import get_free_port, get_local_ip_by_remote

logger = get_colorful_logger(__name__)


@dataclass
class PrefillParallelInfo:
    tp_size: int
    dp_size: int
    enable_mla_l1_5_cache: bool
    kv_item_lens: tuple[int, ...] = ()
    kv_unit_lens: tuple[int, ...] = ()
    state_item_lens: tuple[int, ...] = ()
    state_unit_lens: tuple[int, ...] = ()

    @property
    def prefill_tp_size_per_dp_rank(self):
        return self.tp_size // self.dp_size


def parse_prefill_status_message(
    parts: list[bytes],
) -> tuple[int, int, int, int, list[int] | None]:
    bootstrap_room = int(parts[0].decode("ascii"))
    status = int(parts[1].decode("ascii"))
    prefill_rank = int(parts[2].decode("ascii"))
    bootstrap_token = int(parts[3].decode("ascii")) if len(parts) > 3 else -1
    spec_candidate_ids = None
    if len(parts) > 4 and parts[4] != b"":
        spec_candidate_ids = np.frombuffer(parts[4], dtype=np.int32).copy().tolist()
    return (
        bootstrap_room,
        status,
        prefill_rank,
        bootstrap_token,
        spec_candidate_ids,
    )


class MooncakeKVManagerDecode(MooncakeKVManagerBase):
    def __init__(
        self,
        args: KVManagerArgs,
        kv_args: KVArgs,
    ):
        super().__init__(args, kv_args, DisaggregationMode.DECODE)

        self.heartbeat_failures = {}
        self.session_pool = defaultdict(requests.Session)
        self.session_pool_lock = threading.Lock()
        self.addr_to_rooms_tracker = defaultdict(set)
        self.connection_lock = threading.Lock()
        # Heartbeat interval should be at least 2 seconds
        self.heartbeat_interval = max(
            envs.TOKENSPEED_DISAGGREGATION_HEARTBEAT_INTERVAL.get(),
            2.0,
        )
        # Heartbeat failure should be at least 1
        self.max_failures = max(
            envs.TOKENSPEED_DISAGGREGATION_HEARTBEAT_MAX_FAILURE.get(), 1
        )
        self.start_decode_thread()
        self.connection_pool: dict[str, dict[str, str | int]] = {}
        self.required_prefill_response_num_table: dict[int, int] = {}
        self.prefill_response_tracker: dict[int, set[int]] = defaultdict(set)

        self.prefill_parallel_info: dict[str, PrefillParallelInfo] = {}

        # 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 = envs.TOKENSPEED_DISAGGREGATION_WAITING_TIMEOUT.get()

    def start_decode_thread(self):
        self.rank_port = get_free_port()
        self.server_socket.bind(f"tcp://{get_local_ip_by_remote()}:{self.rank_port}")
        # Maps bootstrap_room -> bootstrap_token (first output token from prefill).
        # Populated by decode_thread when a Success message carries a valid token,
        # consumed by DisaggDecodeExecutor.generate_events() via pop_bootstrap_token().
        self.bootstrap_token_table: dict[int, int] = {}
        self.spec_candidate_ids_table: dict[int, list[int]] = {}
        self._pending_bootstrap_token_table: dict[int, int] = {}
        self._pending_spec_candidate_ids_table: dict[int, list[int]] = {}

        def decode_thread():
            while True:
                parts = self.server_socket.recv_multipart()
                self._handle_prefill_status(*parse_prefill_status_message(parts))

        def heartbeat_checker():
            while True:
                time.sleep(self.heartbeat_interval)
                with self.connection_lock:
                    addresses = list(self.prefill_parallel_info.keys())

                for bootstrap_addr in addresses:
                    session = None
                    try:
                        with self.session_pool_lock:
                            session = self.session_pool[bootstrap_addr]
                        response = session.get(
                            f"http://{bootstrap_addr}/health",
                            timeout=(2, 3),
                            headers={"Connection": "keep-alive"},
                        )
                        if response.status_code == 200:
                            self.heartbeat_failures[bootstrap_addr] = 0

                            current_rooms = self.addr_to_rooms_tracker[
                                bootstrap_addr
                            ].copy()

                            for bootstrap_room in current_rooms:
                                # Remove TransferPoll.Success requests from the map
                                if bootstrap_room not in self.request_status:
                                    self.addr_to_rooms_tracker[bootstrap_addr].discard(
                                        bootstrap_room
                                    )
                        else:
                            logger.info(
                                "Attempting to reconnect to %s...", bootstrap_addr
                            )
                            self.heartbeat_failures[bootstrap_addr] = (
                                self.heartbeat_failures.get(bootstrap_addr, 0) + 1
                            )
                            with self.session_pool_lock:
                                if bootstrap_addr in self.session_pool:
                                    del self.session_pool[bootstrap_addr]
                    except Exception:
                        logger.info("Attempting to reconnect to %s...", bootstrap_addr)
                        self.heartbeat_failures[bootstrap_addr] = (
                            self.heartbeat_failures.get(bootstrap_addr, 0) + 1
                        )

                    if (
                        self.heartbeat_failures.get(bootstrap_addr, 0)
                        >= self.max_failures
                    ):
                        self._handle_node_failure(bootstrap_addr)
                        with self.session_pool_lock:
                            if bootstrap_addr in self.session_pool:
                                del self.session_pool[bootstrap_addr]

        threading.Thread(target=decode_thread).start()
        threading.Thread(target=heartbeat_checker).start()

    def _handle_prefill_status(
        self,
        bootstrap_room: int,
        status: int,
        prefill_rank: int,
        bootstrap_token: int,
        spec_candidate_ids: list[int] | None,
    ) -> None:
        if status == TransferPoll.Success:
            if bootstrap_room not in self.request_status:
                return

            self.prefill_response_tracker[bootstrap_room].add(prefill_rank)
            if bootstrap_token != -1:
                self._pending_bootstrap_token_table.setdefault(
                    bootstrap_room, bootstrap_token
                )
            if spec_candidate_ids is not None:
                self._pending_spec_candidate_ids_table.setdefault(
                    bootstrap_room, spec_candidate_ids
                )

            expected_response_num = self.required_prefill_response_num_table.get(
                bootstrap_room, 1
            )
            arrived_response_num = len(self.prefill_response_tracker[bootstrap_room])
            if arrived_response_num < expected_response_num:
                return

            # Store metadata before marking Success so generate_events() can read
            # it atomically in the same iteration. In heterogeneous TP, the rank
            # carrying bootstrap metadata is not guaranteed to be the last rank
            # whose transfer completes, so preserve the first valid value seen.
            if bootstrap_room in self._pending_bootstrap_token_table:
                self.bootstrap_token_table[bootstrap_room] = (
                    self._pending_bootstrap_token_table.pop(bootstrap_room)
                )
            if bootstrap_room in self._pending_spec_candidate_ids_table:
                self.spec_candidate_ids_table[bootstrap_room] = (
                    self._pending_spec_candidate_ids_table.pop(bootstrap_room)
                )
            self.update_status(bootstrap_room, TransferPoll.Success)
            return

        if status == TransferPoll.Failed:
            self.record_failure(
                bootstrap_room,
                "Failed to get kvcache from prefill instance, it might be dead",
            )
            self.update_status(bootstrap_room, TransferPoll.Failed)
            return

        self.update_status(bootstrap_room, status)

    def pop_bootstrap_token(self, bootstrap_room: int) -> int:
        """Pop and return the bootstrap_token for the given room, or -1 if absent."""
        return self.bootstrap_token_table.pop(bootstrap_room, -1)

    def pop_prefill_metadata(self, bootstrap_room: int) -> tuple[int, list[int] | None]:
        return (
            self.bootstrap_token_table.pop(bootstrap_room, -1),
            self.spec_candidate_ids_table.pop(bootstrap_room, None),
        )

    def get_session_id(self):
        return self.engine.get_session_id()

    def _handle_node_failure(self, failed_bootstrap_addr):
        with self.connection_lock:
            keys_to_remove = [
                k for k in self.connection_pool if k.startswith(failed_bootstrap_addr)
            ]
            for k in keys_to_remove:
                del self.connection_pool[k]
            if failed_bootstrap_addr in self.prefill_parallel_info:
                del self.prefill_parallel_info[failed_bootstrap_addr]

            possible_affected_rooms = self.addr_to_rooms_tracker.get(
                failed_bootstrap_addr, []
            )
            if failed_bootstrap_addr in self.addr_to_rooms_tracker:
                del self.addr_to_rooms_tracker[failed_bootstrap_addr]

        # Report the requests associated with the failed bootstrap addr and mark their status as TransferPoll.Failed
        affected_rooms = []
        for room in possible_affected_rooms:
            if (
                room in self.request_status
                and self.check_status(room) != TransferPoll.Success
            ):
                self.record_failure(
                    room,
                    f"Losing connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr})",
                )
                self.update_status(room, TransferPoll.Failed)
                affected_rooms.append(room)
        logger.error(
            "Losing connection with prefill instance (bootstrap_addr: %s), affected %s requests",
            failed_bootstrap_addr,
            len(affected_rooms),
        )