# 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 time import numpy as np import numpy.typing as npt from tokenspeed.runtime.pd.base.status import TransferPoll from tokenspeed.runtime.pd.mooncake.entities import KVTransferError from tokenspeed.runtime.pd.utils import PageTransferMetadata from tokenspeed.runtime.utils import get_colorful_logger logger = get_colorful_logger(__name__) class MooncakeKVSender: def __init__( self, mgr, # MooncakeKVManagerPrefill bootstrap_addr: str, bootstrap_room: int, ): self.kv_mgr = mgr self.bootstrap_server_url = bootstrap_addr self.bootstrap_room = bootstrap_room self.kv_mgr.update_status(bootstrap_room, TransferPoll.Bootstrapping) logger.info( "[MooncakeKVSender.__init__] bootstrap_room=%s bootstrap_addr=%s status=Bootstrapping", bootstrap_room, bootstrap_addr, ) # inner state self.init_time = None self.conclude_state = None self.curr_idx = 0 self._layerwise_transfer_started = False def has_layerwise_transfer(self) -> bool: return self._layerwise_transfer_started def send( self, kv_indices: npt.NDArray[np.int64], aux_index, is_last, mla_l1_5_args: PageTransferMetadata | None = None, bootstrap_token: int = -1, spec_candidate_ids: list[int] | None = None, mamba_indices: npt.NDArray[np.int64] | None = None, ): """ Send the kv cache at the given kv indices to the decoder server mla_l1_5_args: optional (page_transfer_mask, page_local_indices) page_transfer_mask: boolean mask to select decode pages that will receive data from this prefill rank page_local_indices: remapped local page indices that this prefill rank will send bootstrap_token: first output token produced by prefill (shipped via ZMQ status msg). """ index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices)) self.curr_idx += len(kv_indices) logger.info( "[MooncakeKVSender.send] bootstrap_room=%s kv_indices_len=%d is_last=%s curr_idx=%d bootstrap_token=%s", self.bootstrap_room, len(kv_indices), is_last, self.curr_idx, bootstrap_token, ) if not is_last: self.kv_mgr.add_transfer_request( self.bootstrap_room, kv_indices, index_slice, False, mla_l1_5_args=mla_l1_5_args, mamba_indices=mamba_indices, ) else: self.kv_mgr.add_transfer_request( self.bootstrap_room, kv_indices, index_slice, True, aux_index=aux_index, mla_l1_5_args=mla_l1_5_args, bootstrap_token=bootstrap_token, spec_candidate_ids=spec_candidate_ids, mamba_indices=mamba_indices, ) def send_layerwise( self, kv_indices: npt.NDArray[np.int64], index_slice: slice, aux_index, is_last, begin_cache_step: int, layerwise_interval: int, mla_l1_5_args: PageTransferMetadata | None = None, bootstrap_token: int = -1, wait_for_bootstrap_token: bool = False, spec_candidate_ids: list[int] | None = None, mamba_indices: npt.NDArray[np.int64] | None = None, ): self._layerwise_transfer_started = True self.curr_idx = max(self.curr_idx, index_slice.stop) if len(kv_indices) == 0 and not is_last: return logger.info( "[MooncakeKVSender.send_layerwise] bootstrap_room=%s kv_indices_len=%d " "slice=(%s,%s) is_last=%s begin_cache_step=%s interval=%s", self.bootstrap_room, len(kv_indices), index_slice.start, index_slice.stop, is_last, begin_cache_step, layerwise_interval, ) self.kv_mgr.add_transfer_request( self.bootstrap_room, kv_indices, index_slice, is_last, aux_index=aux_index if is_last else None, mla_l1_5_args=mla_l1_5_args, bootstrap_token=bootstrap_token, begin_cache_step=begin_cache_step, layerwise_interval=layerwise_interval, wait_for_bootstrap_token=wait_for_bootstrap_token, spec_candidate_ids=spec_candidate_ids, mamba_indices=mamba_indices, ) 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.Bootstrapping: if self.init_time is not None: now = time.time() elapsed = now - self.init_time if elapsed >= self.kv_mgr.bootstrap_time_out: logger.warning_once( "Some requests timed out when bootstrapping, " "which means prefill instances fail to receive the KV indices from the decode instance of this request. " "If a greater mean TTFT is acceptable, you can 'export TOKENSPEED_DISAGGREGATION_BOOTSTRAP_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.Bootstrapping", ) self.conclude_state = TransferPoll.Failed return TransferPoll.Failed 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) 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_server_url )