# 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 concurrent.futures import os import socket import threading import time from collections import defaultdict import numpy as np import numpy.typing as npt 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, KVArgsRegisterInfo, KVManagerArgs, TransferIndexResolution, TransferInfo, TransferKVChunk, ) from tokenspeed.runtime.pd.transfer_plan import ( BufferKind, TransferFragment, ) from tokenspeed.runtime.pd.utils import ( DisaggregationMode, FastQueue, PageTransferMetadata, StepCounter, group_concurrent_contiguous, ) 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_ip, get_local_ip_by_remote, ) logger = get_colorful_logger(__name__) class MooncakeKVManagerPrefill(MooncakeKVManagerBase): def __init__( self, args: KVManagerArgs, kv_args: KVArgs, ): super().__init__(args, kv_args, 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: dict[str, float] = {} self.failed_session_ttl = max( envs.TOKENSPEED_DISAGGREGATION_FAILED_SESSION_TTL.get(), 0 ) self.session_lock = threading.Lock() self.kv_layer_ids = list( getattr(self.kv_args, "kv_layer_ids", None) or range(len(self.kv_args.offsets)) ) self.state_layer_ids = list(getattr(self.kv_args, "state_layer_ids", []) or []) layer_ids = self.kv_layer_ids + self.state_layer_ids self.layer_num = ( (max(layer_ids) + 1) if layer_ids else len(self.kv_args.offsets) ) self._kv_layer_to_index = { layer_id: i for i, layer_id in enumerate(self.kv_layer_ids[: len(self.kv_args.offsets)]) } self.layerwise_interval = 1 self.layerwise_debug = envs.TOKENSPEED_PD_LAYERWISE_DEBUG.get() self.step_counter = None # room -> (bootstrap_token, spec_candidate_ids). Published after the prefill # forward; the transfer thread reads it on the wait_for_bootstrap_token path. self.prefill_metadata: dict[int, tuple[int, list[int] | None]] = {} self.expired_prefill_metadata_rooms: set[int] = set() self.bootstrap_token_cond = threading.Condition() # Determine the number of threads to use for kv sender cpu_count = os.cpu_count() transfer_thread_pool_size = ( envs.TOKENSPEED_DISAGGREGATION_THREAD_POOL_SIZE.get_set_value_or( min(max(4, int(0.75 * cpu_count) // 8), 12) ) ) transfer_queue_size = envs.TOKENSPEED_DISAGGREGATION_QUEUE_SIZE.get() if transfer_thread_pool_size < transfer_queue_size: raise ValueError( "TOKENSPEED_DISAGGREGATION_THREAD_POOL_SIZE=" f"{transfer_thread_pool_size} must be greater than or equal to " f"TOKENSPEED_DISAGGREGATION_QUEUE_SIZE={transfer_queue_size}." ) self.start_transfer_thread(transfer_thread_pool_size, transfer_queue_size) self.bootstrap_time_out = envs.TOKENSPEED_DISAGGREGATION_BOOTSTRAP_TIMEOUT.get() def register_layerwise_step_counter( self, step_counter: StepCounter, interval: int ) -> None: self.step_counter = step_counter self.layerwise_interval = max(int(interval), 1) def reserve_layerwise_cache_steps(self) -> int: if self.step_counter is None: return 0 cache_step, _ = self.step_counter.current_step() self.step_counter.advance_step( delta_cache_step=self.layer_num, delta_aux_step=0, ) return cache_step def set_prefill_metadata( self, room: int, token: int, spec_candidate_ids: list[int] | None = None, ) -> None: with self.bootstrap_token_cond: if room in self.expired_prefill_metadata_rooms: self.expired_prefill_metadata_rooms.discard(room) logger.warning( "Dropping late prefill metadata for expired bootstrap_room=%s", room, ) return self.prefill_metadata[room] = ( token, spec_candidate_ids, ) self.bootstrap_token_cond.notify_all() def discard_expired_metadata_room(self, room: int) -> None: """Best-effort cleanup of per-room metadata and expiry markers.""" with self.bootstrap_token_cond: self.expired_prefill_metadata_rooms.discard(room) self.prefill_metadata.pop(room, None) def _wait_prefill_metadata( self, room: int | None, fallback_token: int, fallback_candidate_ids: list[int] | None, ) -> tuple[int, list[int] | None]: if room is None or fallback_token != -1: return fallback_token, fallback_candidate_ids wait_log_interval = max(envs.TOKENSPEED_PD_PREFILL_METADATA_TIMEOUT.get(), 0.01) start_time = time.monotonic() next_log_time = start_time + wait_log_interval with self.bootstrap_token_cond: while room not in self.prefill_metadata: if self.request_status.get(room) == TransferPoll.Failed: self.expired_prefill_metadata_rooms.add(room) logger.warning( "Prefill metadata unavailable for failed bootstrap_room=%s; using fallback=%s", room, fallback_token, ) return ( fallback_token, fallback_candidate_ids, ) now = time.monotonic() if now >= next_log_time: logger.debug( "Still waiting for prefill metadata for bootstrap_room=%s after %.2fs", room, now - start_time, ) next_log_time = now + wait_log_interval self.bootstrap_token_cond.wait(timeout=0.01) return self.prefill_metadata[room] def _is_session_failed(self, mooncake_session_id: str) -> bool: if self.failed_session_ttl <= 0: return False failed_at = self.failed_sessions.get(mooncake_session_id) if failed_at is None: return False elapsed = time.monotonic() - failed_at logger.info( "Session %s failed for %.2fs (TTL=%ds).", mooncake_session_id, elapsed, self.failed_session_ttl, ) if elapsed < self.failed_session_ttl: return True del self.failed_sessions[mooncake_session_id] logger.info( "Session %s failed TTL expired (%.2fs >= %ds), reset.", mooncake_session_id, elapsed, self.failed_session_ttl, ) return False def _mark_session_failed( self, mooncake_session_id: str, reason: str = "transfer_failed" ) -> None: if self.failed_session_ttl <= 0: return self.failed_sessions[mooncake_session_id] = time.monotonic() logger.warning( "Session %s marked failed (reason=%s, ttl=%ds).", mooncake_session_id, reason, self.failed_session_ttl, ) def _clear_failed_session(self, mooncake_session_id: str) -> None: if mooncake_session_id in self.failed_sessions: del self.failed_sessions[mooncake_session_id] logger.info( "Session %s failed state cleared due to KVArgs registration.", mooncake_session_id, ) if mooncake_session_id in self.session_failures: del self.session_failures[mooncake_session_id] def resolve_transfer_indices( self, kv_chunk: TransferKVChunk, req: TransferInfo, ) -> TransferIndexResolution: src_indices = kv_chunk.prefill_kv_indices dst_indices = req.dst_kv_indices[kv_chunk.index_slice] valid_len = min(len(src_indices), len(dst_indices)) # Fast path: empty transfer chunk. Avoid MLA assertions/index ops on empty payload. if valid_len == 0: empty = np.array([], dtype=np.int64) return TransferIndexResolution(src_indices=empty, dst_indices=empty) if valid_len < len(src_indices) or valid_len < len(dst_indices): logger.warning( "Mismatched transfer indices, truncating to %s (src=%s, dst=%s)", valid_len, len(src_indices), len(dst_indices), ) src_indices = src_indices[:valid_len] dst_indices = dst_indices[:valid_len] src_mode = self.src_mode dst_mode = "ON" if req.dst_indices_are_local else "OFF" src_args = kv_chunk.mla_l1_5_args # Prefill OFF/Decode OFF: use original indices. if src_mode == "OFF" and dst_mode == "OFF": return TransferIndexResolution(src_indices, dst_indices) # Prefill ON/Decode OFF: prefill-side only has partial kv cache, only send local part. if src_mode == "ON" and dst_mode == "OFF": if src_args is None: raise RuntimeError( "Prefill MLA L1.5 cache is enabled but no transfer " "metadata provided" ) src_mask = src_args.page_transfer_mask src_local = src_args.page_local_indices return TransferIndexResolution( src_indices=src_local, dst_indices=dst_indices[src_mask], ) # Prefill OFF/Decode ON: decode-side only has partial kv cache, only send requested part. if src_mode == "OFF" and dst_mode == "ON": if req.dst_page_transfer_mask is None: raise RuntimeError( "OFF/ON expects decode ownership mask when destination " "reports local indices." ) dst_mapping = req.dst_page_indices_mapping[kv_chunk.index_slice] dst_mask = req.dst_page_transfer_mask[kv_chunk.index_slice] dst_mapping = dst_mapping[dst_mask] dst_local = req.dst_page_local_indices[dst_mapping] return TransferIndexResolution( src_indices=src_indices[dst_mask], dst_indices=dst_local, ) # Prefill ON/Decode ON: both sides hold partial kv cache, find the intersection part. if src_args is None: raise RuntimeError( "ON/ON expects prefill metadata " "(page_transfer_mask/page_local_indices)." ) if req.dst_page_transfer_mask is None: raise RuntimeError( "ON/ON expects decode ownership mask when destination uses local " "index space." ) # src_args.page_transfer_mask is generated from current chunk page_indices, # so it is already in chunk-local coordinates. src_mask = src_args.page_transfer_mask[:valid_len] src_local = src_args.page_local_indices # req.dst_page_transfer_mask is request-global and should be sliced by chunk. dst_mask = req.dst_page_transfer_mask[kv_chunk.index_slice][:valid_len] # Only positions owned by both sides should be transferred. common_mask = src_mask & dst_mask dst_mapping = req.dst_page_indices_mapping[kv_chunk.index_slice][:valid_len] dst_mapping = dst_mapping[common_mask] dst_local = req.dst_page_local_indices[dst_mapping] src_mapping = np.cumsum(src_args.page_transfer_mask) - 1 src_mapping = src_mapping[:valid_len] src_mapping = src_mapping[common_mask] src_local = src_args.page_local_indices[src_mapping] return TransferIndexResolution( src_indices=src_local, dst_indices=dst_local, ) def _transfer_data(self, mooncake_session_id, transfer_blocks): if not transfer_blocks: return 0 src_addrs, dst_addrs, lengths = zip(*transfer_blocks) return self.engine.batch_transfer_sync( mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths) ) def send_kvcache( self, mooncake_session_id: str, prefill_kv_indices: npt.NDArray[np.int64], dst_kv_ptrs: list[int], dst_kv_indices: npt.NDArray[np.int64], executor: concurrent.futures.ThreadPoolExecutor, transfer_fragments: tuple[TransferFragment, ...] = (), ): # Group by indices prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( prefill_kv_indices, dst_kv_indices ) # ``_layer_transfer_blocks`` indexes ``kv_args.offsets`` by its loop var, # and ``offsets`` is keyed by KV-LAYER INDEX (one entry per attention/KV # layer), not by global layer id. For hybrid models (e.g. Qwen3.5 GDN + # attention) only the attention layers carry KV, so len(offsets) is the # KV-layer count while ``self.layer_num`` is the (larger) global layer # count -- using ``layer_num`` here over-runs ``offsets`` and IndexErrors. # Iterate the full KV-index space instead (the layerwise path already # maps global->KV index via ``_kv_layer_to_index`` before calling this). transfer_blocks = self._layer_transfer_blocks( dst_ptrs=dst_kv_ptrs, src_blocks=prefill_kv_blocks, dst_blocks=dst_kv_blocks, begin_layer_id=0, end_layer_id=len(self.kv_args.offsets), transfer_fragments=transfer_fragments, ) return self._transfer_data(mooncake_session_id, transfer_blocks) def _layer_transfer_blocks( self, dst_ptrs: list[int], src_blocks, dst_blocks, begin_layer_id: int, end_layer_id: int, transfer_fragments: tuple[TransferFragment, ...] = (), ) -> list[tuple[int, int, int]]: transfer_blocks = [] fragments_by_buffer: dict[int, list[TransferFragment]] = defaultdict(list) for fragment in transfer_fragments: if fragment.buffer_kind != BufferKind.MAMBA_STATE: fragments_by_buffer[fragment.buffer_index].append(fragment) has_fragment_plan = bool(transfer_fragments) for layer_id in range(begin_layer_id, end_layer_id): for ptr_offset in self.kv_args.offsets[layer_id]: src_ptr = self.kv_args.kv_data_ptrs[ptr_offset] dst_ptr = dst_ptrs[ptr_offset] item_len = self.kv_args.kv_item_lens[ptr_offset] buffer_fragments = fragments_by_buffer.get(ptr_offset, ()) if has_fragment_plan: for fragment in buffer_fragments: for prefill_index, decode_index in zip(src_blocks, dst_blocks): page_count = min(len(prefill_index), len(decode_index)) if fragment.page_count is not None: page_count = min(page_count, fragment.page_count) for page_offset in range(page_count): src_addr = ( src_ptr + int(prefill_index[page_offset]) * fragment.src_page_stride_bytes + fragment.src_byte_offset ) dst_addr = ( dst_ptr + int(decode_index[page_offset]) * fragment.dst_page_stride_bytes + fragment.dst_byte_offset ) transfer_blocks.append( (src_addr, dst_addr, fragment.bytes_per_page) ) continue for prefill_index, decode_index in zip(src_blocks, dst_blocks): src_addr = src_ptr + int(prefill_index[0]) * item_len dst_addr = dst_ptr + int(decode_index[0]) * item_len length = item_len * len(prefill_index) transfer_blocks.append((src_addr, dst_addr, length)) return transfer_blocks def _wait_until_cache_step(self, target_step: int) -> None: if self.step_counter is None: return while True: ready_step = self.step_counter.query_ready_cache_step() if StepCounter.is_step_ready(ready_step, target_step): return time.sleep(1e-4) def send_mamba_cache( self, mooncake_session_id: str, prefill_mamba_indices: npt.NDArray[np.int64] | None, dst_state_data_ptrs: list[int], dst_mamba_indices: npt.NDArray[np.int64] | None, begin_layer_id: int | None = None, end_layer_id: int | None = None, transfer_fragments: tuple[TransferFragment, ...] = (), ) -> int: if self.kv_args.state_type != "mamba": return 0 state_ptrs = self.kv_args.state_data_ptrs state_item_lens = self.kv_args.state_item_lens if ( not state_ptrs or not dst_state_data_ptrs or prefill_mamba_indices is None or dst_mamba_indices is None ): return 0 if len(state_ptrs) != len(dst_state_data_ptrs): logger.error( "Mamba state tensor count mismatch: prefill=%d decode=%d", len(state_ptrs), len(dst_state_data_ptrs), ) return -1 if prefill_mamba_indices.shape != dst_mamba_indices.shape: if prefill_mamba_indices.size == 1 and dst_mamba_indices.size > 1: prefill_mamba_indices = np.full( dst_mamba_indices.shape, int(prefill_mamba_indices[0]), dtype=np.int64, ) else: logger.error( "Mamba state slot count mismatch: prefill=%s decode=%s", prefill_mamba_indices.tolist(), dst_mamba_indices.tolist(), ) return -1 state_items = list( enumerate(zip(state_ptrs, dst_state_data_ptrs, state_item_lens)) ) if begin_layer_id is not None or end_layer_id is not None: begin = 0 if begin_layer_id is None else begin_layer_id end = self.layer_num if end_layer_id is None else end_layer_id if len(self.state_layer_ids) != len(state_items): logger.error( "Mamba state layer id count mismatch: ids=%d tensors=%d", len(self.state_layer_ids), len(state_items), ) return -1 state_items = [ item for item, layer_id in zip(state_items, self.state_layer_ids) if begin <= layer_id < end ] if not state_items: return 0 valid = (prefill_mamba_indices >= 0) & (dst_mamba_indices >= 0) log_layerwise = getattr(self, "layerwise_debug", False) if log_layerwise and begin_layer_id is not None and end_layer_id is not None: logger.info( "[layerwise_transfer] session=%s layers=[%d,%d) " "send mamba tensors=%d bytes=%d", mooncake_session_id, begin_layer_id, end_layer_id, len(state_items), sum(item_len for _, (_, _, item_len) in state_items) * int(valid.sum()), ) if not valid.any(): return 0 src_indices = prefill_mamba_indices[valid] dst_indices = dst_mamba_indices[valid] src_blocks, dst_blocks = group_concurrent_contiguous(src_indices, dst_indices) transfer_blocks = [] state_fragments_by_buffer: dict[int, list[TransferFragment]] = defaultdict(list) for fragment in transfer_fragments: if fragment.buffer_kind == BufferKind.MAMBA_STATE: state_fragments_by_buffer[fragment.buffer_index].append(fragment) has_state_fragment_plan = bool(state_fragments_by_buffer) for state_index, (src_ptr, dst_ptr, item_len) in state_items: buffer_fragments = state_fragments_by_buffer.get(state_index, ()) if has_state_fragment_plan: for fragment in buffer_fragments: for prefill_index, decode_index in zip(src_blocks, dst_blocks): page_count = min(len(prefill_index), len(decode_index)) if fragment.page_count is not None: page_count = min(page_count, fragment.page_count) for page_offset in range(page_count): src_addr = ( src_ptr + int(prefill_index[page_offset]) * fragment.src_page_stride_bytes + fragment.src_byte_offset ) dst_addr = ( dst_ptr + int(decode_index[page_offset]) * fragment.dst_page_stride_bytes + fragment.dst_byte_offset ) transfer_blocks.append( (src_addr, dst_addr, fragment.bytes_per_page) ) continue for prefill_index, decode_index in zip(src_blocks, dst_blocks): src_addr = src_ptr + int(prefill_index[0]) * item_len dst_addr = dst_ptr + int(decode_index[0]) * item_len length = item_len * len(prefill_index) transfer_blocks.append((src_addr, dst_addr, length)) total_bytes = sum(length for _, _, length in transfer_blocks) ret = self._transfer_data(mooncake_session_id, transfer_blocks) logger.debug( "Transferred mamba cache for session=%s slots=%s blocks=%d bytes=%d ret=%s", mooncake_session_id, src_indices.tolist(), len(transfer_blocks), total_bytes, ret, ) return ret def send_kvcache_layerwise( self, mooncake_session_id: str, prefill_kv_indices: npt.NDArray[np.int64], dst_kv_ptrs: list[int], dst_kv_indices: npt.NDArray[np.int64], begin_cache_step: int, interval: int, dst_state_data_ptrs: list[int] | None = None, prefill_mamba_indices: npt.NDArray[np.int64] | None = None, dst_mamba_indices: npt.NDArray[np.int64] | None = None, transfer_fragments: tuple[TransferFragment, ...] = (), ) -> int: prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( prefill_kv_indices, dst_kv_indices ) interval = max(int(interval), 1) log_layerwise = getattr(self, "layerwise_debug", False) for begin_layer_id in range(0, self.layer_num, interval): end_layer_id = min(begin_layer_id + interval, self.layer_num) target_step = begin_cache_step + end_layer_id - 1 if log_layerwise: logger.info( "[layerwise_transfer] session=%s layers=[%d,%d) wait_cache_step=%d pages=%d", mooncake_session_id, begin_layer_id, end_layer_id, target_step, len(prefill_kv_indices), ) self._wait_until_cache_step(target_step) transfer_blocks = [] if prefill_kv_blocks: for global_layer_id in range(begin_layer_id, end_layer_id): kv_layer_index = self._kv_layer_to_index.get(global_layer_id) if kv_layer_index is None: continue transfer_blocks.extend( self._layer_transfer_blocks( dst_ptrs=dst_kv_ptrs, src_blocks=prefill_kv_blocks, dst_blocks=dst_kv_blocks, begin_layer_id=kv_layer_index, end_layer_id=kv_layer_index + 1, transfer_fragments=transfer_fragments, ) ) if transfer_blocks: if log_layerwise: total_bytes = sum(length for _, _, length in transfer_blocks) logger.info( "[layerwise_transfer] session=%s layers=[%d,%d) send kv blocks=%d bytes=%d", mooncake_session_id, begin_layer_id, end_layer_id, len(transfer_blocks), total_bytes, ) ret = self._transfer_data(mooncake_session_id, transfer_blocks) if ret != 0: return ret ret = self.send_mamba_cache( mooncake_session_id, prefill_mamba_indices, dst_state_data_ptrs or [], dst_mamba_indices, begin_layer_id=begin_layer_id, end_layer_id=end_layer_id, transfer_fragments=transfer_fragments, ) if ret != 0: return ret if log_layerwise: logger.info( "[layerwise_transfer] session=%s layers=[%d,%d) done", mooncake_session_id, begin_layer_id, end_layer_id, ) return 0 def sync_status_to_decode_endpoint( self, remote: str, dst_port: int, room: int, status: int, prefill_rank: int, bootstrap_token: int = -1, spec_candidate_ids: list[int] | None = None, ): if ":" in remote: remote = remote.split(":")[0] spec_candidate_payload = ( np.asarray(spec_candidate_ids, dtype=np.int32).tobytes() if spec_candidate_ids is not None else b"" ) self._connect("tcp://" + remote + ":" + str(dst_port)).send_multipart( [ str(room).encode("ascii"), str(status).encode("ascii"), str(prefill_rank).encode("ascii"), str(bootstrap_token).encode("ascii"), spec_candidate_payload, ] ) def abort_room(self, room: int, reason: str) -> None: """Notify the decode that a room failed before any KV transfer. EPD: when the prefill aborts a request on embedding-receive timeout it never sends KV, so the decode's dual-dispatched KV receiver would wait indefinitely -- its heartbeat only trips if the prefill /health dies, and the receiver waiting_timeout only covers the WaitingForInput state (a receiver whose prefill never registered a sender is stuck earlier). Push a Failed status to every decode endpoint that already pre-allocated for this room (mirrors the in-transfer failure path), so the decode raises a FailedEvent and the client gets an error instead of hanging. A room whose decode has not pre-allocated yet is only marked Failed locally (no endpoint to notify). """ self.record_failure(room, reason) self.update_status(room, TransferPoll.Failed) for req in list(self.transfer_infos.get(room, {}).values()): if not req.is_dummy: self.sync_status_to_decode_endpoint( req.endpoint, req.dst_port, req.room, TransferPoll.Failed, self.attn_tp_rank, ) def transfer_worker( self, queue: FastQueue, executor: concurrent.futures.ThreadPoolExecutor ): while True: try: kv_chunk = queue.get() logger.debug( "[TRANSFER_WORKER] Got transfer request for room %s, is_last=%s, kv_indices_len=%s", kv_chunk.room, kv_chunk.is_last, len(kv_chunk.prefill_kv_indices), ) reqs_to_be_processed = ( self.transfer_infos[kv_chunk.room].values() if kv_chunk.room in self.transfer_infos else [] ) polls = [] dst_ranks_infos = [] for req in reqs_to_be_processed: if not req.is_dummy: # Early exit if the request has failed with self.session_lock: if self._is_session_failed(req.mooncake_session_id): logger.info( "Blocked transfer due to failed session (room=%s, session=%s).", kv_chunk.room, req.mooncake_session_id, ) self.record_failure( kv_chunk.room, f"Decode instance could be dead, remote mooncake session {req.mooncake_session_id} is not alive", ) self.update_status(kv_chunk.room, TransferPoll.Failed) self.sync_status_to_decode_endpoint( req.endpoint, req.dst_port, req.room, TransferPoll.Failed, self.attn_tp_rank, ) break resolved = self.resolve_transfer_indices(kv_chunk, req) logger.debug( "[TRANSFER_WORKER] Calling send_kvcache for room %s, session %s", kv_chunk.room, req.mooncake_session_id, ) tm_start = time.monotonic() prefill_metadata = None dst_kv_ptrs = self.decode_kv_args_table[ req.mooncake_session_id ].dst_kv_ptrs dst_state_data_ptrs = self.decode_kv_args_table[ req.mooncake_session_id ].dst_state_data_ptrs if kv_chunk.begin_cache_step is None: ret = self.send_kvcache( req.mooncake_session_id, resolved.src_indices, dst_kv_ptrs, resolved.dst_indices, executor, req.transfer_fragments, ) else: ret = self.send_kvcache_layerwise( req.mooncake_session_id, resolved.src_indices, dst_kv_ptrs, resolved.dst_indices, kv_chunk.begin_cache_step, kv_chunk.layerwise_interval, dst_state_data_ptrs, kv_chunk.prefill_mamba_indices, req.dst_mamba_indices, req.transfer_fragments, ) if ret == 0 and kv_chunk.is_last: if kv_chunk.wait_for_bootstrap_token: # The first decode/target-verify step consumes prefill's # sampled bootstrap token and spec candidates; publish # Success only after that metadata is ready. prefill_metadata = self._wait_prefill_metadata( kv_chunk.room, kv_chunk.bootstrap_token, kv_chunk.spec_candidate_ids, ) if kv_chunk.begin_cache_step is None: ret = self.send_mamba_cache( req.mooncake_session_id, kv_chunk.prefill_mamba_indices, dst_state_data_ptrs, req.dst_mamba_indices, transfer_fragments=req.transfer_fragments, ) logger.debug( "[TRANSFER_WORKER] send_kvcache returned %s for room %s", ret, kv_chunk.room, ) if ret != 0: with self.session_lock: self.session_failures[req.mooncake_session_id] += 1 # Failures should never happen if the session is not dead, if the session fails once, mark it as failed if self.session_failures[req.mooncake_session_id] >= 1: self._mark_session_failed( req.mooncake_session_id, reason="send_kvcache" ) logger.error( "Session %s failed.", req.mooncake_session_id ) self.record_failure( kv_chunk.room, f"Failed to send kv chunk of {kv_chunk.room} to {req.endpoint}:{req.dst_port}", ) self.update_status(kv_chunk.room, TransferPoll.Failed) self.sync_status_to_decode_endpoint( req.endpoint, req.dst_port, req.room, TransferPoll.Failed, self.attn_tp_rank, ) break if kv_chunk.is_last: polls.append(True) dst_ranks_infos.append( (req.endpoint, req.dst_port, req.room) ) # Only sync status when all the dst ranks have received the kvcache if len(polls) == req.required_dst_info_num: status = ( TransferPoll.Success if all(polls) else TransferPoll.Failed ) self.update_status(req.room, status) # bootstrap_token is carried directly in the chunk (set by # DisaggPrefillExecutor._decode after prefill forward). if kv_chunk.wait_for_bootstrap_token: if prefill_metadata is None: prefill_metadata = self._wait_prefill_metadata( kv_chunk.room, kv_chunk.bootstrap_token, kv_chunk.spec_candidate_ids, ) bootstrap_token, spec_candidate_ids = ( prefill_metadata ) else: bootstrap_token, spec_candidate_ids = ( kv_chunk.bootstrap_token, kv_chunk.spec_candidate_ids, ) if self.check_status(req.room) == TransferPoll.Failed: status = TransferPoll.Failed for endpoint, dst_port, room in dst_ranks_infos: self.sync_status_to_decode_endpoint( endpoint, dst_port, room, status, self.attn_tp_rank, bootstrap_token=bootstrap_token, spec_candidate_ids=spec_candidate_ids, ) elapsed_seconds = time.monotonic() - tm_start if self.kv_transfer_metrics: self.kv_transfer_metrics.observe_kv_transfer_latency( elapsed_seconds ) else: # Dummy request means the decode instance is not used, so its status can be marked as success directly # Dummy request does not need to sync status to decode endpoint if kv_chunk.is_last and req.room in self.request_status: self.update_status(req.room, TransferPoll.Success) if ( kv_chunk.room not in self.request_status or self.check_status(kv_chunk.room) == TransferPoll.Success ): if kv_chunk.room in self.transfer_infos: self.transfer_infos.pop(kv_chunk.room) except Exception as exc: raise RuntimeError( f"Transfer thread failed because of {exc}. Prefill instance " f"with bootstrap_port={self.bootstrap_port} is dead." ) from exc def start_prefill_thread(self): self.rank_port = get_free_port() self.server_socket.bind(f"tcp://{get_local_ip_by_remote()}:{self.rank_port}") def bootstrap_thread(): """This thread recvs pre-alloc notification from the decode engine""" # TransferPoll.Bootstrapping -> TransferPoll.WaitingForInput while True: waiting_req_bytes = self.server_socket.recv_multipart() room = waiting_req_bytes[0].decode("ascii") mooncake_session_id = waiting_req_bytes[3].decode("ascii") logger.info( "[Prefill bootstrap_thread] recv multipart: room=%s session_id=%s", room, mooncake_session_id, ) if room == "None": self.decode_kv_args_table[mooncake_session_id] = ( KVArgsRegisterInfo.from_zmq(waiting_req_bytes) ) with self.session_lock: self._clear_failed_session(mooncake_session_id) logger.info( "[Prefill bootstrap_thread] registered kv_args from decode session=%s", mooncake_session_id, ) continue else: required_dst_info_num = int(waiting_req_bytes[6].decode("ascii")) room = int(room) if room not in self.transfer_infos: self.transfer_infos[room] = {} self.transfer_infos[room][mooncake_session_id] = ( TransferInfo.from_zmq(waiting_req_bytes) ) logger.info( "[Prefill bootstrap_thread] pre-alloc received: room=%d session=%s got=%d/%d, status -> %s", room, mooncake_session_id, len(self.transfer_infos[room]), required_dst_info_num, ( "Bootstrapped" if len(self.transfer_infos[room]) == required_dst_info_num else "waiting more" ), ) if len(self.transfer_infos[room]) == required_dst_info_num: self.update_status(room, TransferPoll.Bootstrapped) threading.Thread(target=bootstrap_thread).start() def start_transfer_thread( self, transfer_thread_pool_size: int, transfer_queue_size: int ): self.transfer_queues: list[FastQueue] = [ FastQueue() for _ in range(transfer_queue_size) ] self.executors = [ concurrent.futures.ThreadPoolExecutor( transfer_thread_pool_size // transfer_queue_size ) for _ in range(transfer_queue_size) ] for queue, executor in zip(self.transfer_queues, self.executors): threading.Thread( target=self.transfer_worker, args=(queue, executor), daemon=True ).start() def add_transfer_request( self, bootstrap_room: int, kv_indices: npt.NDArray[np.int64], index_slice: slice, is_last: bool, aux_index: int | None = None, mla_l1_5_args: PageTransferMetadata | None = None, bootstrap_token: int = -1, begin_cache_step: int | None = None, layerwise_interval: int = 1, wait_for_bootstrap_token: bool = False, mamba_indices: npt.NDArray[np.int64] | None = None, spec_candidate_ids: list[int] | None = None, ): if self.disaggregation_mode != DisaggregationMode.PREFILL: raise RuntimeError("Transfer requests can only be added in prefill mode.") if is_last and aux_index is None: raise ValueError("aux_index must be set for the last transfer chunk.") if ( bootstrap_room not in self.request_status or self.check_status(bootstrap_room) == TransferPoll.Failed ): logger.debug( "Request with bootstrap_room=%s already failed", bootstrap_room ) return if bootstrap_room not in self.transfer_infos: # This means that the current rank is a dummy rank for this request, # and it has already been marked as success, so there is no need to # add further chunks into the transfer queue. return # sharding according to the dst_infos to make sure # requests with the same dst_sessions will be added into the same # queue, which enables early abort with failed sessions. dst_infos = self.transfer_infos[bootstrap_room].keys() session_port_sum = sum(int(session.split(":")[1]) for session in dst_infos) shard_idx = session_port_sum % len(self.transfer_queues) self.transfer_queues[shard_idx].put( TransferKVChunk( room=bootstrap_room, prefill_kv_indices=kv_indices, index_slice=index_slice, is_last=is_last, prefill_aux_index=aux_index, 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, prefill_mamba_indices=mamba_indices, spec_candidate_ids=spec_candidate_ids, ) ) def receive_decode_prefix_info(self, bootstrap_room: int) -> int: """Receive decode prefix info from decode side""" # In mooncake implementation, decode_prefix_len is handled via ZMQ messages # Check the stored transfer info for this room if bootstrap_room in self.transfer_infos: for transfer_info in self.transfer_infos[bootstrap_room].values(): if ( hasattr(transfer_info, "decode_prefix_len") and transfer_info.decode_prefix_len > 0 ): logger.debug( "Found decode_prefix_len=%s for room %s", transfer_info.decode_prefix_len, bootstrap_room, ) return transfer_info.decode_prefix_len logger.debug("No decode_prefix_len found for room %s, using 0", bootstrap_room) return 0 def _register_to_bootstrap(self): """Register KVSender to bootstrap server via HTTP POST.""" if self.dist_init_addr: ip_address = socket.gethostbyname(self.dist_init_addr.split(":")[0]) else: ip_address = get_ip() bootstrap_server_url = f"{ip_address}:{self.bootstrap_port}" url = f"http://{bootstrap_server_url}/route" payload = { "role": "Prefill", "world_size": self.world_size, "dp_size": self.dp_size, "rank_ip": get_local_ip_by_remote(), "rank_port": self.rank_port, "engine_rank": self.kv_args.engine_rank, "enable_mla_l1_5_cache": self.args.enable_mla_l1_5_cache, "kv_item_lens": self.kv_args.kv_item_lens, "kv_unit_lens": getattr(self.kv_args, "kv_unit_lens", []), "state_item_lens": self.kv_args.state_item_lens, "state_unit_lens": getattr(self.kv_args, "state_unit_lens", []), } 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( "Prefill instance failed to connect to bootstrap server: %s, %s", response.status_code, response.text, ) except Exception as exc: logger.error( "Prefill instance failed to register with bootstrap server: %s", exc ) from tokenspeed.runtime.pd.mooncake.sender import MooncakeKVSender __all__ = ["MooncakeKVManagerPrefill", "MooncakeKVSender"]