diff --git a/lightllm/server/httpserver_for_pd_master/pd_selector/cache_aware.py b/lightllm/server/httpserver_for_pd_master/pd_selector/cache_aware.py index 8bdf9b73b1..ef417679be 100644 --- a/lightllm/server/httpserver_for_pd_master/pd_selector/cache_aware.py +++ b/lightllm/server/httpserver_for_pd_master/pd_selector/cache_aware.py @@ -11,8 +11,8 @@ - 树中的 prefill_node 对应 worker.connection_key，可区分同一服务的不同 DP 连接； - prompt 会按 sample_stride 抽稀后再插入/匹配，降低树的深度与内存； - 根据推理侧返回的平均 prompt cache 命中率，动态调整 cache 亲和与负载均衡的权重； - - 优先使用 dispatched_req_num 为 0 的空闲节点，避免 GPU 闲置； - - 所有节点都忙时，用 worker.dispatched_prompt_chars 做负载均衡。 + - 长且高匹配的前缀按未缓存后缀估算新增工作量，在负载阈值内保留缓存亲和； + - 其余请求优先使用空闲节点；所有节点都忙时，用 dispatched_prompt_chars 做负载均衡。 选点流程见 CacheAwarePolicy.select_worker。 """ @@ -54,6 +54,9 @@ class CacheAwareConfig: evict_node_batch: int = 10_000 # 每隔 sample_stride 个字符抽 1 个作为前缀树 key，降低匹配开销与内存。 sample_stride: int = 512 + # 达到两个门槛的历史前缀可以优先于空闲节点，但仍受 balance_rel_threshold 限制。 + long_prefix_min_chars: int = 65536 + long_prefix_min_hit_rate: float = 0.8 class BalanceRelThresholdController: @@ -128,9 +131,9 @@ async def select_worker( 2) 对 request_text 做前缀匹配，计算 match_rate = matched_char_count / input_char_count； 3) match_rate > cache_threshold 且命中 prefill_node 仍在线 -> 得到 cache 命中节点； - 4) 存在 dispatched_req_num 为 0 的空闲节点 -> 强制从空闲节点中选择； - 多个节点空闲时优先选择 cache 命中节点； - 5) cache 命中节点负载未严重高于最空闲节点 -> 选择 cache 命中节点； + 4) 命中长且高匹配前缀时，按未缓存后缀估算缓存节点的新增工作量； + 其余请求优先选择空闲节点，多个节点空闲时优先选择 cache 命中节点； + 5) cache 命中节点的预计负载未超过阈值 -> 选择 cache 命中节点； 6) 未命中或负载严重失衡 -> 选择最空闲节点； """ if not workers: @@ -139,17 +142,30 @@ async def select_worker( raise ValueError(f"request_text length must be > 1, got {len(request_text)}") cache_worker, cache_match = await self._get_cache_worker(workers, request_text) + request_load = len(request_text) + matched_chars = min(request_load, max(0, cache_match.matched_char_count)) + long_prefix_match = ( + cache_worker is not None + and matched_chars >= self.config.long_prefix_min_chars + and matched_chars / request_load >= self.config.long_prefix_min_hit_rate + ) - # ---- 1. 空闲优先：避免有可用 GPU 闲置 ---- + # 长前缀允许与空闲节点比较预计负载；历史位置仅是缓存提示，不保证 KV 仍驻留。 if cache_worker is not None and cache_worker.dispatched_req_num == 0: idle_worker = cache_worker + elif long_prefix_match: + idle_worker = None else: idle_worker = self._select_idle_worker(workers) if idle_worker is not None: selected_worker = idle_worker else: - # ---- 2. 所有节点都忙时，在 cache 亲和与负载均衡之间权衡 ---- - selected_worker = self._select_worker_by_cache_and_load(workers, cache_worker, len(request_text)) + selected_worker = self._select_worker_by_cache_and_load( + workers=workers, + cache_worker=cache_worker, + request_load=request_load, + cached_chars=matched_chars if long_prefix_match else 0, + ) selection_extra_info = PDSelectionExtraInfo() if cache_match.prefill_node == selected_worker.connection_key and cache_match.input_char_count != 0: @@ -209,14 +225,15 @@ def _select_worker_by_cache_and_load( workers: List[PD_Client_Obj], cache_worker: Optional[PD_Client_Obj], request_load: int, + cached_chars: int = 0, ) -> PD_Client_Obj: - """所有节点都忙时，在 cache 亲和与 prompt 负载之间选择节点。""" + """比较在途 prompt 负载与新增工作量；长前缀可抵扣缓存节点的新增成本。""" least_loaded_worker = min(workers, key=lambda worker: worker.dispatched_prompt_chars) least_projected_load = least_loaded_worker.dispatched_prompt_chars + request_load cache_projected_load = None cache_worker_is_overloaded = False if cache_worker is not None: - cache_projected_load = cache_worker.dispatched_prompt_chars + request_load + cache_projected_load = cache_worker.dispatched_prompt_chars + max(1, request_load - cached_chars) cache_worker_is_overloaded = cache_projected_load > least_projected_load * self.config.balance_rel_threshold if cache_worker is None or cache_worker_is_overloaded: diff --git a/unit_tests/server/test_pd_cache_aware.py b/unit_tests/server/test_pd_cache_aware.py index 1b111122fb..70f2756eb4 100644 --- a/unit_tests/server/test_pd_cache_aware.py +++ b/unit_tests/server/test_pd_cache_aware.py @@ -319,3 +319,85 @@ def test_cache_aware_matches_cache_only_within_idle_workers(): selected_worker, _ = asyncio.run(policy.select_worker([other_idle_worker, busy_worker, cache_idle_worker], prompt)) assert selected_worker is cache_idle_worker + + +@pytest.mark.parametrize("other_requests", [0, 1]) +@pytest.mark.parametrize("cached_load,expect_cache", [(80000, True), (250000, False)]) +def test_long_prefix_selection_accounts_for_suffix_and_still_sheds_load( + other_requests: int, cached_load: int, expect_cache: bool +) -> None: + selector = LoadBalancedCacheAwareSelector(pd_manager=None) + cached = _worker(address="cached", dispatched_prompt_chars=cached_load, dispatched_req_num=1) + other = _worker(address="other", dispatched_prompt_chars=other_requests * 1000, dispatched_req_num=other_requests) + decode = SimpleNamespace(client_ip_port="decode", run_status=SimpleNamespace(total_token_usage_rate=0.0)) + selector.update_nodes(prefill_nodes=[other, cached], decode_nodes=[decode]) + prefix = "a" * 90001 + prompt = prefix + "b" * 10000 + _insert_prompt(policy=selector.policy, prompt=prefix, worker=cached) + + selected, selected_decode, info = asyncio.run( + main=selector.select_p_d_node(prompt=prompt, sampling_params=None, multimodal_params=None) + ) + + assert selected is (cached if expect_cache else other) + assert selected_decode is decode + if expect_cache: + match = _prefix_match(policy=selector.policy, prompt=prompt) + assert info == PDSelectionExtraInfo( + estimated_cache_hit_rate=match.matched_char_count / len(prompt), + cache_last_insert_time=match.last_insert_time, + ) + else: + assert info == PDSelectionExtraInfo() + + +@pytest.mark.parametrize( + "prefix_chars,prompt_chars,cached_load,balance_threshold,expect_cache", + [ + (16, 20, 1, 1.5, True), # Both long-prefix gates are inclusive. + (15, 18, 1, 1.5, False), # High hit rate, but the matched prefix is too short. + (16, 21, 1, 1.5, False), # Long prefix, but the hit rate is too low. + (16, 20, 16, 1.0, True), + (16, 20, 17, 1.0, False), + (16, 20, 26, 1.5, True), + (16, 20, 27, 1.5, False), + (16, 20, 36, 2.0, True), + (16, 20, 37, 2.0, False), + (20, 20, 20, 1.0, False), # A full prefix match still has nonzero work. + ], +) +def test_long_prefix_affinity_respects_hit_and_load_boundaries( + prefix_chars: int, prompt_chars: int, cached_load: int, balance_threshold: float, expect_cache: bool +) -> None: + policy = _local_policy( + config=CacheAwareConfig(sample_stride=1, long_prefix_min_chars=16, balance_rel_threshold=balance_threshold) + ) + cached = _worker(address="cached", dispatched_prompt_chars=cached_load, dispatched_req_num=1) + idle = _worker(address="idle") + prefix = "a" * prefix_chars + _insert_prompt(policy=policy, prompt=prefix, worker=cached) + + selected, _ = asyncio.run( + main=policy.select_worker(workers=[idle, cached], request_text=prefix + "b" * (prompt_chars - prefix_chars)) + ) + + assert selected is (cached if expect_cache else idle) + + +@pytest.mark.parametrize("available", [False, True]) +def test_long_prefix_affinity_requires_an_eligible_online_cache_worker(available: bool) -> None: + policy = _local_policy(config=CacheAwareConfig(sample_stride=1, long_prefix_min_chars=16, cache_threshold=0.9)) + cached = _worker(address="cached", dispatched_prompt_chars=1, dispatched_req_num=1) + idle = _worker(address="idle") + # The offline case has a full hit; the online case fails the configured cache threshold. + prefix = "a" * (20 if not available else 16) + _insert_prompt(policy=policy, prompt=prefix, worker=cached) + + selected, info = asyncio.run( + main=policy.select_worker( + workers=[idle, cached] if available else [idle], request_text=prefix + "b" * (20 - len(prefix)) + ) + ) + + assert selected is idle + assert info == PDSelectionExtraInfo()