perf(disaggregation): compact decode queues in one pass
This commit is contained in:
wxiwnd
@@ -433,14 +433,17 @@ class DecodePreallocQueue:
|
||||
|
||||
# allocate memory
|
||||
resumed_reqs = []
|
||||
indices_to_remove = set()
|
||||
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 = (
|
||||
@@ -449,10 +452,11 @@ class DecodePreallocQueue:
|
||||
+ 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)
|
||||
indices_to_remove.add(i)
|
||||
req.is_retracted = False
|
||||
self._pre_alloc(req)
|
||||
allocatable_tokens -= required_tokens_for_request
|
||||
@@ -460,11 +464,7 @@ class DecodePreallocQueue:
|
||||
# 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 = [
|
||||
entry
|
||||
for i, entry in enumerate(self.retracted_queue)
|
||||
if i not in indices_to_remove
|
||||
]
|
||||
self.retracted_queue = remaining_retracted_queue
|
||||
|
||||
return resumed_reqs
|
||||
|
||||
@@ -605,7 +605,8 @@ class DecodePreallocQueue:
|
||||
|
||||
failed_reqs = []
|
||||
preallocated_reqs = []
|
||||
indices_to_remove = set()
|
||||
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.
|
||||
@@ -616,33 +617,35 @@ class DecodePreallocQueue:
|
||||
allocatable_tokens = self._allocatable_tokens(
|
||||
retractable_tokens=retractable_tokens, count_retracted=True
|
||||
)
|
||||
# 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:
|
||||
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)
|
||||
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:
|
||||
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:
|
||||
break
|
||||
remaining_queue.append(decode_req)
|
||||
prealloc_blocked = True
|
||||
continue
|
||||
|
||||
if self.req_to_metadata_buffer_idx_allocator.available_size() <= 0:
|
||||
break
|
||||
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
|
||||
@@ -663,9 +666,13 @@ class DecodePreallocQueue:
|
||||
)
|
||||
> allocatable_tokens
|
||||
):
|
||||
break
|
||||
remaining_queue.append(decode_req)
|
||||
prealloc_blocked = True
|
||||
continue
|
||||
if required_tokens_for_request > allocatable_tokens:
|
||||
break
|
||||
remaining_queue.append(decode_req)
|
||||
prealloc_blocked = True
|
||||
continue
|
||||
|
||||
allocatable_tokens -= required_tokens_for_request
|
||||
self._pre_alloc(decode_req.req)
|
||||
@@ -727,12 +734,9 @@ class DecodePreallocQueue:
|
||||
page_indices, decode_req.metadata_buffer_index, state_indices
|
||||
)
|
||||
preallocated_reqs.append(decode_req)
|
||||
indices_to_remove.add(i)
|
||||
decode_req.req.time_stats.set_decode_transfer_queue_entry_time()
|
||||
|
||||
self.queue = [
|
||||
entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
|
||||
]
|
||||
self.queue = remaining_queue
|
||||
|
||||
return preallocated_reqs, failed_reqs
|
||||
|
||||
@@ -998,9 +1002,11 @@ class DecodeTransferQueue:
|
||||
)
|
||||
|
||||
transferred_reqs = []
|
||||
indices_to_remove = set()
|
||||
for i, (decode_req, poll) in enumerate(zip(self.queue, polls)):
|
||||
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=}"
|
||||
@@ -1019,14 +1025,14 @@ class DecodeTransferQueue:
|
||||
)
|
||||
# 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)
|
||||
indices_to_remove.add(i)
|
||||
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:
|
||||
indices_to_remove.add(i)
|
||||
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(
|
||||
@@ -1039,23 +1045,26 @@ class DecodeTransferQueue:
|
||||
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,
|
||||
]:
|
||||
pass
|
||||
remaining_queue.append(decode_req)
|
||||
else:
|
||||
raise ValueError(f"Unexpected poll case: {poll}")
|
||||
|
||||
for i in indices_to_remove:
|
||||
idx = self.queue[i].metadata_buffer_index
|
||||
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 = [
|
||||
entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
|
||||
]
|
||||
self.queue = remaining_queue
|
||||
|
||||
return transferred_reqs
|
||||
|
||||
|
||||
Reference in New Issue
Block a user