diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index e3dfdd8b4..2f8a42756 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -36,6 +36,7 @@ from sglang.srt.managers.io_struct import ( TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, WatchLoadUpdateReq, + DpRequestInfoqOutput, ) from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.scheduler import run_scheduler_process @@ -57,6 +58,9 @@ from sglang.srt.utils.network import NetworkAddress, bind_port, get_zmq_socket from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter from sglang.srt.utils.watchdog import Watchdog from sglang.utils import TypeBasedDispatcher, get_exception_traceback +from sglang.srt.disaggregation.utils import ( + DisaggregationMode +) logger = logging.getLogger(__name__) @@ -77,35 +81,137 @@ class LoadBalanceMethod(Enum): except KeyError as exc: raise ValueError(f"Invalid load balance method: {method}") from exc - +import random class DPBudget: - def __init__(self, dp_size: int): + def __init__(self, server_args: ServerArgs): + dp_size = server_args.dp_size self.dp_size = dp_size self.total_requests = [0] * dp_size self.total_tokens = [0] * dp_size + self.ts_tic = [0.0] * dp_size + self.pool_tokens = [0] * dp_size + self.req_sum_tokens = [0] * dp_size + self.consume_token_num_per_rank = 4000 + self.disaggregation_mode = DisaggregationMode( + server_args.disaggregation_mode + ) + + # 每个 rank 一个账本，按 timestamp 排序 + self.reqs_info = [list() for _ in range(dp_size)] + + def _sort_rank_ledger(self, rank: int): + self.reqs_info[rank].sort(key=lambda x: x.timestamp) + + def _append_local_req(self, rank: int, req: Req): + if req is None: + return + + rid = getattr(req, "rid", None) + if rid is None: + return + + num_tokens = len(req.input_ids) if getattr(req, "input_ids", None) is not None else 0 + completion_tokens = ( + len(req.output_ids) if getattr(req, "output_ids", None) is not None else 0 + ) + cached_tokens = getattr(req, "cached_tokens", 0) + + info = DpRequestInfoqOutput( + rid=rid, + timestamp=time.perf_counter(), + dp_rank=rank, + num_tokens=num_tokens, + cached_tokens=cached_tokens, + completion_tokens=completion_tokens, + ) + self.reqs_info[rank].append(info) + + def _merge_rank_reqs(self, rank: int, incoming_reqs: List[DpRequestInfoqOutput]): + if incoming_reqs is None: + return + + if len(incoming_reqs) == 0: + self.reqs_info[rank] = [] + return + + incoming_reqs = sorted(incoming_reqs, key=lambda x: x.timestamp) + earliest_ts = incoming_reqs[0].timestamp + + # 删除本地账本中比 incoming_reqs 中最早 timestamp 更旧的请求 + local_reqs = [x for x in self.reqs_info[rank] if x.timestamp >= earliest_ts] + + # 按 rid 合并 + merged = {x.rid: x for x in local_reqs} + for x in incoming_reqs: + merged[x.rid] = x + + self.reqs_info[rank] = sorted(merged.values(), key=lambda x: x.timestamp) def update_budget(self, load_update: WatchLoadUpdateReq): - """Update the budget.""" for load in load_update.loads: - self.total_requests[load.dp_rank] = load.num_reqs - self.total_tokens[load.dp_rank] = load.num_tokens + if load is None: + continue - def dispatch(self, method: LoadBalanceMethod): + rank = load.dp_rank + if rank is None or rank >= self.dp_size: + continue + + self.total_requests[rank] = load.num_reqs + self.total_tokens[rank] = load.num_tokens + self.ts_tic[rank] = load.ts_tic + self.pool_tokens[rank] = load.pool_tokens + if self.disaggregation_mode == DisaggregationMode.PREFILL: + self._merge_rank_reqs(rank, load.reqs_list) + else: + self._merge_rank_reqs(rank, load.reqs_list) + self.req_sum_tokens[rank] = sum(x.num_tokens for x in self.reqs_info[rank])+self.pool_tokens[rank] + + def dispatch(self, method: LoadBalanceMethod, req: Req = None): if method == LoadBalanceMethod.TOTAL_REQUESTS: - target_rank = self.total_requests.index(min(self.total_requests)) + min_req = min(self.total_requests) + candidates = [i for i, x in enumerate(self.total_requests) if x == min_req] + target_rank = random.choice(candidates) + elif method == LoadBalanceMethod.TOTAL_TOKENS: - # Use total_requests as a tie-breaker when total_tokens are equal - target_rank = min( - range(self.dp_size), - key=lambda i: (self.total_tokens[i], self.total_requests[i]), + if self.disaggregation_mode == DisaggregationMode.PREFILL: + pairs = [ + (self.total_tokens[i], self.total_requests[i]) + for i in range(self.dp_size) + ] + min_pair = min(pairs) + candidates = [ + i for i in range(self.dp_size) + if (self.total_tokens[i], self.total_requests[i]) == min_pair + ] + target_rank = random.choice(candidates) + min_tokens = self.total_tokens[target_rank] + else: + pairs = [ + (self.total_tokens[i], self.total_requests[i]) + for i in range(self.dp_size) + ] + min_pair = min(pairs) + candidates = [ + i for i in range(self.dp_size) + if (self.total_tokens[i], self.total_requests[i]) == min_pair + ] + target_rank = random.choice(candidates) + min_tokens = self.total_tokens[target_rank] + + logger.info( + f"Dispatching to DP rank {target_rank} with rough_num_tokens={min_tokens} and req_sum_tokens={self.req_sum_tokens[target_rank]}" ) + else: return None - # Increment the load of that worker by one as a heuristic self.total_requests[target_rank] += 1 - return target_rank + if req is not None: + self.total_tokens[target_rank] += len(req.input_ids) + self.req_sum_tokens[target_rank] += len(req.input_ids) + self._append_local_req(target_rank, req) + return target_rank class DataParallelController: """A controller that dispatches requests to multiple data parallel workers.""" @@ -145,7 +251,7 @@ class DataParallelController: self.dispatching = dispatch_lookup[self.load_balance_method] # Load balance budget - self.dp_budget = DPBudget(server_args.dp_size) + self.dp_budget = DPBudget(server_args) # To protect changing env vars to set CUDA_VISIBLE_DEVICES. self.env_lock = threading.Lock() @@ -557,7 +663,7 @@ class DataParallelController: def total_tokens_scheduler(self, req: Req): if self.maybe_external_dp_rank_routing(req): return - target_worker = self.dp_budget.dispatch(LoadBalanceMethod.TOTAL_TOKENS) + target_worker = self.dp_budget.dispatch(LoadBalanceMethod.TOTAL_TOKENS,req=req) self.workers[target_worker].send_pyobj(req) def event_loop(self): diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index d9c97a0b4..eea640d50 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1803,6 +1803,13 @@ class BlockReqInput(BaseReq): class GetLoadReqInput(BaseReq): pass +@dataclass +class DpRequestInfoqOutput(BaseReq): + timestamp: float + dp_rank: int + num_tokens: int + cached_tokens: int + completion_tokens: int @dataclass class GetLoadReqOutput(BaseReq): @@ -1811,6 +1818,9 @@ class GetLoadReqOutput(BaseReq): num_waiting_reqs: int num_tokens: int ts_tic: float + reqs_list: List[DpRequestInfoqOutput] + pool_tokens: int + @dataclass diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 2d3e579bb..1075b0fc1 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -641,8 +641,12 @@ class HiRadixCache(RadixCache): assert len(node.host_value) > 0 self.ongoing_write_through[node.id] = node if not write_back: - # no need to lock nodes if write back - self.inc_lock_ref(node) + # Only lock the specific node being written through, not the + # entire path to root. Ancestors cannot be evicted while this + # node holds a device value (_collect_leaves_device skips + # parents with non-evicted children), so path locking is + # unnecessary and would over-reduce evictable_size_. + self.inc_node_lock_ref(node) else: return 0 @@ -710,7 +714,7 @@ class HiRadixCache(RadixCache): finish_event.synchronize() for ack_id in ack_list: backuped_node = self.ongoing_write_through.pop(ack_id) - self.dec_lock_ref(backuped_node) + self.dec_node_lock_ref(backuped_node) if self.enable_storage: self.write_backup_storage(backuped_node) finish_count -= 1 diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 501fc4223..1860fb817 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -648,6 +648,37 @@ class RadixCache(BasePrefixCache): node = node.parent return DecLockRefResult(delta=delta) + def inc_node_lock_ref(self, node: TreeNode): + """Increment lock_ref on a single node without walking to root. + + Used by hicache write-through to protect only the specific node being + backed up, avoiding unnecessary locking of ancestor nodes. Ancestors + cannot be evicted anyway while this node has a device value because + _collect_leaves_device() skips parents with non-evicted children. + """ + if self.disable: + return + if node == self.root_node: + return + if node.lock_ref == 0: + self.evictable_size_ -= len(node.key) + self.protected_size_ += len(node.key) + node.lock_ref += 1 + + def dec_node_lock_ref(self, node: TreeNode): + """Decrement lock_ref on a single node without walking to root. + + Counterpart of inc_node_lock_ref for releasing write-through locks. + """ + if self.disable: + return + if node == self.root_node: + return + if node.lock_ref == 1: + self.evictable_size_ += len(node.key) + self.protected_size_ -= len(node.key) + node.lock_ref -= 1 + def evictable_size(self): return self.evictable_size_ diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index 052c23645..dac256b1a 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -14,6 +14,7 @@ from sglang.srt.managers.io_struct import ( DisaggregationMetrics, GetLoadReqInput, GetLoadReqOutput, + DpRequestInfoqOutput, GetLoadsReqInput, GetLoadsReqOutput, LoRAMetrics, @@ -678,6 +679,40 @@ class SchedulerMetricsMixin: / self.stats.max_running_requests_under_SLO, self.stats.token_usage / 0.9, ) + def _build_request_info(self, req): + prompt_tokens = 0 + if getattr(req, "origin_input_ids_unpadded", None) is not None: + prompt_tokens = len(req.origin_input_ids_unpadded) + elif getattr(req, "origin_input_ids", None) is not None: + prompt_tokens = len(req.origin_input_ids) + + completion_tokens = ( + len(req.output_ids) if getattr(req, "output_ids", None) is not None else 0 + ) + + num_tokens = getattr(req, "seqlen", None) + if num_tokens is None: + # logger.info("num tokens is None change to prmpt token and completion") + num_tokens = prompt_tokens + completion_tokens + + return DpRequestInfoqOutput( + rid=req.rid, + dp_rank=getattr(req, "dp_rank", -1), + num_tokens=num_tokens, + cached_tokens=getattr(req, "cached_tokens", 0), + completion_tokens=completion_tokens, + timestamp=time.perf_counter(), + ) + def _collect_request_infos(self, queues,request_infos): + + # all queues + for queue in queues: + for req in queue: + rid = getattr(req, "rid", None) + if rid is None: + # logger.info("rid is None, skipping request") + continue + request_infos.append(self._build_request_info(req)) def get_load(self: Scheduler, _: GetLoadReqInput = None) -> GetLoadReqOutput: if self.is_hybrid_swa: @@ -687,7 +722,7 @@ class SchedulerMetricsMixin: num_tokens = self._get_mamba_token_info()[0] else: num_tokens = self._get_token_info()[0] - + pool_tokens=num_tokens # Tokens in waiting queue, bootstrap queue, prealloc queue waiting_queues = [self.waiting_queue] if self.disaggregation_mode == DisaggregationMode.PREFILL: @@ -696,16 +731,20 @@ class SchedulerMetricsMixin: waiting_queues.append(self.disagg_decode_prealloc_queue.queue) waiting_queues.append(self.disagg_decode_transfer_queue.queue) waiting_queues.append(self.disagg_decode_prealloc_queue.retracted_queue) - num_tokens += sum(req.seqlen for queue in waiting_queues for req in queue) num_waiting_reqs = sum(len(queue) for queue in waiting_queues) + request_infos=[] + self._collect_request_infos(waiting_queues,request_infos) + # logger.info("current dp rank %d num tokens %d waiting queues %d",self.dp_rank,num_tokens,num_waiting_reqs) return GetLoadReqOutput( dp_rank=self.dp_rank, num_reqs=len(self.running_batch.reqs) + num_waiting_reqs, num_waiting_reqs=num_waiting_reqs, num_tokens=num_tokens, + pool_tokens=pool_tokens, ts_tic=time.perf_counter(), + reqs_list=request_infos, ) def get_loads(self: Scheduler, req: GetLoadsReqInput = None) -> GetLoadsReqOutput: