调度策略和负载均衡
有效的调度策略和负载均衡机制是PD分离系统高效运行的关键。本节详细讨论Prefill和Decode节点的选择策略以及动态负载均衡。
4.8.1 Prefill 节点选择策略
Prefill节点的选择直接影响TTFT和Prefill吞吐量。
选择因素
latex
┌─────────────────────────────────────────────────────────────────┐
│ Prefill 节点选择因素 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 1. 计算能力 (权重: 30%) │
│ • GPU类型 (A100/H100/H800) │
│ • GPU数量 │
│ • Tensor Parallel效率 │
│ │
│ 2. 当前负载 (权重: 25%) │
│ • GPU利用率 │
│ • 内存使用率 │
│ • 正在执行的Prefill数量 │
│ │
│ 3. 队列深度 (权重: 20%) │
│ • 等待队列长度 │
│ • 预估等待时间 │
│ │
│ 4. 网络位置 (权重: 15%) │
│ • 与Decode节点的网络距离 │
│ • 预估KV Cache传输时间 │
│ │
│ 5. 模型缓存 (权重: 10%) │
│ • 模型是否已加载 │
│ • 模型版本匹配 │
│ │
└─────────────────────────────────────────────────────────────────┘选择算法实现
python
class PrefillNodeSelector:
"""
Prefill 节点选择器
"""
def __init__(self, weights: Dict[str, float] = None):
# 默认权重配置
self.weights = weights or {
'compute': 0.30,
'load': 0.25,
'queue': 0.20,
'network': 0.15,
'model': 0.10,
}
def select_node(
self,
request: Request,
prefill_nodes: List[PrefillNode],
) -> PrefillNode:
"""
选择最优Prefill节点
"""
scores = []
for node in prefill_nodes:
if not node.is_healthy:
continue
score = self._compute_score(node, request)
scores.append((node, score))
if not scores:
raise NoAvailableNodeError(" 没有可用的Prefill节点")
# 选择得分最高的节点
best_node = max(scores, key=lambda x: x[1])[0]
return best_node
def _compute_score(
self,
node: PrefillNode,
request: Request,
) -> float:
""" 计算节点得分"""
scores = {
'compute': self._compute_score(node, request),
'load': self._load_score(node),
'queue': self._queue_score(node),
'network': self._network_score(node, request),
'model': self._model_score(node, request),
}
# 加权求和
total_score = sum(
scores[k] * self.weights[k] for k in scores
)
return total_score
def _compute_score(
self,
node: PrefillNode,
request: Request,
) -> float:
""" 计算能力得分"""
# 基于GPU类型和数量
gpu_scores = {
'H100': 1.0,
'H800': 0.95,
'A100-80GB': 0.75,
'A100-40GB': 0.65,
'A10': 0.30,
}
base_score = gpu_scores.get(node.gpu_type, 0.5)
# 考虑Tensor Parallel效率
tp_efficiency = 0.85 if node.tp_size > 1 else 1.0
return base_score * tp_efficiency * math.sqrt(node.gpu_count)
def _load_score(self, node: PrefillNode) -> float:
""" 负载得分 (越低越好)"""
# 基于GPU利用率和内存使用率
gpu_util = node.gpu_utilization
mem_util = node.memory_utilization
# 综合负载
load = 0.6 * gpu_util + 0.4 * mem_util
# 转换为得分 (1 - load)
return 1.0 - load
def _queue_score(self, node: PrefillNode) -> float:
""" 队列得分 (越短越好)"""
queue_depth = node.queue_depth
max_queue = node.max_queue_depth
# 归一化队列深度
normalized = min(queue_depth / max_queue, 1.0)
return 1.0 - normalized
def _network_score(
self,
node: PrefillNode,
request: Request,
) -> float:
""" 网络位置得分"""
# 获取目标Decode节点
decode_node = request.preferred_decode_node
if decode_node is None:
return 0.5 # 默认值
# 网络延迟
latency = self.network_topology.get_latency(node, decode_node)
bandwidth = self.network_topology.get_bandwidth(node, decode_node)
# 预估传输时间
kv_cache_size = self.estimate_kv_cache_size(request)
transfer_time = kv_cache_size / bandwidth + latency
# 转换为得分 (传输时间越短越好)
max_acceptable = 100 # 100ms
score = max(0, 1.0 - transfer_time / max_acceptable)
return score
def _model_score(
self,
node: PrefillNode,
request: Request,
) -> float:
""" 模型缓存得分"""
if node.has_model_cached(request.model_id):
return 1.0
elif node.can_load_model(request.model_id):
return 0.5
else:
return 0.04.8.2 Decode 节点选择策略
Decode节点的选择直接影响ITL和用户体验。
选择因素
latex
┌─────────────────────────────────────────────────────────────────┐
│ Decode 节点选择因素 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 1. 当前ITL (权重: 35%) │
│ • 当前批处理的ITL │
│ • ITL趋势（上升/下降） │
│ • ITL P99 │
│ │
│ 2. 内存容量 (权重: 25%) │
│ • 可用内存 │
│ • KV Cache存储能力 │
│ • 预估内存需求 vs 可用内存 │
│ │
│ 3. 批大小 (权重: 20%) │
│ • 当前批大小 │
│ • 批大小上限 │
│ • 批处理效率 │
│ │
│ 4. 内存带宽 (权重: 15%) │
│ • GPU内存带宽 │
│ • HBM带宽利用率 │
│ │
│ 5. 网络位置 (权重: 5%) │
│ • 与Prefill节点的距离 │
│ │
└─────────────────────────────────────────────────────────────────┘选择算法实现
python
class DecodeNodeSelector:
"""
Decode 节点选择器
"""
def __init__(self, weights: Dict[str, float] = None):
self.weights = weights or {
'itl': 0.35,
'memory': 0.25,
'batch': 0.20,
'bandwidth': 0.15,
'network': 0.05,
}
def select_node(
self,
request: Request,
kv_cache_size: int,
decode_nodes: List[DecodeNode],
) -> DecodeNode:
"""
选择最优Decode节点
"""
candidates = []
for node in decode_nodes:
# 首先检查内存容量
if node.available_memory < kv_cache_size * 1.2: # 20%裕量
continue
score = self._compute_score(node, request, kv_cache_size)
candidates.append((node, score))
if not candidates:
# 没有足够内存的节点，触发扩容
self.trigger_scale_up(kv_cache_size)
raise NoAvailableNodeError("没有可用的Decode节点")
# 选择得分最高的节点
best_node = max(candidates, key=lambda x: x[1])[0]
return best_node
def _compute_score(
self,
node: DecodeNode,
request: Request,
kv_cache_size: int,
) -> float:
""" 计算节点得分"""
scores = {
'itl': self._itl_score(node),
'memory': self._memory_score(node, kv_cache_size),
'batch': self._batch_score(node),
'bandwidth': self._bandwidth_score(node),
'network': self._network_score(node, request),
}
total_score = sum(
scores[k] * self.weights[k] for k in scores
)
return total_score
def _itl_score(self, node: DecodeNode) -> float:
"""ITL得分 (越低越好)"""
current_itl = node.current_itl
target_itl = 50 # 目标ITL 50ms
# 超过目标ITL的惩罚
if current_itl > target_itl:
return max(0, 1.0 - (current_itl - target_itl) / target_itl)
else:
# 低于目标ITL的奖励
return 1.0 + (target_itl - current_itl) / target_itl * 0.2
def _memory_score(
self,
node: DecodeNode,
kv_cache_size: int,
) -> float:
""" 内存得分"""
available = node.available_memory
# 内存充足度
ratio = available / kv_cache_size
if ratio < 1.2:
return 0.0 # 内存不足
elif ratio < 2.0:
return (ratio - 1.2) / 0.8 * 0.5 # 0.0 - 0.5
else:
return 0.5 + min(0.5, (ratio - 2.0) / 4.0 * 0.5) # 0.5 - 1.0
def _batch_score(self, node: DecodeNode) -> float:
""" 批大小得分"""
current = node.current_batch_size
max_size = node.max_batch_size
# 理想批大小: 70% of max
ideal = max_size * 0.7
# 距离理想的差距
diff = abs(current - ideal) / ideal
return max(0, 1.0 - diff)4.8.3 动态负载均衡
动态负载均衡确保系统在高负载下仍能保持稳定的性能。
负载均衡策略
python
class DynamicLoadBalancer:
"""
动态负载均衡器
根据系统负载动态调整请求分布
"""
def __init__(self):
self.prefill_nodes = []
self.decode_nodes = []
self.load_history = deque(maxlen=100)
async def balance_load(self):
"""
执行负载均衡
策略:
1. 监控各节点负载
2. 检测负载不均
3. 触发请求迁移
"""
while True:
# 收集负载信息
prefill_loads = await self._collect_prefill_loads()
decode_loads = await self._collect_decode_loads()
# 检测负载不均
if self._detect_imbalance(prefill_loads):
await self._rebalance_prefill(prefill_loads)
if self._detect_imbalance(decode_loads):
await self._rebalance_decode(decode_loads)
# 等待下一次检查
await asyncio.sleep(5) # 5 秒检查一次
def _detect_imbalance(self, loads: List[float]) -> bool:
""" 检测负载是否不均"""
if len(loads) < 2:
return False
avg_load = sum(loads) / len(loads)
max_load = max(loads)
min_load = min(loads)
# 负载差异超过30%认为不均衡
if max_load - min_load > 0.3:
return True
# 或存在节点过载 (>90%)
if max_load > 0.9:
return True
return False
async def _rebalance_prefill(
self,
loads: List[Tuple[PrefillNode, float]],
):
"""重新均衡Prefill负载"""
# 找到最忙和最闲的节点
sorted_nodes = sorted(loads, key=lambda x: x[1], reverse=True)
busiest = sorted_nodes[0]
idlest = sorted_nodes[-1]
# 如果差异足够大，迁移请求
if busiest[1] - idlest[1] > 0.3:
# 从 busiest 迁移部分请求到 idlest
requests_to_migrate = self._select_migratable_requests(
busiest[0],
count=5
)
for request in requests_to_migrate:
await self._migrate_prefill_request(request, idlest[0])
async def _migrate_prefill_request(
self,
request: Request,
target_node: PrefillNode,
):
""" 迁移Prefill请求"""
# 1. 暂停原节点上的请求
source_node = request.assigned_prefill_node
await source_node.pause_request(request)
# 2. 传输中间状态（如果有）
if request.has_partial_kv_cache:
kv_cache = await source_node.extract_kv_cache(request)
await self.transfer_engine.send(kv_cache, target_node.addr)
# 3. 在目标节点恢复执行
await target_node.resume_request(request)自适应批处理
python
class AdaptiveBatcher:
"""
自适应批处理器
根据当前负载动态调整批处理策略
"""
def __init__(self):
self.target_itl = 50 # 目标ITL
self.target_ttft = 200 # 目标TTFT
def adapt_batch_size(self, node: Node, metrics: Metrics):
"""
自适应调整批大小
策略:
- ITL过高 → 减小批大小
- ITL过低且有积压 → 增大批大小
"""
current_itl = metrics.itl_p99
queue_depth = metrics.queue_depth
if current_itl > self.target_itl * 1.2:
# ITL过高，减小批大小
new_batch_size = int(node.max_batch_size * 0.9)
logger.info(f"Reducing batch size to {new_batch_size} due to high ITL")
elif current_itl < self.target_itl * 0.8 and queue_depth > 10:
# ITL较低且有积压，增大批大小
new_batch_size = min(
int(node.max_batch_size * 1.1),
node.hard_max_batch_size
)
logger.info(f"Increasing batch size to {new_batch_size}")
else:
# 保持当前批大小
new_batch_size = node.max_batch_size
node.max_batch_size = new_batch_size