laoyao0822
99b669f8b9
Prefill CP only needs the local hidden shard for DeepSeek NextN draft extend. The change adds a draft shared-KV path that captures target hidden locally, feeds only the CP-local slice into the draft model, and keeps draft KV writes/transfers on the same shared logical-to-physical page mapping as target KV.\n\nDebug logs are gated behind SGLANG_CP_DRAFT_SHARED_KV_DEBUG and cover scheduler pool selection, KV manager buffer registration, local physical writes, prefill sender filtering, transfer pages, and decode commit metadata so ETE runs can prove draft KV is sharded rather than full-concatenated on a prefill rank.\n\nConstraint: Prefill runs CP while decode remains DP, so prefill must avoid full hidden/KV materialization but decode still receives full logical KV pages.\nRejected: Keep draft extend on full hidden state | preserves correctness but wastes prefill memory and defeats CP shared-KV intent.\nRejected: Transfer draft KV with a separate mapping | target and draft pools share req_to_token logical indices, so duplicating mapping adds risk without benefit.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not remove the debug logs until ETE evidence confirms draft MLA/index writes and transfer pages are CP-sharded on all ranks.\nTested: Remote compileall for changed CP draft, transfer, scheduler, NSA index, MLA write, and EAGLE files.\nNot-tested: Full GLM-5 EAGLE ETE with SGLANG_CP_DRAFT_SHARED_KV_DEBUG=1 after this logging addition; local pytest intentionally not run.
1387 lines
52 KiB
Python
1387 lines
52 KiB
Python
"""
|
|
Life cycle of a request in the decode server
|
|
|
|
1. PreallocQueue:
|
|
a. Initialize a receiver for each request
|
|
b. The request handshakes first, and pre-allocate kv once there is available kv.
|
|
c. Move the request to TransferQueue.
|
|
|
|
2. TransferQueue:
|
|
a. Poll the receiver to check the transfer state
|
|
b. If the transfer has finished, move the request to waiting queue
|
|
|
|
3. WaitingQueue:
|
|
a. Use the requests in the queue to construct a PrebuiltExtendBatch
|
|
b. Skip the prefill forward but only populate metadata
|
|
|
|
4. RunningBatch:
|
|
a. Merge the resolved PrebuiltExtendBatch into running batch to run decoding
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import time
|
|
from collections import deque
|
|
from dataclasses import dataclass
|
|
from http import HTTPStatus
|
|
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
|
|
|
|
import torch
|
|
from torch.distributed import ProcessGroup
|
|
|
|
from sglang.srt.configs.mamba_utils import Mamba2CacheParams
|
|
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
|
|
from sglang.srt.disaggregation.base import KVPoll
|
|
from sglang.srt.disaggregation.common.conn import CommonKVManager, CommonKVReceiver
|
|
from sglang.srt.disaggregation.utils import (
|
|
FAKE_BOOTSTRAP_HOST,
|
|
DisaggregationMode,
|
|
KVClassType,
|
|
MetadataBuffers,
|
|
ReqToMetadataIdxAllocator,
|
|
TransferBackend,
|
|
get_kv_class,
|
|
is_mla_backend,
|
|
poll_and_all_reduce,
|
|
prepare_abort,
|
|
)
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.layers.dp_attention import get_attention_tp_size
|
|
from sglang.srt.managers.schedule_batch import FINISH_ABORT, ScheduleBatch
|
|
from sglang.srt.managers.utils import GenerationBatchResult
|
|
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
|
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
|
|
from sglang.srt.mem_cache.common import release_kv_cache
|
|
from sglang.srt.mem_cache.memory_pool import (
|
|
HybridLinearKVPool,
|
|
HybridReqToTokenPool,
|
|
KVCache,
|
|
NSATokenToKVPool,
|
|
ReqToTokenPool,
|
|
)
|
|
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
|
|
from sglang.srt.observability.req_time_stats import (
|
|
set_schedule_time_batch,
|
|
set_time_batch,
|
|
)
|
|
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.managers.schedule_batch import Req
|
|
from sglang.srt.managers.scheduler import Scheduler
|
|
from sglang.srt.server_args import ServerArgs
|
|
|
|
CLIP_MAX_NEW_TOKEN = envs.SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION.get()
|
|
|
|
|
|
def _cp_draft_shared_kv_debug(message: str, *args) -> None:
|
|
if envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
|
|
logger.info("[CP_DRAFT_SHARED_KV] " + message, *args)
|
|
|
|
|
|
def _seq_summary(values) -> str:
|
|
if values is None:
|
|
return "None"
|
|
try:
|
|
size = len(values)
|
|
except TypeError:
|
|
return str(values)
|
|
if size == 0:
|
|
return "size=0"
|
|
try:
|
|
head = list(values[: min(8, size)])
|
|
except TypeError:
|
|
head = list(values)[: min(8, size)]
|
|
try:
|
|
min_val = min(values)
|
|
max_val = max(values)
|
|
return f"size={size} min={min_val} max={max_val} head={head}"
|
|
except (TypeError, ValueError):
|
|
return f"size={size} head={head}"
|
|
|
|
|
|
def _pool_summary(pool) -> str:
|
|
if pool is None:
|
|
return "None"
|
|
parts = [pool.__class__.__name__]
|
|
for attr in ("size", "page_size", "start_layer", "end_layer", "layer_num"):
|
|
if hasattr(pool, attr):
|
|
parts.append(f"{attr}={getattr(pool, attr)}")
|
|
return " ".join(parts)
|
|
|
|
|
|
def _kv_locs_to_page_indices_cpu(
|
|
kv_locs: torch.Tensor,
|
|
page_size: int,
|
|
*,
|
|
num_tokens: Optional[int] = None,
|
|
):
|
|
"""Return compact int32 page indices for a token-location tensor.
|
|
|
|
PD bootstrap only needs page-level destination indices. Avoid materializing
|
|
a full token-level CPU copy before converting to pages.
|
|
"""
|
|
|
|
if num_tokens is not None:
|
|
kv_locs = kv_locs[:num_tokens]
|
|
|
|
if page_size == 1:
|
|
page_locs = kv_locs
|
|
else:
|
|
page_locs = kv_locs[::page_size] // page_size
|
|
|
|
return page_locs.to(dtype=torch.int32).cpu().numpy()
|
|
|
|
|
|
def _is_fake_transfer(req: Req, server_args: ServerArgs) -> bool:
|
|
return req.bootstrap_host == FAKE_BOOTSTRAP_HOST or (
|
|
req.bootstrap_host is None
|
|
and server_args.disaggregation_transfer_backend == "fake"
|
|
)
|
|
|
|
|
|
def _bootstrap_addr(req: Req) -> str:
|
|
# FIXME: make a property of a req
|
|
return f"{req.bootstrap_host}:{req.bootstrap_port}"
|
|
|
|
|
|
class DecodeReqToTokenPool:
|
|
"""
|
|
The difference of DecodeReqToTokenPool and ReqToTokenPool is that
|
|
DecodeReqToTokenPool subscribes memory for pre-allocated requests.
|
|
|
|
In ReqToTokenPool, if `--max-running-requests` is 8,
|
|
#pre-allocated + #transfer + #running <= 8, but there are in fact more memory can carry pre-allocated requests.
|
|
|
|
In DecodeReqToTokenPool, if `--max-running-requests` is 8,
|
|
#running <= 8, #pre-allocated + #transfer <= pre_alloc_size, so we can use the free memory to pre-allocate requests to unblock prefill.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
size: int,
|
|
max_context_len: int,
|
|
device: str,
|
|
enable_memory_saver: bool,
|
|
pre_alloc_size: int,
|
|
):
|
|
memory_saver_adapter = TorchMemorySaverAdapter.create(
|
|
enable=enable_memory_saver
|
|
)
|
|
|
|
self.size = size
|
|
self.max_context_len = max_context_len
|
|
self.device = device
|
|
self.pre_alloc_size = pre_alloc_size
|
|
with memory_saver_adapter.region(tag=GPU_MEMORY_TYPE_KV_CACHE):
|
|
self.req_to_token = torch.zeros(
|
|
(size + pre_alloc_size, max_context_len),
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
|
|
self.free_slots = deque(range(size + pre_alloc_size))
|
|
|
|
def write(self, indices, values):
|
|
self.req_to_token[indices] = values
|
|
|
|
def available_size(self):
|
|
return len(self.free_slots)
|
|
|
|
def alloc(self, reqs: List["Req"]) -> Optional[List[int]]:
|
|
# Indices of reqs that already have a req_pool_idx and will reuse
|
|
# their existing slot (e.g. chunked prefill continuing across chunks).
|
|
reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None]
|
|
assert len(reusing) <= 1, (
|
|
"only one chunked request may reuse req_pool_idx in a batch"
|
|
)
|
|
assert all(
|
|
reqs[i].is_chunked > 0 or reqs[i].kv_committed_len > 0 for i in reusing
|
|
), "reusing request must be chunked or have committed KV"
|
|
|
|
need_size = len(reqs) - len(reusing)
|
|
if need_size > len(self.free_slots):
|
|
return None
|
|
select_index = [self.free_slots.popleft() for _ in range(need_size)]
|
|
offset = 0
|
|
for r in reqs:
|
|
if r.req_pool_idx is None:
|
|
r.req_pool_idx = select_index[offset]
|
|
offset += 1
|
|
return [r.req_pool_idx for r in reqs]
|
|
|
|
def free(self, req: "Req"):
|
|
assert req.req_pool_idx is not None, "request must have req_pool_idx"
|
|
self.free_slots.append(req.req_pool_idx)
|
|
req.req_pool_idx = None
|
|
|
|
def clear(self):
|
|
self.free_slots = deque(range(self.size + self.pre_alloc_size))
|
|
|
|
|
|
class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
|
|
def __init__(
|
|
self,
|
|
size: int,
|
|
max_context_len: int,
|
|
device: str,
|
|
enable_memory_saver: bool,
|
|
cache_params: "Mamba2CacheParams",
|
|
speculative_num_draft_tokens: int,
|
|
enable_mamba_extra_buffer: bool,
|
|
pre_alloc_size: int,
|
|
enable_overlap_schedule: bool,
|
|
mamba_size: int = None,
|
|
):
|
|
DecodeReqToTokenPool.__init__(
|
|
self,
|
|
size=size,
|
|
max_context_len=max_context_len,
|
|
device=device,
|
|
enable_memory_saver=enable_memory_saver,
|
|
pre_alloc_size=pre_alloc_size,
|
|
)
|
|
|
|
self.mamba_ping_pong_track_buffer_size = 2 if enable_overlap_schedule else 1
|
|
self.enable_mamba_extra_buffer = enable_mamba_extra_buffer
|
|
self.enable_memory_saver = enable_memory_saver
|
|
effective_mamba_size = (
|
|
mamba_size if mamba_size is not None else size
|
|
) + pre_alloc_size
|
|
self._init_mamba_pool(
|
|
size=effective_mamba_size,
|
|
mamba_spec_state_size=size + pre_alloc_size,
|
|
cache_params=cache_params,
|
|
device=device,
|
|
enable_mamba_extra_buffer=self.enable_mamba_extra_buffer,
|
|
speculative_num_draft_tokens=speculative_num_draft_tokens,
|
|
)
|
|
|
|
def clear(self):
|
|
self.free_slots = deque(range(self.size + self.pre_alloc_size))
|
|
self.mamba_pool.clear()
|
|
|
|
|
|
@dataclass
|
|
class DecodeRequest:
|
|
req: Req
|
|
kv_receiver: CommonKVReceiver
|
|
waiting_for_input: bool = False
|
|
metadata_buffer_index: int = -1
|
|
|
|
@property
|
|
def seqlen(self) -> int:
|
|
return self.req.seqlen
|
|
|
|
|
|
class DecodePreallocQueue:
|
|
"""
|
|
Store the requests that are preallocating.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
req_to_token_pool: ReqToTokenPool,
|
|
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
|
|
draft_token_to_kv_pool: Optional[KVCache],
|
|
req_to_metadata_buffer_idx_allocator: ReqToMetadataIdxAllocator,
|
|
metadata_buffers: MetadataBuffers,
|
|
scheduler: Scheduler,
|
|
transfer_queue: DecodeTransferQueue,
|
|
tree_cache: BasePrefixCache,
|
|
gloo_group: ProcessGroup,
|
|
tp_rank: int,
|
|
tp_size: int,
|
|
dp_size: int,
|
|
gpu_id: int,
|
|
bootstrap_port: int,
|
|
max_total_num_tokens: int,
|
|
pp_rank: int,
|
|
num_reserved_decode_tokens: int,
|
|
transfer_backend: TransferBackend,
|
|
):
|
|
self.req_to_token_pool = req_to_token_pool
|
|
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
|
|
self.token_to_kv_pool = token_to_kv_pool_allocator.get_kvcache()
|
|
self.draft_token_to_kv_pool = draft_token_to_kv_pool
|
|
self.is_mla_backend = is_mla_backend(self.token_to_kv_pool)
|
|
self.metadata_buffers = metadata_buffers
|
|
self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator
|
|
self.scheduler = scheduler
|
|
self.transfer_queue = transfer_queue
|
|
self.tree_cache = tree_cache # this is always a chunk cache
|
|
self.gloo_group = gloo_group
|
|
self.tp_rank = tp_rank
|
|
self.tp_size = tp_size
|
|
self.dp_size = dp_size
|
|
self.gpu_id = gpu_id
|
|
self.bootstrap_port = bootstrap_port
|
|
self.max_total_num_tokens = max_total_num_tokens
|
|
self.pp_rank = pp_rank
|
|
self.num_reserved_decode_tokens = num_reserved_decode_tokens
|
|
self.transfer_backend = transfer_backend
|
|
# Queue for requests pending pre-allocation
|
|
self.queue: List[DecodeRequest] = []
|
|
self.retracted_queue: List[Req] = []
|
|
self.pending_reqs: List[Req] = []
|
|
self._ensure_retry_count: Dict[str, int] = {}
|
|
self._max_ensure_retries: int = 20 # scheduling cycles
|
|
self._ensure_last_attempt_time: Dict[str, float] = {}
|
|
self._ensure_retry_interval: float = 1.0 # seconds
|
|
self.kv_manager = self._init_kv_manager()
|
|
self._alloc_extend_prefix_lens: Optional[torch.Tensor] = None
|
|
self._alloc_extend_prefix_lens_cpu: Optional[torch.Tensor] = None
|
|
self._alloc_extend_seq_lens: Optional[torch.Tensor] = None
|
|
self._alloc_extend_seq_lens_cpu: Optional[torch.Tensor] = None
|
|
self._alloc_extend_last_loc: Optional[torch.Tensor] = None
|
|
|
|
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) -> CommonKVManager:
|
|
kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS)
|
|
kv_args = kv_args_class()
|
|
|
|
attn_tp_size = get_attention_tp_size()
|
|
kv_args.engine_rank = self.tp_rank % (attn_tp_size)
|
|
|
|
kv_args.pp_rank = self.pp_rank
|
|
kv_args.system_dp_rank = self.scheduler.dp_rank
|
|
kv_data_ptrs, kv_data_lens, kv_item_lens = (
|
|
self.token_to_kv_pool.get_contiguous_buf_infos()
|
|
)
|
|
target_kv_buffer_count = len(kv_data_ptrs)
|
|
draft_kv_data_lens = []
|
|
draft_kv_item_lens = []
|
|
draft_kv_buffer_count = 0
|
|
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()
|
|
)
|
|
draft_kv_buffer_count = len(draft_kv_data_ptrs)
|
|
kv_data_ptrs += draft_kv_data_ptrs
|
|
kv_data_lens += draft_kv_data_lens
|
|
kv_item_lens += draft_kv_item_lens
|
|
|
|
kv_args.draft_kv_buffer_start = target_kv_buffer_count
|
|
kv_args.draft_kv_buffer_count = draft_kv_buffer_count
|
|
_cp_draft_shared_kv_debug(
|
|
"decode_kv_manager cp_rank=%s target_pool=(%s) draft_pool=(%s) "
|
|
"target_bufs=%s draft_bufs=%s total_bufs=%s target_lens=%s "
|
|
"draft_lens=%s target_item_lens=%s draft_item_lens=%s",
|
|
self.tp_rank,
|
|
_pool_summary(self.token_to_kv_pool),
|
|
_pool_summary(self.draft_token_to_kv_pool),
|
|
target_kv_buffer_count,
|
|
draft_kv_buffer_count,
|
|
len(kv_data_ptrs),
|
|
_seq_summary(kv_data_lens[:target_kv_buffer_count]),
|
|
_seq_summary(draft_kv_data_lens),
|
|
_seq_summary(kv_item_lens[:target_kv_buffer_count]),
|
|
_seq_summary(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
|
|
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()
|
|
)
|
|
|
|
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_args.ib_device = self.scheduler.server_args.disaggregation_ib_device
|
|
kv_args.gpu_id = self.scheduler.gpu_id
|
|
kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)
|
|
kv_manager = kv_manager_class(
|
|
kv_args,
|
|
DisaggregationMode.DECODE,
|
|
self.scheduler.server_args,
|
|
self.is_mla_backend,
|
|
)
|
|
return kv_manager
|
|
|
|
def add(self, req: Req, is_retracted: bool = False) -> None:
|
|
"""Add a request to the pending queue."""
|
|
if self._check_if_req_exceed_kv_capacity(req):
|
|
return
|
|
|
|
if is_retracted:
|
|
req.retraction_mb_id = None
|
|
self.retracted_queue.append(req)
|
|
else:
|
|
# NOTE: fake transfer does not need to resolve prefill dp rank in the pending queue
|
|
if _is_fake_transfer(req, self.scheduler.server_args):
|
|
self._create_receiver_and_enqueue(req, 0)
|
|
return
|
|
|
|
# Fast path: cache-only lookup, no network calls
|
|
prefill_dp_rank = self._resolve_prefill_dp_rank(req)
|
|
if prefill_dp_rank is not None:
|
|
self._create_receiver_and_enqueue(req, prefill_dp_rank)
|
|
else:
|
|
self.pending_reqs.append(req)
|
|
|
|
def _resolve_prefill_dp_rank(self, req: Req) -> Optional[int]:
|
|
if req.disagg_prefill_dp_rank is not None:
|
|
return req.disagg_prefill_dp_rank
|
|
|
|
prefill_info = self.kv_manager.prefill_info_table.get(_bootstrap_addr(req))
|
|
if prefill_info is None:
|
|
return None
|
|
|
|
if prefill_info.dp_size == 1:
|
|
return 0
|
|
|
|
if prefill_info.follow_bootstrap_room:
|
|
return req.bootstrap_room % prefill_info.dp_size
|
|
|
|
return None
|
|
|
|
def _create_receiver_and_enqueue(self, req: Req, prefill_dp_rank: int) -> None:
|
|
backend = (
|
|
TransferBackend.FAKE
|
|
if _is_fake_transfer(req, self.scheduler.server_args)
|
|
else self.transfer_backend
|
|
)
|
|
kv_receiver_class = get_kv_class(backend, KVClassType.RECEIVER)
|
|
|
|
kv_receiver = kv_receiver_class(
|
|
mgr=self.kv_manager,
|
|
bootstrap_addr=_bootstrap_addr(req),
|
|
bootstrap_room=req.bootstrap_room,
|
|
prefill_dp_rank=prefill_dp_rank,
|
|
)
|
|
|
|
self.queue.append(
|
|
DecodeRequest(
|
|
req=req,
|
|
kv_receiver=kv_receiver,
|
|
waiting_for_input=False,
|
|
)
|
|
)
|
|
|
|
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 extend(self, reqs: List[Req], is_retracted: bool = False) -> None:
|
|
"""Add a request to the pending queue."""
|
|
for req in reqs:
|
|
self.add(req, is_retracted=is_retracted)
|
|
|
|
def resume_retracted_reqs(
|
|
self, rids_to_check: Optional[List[str]] = None
|
|
) -> List[Req]:
|
|
# TODO refactor the scheduling part, reuse with the unified engine logic as much as possible
|
|
|
|
# allocate memory
|
|
resumed_reqs = []
|
|
remaining_retracted_queue = []
|
|
allocatable_tokens = self._allocatable_tokens(count_retracted=False)
|
|
|
|
for i, req in enumerate(self.retracted_queue):
|
|
if rids_to_check is not None and req.rid not in rids_to_check:
|
|
remaining_retracted_queue.append(req)
|
|
continue
|
|
|
|
if self.req_to_token_pool.available_size() <= 0:
|
|
remaining_retracted_queue.append(req)
|
|
remaining_retracted_queue.extend(self.retracted_queue[i + 1 :])
|
|
break
|
|
|
|
required_tokens_for_request = (
|
|
len(req.origin_input_ids)
|
|
+ len(req.output_ids)
|
|
+ self.num_reserved_decode_tokens
|
|
)
|
|
if required_tokens_for_request > allocatable_tokens:
|
|
remaining_retracted_queue.append(req)
|
|
remaining_retracted_queue.extend(self.retracted_queue[i + 1 :])
|
|
break
|
|
|
|
resumed_reqs.append(req)
|
|
req.is_retracted = False
|
|
self._pre_alloc(req)
|
|
allocatable_tokens -= required_tokens_for_request
|
|
|
|
# load from cpu, release the cpu copy
|
|
req.load_kv_cache(self.req_to_token_pool, self.token_to_kv_pool_allocator)
|
|
|
|
self.retracted_queue = remaining_retracted_queue
|
|
|
|
return resumed_reqs
|
|
|
|
def _update_handshake_waiters(
|
|
self, rids_to_check: Optional[List[str]] = None
|
|
) -> None:
|
|
if not self.queue:
|
|
return
|
|
|
|
if all(decode_req.waiting_for_input for decode_req in self.queue):
|
|
return
|
|
|
|
polls = poll_and_all_reduce(
|
|
[decode_req.kv_receiver for decode_req in self.queue], self.gloo_group
|
|
)
|
|
|
|
for i, (decode_req, poll) in enumerate(zip(self.queue, polls)):
|
|
if rids_to_check is not None and decode_req.req.rid not in rids_to_check:
|
|
continue
|
|
|
|
if poll == KVPoll.Bootstrapping:
|
|
pass
|
|
elif poll == KVPoll.WaitingForInput:
|
|
decode_req.waiting_for_input = True
|
|
decode_req.req.time_stats.set_bootstrap_done_time()
|
|
elif poll == KVPoll.Failed:
|
|
error_message = f"Decode handshake failed for request rank={self.tp_rank} {decode_req.req.rid=} {decode_req.req.bootstrap_room=}"
|
|
try:
|
|
decode_req.kv_receiver.failure_exception()
|
|
except Exception as e:
|
|
error_message += f" with exception {e}"
|
|
logger.error(error_message)
|
|
prepare_abort(
|
|
decode_req.req,
|
|
error_message,
|
|
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
|
|
)
|
|
if self.scheduler.enable_metrics:
|
|
self.scheduler.metrics_collector.increment_bootstrap_failed_reqs()
|
|
else:
|
|
raise ValueError(f"Unexpected poll case: {poll}")
|
|
|
|
def _ensure_prefill_info(
|
|
self, addr_to_reqs: Dict[str, List[Req]]
|
|
) -> Tuple[Dict[str, List[Req]], List[Req]]:
|
|
"""Non-blocking ensure parallel info for each addr.
|
|
Returns (ready_addrs, remaining_reqs)."""
|
|
ready: Dict[str, List[Req]] = {}
|
|
remaining: List[Req] = []
|
|
|
|
now = time.monotonic()
|
|
for bootstrap_addr, reqs in addr_to_reqs.items():
|
|
last_attempt = self._ensure_last_attempt_time.get(bootstrap_addr)
|
|
if last_attempt is not None and (
|
|
now - last_attempt < self._ensure_retry_interval
|
|
):
|
|
remaining.extend(reqs)
|
|
continue
|
|
|
|
self._ensure_last_attempt_time[bootstrap_addr] = now
|
|
|
|
if self.kv_manager.try_ensure_parallel_info(bootstrap_addr):
|
|
if bootstrap_addr in self._ensure_retry_count:
|
|
del self._ensure_retry_count[bootstrap_addr]
|
|
if bootstrap_addr in self._ensure_last_attempt_time:
|
|
del self._ensure_last_attempt_time[bootstrap_addr]
|
|
ready[bootstrap_addr] = reqs
|
|
continue
|
|
|
|
count = self._ensure_retry_count.get(bootstrap_addr, 0) + 1
|
|
self._ensure_retry_count[bootstrap_addr] = count
|
|
|
|
if count >= self._max_ensure_retries:
|
|
error_msg = f"Could not fetch prefill parallel info from {bootstrap_addr} after {count} attempts"
|
|
logger.error(error_msg)
|
|
for req in reqs:
|
|
prepare_abort(
|
|
req, error_msg, status_code=HTTPStatus.INTERNAL_SERVER_ERROR
|
|
)
|
|
if self.scheduler.enable_metrics:
|
|
self.scheduler.metrics_collector.increment_bootstrap_failed_reqs()
|
|
self.scheduler.stream_output([req], req.return_logprob)
|
|
del self._ensure_retry_count[bootstrap_addr]
|
|
del self._ensure_last_attempt_time[bootstrap_addr]
|
|
else:
|
|
remaining.extend(reqs)
|
|
|
|
return ready, remaining
|
|
|
|
def _resolve_pending_reqs(self) -> None:
|
|
"""Batch-resolve prefill_dp_ranks for pending requests and create receivers."""
|
|
if not self.pending_reqs:
|
|
return
|
|
|
|
# Group pending requests by bootstrap_addr
|
|
addr_to_reqs: Dict[str, List[Req]] = {}
|
|
for req in self.pending_reqs:
|
|
addr = _bootstrap_addr(req)
|
|
addr_to_reqs.setdefault(addr, []).append(req)
|
|
|
|
# Pass 1: ensure parallel info for each addr
|
|
ready_addrs, remaining = self._ensure_prefill_info(addr_to_reqs)
|
|
|
|
# Pass 2: resolve dp rank for addrs whose info is available
|
|
resolved = []
|
|
for bootstrap_addr, reqs in ready_addrs.items():
|
|
need_query: List[Req] = []
|
|
for req in reqs:
|
|
prefill_dp_rank = self._resolve_prefill_dp_rank(req)
|
|
if prefill_dp_rank is not None:
|
|
resolved.append((req, prefill_dp_rank))
|
|
else:
|
|
need_query.append(req)
|
|
|
|
if need_query:
|
|
rooms = [req.bootstrap_room for req in need_query]
|
|
room_to_rank = CommonKVReceiver.query_prefill_dp_ranks(
|
|
bootstrap_addr, rooms
|
|
)
|
|
for req in need_query:
|
|
prefill_dp_rank = room_to_rank.get(str(req.bootstrap_room))
|
|
if prefill_dp_rank is not None:
|
|
resolved.append((req, int(prefill_dp_rank)))
|
|
else:
|
|
remaining.append(req)
|
|
|
|
self.pending_reqs = remaining
|
|
|
|
for req, prefill_dp_rank in resolved:
|
|
self._create_receiver_and_enqueue(req, prefill_dp_rank)
|
|
|
|
def pop_preallocated(
|
|
self, rids_to_check: Optional[List[str]] = None
|
|
) -> Tuple[List[DecodeRequest], List[DecodeRequest]]:
|
|
"""Pop the preallocated requests from the pending queue (FIFO)."""
|
|
self._resolve_pending_reqs()
|
|
self._update_handshake_waiters(rids_to_check)
|
|
|
|
failed_reqs = []
|
|
preallocated_reqs = []
|
|
remaining_queue = []
|
|
prealloc_blocked = False
|
|
|
|
# We need to make sure that the sum of inflight tokens and allocatable tokens is greater than maximum input+output length of each inflight request
|
|
# Otherwise it is possible for one request running decode out of memory, while all other requests are in the transfer queue that cannot be retracted.
|
|
retractable_tokens = sum(
|
|
len(r.origin_input_ids) + len(r.output_ids)
|
|
for r in self.scheduler.running_batch.reqs
|
|
)
|
|
allocatable_tokens = self._allocatable_tokens(
|
|
retractable_tokens=retractable_tokens, count_retracted=True
|
|
)
|
|
for i, decode_req in enumerate(self.queue):
|
|
if rids_to_check is not None and decode_req.req.rid not in rids_to_check:
|
|
remaining_queue.append(decode_req)
|
|
continue
|
|
|
|
if isinstance(decode_req.req.finished_reason, FINISH_ABORT):
|
|
self.scheduler.stream_output(
|
|
[decode_req.req], decode_req.req.return_logprob
|
|
)
|
|
failed_reqs.append(decode_req)
|
|
continue
|
|
|
|
if prealloc_blocked:
|
|
remaining_queue.append(decode_req)
|
|
continue
|
|
|
|
if not decode_req.waiting_for_input:
|
|
remaining_queue.append(decode_req)
|
|
continue
|
|
|
|
if self.req_to_token_pool.available_size() <= 0:
|
|
remaining_queue.append(decode_req)
|
|
prealloc_blocked = True
|
|
continue
|
|
|
|
if self.req_to_metadata_buffer_idx_allocator.available_size() <= 0:
|
|
remaining_queue.append(decode_req)
|
|
prealloc_blocked = True
|
|
continue
|
|
|
|
# Memory estimation: don't add if the projected memory cannot be met
|
|
# TODO: add new_token ratio
|
|
origin_input_len = len(decode_req.req.origin_input_ids)
|
|
required_tokens_for_request = (
|
|
origin_input_len + self.num_reserved_decode_tokens
|
|
)
|
|
|
|
if (
|
|
max(
|
|
required_tokens_for_request,
|
|
origin_input_len
|
|
+ min(
|
|
decode_req.req.sampling_params.max_new_tokens,
|
|
CLIP_MAX_NEW_TOKEN,
|
|
)
|
|
- retractable_tokens,
|
|
)
|
|
> allocatable_tokens
|
|
):
|
|
remaining_queue.append(decode_req)
|
|
prealloc_blocked = True
|
|
continue
|
|
if required_tokens_for_request > allocatable_tokens:
|
|
remaining_queue.append(decode_req)
|
|
prealloc_blocked = True
|
|
continue
|
|
|
|
allocatable_tokens -= required_tokens_for_request
|
|
kv_loc = self._pre_alloc(
|
|
decode_req.req,
|
|
write_req_to_token=False,
|
|
)
|
|
|
|
page_size = self.token_to_kv_pool_allocator.page_size
|
|
page_indices = _kv_locs_to_page_indices_cpu(
|
|
kv_loc,
|
|
page_size,
|
|
num_tokens=origin_input_len,
|
|
)
|
|
|
|
# Prepare extra pool indices for hybrid models
|
|
if isinstance(self.token_to_kv_pool, HybridLinearKVPool):
|
|
# Mamba hybrid model: single mamba state index
|
|
state_indices = [
|
|
self.req_to_token_pool.req_index_to_mamba_index_mapping[
|
|
decode_req.req.req_pool_idx
|
|
]
|
|
.cpu()
|
|
.numpy()
|
|
]
|
|
elif isinstance(self.token_to_kv_pool, SWAKVPool):
|
|
# SWA hybrid model: send decode-side SWA window indices
|
|
seq_len = len(decode_req.req.origin_input_ids)
|
|
window_size = self.scheduler.sliding_window_size
|
|
|
|
window_start = max(0, seq_len - window_size)
|
|
window_start = (window_start // page_size) * page_size
|
|
window_kv_indices_full = kv_loc[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 = _kv_locs_to_page_indices_cpu(
|
|
window_kv_indices_swa,
|
|
page_size,
|
|
)
|
|
elif isinstance(self.token_to_kv_pool, NSATokenToKVPool):
|
|
state_indices = page_indices
|
|
else:
|
|
state_indices = None
|
|
|
|
self.req_to_token_pool.write(
|
|
(decode_req.req.req_pool_idx, slice(0, len(kv_loc))), kv_loc
|
|
)
|
|
|
|
decode_req.metadata_buffer_index = (
|
|
self.req_to_metadata_buffer_idx_allocator.alloc()
|
|
)
|
|
assert decode_req.metadata_buffer_index is not None
|
|
_cp_draft_shared_kv_debug(
|
|
"decode_prealloc rid=%s room=%s origin_tokens=%s fill_tokens=%s "
|
|
"page_size=%s pages=%s state_pages=%s metadata_idx=%s has_draft_pool=%s",
|
|
decode_req.req.rid,
|
|
decode_req.req.bootstrap_room,
|
|
origin_input_len,
|
|
len(kv_loc),
|
|
page_size,
|
|
_seq_summary(page_indices),
|
|
_seq_summary(state_indices),
|
|
decode_req.metadata_buffer_index,
|
|
self.draft_token_to_kv_pool is not None,
|
|
)
|
|
decode_req.kv_receiver.init(
|
|
page_indices, decode_req.metadata_buffer_index, state_indices
|
|
)
|
|
preallocated_reqs.append(decode_req)
|
|
decode_req.req.time_stats.set_decode_transfer_queue_entry_time()
|
|
|
|
self.queue = remaining_queue
|
|
|
|
return preallocated_reqs, failed_reqs
|
|
|
|
@property
|
|
def num_tokens_pre_allocated(self):
|
|
return sum(
|
|
len(decode_req.req.fill_ids) for decode_req in self.transfer_queue.queue
|
|
)
|
|
|
|
def _allocatable_tokens(
|
|
self, retractable_tokens: Optional[int] = None, count_retracted: bool = True
|
|
) -> int:
|
|
need_space_for_single_req = (
|
|
max(
|
|
[
|
|
min(x.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKEN)
|
|
+ len(x.origin_input_ids)
|
|
- retractable_tokens
|
|
for x in self.scheduler.running_batch.reqs
|
|
]
|
|
)
|
|
if retractable_tokens is not None
|
|
and len(self.scheduler.running_batch.reqs) > 0
|
|
else 0
|
|
)
|
|
available_size = self.token_to_kv_pool_allocator.available_size()
|
|
allocatable_tokens = available_size - max(
|
|
# preserve some space for future decode
|
|
self.num_reserved_decode_tokens
|
|
* (
|
|
len(self.scheduler.running_batch.reqs)
|
|
+ len(self.transfer_queue.queue)
|
|
+ len(self.scheduler.waiting_queue)
|
|
),
|
|
# make sure each request can finish if reach max_tokens with all other requests retracted
|
|
need_space_for_single_req,
|
|
)
|
|
|
|
# Note: if the last prebuilt extend just finishes, and we enter `pop_preallocated` immediately in the next iteration
|
|
# the extend batch is not in any queue, so we need to explicitly add the tokens slots here
|
|
if (
|
|
self.scheduler.last_batch
|
|
and self.scheduler.last_batch.forward_mode.is_prebuilt()
|
|
):
|
|
allocatable_tokens -= self.num_reserved_decode_tokens * len(
|
|
self.scheduler.last_batch.reqs
|
|
)
|
|
|
|
if count_retracted:
|
|
allocatable_tokens -= sum(
|
|
[
|
|
len(req.origin_input_ids)
|
|
+ len(req.output_ids)
|
|
+ self.num_reserved_decode_tokens
|
|
for req in self.retracted_queue
|
|
]
|
|
)
|
|
return allocatable_tokens
|
|
|
|
def _get_alloc_extend_args(
|
|
self, fill_len: int
|
|
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
device = self.token_to_kv_pool_allocator.device
|
|
|
|
if self._alloc_extend_prefix_lens is None:
|
|
self._alloc_extend_prefix_lens = torch.zeros(
|
|
1, dtype=torch.int64, device=device
|
|
)
|
|
self._alloc_extend_prefix_lens_cpu = torch.zeros(1, dtype=torch.int64)
|
|
self._alloc_extend_seq_lens = torch.empty(
|
|
1, dtype=torch.int64, device=device
|
|
)
|
|
self._alloc_extend_seq_lens_cpu = torch.empty(1, dtype=torch.int64)
|
|
self._alloc_extend_last_loc = torch.empty(
|
|
1, dtype=torch.int64, device=device
|
|
)
|
|
|
|
prefix_lens = self._alloc_extend_prefix_lens
|
|
prefix_lens_cpu = self._alloc_extend_prefix_lens_cpu
|
|
seq_lens = self._alloc_extend_seq_lens
|
|
seq_lens_cpu = self._alloc_extend_seq_lens_cpu
|
|
last_loc = self._alloc_extend_last_loc
|
|
|
|
assert prefix_lens is not None
|
|
assert prefix_lens_cpu is not None
|
|
assert seq_lens is not None
|
|
assert seq_lens_cpu is not None
|
|
assert last_loc is not None
|
|
|
|
prefix_lens.zero_()
|
|
prefix_lens_cpu.zero_()
|
|
seq_lens.fill_(fill_len)
|
|
seq_lens_cpu.fill_(fill_len)
|
|
last_loc.fill_(-1)
|
|
|
|
return (
|
|
prefix_lens,
|
|
prefix_lens_cpu,
|
|
seq_lens,
|
|
seq_lens_cpu,
|
|
last_loc,
|
|
)
|
|
|
|
def _pre_alloc(
|
|
self,
|
|
req: Req,
|
|
*,
|
|
write_req_to_token: bool = True,
|
|
) -> torch.Tensor:
|
|
"""Pre-allocate the memory for req_to_token and token_kv_pool"""
|
|
req_pool_indices = self.req_to_token_pool.alloc([req])
|
|
|
|
assert req_pool_indices is not None, (
|
|
"req_pool_indices is full! There is a bug in memory estimation."
|
|
)
|
|
|
|
# Alloc all tokens for the prebuilt req (except for the reserved input token for decoding)
|
|
fill_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
|
|
req.kv_allocated_len = fill_len
|
|
req.kv_committed_len = fill_len
|
|
if self.token_to_kv_pool_allocator.page_size == 1:
|
|
kv_loc = self.token_to_kv_pool_allocator.alloc(fill_len)
|
|
else:
|
|
(
|
|
prefix_lens,
|
|
prefix_lens_cpu,
|
|
seq_lens,
|
|
seq_lens_cpu,
|
|
last_loc,
|
|
) = self._get_alloc_extend_args(fill_len)
|
|
kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
|
|
prefix_lens=prefix_lens,
|
|
prefix_lens_cpu=prefix_lens_cpu,
|
|
seq_lens=seq_lens,
|
|
seq_lens_cpu=seq_lens_cpu,
|
|
last_loc=last_loc,
|
|
extend_num_tokens=fill_len,
|
|
)
|
|
|
|
assert kv_loc is not None, (
|
|
"KV cache is full! There is a bug in memory estimation."
|
|
)
|
|
|
|
if write_req_to_token:
|
|
self.req_to_token_pool.write(
|
|
(req.req_pool_idx, slice(0, len(kv_loc))), kv_loc
|
|
)
|
|
|
|
# populate metadata
|
|
req.fill_ids = req.origin_input_ids + req.output_ids
|
|
req.set_extend_input_len(len(req.fill_ids))
|
|
|
|
return kv_loc
|
|
|
|
|
|
class DecodeTransferQueue:
|
|
"""
|
|
Store the requests that is polling kv
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
gloo_group: ProcessGroup,
|
|
req_to_metadata_buffer_idx_allocator: ReqToMetadataIdxAllocator,
|
|
tp_rank: int,
|
|
metadata_buffers: MetadataBuffers,
|
|
scheduler: Scheduler,
|
|
tree_cache: BasePrefixCache,
|
|
):
|
|
self.queue: List[DecodeRequest] = []
|
|
self.gloo_group = gloo_group
|
|
self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator
|
|
self.tp_rank = tp_rank
|
|
self.metadata_buffers = metadata_buffers
|
|
self.scheduler = scheduler
|
|
self.tree_cache = tree_cache
|
|
self.spec_algorithm = scheduler.spec_algorithm
|
|
|
|
def add(self, decode_req: DecodeRequest) -> None:
|
|
self.queue.append(decode_req)
|
|
|
|
def extend(self, decode_reqs: List[DecodeRequest]) -> None:
|
|
self.queue.extend(decode_reqs)
|
|
|
|
def _commit_transfer_to_req(self, decode_req: DecodeRequest) -> bool:
|
|
"""
|
|
Returns:
|
|
True if the request should be removed from the queue (success or corruption)
|
|
False if metadata not ready yet (keep in queue for next poll)
|
|
"""
|
|
idx = decode_req.metadata_buffer_index
|
|
(
|
|
output_id,
|
|
cached_tokens,
|
|
output_token_logprobs_val,
|
|
output_token_logprobs_idx,
|
|
output_top_logprobs_val,
|
|
output_top_logprobs_idx,
|
|
output_topk_p,
|
|
output_topk_index,
|
|
output_hidden_states,
|
|
output_bootstrap_room,
|
|
) = self.metadata_buffers.get_buf(idx)
|
|
|
|
# Validate bootstrap_room to detect context corruption
|
|
actual_room = output_bootstrap_room[0].item()
|
|
expected_room = (
|
|
decode_req.req.bootstrap_room
|
|
if decode_req.req.bootstrap_room is not None
|
|
else 0
|
|
)
|
|
|
|
if _is_fake_transfer(decode_req.req, self.scheduler.server_args):
|
|
pass
|
|
elif actual_room == 0:
|
|
# Case 1: Metadata not ready yet (actual_room == 0)
|
|
# Keep request in queue and wait for next poll
|
|
return False
|
|
elif actual_room != expected_room:
|
|
# Case 2: Real corruption detected (mismatch)
|
|
# Abort the request and remove from the queue
|
|
error_msg = (
|
|
f"Context corruption detected: Request {decode_req.req.rid} "
|
|
f"(bootstrap_room={expected_room}) received metadata from "
|
|
f"bootstrap_room={actual_room}. "
|
|
f"Metadata buffer index: {idx}. "
|
|
f"This indicates metadata buffer index collision."
|
|
)
|
|
logger.error(error_msg)
|
|
prepare_abort(
|
|
decode_req.req,
|
|
"Metadata corruption detected - bootstrap_room mismatch",
|
|
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
|
|
)
|
|
decode_req.kv_receiver.clear()
|
|
decode_req.kv_receiver = None
|
|
return True
|
|
|
|
# Case 3: Success - commit the transfer
|
|
decode_req.req.output_ids.append(output_id[0].item())
|
|
decode_req.req.cached_tokens = cached_tokens[0].item()
|
|
if not self.spec_algorithm.is_none():
|
|
decode_req.req.output_topk_p = output_topk_p
|
|
decode_req.req.output_topk_index = output_topk_index
|
|
decode_req.req.hidden_states_tensor = output_hidden_states
|
|
|
|
_cp_draft_shared_kv_debug(
|
|
"decode_transfer_commit rid=%s room=%s metadata_idx=%s cached_tokens=%s "
|
|
"topk_p_shape=%s topk_index_shape=%s hidden_shape=%s",
|
|
decode_req.req.rid,
|
|
decode_req.req.bootstrap_room,
|
|
idx,
|
|
decode_req.req.cached_tokens,
|
|
tuple(output_topk_p.shape) if output_topk_p is not None else None,
|
|
tuple(output_topk_index.shape) if output_topk_index is not None else None,
|
|
tuple(output_hidden_states.shape) if output_hidden_states is not None else None,
|
|
)
|
|
|
|
if decode_req.req.return_logprob:
|
|
decode_req.req.output_token_logprobs_val.append(
|
|
output_token_logprobs_val[0].item()
|
|
)
|
|
decode_req.req.output_token_logprobs_idx.append(
|
|
output_token_logprobs_idx[0].item()
|
|
)
|
|
decode_req.req.output_top_logprobs_val.append(
|
|
output_top_logprobs_val[: decode_req.req.top_logprobs_num].tolist()
|
|
)
|
|
decode_req.req.output_top_logprobs_idx.append(
|
|
output_top_logprobs_idx[: decode_req.req.top_logprobs_num].tolist()
|
|
)
|
|
|
|
decode_req.kv_receiver.clear()
|
|
decode_req.kv_receiver = None
|
|
decode_req.req.time_stats.set_wait_queue_entry_time()
|
|
return True
|
|
|
|
def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req]:
|
|
if not self.queue:
|
|
return []
|
|
polls = poll_and_all_reduce(
|
|
[decode_req.kv_receiver for decode_req in self.queue], self.gloo_group
|
|
)
|
|
|
|
transferred_reqs = []
|
|
completed_decode_reqs = []
|
|
remaining_queue = []
|
|
for decode_req, poll in zip(self.queue, polls):
|
|
if rids_to_check is not None and decode_req.req.rid not in rids_to_check:
|
|
remaining_queue.append(decode_req)
|
|
continue
|
|
if poll == KVPoll.Failed:
|
|
error_message = f"Decode transfer failed for request rank={self.tp_rank} {decode_req.req.rid=} {decode_req.req.bootstrap_room=}"
|
|
try:
|
|
decode_req.kv_receiver.failure_exception()
|
|
except Exception as e:
|
|
error_message += f" with exception {e}"
|
|
logger.error(error_message)
|
|
prepare_abort(
|
|
decode_req.req,
|
|
error_message,
|
|
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
|
|
)
|
|
self.scheduler.stream_output(
|
|
[decode_req.req], decode_req.req.return_logprob
|
|
)
|
|
# release pre-allocated kv cache, but don't insert into the tree since it's failed
|
|
release_kv_cache(decode_req.req, self.tree_cache, is_insert=False)
|
|
completed_decode_reqs.append(decode_req)
|
|
if self.scheduler.enable_metrics:
|
|
self.scheduler.metrics_collector.increment_transfer_failed_reqs()
|
|
continue
|
|
elif poll == KVPoll.Success:
|
|
should_remove = self._commit_transfer_to_req(decode_req)
|
|
if should_remove:
|
|
completed_decode_reqs.append(decode_req)
|
|
# Check if request was aborted due to corruption
|
|
if isinstance(decode_req.req.finished_reason, FINISH_ABORT):
|
|
self.scheduler.stream_output(
|
|
[decode_req.req], decode_req.req.return_logprob
|
|
)
|
|
release_kv_cache(
|
|
decode_req.req, self.tree_cache, is_insert=False
|
|
)
|
|
if self.scheduler.enable_metrics:
|
|
self.scheduler.metrics_collector.increment_transfer_failed_reqs()
|
|
else:
|
|
transferred_reqs.append(decode_req.req)
|
|
else:
|
|
remaining_queue.append(decode_req)
|
|
elif poll in [
|
|
KVPoll.Bootstrapping,
|
|
KVPoll.WaitingForInput,
|
|
KVPoll.Transferring,
|
|
]:
|
|
remaining_queue.append(decode_req)
|
|
else:
|
|
raise ValueError(f"Unexpected poll case: {poll}")
|
|
|
|
if len(polls) < len(self.queue):
|
|
remaining_queue.extend(self.queue[len(polls) :])
|
|
|
|
for decode_req in completed_decode_reqs:
|
|
idx = decode_req.metadata_buffer_index
|
|
assert idx != -1
|
|
self.req_to_metadata_buffer_idx_allocator.free(idx)
|
|
|
|
self.queue = remaining_queue
|
|
|
|
return transferred_reqs
|
|
|
|
|
|
class SchedulerDisaggregationDecodeMixin:
|
|
@torch.no_grad()
|
|
def event_loop_normal_disagg_decode(self: Scheduler):
|
|
"""A normal scheduler loop for decode worker in disaggregation mode."""
|
|
|
|
while True:
|
|
# Receive requests
|
|
recv_reqs = self.recv_requests()
|
|
self.process_input_requests(recv_reqs)
|
|
# polling and allocating kv cache
|
|
self.process_decode_queue()
|
|
|
|
# Get the next batch to run
|
|
batch = self.get_next_disagg_decode_batch_to_run()
|
|
self.cur_batch = batch
|
|
disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch)
|
|
|
|
def pop_and_process():
|
|
# Process the results of the last batch
|
|
tmp_batch, tmp_result = self.result_queue.popleft()
|
|
self.process_batch_result(tmp_batch, tmp_result)
|
|
|
|
# If we need grammar sync (spec + grammar), process the last batch
|
|
# results first so that the grammar state is up-to-date before
|
|
# generating the bitmask for the current batch.
|
|
if disable_overlap_for_batch:
|
|
pop_and_process()
|
|
|
|
# Launch the current batch
|
|
if batch:
|
|
result = self.run_batch(batch)
|
|
self.process_batch_result(batch, result)
|
|
else:
|
|
# When the server is idle, do self-check and re-init some states
|
|
self.self_check_during_idle()
|
|
|
|
# Update last_batch
|
|
self.last_batch = batch
|
|
|
|
@torch.no_grad()
|
|
def event_loop_overlap_disagg_decode(self: Scheduler):
|
|
self.result_queue = deque()
|
|
self.last_batch: Optional[ScheduleBatch] = None
|
|
|
|
def pop_and_process():
|
|
# Process the results of the last batch
|
|
tmp_batch, tmp_result = self.result_queue.popleft()
|
|
self.process_batch_result(tmp_batch, tmp_result)
|
|
|
|
while True:
|
|
# Receive requests
|
|
recv_reqs = self.recv_requests()
|
|
self.process_input_requests(recv_reqs)
|
|
# polling and allocating kv cache
|
|
self.process_decode_queue()
|
|
|
|
# Get the next batch to run
|
|
batch = self.get_next_disagg_decode_batch_to_run()
|
|
self.cur_batch = batch
|
|
disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch)
|
|
|
|
# If we need grammar sync (spec + grammar), process the last batch
|
|
# results first so that the grammar state is up-to-date before
|
|
# generating the bitmask for the current batch.
|
|
if self.last_batch and disable_overlap_for_batch:
|
|
pop_and_process()
|
|
|
|
# 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:
|
|
if not disable_overlap_for_batch:
|
|
pop_and_process()
|
|
elif batch is None:
|
|
self.self_check_during_idle()
|
|
|
|
# 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 _run_batch_prebuilt(
|
|
self: Scheduler, batch: ScheduleBatch
|
|
) -> GenerationBatchResult:
|
|
if batch.inner_idle_batch is not None:
|
|
idle_batch = batch.inner_idle_batch
|
|
# Reset the inner idle batch to avoid reusing it.
|
|
batch.inner_idle_batch = None
|
|
return self.run_batch(idle_batch)
|
|
|
|
return GenerationBatchResult()
|
|
|
|
def get_next_disagg_decode_batch_to_run(
|
|
self: Scheduler,
|
|
) -> Optional[ScheduleBatch]:
|
|
"""Process prebuilt batch and schedule the next decode batch."""
|
|
# Process pending prebuilt batch: output processing + filter + merge
|
|
new_prebuilt_batch = self.get_new_prebuilt_batch()
|
|
if new_prebuilt_batch:
|
|
assert self.chunked_req is None
|
|
self.process_batch_result_prebuilt(new_prebuilt_batch)
|
|
new_prebuilt_batch.filter_batch()
|
|
if not new_prebuilt_batch.is_empty():
|
|
if self.running_batch.is_empty():
|
|
self.running_batch = new_prebuilt_batch
|
|
else:
|
|
self.running_batch.merge_batch(new_prebuilt_batch)
|
|
|
|
# Schedule decode batch
|
|
if self.running_batch.is_empty():
|
|
ret = None
|
|
else:
|
|
self.running_batch = self.update_running_batch(self.running_batch)
|
|
ret = self.running_batch if not self.running_batch.is_empty() else None
|
|
|
|
ret = self.maybe_prepare_mlp_sync_batch(ret)
|
|
if ret:
|
|
set_schedule_time_batch(ret)
|
|
return ret
|
|
|
|
def get_new_prebuilt_batch(self: Scheduler) -> Optional[ScheduleBatch]:
|
|
"""Create a schedulebatch for fake completed prefill"""
|
|
if self.grammar_manager.has_waiting_grammars():
|
|
ready_grammar_requests = self.grammar_manager.get_ready_grammar_requests()
|
|
for req in ready_grammar_requests:
|
|
self._add_request_to_queue(req)
|
|
|
|
if len(self.waiting_queue) == 0:
|
|
return None
|
|
|
|
curr_batch_size = self.running_batch.batch_size()
|
|
|
|
batch_size = min(self.req_to_token_pool.size, self.max_running_requests)
|
|
|
|
num_not_used_batch = batch_size - curr_batch_size
|
|
|
|
num_new_reqs = min(num_not_used_batch, len(self.waiting_queue))
|
|
if num_new_reqs <= 0:
|
|
return None
|
|
|
|
can_run_list = self.waiting_queue[:num_new_reqs]
|
|
self.waiting_queue = self.waiting_queue[num_new_reqs:]
|
|
|
|
for req in can_run_list:
|
|
req.init_next_round_input(self.tree_cache)
|
|
|
|
set_time_batch(can_run_list, "set_forward_entry_time")
|
|
|
|
# construct a schedule batch with those requests and mark as decode
|
|
new_batch = ScheduleBatch.init_new(
|
|
can_run_list,
|
|
self.req_to_token_pool,
|
|
self.token_to_kv_pool_allocator,
|
|
self.tree_cache,
|
|
self.model_config,
|
|
self.enable_overlap,
|
|
self.spec_algorithm,
|
|
)
|
|
|
|
# construct fake completed prefill
|
|
new_batch.prepare_for_prebuilt()
|
|
new_batch.process_prebuilt(self.server_args, self.future_map)
|
|
|
|
return new_batch
|
|
|
|
def process_decode_queue(self: Scheduler):
|
|
if self.server_args.disaggregation_decode_enable_offload_kvcache:
|
|
self.decode_offload_manager.check_offload_progress()
|
|
|
|
# try to resume retracted requests if there are enough space for another `num_reserved_decode_tokens` decode steps
|
|
resumed_reqs = self.disagg_decode_prealloc_queue.resume_retracted_reqs()
|
|
self.waiting_queue.extend(resumed_reqs)
|
|
if len(self.disagg_decode_prealloc_queue.retracted_queue) > 0:
|
|
# if there are still retracted requests, we do not allocate new requests
|
|
return
|
|
|
|
if not hasattr(self, "polling_count"):
|
|
self.polling_count = 0
|
|
self.polling_interval = (
|
|
self.server_args.disaggregation_decode_polling_interval
|
|
)
|
|
|
|
self.polling_count = (self.polling_count + 1) % self.polling_interval
|
|
|
|
if self.polling_count % self.polling_interval == 0:
|
|
req_conns, _ = self.disagg_decode_prealloc_queue.pop_preallocated()
|
|
self.disagg_decode_transfer_queue.extend(req_conns)
|
|
transferred_reqs = (
|
|
self.disagg_decode_transfer_queue.pop_transferred()
|
|
) # the requests which kv has arrived
|
|
self.waiting_queue.extend(transferred_reqs)
|