diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 04e0187452..e898e439c9 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -57,7 +57,7 @@ def make_argument_parser() -> argparse.ArgumentParser: "--select_p_d_node_strategy", type=str, default="round_robin", - choices=["random", "round_robin", "adaptive_load"], + choices=["random", "round_robin", "adaptive_load", "cache_aware"], help="pd master use this strategy to select p d node, can be round_robin, random or adaptive_load", ) parser.add_argument( diff --git a/lightllm/server/httpserver_for_pd_master/pd_selector/__init__.py b/lightllm/server/httpserver_for_pd_master/pd_selector/__init__.py index a48024a39d..d2c9506396 100644 --- a/lightllm/server/httpserver_for_pd_master/pd_selector/__init__.py +++ b/lightllm/server/httpserver_for_pd_master/pd_selector/__init__.py @@ -1,4 +1,10 @@ -from .pd_selector import PDSelector, RandomSelector, RoundRobinSelector, AdaptiveLoadSelector +from .pd_selector import ( + PDSelector, + RandomSelector, + RoundRobinSelector, + AdaptiveLoadSelector, + LoadBalancedCacheAwareSelector, +) def create_selector(selector_type: str, pd_manager) -> PDSelector: @@ -8,5 +14,7 @@ def create_selector(selector_type: str, pd_manager) -> PDSelector: return RoundRobinSelector(pd_manager) elif selector_type == "adaptive_load": return AdaptiveLoadSelector(pd_manager) + elif selector_type == "cache_aware": + return LoadBalancedCacheAwareSelector(pd_manager) else: raise ValueError(f"Invalid selector type: {selector_type}") 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 new file mode 100644 index 0000000000..945eeb3284 --- /dev/null +++ b/lightllm/server/httpserver_for_pd_master/pd_selector/cache_aware.py @@ -0,0 +1,202 @@ +""" +PD Master 的 cache-aware prefill 选点策略。 + +目标: + 在多 prefill 节点场景下,尽量把 prompt 前缀相近的请求打到同一 P 节点, + 以提高该节点上的前缀 KV cache 命中率;同时在节点负载差距过大时优先做 + 负载均衡,避免热点。 + +实现要点: + - 用前缀树(见 PromptCacheTree)记录「历史 prompt -> 处理它的 worker」; + - 树中的 tenant 对应 worker.client_ip_port; + - prompt 会按 sample_stride 抽稀后再插入/匹配,降低树的深度与内存; + - 用 worker.dispatched_prompt_chars(累计派发的 prompt 字符数)做粗粒度均衡。 + +选点流程见 CacheAwarePolicy.select_worker。 +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from typing import List, Optional + +from lightllm.server.pd_io_struct import PD_Client_Obj +from lightllm.utils.log_utils import init_logger + +from .prompt_cache_tree import DEFAULT_SAMPLE_STRIDE, PromptCacheTree + + +logger = init_logger(__name__) + + +@dataclass(slots=True) +class CacheAwareConfig: + """cache-aware 策略超参。""" + + # 前缀匹配成功率阈值:matched_key_len / input_key_len 超过该值才路由到命中节点。 + cache_threshold: float = 0.5 + # 派发量不均衡判定:max > min * balance_rel_threshold 时强制选派发量最少的节点。 + balance_rel_threshold: float = 1.2 + # 后台按间隔驱逐前缀树中过大的 tenant 缓存;<=0 表示不启动驱逐线程。 + eviction_interval_secs: int = 30 + # 单个 tenant 在前缀树中允许占用的最大 key 字符量(抽稀后的长度口径)。 + max_tree_size: int = 1000000 + # 每隔 sample_stride 个字符抽 1 个作为前缀树 key,降低匹配开销与内存。 + sample_stride: int = DEFAULT_SAMPLE_STRIDE + + +class CacheAwarePolicy: + """ + 维护 prompt 前缀树,并据此为请求选择 prefill worker。 + + 树生命周期: + - 选中 worker 后会把当前 prompt 插入该 worker 对应的 tenant; + - worker 下线时可通过 remove_worker 清理其 tenant; + - 后台线程周期性按 max_tree_size 做 LRU 风格驱逐,防止树无限增长。 + """ + + def __init__(self, config: Optional[CacheAwareConfig] = None) -> None: + self.config = config or CacheAwareConfig() + self.prompt_cache_tree: PromptCacheTree = PromptCacheTree(sample_stride=self.config.sample_stride) + self._stop_eviction = threading.Event() + self._eviction_thread: Optional[threading.Thread] = None + if self.config.eviction_interval_secs > 0: + self._eviction_thread = threading.Thread( + target=self._run_eviction_loop, name="cache-aware-eviction", daemon=True + ) + self._eviction_thread.start() + + def _run_eviction_loop(self) -> None: + """周期性裁剪前缀树,控制每个 tenant 的缓存规模。""" + while not self._stop_eviction.wait(self.config.eviction_interval_secs): + logger.info("Running cache eviction...") + self.evict_cache(self.config.max_tree_size) + logger.info(f"Cache eviction completed.: {self.prompt_cache_tree.get_used_size_per_tenant()}") + + def close(self) -> None: + """停止后台驱逐线程。""" + self._stop_eviction.set() + if self._eviction_thread is not None and self._eviction_thread.is_alive(): + self._eviction_thread.join(timeout=1.0) + + def init_workers(self, workers: List[PD_Client_Obj]) -> None: + """在树中注册一批 worker(插入空前缀,仅建立 tenant)。""" + for worker in workers: + self.prompt_cache_tree.insert("", worker.client_ip_port) + + def add_worker(self, worker: PD_Client_Obj) -> None: + """注册单个新上线的 worker。""" + self.prompt_cache_tree.insert("", worker.client_ip_port) + + def remove_worker(self, worker: PD_Client_Obj) -> None: + """移除下线 worker 在前缀树中的全部记录。""" + self.prompt_cache_tree.remove_tenant(worker.client_ip_port) + + def remove_worker_by_url(self, url: str) -> None: + """按 client_ip_port 移除 worker 对应 tenant。""" + self.prompt_cache_tree.remove_tenant(url) + + def evict_cache(self, max_size: int) -> None: + """将各 tenant 占用压缩到 max_size 以下。""" + self.prompt_cache_tree.evict_tenant_by_size(max_size) + + def _select_worker_min_dispatched( + self, + workers: List[PD_Client_Obj], + request_text: Optional[str], + ) -> Optional[PD_Client_Obj]: + """ + 派发量优先兜底:选择累计 dispatched_prompt_chars 最小的 worker。 + 若提供了 request_text,同时把该 prompt 记到该 worker 的前缀树下, + 便于后续相似请求继续命中。 + """ + min_dispatched_worker = min(workers, key=lambda worker: worker.dispatched_prompt_chars) + + if request_text is not None: + self.prompt_cache_tree.insert(request_text, min_dispatched_worker.client_ip_port) + + return min_dispatched_worker + + def select_worker( + self, workers: List[PD_Client_Obj], request_text: Optional[str] = None + ) -> Optional[PD_Client_Obj]: + """ + 为一次请求选择 prefill worker。 + + 决策顺序: + 1) workers 为空 -> 返回 None; + 2) 若 max(dispatched) > min(dispatched) * balance_rel_threshold, + 认为派发不均衡,直接选派发量最少的节点; + 3) 否则对 request_text 做前缀匹配,计算 + match_rate = matched_key_len / input_key_len; + 4) match_rate > cache_threshold 且命中 tenant 仍在线 -> 路由到该节点, + 并更新树;若 tenant 已不在当前 workers 中,则从树中剔除该 tenant; + 5) 未命中阈值或 tenant 失效 -> 回退到派发量最少选择。 + """ + if not workers: + return None + + # ---- 1. 派发均衡门闩:差距过大时不再追求 cache 亲和 ---- + dispatched_chars = [worker.dispatched_prompt_chars for worker in workers] + min_dispatched = min(dispatched_chars) if dispatched_chars else 0 + max_dispatched = max(dispatched_chars) if dispatched_chars else 0 + + is_imbalanced = max_dispatched > (min_dispatched * self.config.balance_rel_threshold) + + logger.info( + f"CacheAwarePolicy: min_dispatched={min_dispatched}, max_dispatched={max_dispatched}, " + f"balance_rel_threshold={self.config.balance_rel_threshold:.4f}, " + f"is_imbalanced={is_imbalanced}" + ) + + if is_imbalanced: + return self._select_worker_min_dispatched( + workers=workers, + request_text=request_text, + ) + + # ---- 2. 前缀匹配:估计当前请求与历史请求的 cache 复用潜力 ---- + text = request_text or "" + + result = self.prompt_cache_tree.prefix_match_with_counts(text) + # matched/input 均基于抽稀后的 key 长度,比值近似原始前缀重合比例。 + match_rate = 0.0 if result.input_char_count == 0 else result.matched_char_count / result.input_char_count + + logger.info( + f"CacheAwarePolicy: matched_char_count={result.matched_char_count}, " + f"input_char_count={result.input_char_count}, match_rate={match_rate:.4f}, " + f"cache_threshold={self.config.cache_threshold:.4f}" + ) + + selected_worker: Optional[PD_Client_Obj] = None + if match_rate > self.config.cache_threshold: + # 树中的 tenant 是 client_ip_port,需要映射回当前在线 worker 对象。 + for worker in workers: + if worker.client_ip_port == result.tenant: + selected_worker = worker + break + + if selected_worker is None: + # 命中了已下线/不在列表中的 tenant,清理脏数据后走派发量兜底。 + logger.info(f"Evicting tenant: {result.tenant}") + self.prompt_cache_tree.remove_tenant(result.tenant) + + logger.info( + f"CacheAwarePolicy: selected_worker=" + f"{selected_worker.client_ip_port if selected_worker else None}, " + f"match_rate={match_rate:.4f}, cache_threshold={self.config.cache_threshold:.4f}" + ) + + # ---- 3. 命中则更新树;未命中则派发量兜底并写入树 ---- + if selected_worker is not None: + self.prompt_cache_tree.insert(text, selected_worker.client_ip_port) + return selected_worker + else: + return self._select_worker_min_dispatched( + workers=workers, + request_text=request_text, + ) + + def __del__(self) -> None: + self.close() diff --git a/lightllm/server/httpserver_for_pd_master/pd_selector/pd_selector.py b/lightllm/server/httpserver_for_pd_master/pd_selector/pd_selector.py index 2a728fb8c7..989b3827e3 100644 --- a/lightllm/server/httpserver_for_pd_master/pd_selector/pd_selector.py +++ b/lightllm/server/httpserver_for_pd_master/pd_selector/pd_selector.py @@ -3,6 +3,12 @@ from lightllm.server.pd_io_struct import PD_Client_Obj from lightllm.server.core.objs import SamplingParams from lightllm.server.multimodal_params import MultimodalParams +from lightllm.utils.log_utils import init_logger + +from .cache_aware import CacheAwarePolicy, CacheAwareConfig + + +logger = init_logger(__name__) class PDSelector: @@ -64,4 +70,44 @@ def select_p_d_node( return p_node, d_node def _importance_sampling(self, nodes: List[PD_Client_Obj]): - return random.choices(nodes, weights=[max(1.0 - e.run_status.total_token_usage_rate, 0.02) for e in nodes]) + return random.choices(nodes, weights=[max(1.0 - e.run_status.total_token_usage_rate, 0.02) for e in nodes])[0] + + +class LoadBalancedCacheAwareSelector(AdaptiveLoadSelector): + """Cache-aware prefill 选点:按抽稀后的 prompt 前缀匹配 + 负载均衡。""" + + def __init__(self, pd_manager): + super().__init__(pd_manager) + self.policy = CacheAwarePolicy(CacheAwareConfig()) + self.tree_workers = [] + + def update_nodes(self, prefill_nodes, decode_nodes): + super().update_nodes(prefill_nodes, decode_nodes) + + add_workers = set(self.prefill_nodes) - set(self.tree_workers) + remove_workers = set(self.tree_workers) - set(self.prefill_nodes) + + for worker in add_workers: + self.tree_workers.append(worker) + + for worker in remove_workers: + self.tree_workers.remove(worker) + + self.tree_workers = self.prefill_nodes + + def select_p_d_node( + self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams + ) -> Tuple[PD_Client_Obj, PD_Client_Obj]: + assert isinstance(prompt, str), "prompt must be a string for cache-aware selection" + p_node = self.policy.select_worker(self.prefill_nodes, request_text=prompt) + d_node = self._importance_sampling(self.decode_nodes) + + # 累计派发字符数,供后续 cache-aware 做派发均衡判断。 + p_node.dispatched_prompt_chars += len(prompt) + + logger.info( + f"LoadBalancedCacheAwareSelector: selected p_node={p_node.client_ip_port}, " + f"d_node={d_node.client_ip_port}" + ) + + return p_node, d_node diff --git a/lightllm/server/httpserver_for_pd_master/pd_selector/prompt_cache_tree.py b/lightllm/server/httpserver_for_pd_master/pd_selector/prompt_cache_tree.py new file mode 100644 index 0000000000..ea7f9270da --- /dev/null +++ b/lightllm/server/httpserver_for_pd_master/pd_selector/prompt_cache_tree.py @@ -0,0 +1,322 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from heapq import heappop, heappush +from itertools import count +from threading import RLock +from typing import Dict, List, Optional, Tuple + + +_EPOCH_COUNTER = count() + +# 每隔 sample_stride 个字符取 1 个作为前缀树 key,压缩树深度与内存。 +DEFAULT_SAMPLE_STRIDE = 16 + + +def _get_epoch() -> int: + return next(_EPOCH_COUNTER) + + +def _shared_prefix_count(a: str, b: str) -> int: + matched = 0 + for ca, cb in zip(a, b): + if ca != cb: + break + matched += 1 + return matched + + +def sample_prompt_cache_key(text: str, sample_stride: int) -> str: + """从 prompt 中按固定步长抽字符,作为 prompt cache 前缀匹配 key。""" + if sample_stride <= 1 or not text: + return text + return text[::sample_stride] + + +@dataclass(slots=True) +class PromptCacheMatchResult: + tenant: str + matched_char_count: int + input_char_count: int + + +@dataclass(slots=True) +class _PromptCacheNode: + text: str + children: Dict[str, "_PromptCacheNode"] = field(default_factory=dict) + tenant_last_access_time: Dict[str, int] = field(default_factory=dict) + parent: Optional["_PromptCacheNode"] = None + last_tenant: Optional[str] = None + + +class PromptCacheTree: + """ + 用于 cache-aware 选点的 prompt 前缀缓存树。 + + 对原始 prompt 按 sample_stride 抽稀后再建树/匹配,key 长度约为 + len(prompt) / sample_stride,从而降低匹配开销与树节点内存。 + matched_char_count / input_char_count 均基于抽稀后的 key 长度计算, + 比值仍可近似反映原始前缀重合比例。 + """ + + def __init__(self, sample_stride: int = DEFAULT_SAMPLE_STRIDE) -> None: + if sample_stride < 1: + raise ValueError(f"sample_stride must be >= 1, got {sample_stride}") + self.sample_stride = sample_stride + self.root = _PromptCacheNode(text="") + self.tenant_char_count: Dict[str, int] = {} + self._lock = RLock() + + def _to_key(self, text: str) -> str: + return sample_prompt_cache_key(text, self.sample_stride) + + def insert(self, text: str, tenant: str) -> None: + key = self._to_key(text) + with self._lock: + self.root.tenant_last_access_time.setdefault(tenant, 0) + self.tenant_char_count.setdefault(tenant, 0) + + remaining = key + prev = self.root + + while remaining: + first_char = remaining[0] + child = prev.children.get(first_char) + + if child is None: + remaining_char_count = len(remaining) + epoch = _get_epoch() + new_node = _PromptCacheNode( + text=remaining, + tenant_last_access_time={tenant: epoch}, + parent=prev, + last_tenant=tenant, + ) + self.tenant_char_count[tenant] = self.tenant_char_count.get(tenant, 0) + remaining_char_count + prev.children[first_char] = new_node + return + + shared_count = _shared_prefix_count(remaining, child.text) + child_len = len(child.text) + + if shared_count < child_len: + matched_text = child.text[:shared_count] + contracted_text = child.text[shared_count:] + matched_text_count = shared_count + + new_node = _PromptCacheNode( + text=matched_text, + tenant_last_access_time=dict(child.tenant_last_access_time), + parent=prev, + last_tenant=child.last_tenant, + ) + new_node.children[contracted_text[0]] = child + + child.text = contracted_text + child.parent = new_node + prev.children[first_char] = new_node + + if tenant not in new_node.tenant_last_access_time: + self.tenant_char_count[tenant] = self.tenant_char_count.get(tenant, 0) + matched_text_count + new_node.tenant_last_access_time[tenant] = 0 + + prev = new_node + remaining = remaining[shared_count:] + else: + if tenant not in child.tenant_last_access_time: + self.tenant_char_count[tenant] = self.tenant_char_count.get(tenant, 0) + child_len + child.tenant_last_access_time[tenant] = 0 + prev = child + remaining = remaining[shared_count:] + + epoch = _get_epoch() + prev.tenant_last_access_time[tenant] = epoch + prev.last_tenant = tenant + + def prefix_match_with_counts(self, text: str) -> PromptCacheMatchResult: + key = self._to_key(text) + with self._lock: + remaining = key + matched_chars = 0 + prev = self.root + + while remaining: + first_char = remaining[0] + child = prev.children.get(first_char) + if child is None: + break + + shared_count = _shared_prefix_count(remaining, child.text) + child_len = len(child.text) + + if shared_count == child_len: + matched_chars += shared_count + remaining = remaining[shared_count:] + prev = child + else: + matched_chars += shared_count + prev = child + break + + curr = prev + + if curr.last_tenant and curr.last_tenant in curr.tenant_last_access_time: + tenant = curr.last_tenant + else: + tenant = next(iter(curr.tenant_last_access_time), "empty") + curr.last_tenant = tenant + + if tenant != "empty": + curr.tenant_last_access_time[tenant] = _get_epoch() + + return PromptCacheMatchResult( + tenant=tenant, + matched_char_count=matched_chars, + input_char_count=len(key), + ) + + def prefix_match(self, text: str) -> Tuple[str, str]: + key = self._to_key(text) + result = self.prefix_match_with_counts(text) + return key[: result.matched_char_count], result.tenant + + def prefix_match_tenant(self, text: str, tenant: str) -> str: + key = self._to_key(text) + with self._lock: + remaining = key + matched_chars = 0 + prev = self.root + + while remaining: + first_char = remaining[0] + child = prev.children.get(first_char) + if child is None: + break + if tenant not in child.tenant_last_access_time: + break + + shared_count = _shared_prefix_count(remaining, child.text) + child_len = len(child.text) + + if shared_count == child_len: + matched_chars += shared_count + remaining = remaining[shared_count:] + prev = child + else: + matched_chars += shared_count + prev = child + break + + if tenant in prev.tenant_last_access_time: + prev.tenant_last_access_time[tenant] = _get_epoch() + prev.last_tenant = tenant + + return key[:matched_chars] + + @staticmethod + def _leaf_of(node: _PromptCacheNode) -> List[str]: + candidates: Dict[str, bool] = {tenant: True for tenant in node.tenant_last_access_time} + for child in node.children.values(): + for tenant in child.tenant_last_access_time: + candidates[tenant] = False + return [tenant for tenant, is_leaf in candidates.items() if is_leaf] + + def evict_tenant_by_size(self, max_size: int) -> None: + with self._lock: + stack = [self.root] + pq: List[Tuple[int, str, _PromptCacheNode]] = [] + + while stack: + curr = stack.pop() + stack.extend(curr.children.values()) + for tenant in self._leaf_of(curr): + ts = curr.tenant_last_access_time.get(tenant) + if ts is not None: + heappush(pq, (ts, tenant, curr)) + + while pq: + _, tenant, node = heappop(pq) + used_size = self.tenant_char_count.get(tenant, 0) + if used_size <= max_size: + continue + + if tenant not in node.tenant_last_access_time: + continue + if any(tenant in child.tenant_last_access_time for child in node.children.values()): + continue + + node_len = len(node.text) + self.tenant_char_count[tenant] = max(0, self.tenant_char_count.get(tenant, 0) - node_len) + + node.tenant_last_access_time.pop(tenant, None) + if node.last_tenant == tenant: + node.last_tenant = next(iter(node.tenant_last_access_time), None) + + parent = node.parent + if not node.children and not node.tenant_last_access_time and parent is not None: + if node.text: + parent.children.pop(node.text[0], None) + + if parent is not None and tenant in parent.tenant_last_access_time: + has_child_with_tenant = any( + tenant in child.tenant_last_access_time for child in parent.children.values() + ) + if not has_child_with_tenant: + ts = parent.tenant_last_access_time.get(tenant) + if ts is not None: + heappush(pq, (ts, tenant, parent)) + + if self.tenant_char_count.get(tenant, 0) == 0: + self.tenant_char_count.pop(tenant, None) + + def remove_tenant(self, tenant: str) -> None: + with self._lock: + stack = [self.root] + queue: List[_PromptCacheNode] = [] + + while stack: + curr = stack.pop() + stack.extend(curr.children.values()) + + if tenant in curr.tenant_last_access_time: + has_child_with_tenant = any( + tenant in child.tenant_last_access_time for child in curr.children.values() + ) + if not has_child_with_tenant: + queue.append(curr) + + while queue: + curr = queue.pop(0) + curr.tenant_last_access_time.pop(tenant, None) + if curr.last_tenant == tenant: + curr.last_tenant = next(iter(curr.tenant_last_access_time), None) + + parent = curr.parent + if not curr.children and not curr.tenant_last_access_time and parent is not None: + if curr.text: + parent.children.pop(curr.text[0], None) + + if parent is not None and tenant in parent.tenant_last_access_time: + has_child_with_tenant = any( + tenant in child.tenant_last_access_time for child in parent.children.values() + ) + if not has_child_with_tenant: + queue.append(parent) + + self.tenant_char_count.pop(tenant, None) + + def get_tenant_char_count(self) -> Dict[str, int]: + with self._lock: + return dict(self.tenant_char_count) + + def get_used_size_per_tenant(self) -> Dict[str, int]: + with self._lock: + used_size_per_tenant: Dict[str, int] = {} + stack = [self.root] + while stack: + curr = stack.pop() + text_count = len(curr.text) + for tenant in curr.tenant_last_access_time: + used_size_per_tenant[tenant] = used_size_per_tenant.get(tenant, 0) + text_count + stack.extend(curr.children.values()) + return used_size_per_tenant diff --git a/lightllm/server/pd_io_struct.py b/lightllm/server/pd_io_struct.py index 1d68f81a9e..34e6923f7d 100644 --- a/lightllm/server/pd_io_struct.py +++ b/lightllm/server/pd_io_struct.py @@ -54,6 +54,8 @@ class PD_Client_Obj: start_args: object # 节点的启动参数信息,用于做匹配性的校验,防止运行过程中出现问题。 websocket: WebSocket = None # 用于通信的 websocket 连接对象 run_status: _PD_Client_RunStatus = field(default_factory=_PD_Client_RunStatus) + # cache-aware 选点用:累计派发到该节点的 prompt 字符数(只增不减,非实时负载)。 + dispatched_prompt_chars: int = 0 def __post_init__(self): if self.mode not in ["prefill", "decode"]: diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/up_status.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/up_status.py index bc1d00f384..b8fbd3e276 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/up_status.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/up_status.py @@ -94,8 +94,10 @@ async def up_kv_status_task(self, pd_master_obj: PD_Master_Obj): logger.info(f"up kv status: {upkv_status}") else: await asyncio.sleep(3) + except BaseException as e: logger.error(str(e)) + await task_queue.put(upkv_status) raise e except asyncio.CancelledError: logger.info(f"up_kv_status_task {pd_master_obj} cancelled")