chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

fix(dp_ctrl): fix imbalanced DP rank assignment

Browse Source 
1. feat:update total tokens immediately

fea:record request in da controller local
This commit is contained in:
laoyao0822
2026-03-26 00:12:46 +08:00
committed by wxiwnd
parent cc11cac77c
commit d9b6e90b35
5 changed files with 210 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

136
python/sglang/srt/managers/data_parallel_controller.py
View File

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

10
python/sglang/srt/managers/io_struct.py
View File

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

10
python/sglang/srt/mem_cache/hiradix_cache.py
View File

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

31
python/sglang/srt/mem_cache/radix_cache.py
View File

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

43
python/sglang/srt/observability/scheduler_metrics_mixin.py
View File

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