[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:
Kevin Li
2025-12-12 04:15:00 -08:00
committed by GitHub
parent c8cf1cafdb
commit 8fa8d9d7e8
5 changed files with 414 additions and 37 deletions

View File

@@ -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()

View File

@@ -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)

View File

@@ -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

View File

@@ -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()

View File

@@ -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()