feat(pd): prioritize idle prefill nodes

lightllm/server/api_start.py

@@ -256,6 +256,12 @@ def _launch_subprocesses(args: StartArgs):
         per_dp_cache_size = max(1, math.ceil(args.running_max_req_size / dp_size_in_node) * 2)
         args.linear_att_cache_size = min(default_cache_size, per_dp_cache_size)
 
+    if args.run_mode == "decode":
+        # PD Decode 节点只接收 prompt 末尾位置的 linear attention state，不具备
+        # 中间大页边界对应的 state。因此 Decode 节点必须使用默认值关闭大页功能，
+        # 避免请求释放时将不完整的大页 state 写入 radix cache 并触发断言。
+        args.linear_att_page_block_num = 10000000
+
     if args.enable_cpu_cache and is_linear_att_mixed_model(args.model_dir):
         args.cpu_cache_token_page_size = args.linear_att_hash_page_size * args.linear_att_page_block_num
         logger.info(f"set cpu_cache_token_page_size to {args.cpu_cache_token_page_size} for linear hybrid att model")

lightllm/server/httpserver_for_pd_master/manager.py

@@ -211,19 +211,15 @@ async def _generate_one(
         block_group_request_id = origin_request_id
         p_node = None
         d_node = None
-        prefill_load_released = False
+        pending_prefill_load_chars = None
 
         try:
             p_node, d_node = await self.select_p_d_node(prompt, origin_sampling_params, multimodal_params)
-            # 记录当前 P 节点的在途 prompt 负载，首 token 生成后即可释放。
-            p_node.dispatched_prompt_chars += len(prompt)
-
-            history_gen_token_strs = []
-
             if not p_node or not d_node:
                 logger.error(f"{origin_request_id}: No p_node or d_node found")
                 raise Exception(f"{origin_request_id}: No p_node or d_node found")
 
+            history_gen_token_strs = []
             origin_prompt_cache_len = None
 
             for iter_index, block_max_new_tokens in enumerate(max_new_tokens_list):
@@ -233,11 +229,17 @@ async def _generate_one(
                 logger.info(f"pd log gen sub req id {block_group_request_id} for main req id {origin_request_id}")
                 sampling_params.max_new_tokens = block_max_new_tokens
 
+                # 分段请求始终复用循环外选定的 P 节点；这里只按每段实际发送的
+                # prompt 更新该节点的在途 prefill 负载，不会重新选点。
+                block_prompt = prompt + "".join(history_gen_token_strs)
+                pending_prefill_load_chars = len(block_prompt)
+                p_node.dispatched_prompt_chars += pending_prefill_load_chars
+                p_node.dispatched_req_num += 1
                 results_generator = self._wait_to_token_package(
                     p_node,
                     d_node,
                     start_time,
-                    prompt + "".join(history_gen_token_strs),
+                    block_prompt,
                     sampling_params,
                     multimodal_params,
                     request,
@@ -254,10 +256,15 @@ async def _generate_one(
                     metadata["prompt_tokens"] = prompt_tokens
                     if iter_index == 0 and origin_prompt_cache_len is None:
                         origin_prompt_cache_len = metadata.get("prompt_cache_len", 0)
+                        prompt_cache_hit_rate = origin_prompt_cache_len / max(prompt_tokens, 1)
+                        self.pd_manager.selector.record_prompt_cache_hit_rate(prompt_cache_hit_rate)
                     metadata["prompt_cache_len"] = origin_prompt_cache_len or 0
-                    if not prefill_load_released:
-                        p_node.dispatched_prompt_chars = max(0, p_node.dispatched_prompt_chars - len(prompt))
-                        prefill_load_released = True
+                    if pending_prefill_load_chars is not None:
+                        p_node.dispatched_prompt_chars = max(
+                            0, p_node.dispatched_prompt_chars - pending_prefill_load_chars
+                        )
+                        p_node.dispatched_req_num = max(0, p_node.dispatched_req_num - 1)
+                        pending_prefill_load_chars = None
                     yield origin_request_id, request_output, metadata, finish_status
 
                 await self.remove_req(group_request_id=block_group_request_id)
@@ -277,8 +284,9 @@ async def _generate_one(
             raise e
 
         finally:
-            if p_node is not None and not prefill_load_released:
-                p_node.dispatched_prompt_chars = max(0, p_node.dispatched_prompt_chars - len(prompt))
+            if p_node is not None and pending_prefill_load_chars is not None:
+                p_node.dispatched_prompt_chars = max(0, p_node.dispatched_prompt_chars - pending_prefill_load_chars)
+                p_node.dispatched_req_num = max(0, p_node.dispatched_req_num - 1)
             await self.remove_req(block_group_request_id)
         return
 

lightllm/server/httpserver_for_pd_master/pd_selector/cache_aware.py

@@ -10,7 +10,9 @@
   - 用前缀树（见 PromptCacheTree）记录「历史 prompt -> 处理它的 worker」；
   - 树中的 prefill_node 对应 worker.client_ip_port；
   - prompt 会按 sample_stride 抽稀后再插入/匹配，降低树的深度与内存；
-  - 用 worker.dispatched_prompt_chars（当前尚未产出首 token 的 prompt 字符数）做负载均衡。
+  - 根据推理侧返回的平均 prompt cache 命中率，动态调整 cache 亲和与负载均衡的权重；
+  - 优先使用 dispatched_req_num 为 0 的空闲节点，避免 GPU 闲置；
+  - 所有节点都忙时，用 worker.dispatched_prompt_chars 做负载均衡。
 
 选点流程见 CacheAwarePolicy.select_worker。
 """
@@ -36,7 +38,14 @@ class CacheAwareConfig:
     # 前缀匹配成功率阈值：matched_char_count / input_char_count 超过该值才路由到命中节点。
     cache_threshold: float = 0.5
     # cache 命中节点的在途量超过最空闲节点该倍数时，优先选择最空闲节点。
-    balance_rel_threshold: float = 1.8
+    balance_rel_threshold: float = 1.5
+    # 动态调整 balance_rel_threshold 时允许的上下限。
+    min_balance_rel_threshold: float = 1.0
+    max_balance_rel_threshold: float = 2.0
+    # 每轮用于统计平均 prompt cache 命中率的请求数。
+    cache_hit_rate_window_size: int = 1000
+    # 相邻统计周期命中率变化时，负载均衡阈值的调整步长。
+    balance_rel_threshold_step: float = 0.05
     # 前缀树允许的最大节点数（不含 root）。
     max_node_count: int = 1_000_000
     # 每次 LRU 驱逐的叶节点数量。
@@ -47,6 +56,41 @@ class CacheAwareConfig:
     recursion_limit: int = 4000
 
 
+class BalanceRelThresholdController:
+    """根据最近请求的 prompt cache 命中率动态调整负载均衡阈值。"""
+
+    def __init__(self) -> None:
+        self._cache_hit_rates = []
+        self._last_average_cache_hit_rate = None
+
+    def append(self, cache_hit_rate: float) -> None:
+        """追加一次真实 prompt cache 命中率。"""
+        cache_hit_rate = min(max(cache_hit_rate, 0.0), 1.0)
+        self._cache_hit_rates.append(cache_hit_rate)
+
+    def update_config(self, config: CacheAwareConfig) -> None:
+        """每收集一个统计窗口，根据命中率趋势调整负载均衡阈值。"""
+        if len(self._cache_hit_rates) < config.cache_hit_rate_window_size:
+            return
+
+        average_cache_hit_rate = (
+            sum(self._cache_hit_rates[-config.cache_hit_rate_window_size :]) / config.cache_hit_rate_window_size
+        )
+        self._cache_hit_rates.clear()
+
+        if self._last_average_cache_hit_rate is not None:
+            if average_cache_hit_rate > self._last_average_cache_hit_rate:
+                config.balance_rel_threshold += config.balance_rel_threshold_step
+            elif average_cache_hit_rate < self._last_average_cache_hit_rate:
+                config.balance_rel_threshold -= config.balance_rel_threshold_step
+            config.balance_rel_threshold = min(
+                max(config.balance_rel_threshold, config.min_balance_rel_threshold),
+                config.max_balance_rel_threshold,
+            )
+
+        self._last_average_cache_hit_rate = average_cache_hit_rate
+
+
 class CacheAwarePolicy:
     """
     维护 prompt 前缀树，并据此为请求选择 prefill worker。
@@ -65,6 +109,7 @@ def __init__(self, config: Optional[CacheAwareConfig] = None) -> None:
             evict_node_batch=self.config.evict_node_batch,
             recursion_limit=self.config.recursion_limit,
         )
+        self.balance_rel_threshold_controller = BalanceRelThresholdController()
 
     def select_worker(self, workers: List[PD_Client_Obj], request_text: str) -> Optional[PD_Client_Obj]:
         """
@@ -79,19 +124,40 @@ def select_worker(self, workers: List[PD_Client_Obj], request_text: str) -> Opti
 
         决策顺序：
           1) workers 为空 -> 返回 None；
-          2) 对 request_text 做前缀匹配，计算
+          2) 存在 dispatched_req_num 为 0 的空闲节点 -> 强制从空闲节点中选择；
+             多个节点空闲时优先选择 cache 命中节点；
+          3) 对 request_text 做前缀匹配，计算
              match_rate = matched_char_count / input_char_count；
-          3) match_rate > cache_threshold 且命中 prefill_node 仍在线 -> 得到 cache 命中节点；
-          4) cache 命中节点负载未严重高于最空闲节点 -> 选择 cache 命中节点；
-          5) 未命中或负载严重失衡 -> 选择最空闲节点；
-          6) 将当前 prompt 与最终选中的节点写入前缀树。
+          4) match_rate > cache_threshold 且命中 prefill_node 仍在线 -> 得到 cache 命中节点；
+          5) cache 命中节点负载未严重高于最空闲节点 -> 选择 cache 命中节点；
+          6) 未命中或负载严重失衡 -> 选择最空闲节点；
+          7) 将当前 prompt 与最终选中的节点写入前缀树。
         """
         if not workers:
             return None
         if len(request_text) <= 1:
             raise ValueError(f"request_text length must be > 1, got {len(request_text)}")
 
-        # ---- 1. 前缀匹配：估计当前请求与历史请求的 cache 复用潜力 ----
+        # ---- 1. 空闲优先：避免有可用 GPU 闲置 ----
+        idle_worker = self._select_idle_worker(workers, request_text)
+        if idle_worker is not None:
+            self.prompt_cache_tree.insert(request_text, idle_worker.client_ip_port)
+            return idle_worker
+
+        # ---- 2. 所有节点都忙时，在 cache 亲和与负载均衡之间权衡 ----
+        cache_worker = self._get_cache_worker(workers, request_text)
+        selected_worker = self._select_worker_by_cache_and_load(workers, cache_worker, len(request_text))
+
+        self.prompt_cache_tree.insert(request_text, selected_worker.client_ip_port)
+        return selected_worker
+
+    def record_prompt_cache_hit_rate(self, cache_hit_rate: float) -> None:
+        """记录推理侧上报的真实 cache 命中率，并更新动态负载阈值。"""
+        self.balance_rel_threshold_controller.append(cache_hit_rate)
+        self.balance_rel_threshold_controller.update_config(self.config)
+
+    def _get_cache_worker(self, workers: List[PD_Client_Obj], request_text: str) -> Optional[PD_Client_Obj]:
+        """在指定候选节点中返回达到匹配阈值的 cache 节点。"""
         result = self.prompt_cache_tree.prefix_match(request_text)
         match_rate = 0.0 if result.input_char_count == 0 else result.matched_char_count / result.input_char_count
 
@@ -102,16 +168,41 @@ def select_worker(self, workers: List[PD_Client_Obj], request_text: str) -> Opti
             f"prefill_node={result.prefill_node}"
         )
 
-        cache_worker: Optional[PD_Client_Obj] = None
-        if match_rate > self.config.cache_threshold and result.prefill_node is not None:
-            for worker in workers:
-                if worker.client_ip_port == result.prefill_node:
-                    cache_worker = worker
-                    break
+        if match_rate <= self.config.cache_threshold or result.prefill_node is None:
+            return None
+
+        for worker in workers:
+            if worker.client_ip_port == result.prefill_node:
+                return worker
+        return None
+
+    def _select_idle_worker(self, workers: List[PD_Client_Obj], request_text: str) -> Optional[PD_Client_Obj]:
+        """优先选择空闲节点；多个空闲节点之间优先复用 cache。"""
+        idle_workers = [worker for worker in workers if worker.dispatched_req_num == 0]
+        if not idle_workers:
+            return None
+
+        cache_worker = self._get_cache_worker(idle_workers, request_text) if len(idle_workers) > 1 else None
+        selected_worker = cache_worker or min(
+            idle_workers,
+            key=lambda worker: (worker.dispatched_prompt_chars, worker.client_ip_port),
+        )
+        logger.info(
+            f"CacheAwarePolicy: select idle worker, idle_worker_num={len(idle_workers)}, "
+            f"cache_worker={cache_worker.client_ip_port if cache_worker else None}, "
+            f"balance_rel_threshold={self.config.balance_rel_threshold:.4f}, "
+            f"selected_worker={selected_worker.client_ip_port}"
+        )
+        return selected_worker
 
-        # ---- 2. 负载配平：cache 节点过载时选择当前在途量最少的节点 ----
+    def _select_worker_by_cache_and_load(
+        self,
+        workers: List[PD_Client_Obj],
+        cache_worker: Optional[PD_Client_Obj],
+        request_load: int,
+    ) -> PD_Client_Obj:
+        """所有节点都忙时，在 cache 亲和与 prompt 负载之间选择节点。"""
         least_loaded_worker = min(workers, key=lambda worker: worker.dispatched_prompt_chars)
-        request_load = len(request_text)
         least_projected_load = least_loaded_worker.dispatched_prompt_chars + request_load
         cache_projected_load = None
         cache_worker_is_overloaded = False
@@ -135,6 +226,4 @@ def select_worker(self, workers: List[PD_Client_Obj], request_text: str) -> Opti
             f"cache_worker_is_overloaded={cache_worker_is_overloaded}, "
             f"selected_worker={selected_worker.client_ip_port}"
         )
-
-        self.prompt_cache_tree.insert(request_text, selected_worker.client_ip_port)
         return selected_worker

lightllm/server/httpserver_for_pd_master/pd_selector/pd_selector.py

@@ -26,6 +26,10 @@ def select_p_d_node(
     ) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
         raise NotImplementedError("Subclass must implement this method")
 
+    def record_prompt_cache_hit_rate(self, cache_hit_rate: float) -> None:
+        """记录推理侧返回的 prompt cache 命中率；非 cache-aware 策略无需处理。"""
+        return
+
 
 class RandomSelector(PDSelector):
     """随机选择器"""
@@ -93,3 +97,6 @@ def select_p_d_node(
         )
 
         return p_node, d_node
+
+    def record_prompt_cache_hit_rate(self, cache_hit_rate: float) -> None:
+        self.policy.record_prompt_cache_hit_rate(cache_hit_rate)

lightllm/server/pd_io_struct.py

@@ -58,6 +58,8 @@ class PD_Client_Obj:
     run_status: _PD_Client_RunStatus = field(default_factory=_PD_Client_RunStatus)
     # cache-aware 选点用：当前派发到该节点且尚未产出首 token 的 prompt 字符数。
     dispatched_prompt_chars: int = 0
+    # 当前派发到该节点且尚未产出首 token 的请求数。
+    dispatched_req_num: int = 0
 
     def __post_init__(self):
         if self.mode not in ["prefill", "decode"]: