""" Life cycle of a request in the prefill server 1. Bootstrap Queue a. Initialize a sender for each request b. Use the queue to store requests whose bootstrap (handshake and preallocation) has not finished c. Poll senders to check bootstrap state d. Once bootstrap is complete, move request to Waiting Queue 2. Waiting Queue a. Use PrefillAdder to pop requests b. Run forward c. Add the request to Inflight Queue 3. Inflight Queue a. Poll (non-blocking) the sender of the request b. Once the transfer has finished, return the request """ from __future__ import annotations import logging import time from collections import deque from http import HTTPStatus from typing import TYPE_CHECKING, List, Optional, Type import torch from sglang.srt.disaggregation.base import BaseKVManager, KVPoll from sglang.srt.disaggregation.utils import ( FAKE_BOOTSTRAP_HOST, DisaggregationMode, KVClassType, MetadataBuffers, ReqToMetadataIdxAllocator, TransferBackend, get_kv_class, is_mla_backend, kv_to_page_indices, kv_to_page_num, poll_and_all_reduce, prepare_abort, ) from sglang.srt.managers.schedule_batch import ( FINISH_LENGTH, Req, RequestStage, ScheduleBatch, ) from sglang.srt.mem_cache.common import release_kv_cache from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, NSATokenToKVPool from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.tracing.trace import trace_event_batch, trace_slice, trace_slice_end if TYPE_CHECKING: from torch.distributed import ProcessGroup from sglang.srt.managers.scheduler import GenerationBatchResult, Scheduler from sglang.srt.mem_cache.memory_pool import KVCache logger = logging.getLogger(__name__) def release_req_to_metadata_buffer( req: Req, allocator: ReqToMetadataIdxAllocator ) -> None: """ Release the metadata buffer index allocated for a request in prefill disaggregation mode. This function safely releases the metadata buffer index if it was allocated. Args: req: The request object that may have a metadata_buffer_index allocated allocator: The ReqToMetadataIdxAllocator instance to free the index """ if ( hasattr(req, "metadata_buffer_index") and req.metadata_buffer_index is not None and req.metadata_buffer_index >= 0 ): allocator.free(req.metadata_buffer_index) req.metadata_buffer_index = -1 class PrefillBootstrapQueue: """ Store the requests in bootstrapping """ def __init__( self, token_to_kv_pool: KVCache, draft_token_to_kv_pool: Optional[KVCache], req_to_metadata_buffer_idx_allocator: ReqToMetadataIdxAllocator, metadata_buffers: MetadataBuffers, tp_rank: int, tp_size: int, gpu_id: int, bootstrap_port: int, gloo_group: ProcessGroup, max_total_num_tokens: int, decode_tp_size: int, decode_dp_size: int, scheduler: Scheduler, pp_rank: int, pp_size: int, transfer_backend: TransferBackend, ): self.token_to_kv_pool = token_to_kv_pool self.draft_token_to_kv_pool = draft_token_to_kv_pool self.is_mla_backend = is_mla_backend(token_to_kv_pool) self.metadata_buffers = metadata_buffers self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator self.tp_rank = tp_rank self.tp_size = tp_size self.decode_tp_size = decode_tp_size self.decode_dp_size = decode_dp_size self.pp_rank = pp_rank self.pp_size = pp_size self.gpu_id = gpu_id self.bootstrap_port = bootstrap_port self.queue: List[Req] = [] self.gloo_group = gloo_group self.max_total_num_tokens = max_total_num_tokens self.scheduler = scheduler self.transfer_backend = transfer_backend self.kv_manager = self._init_kv_manager() if self.scheduler.tp_worker.is_hybrid_swa: # FIXME: current SWA allocation allocate full kv cache size in prefill self.max_total_num_tokens = min( self.max_total_num_tokens, self.scheduler.tp_worker.model_runner.swa_max_total_num_tokens, ) def _init_kv_manager(self) -> BaseKVManager: kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS) kv_args = kv_args_class() kv_args.engine_rank = self.tp_rank kv_args.pp_rank = self.pp_rank kv_args.system_dp_rank = self.scheduler.dp_rank kv_args.decode_tp_size = self.decode_tp_size // self.decode_dp_size kv_args.prefill_pp_size = self.pp_size kv_args.prefill_start_layer = self.token_to_kv_pool.start_layer kv_data_ptrs, kv_data_lens, kv_item_lens = ( self.token_to_kv_pool.get_contiguous_buf_infos() ) if self.draft_token_to_kv_pool is not None: # We should also transfer draft model kv cache. The indices are # always shared with a target model. draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = ( self.draft_token_to_kv_pool.get_contiguous_buf_infos() ) kv_data_ptrs += draft_kv_data_ptrs kv_data_lens += draft_kv_data_lens kv_item_lens += draft_kv_item_lens kv_args.kv_data_ptrs = kv_data_ptrs kv_args.kv_data_lens = kv_data_lens kv_args.kv_item_lens = kv_item_lens if not self.is_mla_backend: kv_args.kv_head_num = self.token_to_kv_pool.head_num kv_args.page_size = self.token_to_kv_pool.page_size kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = ( self.metadata_buffers.get_buf_infos() ) kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device kv_args.gpu_id = self.scheduler.gpu_id if hasattr(self.token_to_kv_pool, "get_state_buf_infos"): state_data_ptrs, state_data_lens, state_item_lens = ( self.token_to_kv_pool.get_state_buf_infos() ) kv_args.state_data_ptrs = state_data_ptrs kv_args.state_data_lens = state_data_lens kv_args.state_item_lens = state_item_lens if isinstance(self.token_to_kv_pool, SWAKVPool): kv_args.state_type = "swa" elif isinstance(self.token_to_kv_pool, HybridLinearKVPool): kv_args.state_type = "mamba" # Get state dimension info for cross-TP slice transfer if hasattr(self.token_to_kv_pool, "get_state_dim_per_tensor"): kv_args.state_dim_per_tensor = ( self.token_to_kv_pool.get_state_dim_per_tensor() ) elif isinstance(self.token_to_kv_pool, NSATokenToKVPool): kv_args.state_type = "nsa" else: kv_args.state_type = "none" else: kv_args.state_data_ptrs = [] kv_args.state_data_lens = [] kv_args.state_item_lens = [] kv_args.state_type = "none" kv_manager_class: Type[BaseKVManager] = get_kv_class( self.transfer_backend, KVClassType.MANAGER ) kv_manager: BaseKVManager = kv_manager_class( kv_args, DisaggregationMode.PREFILL, self.scheduler.server_args, self.is_mla_backend, ) return kv_manager def add(self, req: Req, num_kv_heads: int) -> None: if self._check_if_req_exceed_kv_capacity(req): return if req.bootstrap_host == FAKE_BOOTSTRAP_HOST: kv_sender_class = get_kv_class(TransferBackend.FAKE, KVClassType.SENDER) else: kv_sender_class = get_kv_class(self.transfer_backend, KVClassType.SENDER) dest_tp_ranks = [self.tp_rank] req.disagg_kv_sender = kv_sender_class( mgr=self.kv_manager, bootstrap_addr=f"{req.bootstrap_host}:{self.bootstrap_port}", bootstrap_room=req.bootstrap_room, dest_tp_ranks=dest_tp_ranks, pp_rank=self.pp_rank, ) self._process_req(req) req.add_latency(RequestStage.PREFILL_PREPARE) self.queue.append(req) trace_slice_end(RequestStage.PREFILL_PREPARE, req.rid, auto_next_anon=True) def extend(self, reqs: List[Req], num_kv_heads: int) -> None: for req in reqs: self.add(req, num_kv_heads) def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: if len(req.origin_input_ids) > self.max_total_num_tokens: message = f"Request {req.rid} exceeds the maximum number of tokens: {len(req.origin_input_ids)} > {self.max_total_num_tokens}" logger.error(message) prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) self.scheduler.stream_output([req], req.return_logprob) return True return False def _process_req(self, req: Req) -> None: """ Set max_new_tokens = 1, so PrefillAdder memory estimation is accurate """ req.sampling_params.max_new_tokens = 1 def pop_bootstrapped( self, return_failed_reqs: bool = False, rids_to_check: Optional[List[str]] = None, ) -> List[Req]: """ pop the reqs which has finished bootstrapping return_failed_reqs: For PP, on rank 0, also return the failed reqs to notify the next rank rids_to_check: For PP, on rank > 0, check the rids from the previous rank has consensus with the current rank. """ bootstrapped_reqs = [] failed_reqs = [] indices_to_remove = set() if len(self.queue) == 0: if return_failed_reqs is False: return [] else: return [], [] polls = poll_and_all_reduce( [req.disagg_kv_sender for req in self.queue], self.gloo_group ) for i, (req, poll) in enumerate(zip(self.queue, polls)): if rids_to_check is not None: # if req not in reqs_info_to_check, skip if req.rid not in rids_to_check: continue if poll == KVPoll.Bootstrapping: continue elif poll == KVPoll.Failed: error_message = f"Prefill bootstrap failed for request rank={self.tp_rank} {req.rid=} {req.bootstrap_room=}" try: req.disagg_kv_sender.failure_exception() except Exception as e: error_message += f" with exception {e}" logger.error(error_message) prepare_abort( req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR ) self.scheduler.stream_output([req], req.return_logprob) indices_to_remove.add(i) failed_reqs.append(req) if self.scheduler.enable_metrics: self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() if self.scheduler.enable_hicache_storage: # to release prefetch events associated with the request self.scheduler.tree_cache.release_aborted_request(req.rid) continue # KV.WaitingForInput - init here num_kv_indices = len(req.origin_input_ids) if self.req_to_metadata_buffer_idx_allocator.available_size() == 0: break req.metadata_buffer_index = ( self.req_to_metadata_buffer_idx_allocator.alloc() ) assert req.metadata_buffer_index is not None num_pages = kv_to_page_num(num_kv_indices, self.token_to_kv_pool.page_size) req.disagg_kv_sender.init(num_pages, req.metadata_buffer_index) bootstrapped_reqs.append(req) indices_to_remove.add(i) req.time_stats.wait_queue_entry_time = time.perf_counter() req.add_latency(RequestStage.PREFILL_BOOTSTRAP) trace_slice_end( RequestStage.PREFILL_BOOTSTRAP, req.rid, auto_next_anon=True ) self.queue = [ entry for i, entry in enumerate(self.queue) if i not in indices_to_remove ] if return_failed_reqs is False: return bootstrapped_reqs else: return bootstrapped_reqs, failed_reqs class SchedulerDisaggregationPrefillMixin: """ Mixin for Scheduler to handle disaggregation prefill """ def get_next_disagg_prefill_batch_to_run( self: Scheduler, ) -> Optional[ScheduleBatch]: # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it # Otherwise, it hangs under high concurrency self.running_batch.batch_is_full = False self.process_prefill_chunk() batch = self.get_new_batch_prefill() batch = self.maybe_prepare_mlp_sync_batch(batch) if batch: trace_event_batch("schedule", batch.reqs) return batch @torch.no_grad() def event_loop_normal_disagg_prefill(self: Scheduler) -> None: """A normal scheduler loop for prefill worker in disaggregation mode.""" while True: # Receive requests recv_reqs = self.recv_requests() self.process_input_requests(recv_reqs) self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() ) # Get the next batch to run batch = self.get_next_disagg_prefill_batch_to_run() self.cur_batch = batch # Launch the current batch if batch: result = self.run_batch(batch) self.process_batch_result_disagg_prefill(batch, result) else: self.self_check_during_idle() self.process_disagg_prefill_inflight_queue() # Update last_batch self.last_batch = batch @torch.no_grad() def event_loop_overlap_disagg_prefill(self: Scheduler) -> None: self.result_queue = deque() while True: # Receive requests recv_reqs = self.recv_requests() self.process_input_requests(recv_reqs) self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() ) # Get the next batch to run batch = self.get_next_disagg_prefill_batch_to_run() self.cur_batch = batch # Launch the current batch if batch: batch_result = self.run_batch(batch) self.result_queue.append((batch.copy(), batch_result)) else: batch_result = None # Process the last batch if self.last_batch: tmp_batch, tmp_result = self.result_queue.popleft() self.process_batch_result_disagg_prefill(tmp_batch, tmp_result) elif batch is None: # When the server is idle, do self-check and re-init some states self.self_check_during_idle() self.process_disagg_prefill_inflight_queue() # Run sample of the current batch # It depends on the result of the last batch (e.g., grammar), so we run it after the last batch is processed. self.launch_batch_sample_if_needed(batch_result) # Update last_batch self.last_batch = batch def process_batch_result_disagg_prefill( self: Scheduler, batch: ScheduleBatch, result: GenerationBatchResult, ) -> None: """ Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue Adapted from process_batch_result_prefill """ ( logits_output, next_token_ids, extend_input_len_per_req, extend_logprob_start_len_per_req, copy_done, ) = ( result.logits_output, result.next_token_ids, result.extend_input_len_per_req, result.extend_logprob_start_len_per_req, result.copy_done, ) if copy_done is not None: copy_done.synchronize() logprob_pt = 0 # Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue next_token_ids = result.next_token_ids.tolist() if batch.return_logprob: if logits_output.next_token_logprobs is not None: logits_output.next_token_logprobs = ( logits_output.next_token_logprobs.tolist() ) if logits_output.input_token_logprobs is not None: logits_output.input_token_logprobs = tuple( logits_output.input_token_logprobs.tolist() ) for i, (req, next_token_id) in enumerate( zip(batch.reqs, next_token_ids, strict=True) ): if req.is_chunked <= 0: if req.time_stats.prefill_finished_ts == 0.0: req.time_stats.prefill_finished_ts = time.time() # There is no output_ids for prefill req.output_ids.append(next_token_id) self.tree_cache.cache_unfinished_req(req) # update the tree and lock req.add_latency(RequestStage.PREFILL_FORWARD) trace_slice(RequestStage.PREFILL_FORWARD, req.rid, auto_next_anon=True) self.disagg_prefill_inflight_queue.append(req) if self.spec_algorithm.is_eagle() and batch.spec_info is not None: req.output_topk_p = batch.spec_info.topk_p[i] req.output_topk_index = batch.spec_info.topk_index[i] req.hidden_states_tensor = ( batch.spec_info.hidden_states[i].cpu().clone() ) else: req.hidden_states_tensor = None if req.return_logprob: assert extend_logprob_start_len_per_req is not None assert extend_input_len_per_req is not None extend_logprob_start_len = extend_logprob_start_len_per_req[i] extend_input_len = extend_input_len_per_req[i] num_input_logprobs = extend_input_len - extend_logprob_start_len self.add_logprob_return_values( i, req, logprob_pt, next_token_ids, num_input_logprobs, logits_output, ) logprob_pt += num_input_logprobs self.send_kv_chunk(req, last_chunk=True) req.time_stats.prefill_transfer_queue_entry_time = time.perf_counter() if req.grammar is not None: # FIXME: this try-except block is for handling unexpected xgrammar issue. try: req.grammar.accept_token(next_token_id) except ValueError as e: # Grammar accept_token can raise ValueError if the token is not in the grammar. # This can happen if the grammar is not set correctly or the token is invalid. error_message = f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}" release_kv_cache(req, self.tree_cache) prepare_abort( req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR, ) req.grammar.finished = req.finished() else: # being chunked reqs' prefill is not finished req.is_chunked -= 1 if req.return_logprob: extend_logprob_start_len = extend_logprob_start_len_per_req[i] extend_input_len = extend_input_len_per_req[i] if extend_logprob_start_len < extend_input_len: # Update input logprobs. num_input_logprobs = extend_input_len - extend_logprob_start_len self.add_input_logprob_return_values( i, req, logits_output, logprob_pt, num_input_logprobs, last_prefill_chunk=False, ) logprob_pt += num_input_logprobs if self.enable_overlap: self.send_kv_chunk(req, last_chunk=False, end_idx=req.tmp_end_idx) trace_slice( RequestStage.PREFILL_CHUNKED_FORWARD, req.rid, auto_next_anon=True ) self.maybe_send_health_check_signal() def process_disagg_prefill_inflight_queue( self: Scheduler, rids_to_check: Optional[List[str]] = None ) -> List[Req]: """ Poll the requests in the middle of transfer. If done, return the request. rids_to_check: For PP, on rank > 0, check the rids from the previous rank has consensus with the current rank. """ if len(self.disagg_prefill_inflight_queue) == 0: return [] done_reqs = [] polls = poll_and_all_reduce( [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue], self.attn_tp_cpu_group, ) undone_reqs: List[Req] = [] # Check .poll() for the reqs in disagg_prefill_inflight_queue. If Success, respond to the client and remove it from the queue for req, poll in zip(self.disagg_prefill_inflight_queue, polls): if rids_to_check is not None: if req.rid not in rids_to_check: undone_reqs.append(req) continue assert poll == KVPoll.Success or poll == KVPoll.Failed if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]: undone_reqs.append(req) elif poll == KVPoll.Success: # transfer done release_kv_cache(req, self.tree_cache) # unlock the tree req.finished_reason = FINISH_LENGTH(length=0) # FIXME: clean up req's data in transfer engine if hasattr(req.disagg_kv_sender, "clear"): req.disagg_kv_sender.clear() done_reqs.append(req) elif poll == KVPoll.Failed: error_message = f"Prefill transfer failed for request rank={self.tp_rank} {req.rid=} {req.bootstrap_room=}" try: req.disagg_kv_sender.failure_exception() except Exception as e: error_message += f" with exception {e}" logger.warning(error_message) release_kv_cache(req, self.tree_cache) # unlock the tree prepare_abort( req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR ) done_reqs.append(req) if self.enable_metrics: self.metrics_collector.increment_transfer_failed_reqs() else: assert False, f"Unexpected polling state {poll=}" for req in done_reqs: req.time_stats.completion_time = time.perf_counter() # Stream requests which have finished transfer self.stream_output( done_reqs, any(req.return_logprob for req in done_reqs), None, ) for req in done_reqs: req: Req req.add_latency(RequestStage.PREFILL_TRANSFER_KV_CACHE) release_req_to_metadata_buffer( req, self.req_to_metadata_buffer_idx_allocator ) trace_slice( RequestStage.PREFILL_TRANSFER_KV_CACHE, req.rid, thread_finish_flag=True ) self.disagg_prefill_inflight_queue = undone_reqs return done_reqs def get_transferred_rids(self: Scheduler) -> List[str]: """ Used by PP, get the transferred rids but **do not pop** """ polls = poll_and_all_reduce( [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue], self.attn_tp_cpu_group, ) transferred_rids: List[str] = [] for req, poll in zip(self.disagg_prefill_inflight_queue, polls): if poll == KVPoll.Success or poll == KVPoll.Failed: transferred_rids.append(req.rid) return transferred_rids def process_prefill_chunk(self: Scheduler) -> None: chunked_req_to_exclude = set() if self.chunked_req: chunked_req_to_exclude.add(self.chunked_req) self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True) if self.enable_overlap: # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved self.chunked_req.tmp_end_idx = min( len(self.chunked_req.fill_ids), len(self.chunked_req.origin_input_ids), ) else: self.send_kv_chunk(self.chunked_req) self.running_batch.batch_is_full = False if self.last_batch and self.last_batch.forward_mode.is_extend(): if self.last_batch.chunked_req: # In the context pipeline parallelism, after the last chunk, the current microbatch still track outdated chunked_req. # We need to discard it. chunked_req_to_exclude.add(self.last_batch.chunked_req) last_bs = self.last_batch.batch_size() self.last_batch.filter_batch( chunked_req_to_exclude=list(chunked_req_to_exclude) ) if self.last_batch.batch_size() < last_bs: self.running_batch.batch_is_full = False def send_kv_chunk( self: Scheduler, req: Req, last_chunk: bool = False, end_idx: Optional[int] = None, ) -> None: """ Send a prefilled chunk to the decode server """ page_size = self.token_to_kv_pool_allocator.page_size start_idx = req.start_send_idx end_idx = ( end_idx if end_idx is not None else min(len(req.fill_ids), len(req.origin_input_ids)) ) if not last_chunk: # if not the last chunk and the last page is partial, delay the last partial page to the next send end_idx = end_idx - end_idx % page_size kv_indices = ( self.req_to_token_pool.req_to_token[req.req_pool_idx, start_idx:end_idx] .cpu() .numpy() ) req.start_send_idx = end_idx state_indices = None if last_chunk: self.disagg_metadata_buffers.set_buf(req) # Prepare extra pool indices for hybrid models if isinstance( self.token_to_kv_pool_allocator.get_kvcache(), HybridLinearKVPool ): # Mamba hybrid model: send single mamba state index state_indices = [ self.req_to_token_pool.req_index_to_mamba_index_mapping[ req.req_pool_idx ] .cpu() .numpy() ] elif isinstance(self.token_to_kv_pool_allocator.get_kvcache(), SWAKVPool): # SWA hybrid model: send last window KV indices seq_len = len(req.fill_ids) window_size = self.sliding_window_size window_start = max(0, seq_len - window_size) window_start = (window_start // page_size) * page_size window_kv_indices_full = self.req_to_token_pool.req_to_token[ req.req_pool_idx, window_start:seq_len ] # Translate to SWA pool indices window_kv_indices_swa = ( self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa( window_kv_indices_full ) ) state_indices = window_kv_indices_swa.cpu().numpy() state_indices = kv_to_page_indices(state_indices, page_size) elif isinstance( self.token_to_kv_pool_allocator.get_kvcache(), NSATokenToKVPool ): seq_len = len(req.fill_ids) kv_indices_full = self.req_to_token_pool.req_to_token[ req.req_pool_idx, :seq_len ] state_indices = kv_indices_full.cpu().numpy() state_indices = kv_to_page_indices(state_indices, page_size) page_indices = kv_to_page_indices(kv_indices, page_size) if len(page_indices) == 0: logger.info( f"Skip sending kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty" ) return req.disagg_kv_sender.send(page_indices, state_indices)