概述 (Overview)

Dynamo 的 KV 路由器 (KV Router) 监听来自工作节点的 KV 事件,以构建一个全局的 KV 缓存前缀树 (global prefix tree of KV caches)。这使得路由器能够预测每个工作节点对传入请求的 KV 命中率(重叠分数,overlap score),并做出智能的路由决策。优化 KV 路由器对于在分布式推理设置中实现最大吞吐量和最低延迟至关重要。

KV 路由器架构 (KV Router Architecture)

KV 路由器维护着所有工作节点上 KV 缓存的全局视图。当一个新的请求到达时:

1.路由器通过计算请求与工作节点前缀缓存之间的匹配块,为每个工作节点计算一个重叠分数 (overlap score)

2.它从每个工作节点收集运行时指标,包括 KV 缓存使用率等待请求数量

3.它应用一个代价函数 (cost function) 来确定处理该请求的最优工作节点。

代价函数调优 (Cost Function Tuning)

KV 路由器的决策主要由其代价函数控制,该函数可在 kv_router.py 中进行自定义。默认的代价函数是:

worker_logits[worker_id] = 2 * score - gpu_cache_usage - normalized_waiting

其中:

score: 归一化的重叠分数 (匹配块数 × 块大小 / 令牌长度)

gpu_cache_usage: 当前已使用的 GPU KV 缓存百分比

normalized_waiting: 等待请求数量,按所有工作节点中的最大等待请求数进行归一化

路由器选择具有最高 logit 值的工作节点。如果出现平局,则在得分最高的工作节点中随机选择。或者,也可以对 logits 应用 softmax 函数,并根据得到的概率进行采样,这可以为路由过程引入随机性 (stochasticity)。这种概率性方法 (probabilistic approach) 有助于防止出现一种故障模式:某个工作节点接收了不成比例的大量请求,导致其前缀缓存饱和。这种饱和会形成一个反馈循环 (feedback loop) —— 缓存丰富的工作节点持续被选中 —— 使得难以确定性地 (deterministically) 打破这种循环。

关键调优参数 (Key Tuning Parameters)

重叠分数权重 (Overlap Score Weight) (默认: 2.0)

◦值越高,越优先考虑 KV 缓存重用

◦值越低,越允许请求更均匀地分布

GPU 缓存使用率权重 (GPU Cache Usage Weight) (默认: 1.0)

◦值越高,越避免选择 KV 缓存接近满的工作节点

◦值越低,越忽略 KV 缓存利用率

等待请求权重 (Waiting Requests Weight) (默认: 1.0)

◦值越高,越避免选择有请求排队的工作节点

◦值越低,越忽略队列长度

调优指南 (Tuning Guidelines)

目前,要优化使用我们的 KV 路由器,需要理解后端引擎的容量和数据的前缀结构 (prefix structure)。我们在 benchmarks 目录中提供了用于此目的的分析工具。未来,我们计划利用工作节点反馈指标和对数据前缀结构的动态分析(进行中,WIP),实现 KV 路由器的自动调优(通过规划器,Planner)。以下是我们建议遵循的几个提示:

1.考虑总的 KV 块分配 (Consider Total KV Block Allocation):

•检查为后端引擎分配的 KV 块总数。对于较小的模型(例如 80 亿参数),这可能超过一百万个块。

•在这种情况下:

◦降低对 KV 缓存使用率 (gpu_cache_usage) 的权重,因为耗尽 KV 缓存的可能性较小。

◦更专注于重叠分数 (score) 和等待请求 (normalized_waiting)。

2.分析数据集的理论命中率 (Analyze Your Dataset’s Theoretical Hit Rate):

考虑数据集预期的理论命中率(假设缓存完美):

•更正式地说,考虑核心前缀树的深度(至少被访问两次的节点)。

•对于命中率较低的情况,或者如果核心前缀树深度相对于输入序列长度 (ISL) 较短,则降低重叠分数权重 (score)。

•或者,用输入序列长度 (ISL) 对重叠分数进行归一化。

3.考虑前缀树的广度 (Consider Prefix Tree Breadth):

前缀树的广度可以通过预期的唯一上下文提示 (context prompts) 数量来近似代表。

◦如果可以识别出不同的相似提示“桶”(buckets),每个桶包含大致相等数量的提示:

◦考虑使用多个 KV 路由器,每个桶一个。

◦使用一个元路由器 (meta-router) 将请求导向处理不同提示类别的专用 KV 路由器。

◦对于上下文提示非常多样化的情况,重叠分数 (score) 很可能应该被优先考虑。

4.平衡延迟与吞吐量 (Balance Latency vs. Throughput):

权重直接影响服务级别目标 (SLO):

▪提高等待请求 (normalized_waiting) 的权重可以改善延迟,但可能降低吞吐量。

▪提高重叠分数 (score) 的权重可能提高吞吐量,但可能增加尾部延迟 (tail latency)。

替代路由策略 (Alternative Routing Strategies)

默认策略使用贪婪选择 (greedy selection)(最高 logit 胜出),但也可以实现其他方法:

Softmax 采样 (Softmax Sampling): 将 logits 转换为概率,并概率性地采样工作节点。

基于温度的采样 (Temperature-Based Sampling): 添加一个温度参数 (temperature) 来控制采样的随机性。

两阶段路由 (Two-Stage Routing): 例如,使用轮询 (round-robin) 作为元路由器,将请求路由到多个 KV 路由器。

监控与优化 (Monitoring and Refinement)

为了有效调优 KV 路由器:

监控路由器日志,查看每个工作节点的实际 logit 计算值。

追踪命中率、延迟和吞吐量指标。

•根据观察到的性能迭代调整权重。

•考虑根据当前负载状况动态调整权重。