[PD] Add decode PP event loop for PD disaggregation (#14945)
Signed-off-by: Shangming Cai <csmthu@gmail.com> Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
@@ -25,7 +25,7 @@ import time
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from http import HTTPStatus
|
||||
from typing import TYPE_CHECKING, List, Optional, Type, Union
|
||||
from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union
|
||||
|
||||
import torch
|
||||
from torch.distributed import ProcessGroup
|
||||
@@ -199,6 +199,7 @@ class DecodePreallocQueue:
|
||||
bootstrap_port: int,
|
||||
max_total_num_tokens: int,
|
||||
prefill_pp_size: int,
|
||||
pp_rank: int,
|
||||
num_reserved_decode_tokens: int,
|
||||
transfer_backend: TransferBackend,
|
||||
):
|
||||
@@ -220,6 +221,7 @@ class DecodePreallocQueue:
|
||||
self.bootstrap_port = bootstrap_port
|
||||
self.max_total_num_tokens = max_total_num_tokens
|
||||
self.prefill_pp_size = prefill_pp_size
|
||||
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
|
||||
@@ -236,8 +238,7 @@ class DecodePreallocQueue:
|
||||
kv_args.engine_rank = self.tp_rank % (attn_tp_size)
|
||||
|
||||
kv_args.decode_tp_size = attn_tp_size
|
||||
# Note(shangming): pp is not supported on the decode side yet, so its rank is fixed to 0
|
||||
kv_args.pp_rank = 0
|
||||
kv_args.pp_rank = self.pp_rank
|
||||
kv_args.system_dp_rank = self.scheduler.dp_rank
|
||||
kv_args.prefill_pp_size = self.prefill_pp_size
|
||||
kv_data_ptrs, kv_data_lens, kv_item_lens = (
|
||||
@@ -302,6 +303,7 @@ class DecodePreallocQueue:
|
||||
return
|
||||
|
||||
if is_retracted:
|
||||
req.retraction_mb_id = None
|
||||
self.retracted_queue.append(req)
|
||||
else:
|
||||
if req.bootstrap_host == FAKE_BOOTSTRAP_HOST:
|
||||
@@ -340,7 +342,9 @@ class DecodePreallocQueue:
|
||||
for req in reqs:
|
||||
self.add(req, is_retracted=is_retracted)
|
||||
|
||||
def resume_retracted_reqs(self) -> List[Req]:
|
||||
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
|
||||
@@ -349,6 +353,9 @@ class DecodePreallocQueue:
|
||||
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:
|
||||
continue
|
||||
|
||||
if self.req_to_token_pool.available_size() <= 0:
|
||||
break
|
||||
|
||||
@@ -377,7 +384,9 @@ class DecodePreallocQueue:
|
||||
|
||||
return resumed_reqs
|
||||
|
||||
def _update_handshake_waiters(self) -> None:
|
||||
def _update_handshake_waiters(
|
||||
self, rids_to_check: Optional[List[str]] = None
|
||||
) -> None:
|
||||
if not self.queue:
|
||||
return
|
||||
|
||||
@@ -389,6 +398,9 @@ class DecodePreallocQueue:
|
||||
)
|
||||
|
||||
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:
|
||||
@@ -410,10 +422,13 @@ class DecodePreallocQueue:
|
||||
else:
|
||||
raise ValueError(f"Unexpected poll case: {poll}")
|
||||
|
||||
def pop_preallocated(self) -> List[DecodeRequest]:
|
||||
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._update_handshake_waiters()
|
||||
self._update_handshake_waiters(rids_to_check)
|
||||
|
||||
failed_reqs = []
|
||||
preallocated_reqs = []
|
||||
indices_to_remove = set()
|
||||
|
||||
@@ -428,14 +443,20 @@ class DecodePreallocQueue:
|
||||
)
|
||||
# First, remove all failed requests from the queue
|
||||
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:
|
||||
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)
|
||||
indices_to_remove.add(i)
|
||||
|
||||
# Then, preallocate the remaining requests if possible
|
||||
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:
|
||||
continue
|
||||
|
||||
if i in indices_to_remove:
|
||||
continue
|
||||
|
||||
@@ -544,7 +565,7 @@ class DecodePreallocQueue:
|
||||
entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
|
||||
]
|
||||
|
||||
return preallocated_reqs
|
||||
return preallocated_reqs, failed_reqs
|
||||
|
||||
@property
|
||||
def num_tokens_pre_allocated(self):
|
||||
@@ -726,7 +747,7 @@ class DecodeTransferQueue:
|
||||
)
|
||||
decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter()
|
||||
|
||||
def pop_transferred(self) -> List[Req]:
|
||||
def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req]:
|
||||
if not self.queue:
|
||||
return []
|
||||
polls = poll_and_all_reduce(
|
||||
@@ -736,6 +757,8 @@ class DecodeTransferQueue:
|
||||
transferred_reqs = []
|
||||
indices_to_remove = set()
|
||||
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.Failed:
|
||||
error_message = f"Decode transfer failed for request rank={self.tp_rank} {decode_req.req.rid=} {decode_req.req.bootstrap_room=}"
|
||||
try:
|
||||
@@ -958,7 +981,7 @@ class SchedulerDisaggregationDecodeMixin:
|
||||
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()
|
||||
req_conns, _ = self.disagg_decode_prealloc_queue.pop_preallocated()
|
||||
self.disagg_decode_transfer_queue.extend(req_conns)
|
||||
alloc_reqs = (
|
||||
self.disagg_decode_transfer_queue.pop_transferred()
|
||||
|
||||
@@ -807,13 +807,12 @@ class MooncakeKVManager(CommonKVManager):
|
||||
executor,
|
||||
)
|
||||
|
||||
if self.pp_group.is_last_rank:
|
||||
# Only the last chunk we need to send the aux data
|
||||
ret = self.send_aux(
|
||||
req,
|
||||
kv_chunk.prefill_aux_index,
|
||||
target_rank_registration_info.dst_aux_ptrs,
|
||||
)
|
||||
# Only the last chunk we need to send the aux data
|
||||
ret = self.send_aux(
|
||||
req,
|
||||
kv_chunk.prefill_aux_index,
|
||||
target_rank_registration_info.dst_aux_ptrs,
|
||||
)
|
||||
polls.append(True if ret == 0 else False)
|
||||
dst_ranks_infos.append(
|
||||
(req.endpoint, req.dst_port, req.room)
|
||||
|
||||
@@ -677,6 +677,7 @@ class Req:
|
||||
|
||||
# The number of times this request has been retracted / preempted.
|
||||
self.retraction_count = 0
|
||||
self.retraction_mb_id = None
|
||||
|
||||
# For metrics
|
||||
self.metrics_collector = metrics_collector
|
||||
|
||||
@@ -901,6 +901,7 @@ class Scheduler(
|
||||
bootstrap_port=self.server_args.disaggregation_bootstrap_port,
|
||||
max_total_num_tokens=self.max_total_num_tokens,
|
||||
prefill_pp_size=self.server_args.disaggregation_prefill_pp,
|
||||
pp_rank=self.pp_rank,
|
||||
num_reserved_decode_tokens=self.server_args.num_reserved_decode_tokens,
|
||||
transfer_backend=self.transfer_backend,
|
||||
)
|
||||
@@ -2771,7 +2772,10 @@ def run_scheduler_process(
|
||||
if scheduler.enable_overlap:
|
||||
scheduler.event_loop_overlap_disagg_decode()
|
||||
else:
|
||||
scheduler.event_loop_normal_disagg_decode()
|
||||
if server_args.pp_size > 1:
|
||||
scheduler.event_loop_pp_disagg_decode()
|
||||
else:
|
||||
scheduler.event_loop_normal_disagg_decode()
|
||||
|
||||
except Exception:
|
||||
traceback = get_exception_traceback()
|
||||
|
||||
@@ -512,12 +512,13 @@ class SchedulerPPMixin:
|
||||
# send ready PP output to rank 0
|
||||
if mbs[next_first_rank_mb_id] is not None:
|
||||
q_event, pp_outputs_to_send = last_rank_comm_queue.popleft()
|
||||
torch.cuda.current_stream().wait_event(q_event)
|
||||
with torch.profiler.record_function("send_res_dict_to_next_stage"):
|
||||
send_output_work = self._pp_send_dict_to_next_stage(
|
||||
pp_outputs_to_send.tensors,
|
||||
async_send=True,
|
||||
)
|
||||
if not mbs[next_first_rank_mb_id].forward_mode.is_prebuilt():
|
||||
torch.cuda.current_stream().wait_event(q_event)
|
||||
with torch.profiler.record_function("send_res_dict_to_next_stage"):
|
||||
send_output_work = self._pp_send_dict_to_next_stage(
|
||||
pp_outputs_to_send.tensors,
|
||||
async_send=True,
|
||||
)
|
||||
# send the outputs from the last round to let the next stage worker run post processing
|
||||
if not self.pp_group.is_last_rank:
|
||||
if pp_outputs:
|
||||
@@ -549,14 +550,19 @@ class SchedulerPPMixin:
|
||||
|
||||
if mbs[next_mb_id] is not None:
|
||||
with torch.profiler.record_function("recv_res_dict_from_prev_stage"):
|
||||
next_pp_outputs = PPProxyTensors(self._pp_recv_dict_from_prev_stage())
|
||||
with self.copy_stream_ctx:
|
||||
self.copy_stream.wait_stream(self.default_stream)
|
||||
batch_result = self._pp_prep_batch_result(
|
||||
mbs[next_mb_id], mb_metadata[next_mb_id], next_pp_outputs
|
||||
)
|
||||
d2h_event = torch.cuda.Event()
|
||||
d2h_event.record(torch.cuda.current_stream())
|
||||
next_pp_outputs = None
|
||||
if not mbs[next_mb_id].forward_mode.is_prebuilt():
|
||||
next_pp_outputs = PPProxyTensors(
|
||||
self._pp_recv_dict_from_prev_stage()
|
||||
)
|
||||
if not mbs[next_mb_id].forward_mode.is_prebuilt():
|
||||
with self.copy_stream_ctx:
|
||||
self.copy_stream.wait_stream(self.default_stream)
|
||||
batch_result = self._pp_prep_batch_result(
|
||||
mbs[next_mb_id], mb_metadata[next_mb_id], next_pp_outputs
|
||||
)
|
||||
d2h_event = torch.cuda.Event()
|
||||
d2h_event.record(torch.cuda.current_stream())
|
||||
|
||||
return next_pp_outputs, batch_result, d2h_event, send_output_work
|
||||
|
||||
@@ -588,19 +594,21 @@ class SchedulerPPMixin:
|
||||
)
|
||||
return result, event
|
||||
|
||||
def get_rids(self: Scheduler, req_queue: List[Req], *poll_statuses_group):
|
||||
def get_rids(
|
||||
self: Scheduler, req_queue: List[Req], is_send: bool, *poll_statuses_group
|
||||
):
|
||||
"""
|
||||
Used by PP, get the required rids with the given poll statuses.
|
||||
"""
|
||||
polls = poll_and_all_reduce(
|
||||
[req.disagg_kv_sender for req in req_queue],
|
||||
[req.disagg_kv_sender if is_send else req.kv_receiver for req in req_queue],
|
||||
self.tp_worker.get_attention_tp_cpu_group(),
|
||||
)
|
||||
rids: List = []
|
||||
for poll_statuses in poll_statuses_group:
|
||||
rids.append(
|
||||
[
|
||||
req.rid
|
||||
req.rid if is_send else req.req.rid
|
||||
for req, poll in zip(req_queue, polls)
|
||||
if poll in poll_statuses
|
||||
]
|
||||
@@ -760,6 +768,7 @@ class SchedulerPPMixin:
|
||||
# First rank, pop the bootstrap reqs from the bootstrap queue
|
||||
good_bootstrapped_rids, bad_bootstrapped_rids = self.get_rids(
|
||||
self.disagg_prefill_bootstrap_queue.queue,
|
||||
True,
|
||||
[KVPoll.WaitingForInput],
|
||||
[KVPoll.Failed],
|
||||
)
|
||||
@@ -771,6 +780,7 @@ class SchedulerPPMixin:
|
||||
)
|
||||
curr_good_bootstrapped_rids, curr_bad_bootstrapped_rids = self.get_rids(
|
||||
self.disagg_prefill_bootstrap_queue.queue,
|
||||
True,
|
||||
[KVPoll.WaitingForInput],
|
||||
[KVPoll.Failed],
|
||||
)
|
||||
@@ -782,11 +792,12 @@ class SchedulerPPMixin:
|
||||
)
|
||||
return [good_bootstrapped_rids, bad_bootstrapped_rids]
|
||||
|
||||
def _pp_pd_get_transferred_ids(self: Scheduler):
|
||||
def _pp_pd_get_prefill_transferred_ids(self: Scheduler):
|
||||
# get the current stage transfer success
|
||||
if self.pp_group.is_first_rank:
|
||||
transferred_rids = self.get_rids(
|
||||
self.disagg_prefill_inflight_queue,
|
||||
True,
|
||||
[KVPoll.Success, KVPoll.Failed],
|
||||
)
|
||||
# if other ranks, do intersection with the previous rank's transferred rids
|
||||
@@ -797,6 +808,7 @@ class SchedulerPPMixin:
|
||||
# 2. get the current stage's transferred reqs info
|
||||
curr_transferred_rids = self.get_rids(
|
||||
self.disagg_prefill_inflight_queue,
|
||||
True,
|
||||
[KVPoll.Success, KVPoll.Failed],
|
||||
)
|
||||
# 3. new consensus rids = intersection(previous consensus rids, transfer finished rids)
|
||||
@@ -938,7 +950,7 @@ class SchedulerPPMixin:
|
||||
bmbs[mb_id] = bootstrapped_rids
|
||||
self._pp_commit_comm_work(send_bootstrapped_work)
|
||||
|
||||
transferred_rids = self._pp_pd_get_transferred_ids()
|
||||
transferred_rids = self._pp_pd_get_prefill_transferred_ids()
|
||||
self._pp_commit_comm_work(send_transfer_work)
|
||||
tmbs[mb_id] = transferred_rids
|
||||
|
||||
@@ -1052,3 +1064,341 @@ class SchedulerPPMixin:
|
||||
self.check_tree_cache()
|
||||
self.new_token_ratio = self.init_new_token_ratio
|
||||
self.maybe_sleep_on_idle()
|
||||
|
||||
def _pp_pd_get_retract_ids(self: Scheduler, mb_id: int):
|
||||
# communicate pre-consensus retracted reqs
|
||||
for req in self.disagg_decode_prealloc_queue.retracted_queue:
|
||||
# assign retracted reqs to the current microbatch
|
||||
if req.retraction_mb_id is None:
|
||||
req.retraction_mb_id = mb_id
|
||||
curr_retract_rids = [
|
||||
req.rid
|
||||
for req in self.disagg_decode_prealloc_queue.retracted_queue
|
||||
if req.retraction_mb_id == mb_id
|
||||
]
|
||||
if self.pp_group.is_first_rank:
|
||||
# First rank, get all retracted req ids for the microbatch
|
||||
return curr_retract_rids
|
||||
else:
|
||||
# Other ranks, receive the retracted reqs info from the previous rank and ensure the consensus
|
||||
prev_retract_rids = self._pp_recv_pyobj_from_prev_stage()
|
||||
return list(set(prev_retract_rids) & set(curr_retract_rids))
|
||||
|
||||
def _pp_pd_get_prealloc_ids(self: Scheduler):
|
||||
# communicate pre-consensus prealloc reqs
|
||||
if self.pp_group.is_first_rank:
|
||||
# First rank, pop the preallocated reqs from the prealloc queue
|
||||
good_prealloc_rids, bad_prealloc_rids = self.get_rids(
|
||||
self.disagg_decode_prealloc_queue.queue,
|
||||
False,
|
||||
[KVPoll.WaitingForInput],
|
||||
[KVPoll.Failed],
|
||||
)
|
||||
else:
|
||||
# Other ranks, receive the preallocated reqs info from the previous rank and ensure the consensus
|
||||
prev_prealloc_rids = self._pp_recv_pyobj_from_prev_stage()
|
||||
prev_good_prealloc_rids, prev_bad_prealloc_rids = prev_prealloc_rids
|
||||
curr_good_prealloc_rids, curr_bad_prealloc_rids = self.get_rids(
|
||||
self.disagg_decode_prealloc_queue.queue,
|
||||
False,
|
||||
[KVPoll.WaitingForInput],
|
||||
[KVPoll.Failed],
|
||||
)
|
||||
good_prealloc_rids = list(
|
||||
set(prev_good_prealloc_rids) & set(curr_good_prealloc_rids)
|
||||
)
|
||||
bad_prealloc_rids = list(
|
||||
set(prev_bad_prealloc_rids) | set(curr_bad_prealloc_rids)
|
||||
)
|
||||
return [good_prealloc_rids, bad_prealloc_rids]
|
||||
|
||||
def _pp_pd_get_decode_transferred_ids(self: Scheduler):
|
||||
# get the current stage transfer success
|
||||
if self.pp_group.is_first_rank:
|
||||
transferred_rids = self.get_rids(
|
||||
self.disagg_decode_transfer_queue.queue,
|
||||
False,
|
||||
[KVPoll.Success, KVPoll.Failed],
|
||||
)
|
||||
# if other ranks, do intersection with the previous rank's transferred rids
|
||||
else:
|
||||
# 2 (Release): Receive the transferred rids from the previous rank
|
||||
# 1. recv previous stage's transferred reqs info
|
||||
prev_transferred_rids = self._pp_recv_pyobj_from_prev_stage()
|
||||
# 2. get the current stage's transferred reqs info
|
||||
curr_transferred_rids = self.get_rids(
|
||||
self.disagg_decode_transfer_queue.queue,
|
||||
False,
|
||||
[KVPoll.Success, KVPoll.Failed],
|
||||
)
|
||||
# 3. new consensus rids = intersection(previous consensus rids, transfer finished rids)
|
||||
transferred_rids = list(
|
||||
set(prev_transferred_rids) & set(curr_transferred_rids)
|
||||
)
|
||||
return transferred_rids
|
||||
|
||||
def process_retract_queue(self: Scheduler, retract_rids: Optional[List[str]]):
|
||||
if retract_rids is not None:
|
||||
# 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(
|
||||
retract_rids
|
||||
)
|
||||
self.waiting_queue.extend(resumed_reqs)
|
||||
return [req.rid for req in resumed_reqs]
|
||||
return None
|
||||
|
||||
def process_prealloc_queue(self: Scheduler, prealloc_rids: Optional[List[str]]):
|
||||
if len(self.disagg_decode_prealloc_queue.retracted_queue) > 0:
|
||||
# if there are still retracted requests, we do not allocate new requests
|
||||
return [[], []]
|
||||
|
||||
if prealloc_rids is not None:
|
||||
(
|
||||
good_consensus_prealloc_rids,
|
||||
bad_consensus_prealloc_rids,
|
||||
) = prealloc_rids
|
||||
good_reqs, failed_reqs = self.disagg_decode_prealloc_queue.pop_preallocated(
|
||||
rids_to_check=good_consensus_prealloc_rids
|
||||
+ bad_consensus_prealloc_rids,
|
||||
)
|
||||
self.disagg_decode_transfer_queue.extend(good_reqs)
|
||||
return [
|
||||
[req.req.rid for req in good_reqs],
|
||||
[req.req.rid for req in failed_reqs],
|
||||
]
|
||||
return None
|
||||
|
||||
def process_decode_transfer_queue(
|
||||
self: Scheduler, release_rids: Optional[List[str]]
|
||||
):
|
||||
if release_rids is not None:
|
||||
released_reqs = self.disagg_decode_transfer_queue.pop_transferred(
|
||||
release_rids
|
||||
)
|
||||
self.waiting_queue.extend(released_reqs)
|
||||
return [req.rid for req in released_reqs]
|
||||
return None
|
||||
|
||||
@DynamicGradMode()
|
||||
def event_loop_pp_disagg_decode(self: Scheduler):
|
||||
self.pp_loop_size: int = self.pp_size + self.server_args.pp_async_batch_depth
|
||||
mbs = [None] * self.pp_loop_size
|
||||
last_mbs = [None] * self.pp_loop_size
|
||||
self.running_mbs = [
|
||||
ScheduleBatch(reqs=[], batch_is_full=False)
|
||||
for _ in range(self.pp_loop_size)
|
||||
]
|
||||
mb_metadata: List[Optional[PPBatchMetadata]] = [None] * self.pp_loop_size
|
||||
pp_outputs: Optional[PPProxyTensors] = None
|
||||
last_rank_comm_queue: deque[Tuple[torch.cuda.Event, PPProxyTensors]] = deque()
|
||||
|
||||
# PD additional
|
||||
|
||||
# consensus rids
|
||||
consensus_retract_rids: Optional[List[str]] = None
|
||||
consensus_prealloc_rids: Optional[List[str]] = None
|
||||
release_rids: Optional[List[str]] = None # consensus transferred rids
|
||||
|
||||
rmbs = [None] * self.pp_loop_size
|
||||
pmbs = [None] * self.pp_loop_size
|
||||
tmbs = [None] * self.pp_loop_size
|
||||
|
||||
send_req_work = []
|
||||
|
||||
# send info to reach consensus
|
||||
send_retract_work = []
|
||||
send_prealloc_work = []
|
||||
send_transfer_work = []
|
||||
|
||||
# send consensus info
|
||||
send_consensus_retract_work = []
|
||||
send_consensus_prealloc_work = []
|
||||
send_release_work = []
|
||||
|
||||
send_proxy_work = []
|
||||
send_output_work = []
|
||||
|
||||
while True:
|
||||
server_is_idle = True
|
||||
for mb_id in range(self.pp_loop_size):
|
||||
self.running_batch = self.running_mbs[mb_id]
|
||||
self.last_batch = last_mbs[mb_id]
|
||||
next_first_rank_mb_id = (mb_id + self.pp_size) % self.pp_loop_size
|
||||
next_mb_id = (mb_id + 1) % self.pp_loop_size
|
||||
|
||||
next_pp_outputs = None
|
||||
next_consensus_retract_rids = None
|
||||
next_consensus_prealloc_rids = None
|
||||
next_release_rids = None
|
||||
d2h_event = None
|
||||
next_batch_result = None
|
||||
|
||||
recv_reqs = self.recv_requests()
|
||||
self.process_input_requests(recv_reqs)
|
||||
|
||||
if not self.pp_group.is_last_rank:
|
||||
self._pp_commit_comm_work(send_req_work)
|
||||
|
||||
# reaching consensus through PP ranks
|
||||
retract_rids = self._pp_pd_get_retract_ids(mb_id)
|
||||
rmbs[mb_id] = retract_rids
|
||||
self._pp_commit_comm_work(send_retract_work)
|
||||
|
||||
prealloc_rids = self._pp_pd_get_prealloc_ids()
|
||||
pmbs[mb_id] = prealloc_rids
|
||||
self._pp_commit_comm_work(send_prealloc_work)
|
||||
|
||||
transferred_rids = self._pp_pd_get_decode_transferred_ids()
|
||||
tmbs[mb_id] = transferred_rids
|
||||
self._pp_commit_comm_work(send_transfer_work)
|
||||
|
||||
# get batch to run and proxy tensors if needed
|
||||
batch = self.get_next_disagg_decode_batch_to_run()
|
||||
mbs[mb_id] = batch
|
||||
self.running_mbs[mb_id] = self.running_batch
|
||||
|
||||
self.cur_batch: Optional[ScheduleBatch] = mbs[mb_id]
|
||||
if self.cur_batch:
|
||||
server_is_idle = False
|
||||
pp_proxy_tensors = None
|
||||
if not self.cur_batch.forward_mode.is_prebuilt():
|
||||
pp_proxy_tensors = self._pp_recv_proxy_tensors()
|
||||
|
||||
# early send output if possible
|
||||
if self.server_args.pp_async_batch_depth > 0:
|
||||
self._pp_commit_comm_work(work=send_output_work)
|
||||
next_pp_outputs, next_batch_result, d2h_event, send_output_work = (
|
||||
self._pp_send_recv_and_preprocess_output_tensors(
|
||||
next_first_rank_mb_id,
|
||||
next_mb_id,
|
||||
mbs,
|
||||
mb_metadata,
|
||||
last_rank_comm_queue,
|
||||
pp_outputs,
|
||||
)
|
||||
)
|
||||
self._pp_commit_comm_work(send_proxy_work)
|
||||
# run batch
|
||||
if self.cur_batch:
|
||||
result, event = self._pp_launch_batch(
|
||||
mb_id, pp_proxy_tensors, mb_metadata, last_rank_comm_queue
|
||||
)
|
||||
# regular send output
|
||||
if self.server_args.pp_async_batch_depth == 0:
|
||||
self._pp_commit_comm_work(work=send_output_work)
|
||||
next_pp_outputs, next_batch_result, d2h_event, send_output_work = (
|
||||
self._pp_send_recv_and_preprocess_output_tensors(
|
||||
next_first_rank_mb_id,
|
||||
next_mb_id,
|
||||
mbs,
|
||||
mb_metadata,
|
||||
last_rank_comm_queue,
|
||||
pp_outputs,
|
||||
)
|
||||
)
|
||||
|
||||
# reach consensus on last rank and send to PP=0
|
||||
# otherwise, just pass along previous consensus
|
||||
send_consensus_retract_work, consensus_retract_rids = (
|
||||
self._pp_pd_send_consensus_bootstrapped_ids( # reuse the function
|
||||
rmbs,
|
||||
next_first_rank_mb_id,
|
||||
consensus_retract_rids,
|
||||
retract_rids,
|
||||
)
|
||||
)
|
||||
|
||||
send_consensus_prealloc_work, consensus_prealloc_rids = (
|
||||
self._pp_pd_send_consensus_bootstrapped_ids( # reuse the function
|
||||
pmbs,
|
||||
next_first_rank_mb_id,
|
||||
consensus_prealloc_rids,
|
||||
prealloc_rids,
|
||||
)
|
||||
)
|
||||
|
||||
send_release_work, release_rids = (
|
||||
self._pp_pd_send_consensus_release_ids(
|
||||
tmbs, next_first_rank_mb_id, release_rids, transferred_rids
|
||||
)
|
||||
)
|
||||
|
||||
if self.server_args.disaggregation_decode_enable_offload_kvcache:
|
||||
self.decode_offload_manager.check_offload_progress()
|
||||
|
||||
if rmbs[next_mb_id] is not None:
|
||||
next_consensus_retract_rids = self._pp_recv_pyobj_from_prev_stage()
|
||||
next_consensus_retract_rids = self.process_retract_queue(
|
||||
next_consensus_retract_rids
|
||||
)
|
||||
self._pp_commit_comm_work(send_consensus_retract_work)
|
||||
|
||||
if pmbs[next_mb_id] is not None:
|
||||
next_consensus_prealloc_rids = self._pp_recv_pyobj_from_prev_stage()
|
||||
next_consensus_prealloc_rids = self.process_prealloc_queue(
|
||||
next_consensus_prealloc_rids
|
||||
)
|
||||
self._pp_commit_comm_work(send_consensus_prealloc_work)
|
||||
|
||||
if tmbs[next_mb_id] is not None:
|
||||
next_release_rids = self._pp_recv_pyobj_from_prev_stage()
|
||||
next_release_rids = self.process_decode_transfer_queue(
|
||||
next_release_rids
|
||||
)
|
||||
self._pp_commit_comm_work(send_release_work)
|
||||
|
||||
# post-process the coming microbatch
|
||||
if mbs[next_mb_id] is not None:
|
||||
if not mbs[next_mb_id].forward_mode.is_prebuilt():
|
||||
d2h_event.synchronize()
|
||||
self._pp_process_batch_result(
|
||||
mbs[next_mb_id],
|
||||
next_batch_result,
|
||||
)
|
||||
last_mbs[next_mb_id] = mbs[next_mb_id]
|
||||
|
||||
if not self.pp_group.is_last_rank:
|
||||
send_req_work = self._pp_send_pyobj_to_next_stage(
|
||||
recv_reqs, async_send=True
|
||||
)
|
||||
send_retract_work = self._pp_send_pyobj_to_next_stage(
|
||||
retract_rids, async_send=True
|
||||
)
|
||||
send_prealloc_work = self._pp_send_pyobj_to_next_stage(
|
||||
prealloc_rids, async_send=True
|
||||
)
|
||||
send_transfer_work = self._pp_send_pyobj_to_next_stage(
|
||||
transferred_rids, async_send=True
|
||||
)
|
||||
if self.cur_batch and not self.cur_batch.forward_mode.is_prebuilt():
|
||||
torch.cuda.current_stream().wait_event(event)
|
||||
send_proxy_work = self._pp_send_dict_to_next_stage(
|
||||
result.pp_hidden_states_proxy_tensors.tensors,
|
||||
async_send=True,
|
||||
)
|
||||
|
||||
if hasattr(self, "delayed_weight_sync_fn"):
|
||||
self.delayed_weight_sync_fn()
|
||||
self.delayed_weight_sync_fn = None
|
||||
|
||||
pp_outputs = next_pp_outputs
|
||||
release_rids = next_release_rids
|
||||
consensus_retract_rids = next_consensus_retract_rids
|
||||
consensus_prealloc_rids = next_consensus_prealloc_rids
|
||||
|
||||
self.running_batch.batch_is_full = False
|
||||
|
||||
# When the server is idle, self-check and re-init some states
|
||||
queue_size = (
|
||||
len(self.waiting_queue)
|
||||
+ len(self.disagg_decode_transfer_queue.queue)
|
||||
+ len(self.disagg_decode_prealloc_queue.queue)
|
||||
)
|
||||
if self.server_args.disaggregation_decode_enable_offload_kvcache:
|
||||
queue_size += len(self.decode_offload_manager.ongoing_offload)
|
||||
|
||||
if server_is_idle and queue_size == 0:
|
||||
self.check_memory()
|
||||
self.check_tree_cache()
|
||||
self.new_token_ratio = self.init_new_token_ratio
|
||||
self.maybe_sleep_on_idle()
|
||||
|
||||
Reference in New Issue
Block a user