KV Cache 管理策略:显存告急时,优先淘汰哪些请求才损失最小?
KV Cache 管理策略显存告急时优先淘汰哪些请求才损失最小一、OOM 前的最后一道防线KV Cache 的选择困难推理服务的 OOMOut of Memory通常是突发性的。上一秒 GPU 显存利用率 85%下一秒一个超长 context 请求进入显存瞬间吃满CUDA OOM 异常抛出——整个进程崩溃所有在线请求丢失。KV Cache 是大模型推理中最大的显存消耗者。以 Llama-70B 为例单层 Transformer 的 KV Cache 占用 2 × hidden_size × num_layers × num_kv_heads × head_dim × dtype_size。对于 32K context 的请求仅 KV Cache 就超过 10GB。这意味着显存紧张时淘汰 2~3 个长 context 请求就能释放可观的空间。问题的核心在于选择。当前有 8 个活跃请求每个请求的 KV Cache 都已经部分计算完成。显存接近阈值时必须从这 8 个中选择一个驱逐。选谁驱逐错误的请求意味着大量 GPU 计算被浪费——该请求恢复时需要从零重建 KV Cache。一个真实场景4 个短 prompt 请求各消耗 256MB KV Cache3 个中长 prompt 请求各消耗 2GB1 个超长 context 请求消耗 8GB。按驱逐最大的策略直接干掉 8GB 那个。恢复时这个请求需要重新 prefill 全部 32K context耗时约 3.5 秒。在这 3.5 秒内新到达的请求会进一步挤占显存可能触发二次驱逐——连锁反应。二、多维淘汰策略从 FIFO 到价值函数flowchart TD A[活跃请求 8 个] -- B{显存阈值检测} B --|正常| C[继续推理] B --|超过阈值| D[触发驱逐决策] D -- E[收集候选驱逐集] E -- F[计算驱逐价值函数V] F -- G{V w1 × 存活时间\n w2 × 剩余长度\n w3 × Cache大小\n w4 × 优先级} G -- H[按 V 排序取最小 V] H -- I[驱逐该请求 KV Cache] I -- J{显存仍不足?} J --|是| E J --|否| C驱逐决策的核心是定义一个价值函数 V(x)值越低的请求越应该被淘汰存活时间Age已经执行了较长时间的请求其 KV Cache 重建成本更高。存活时间越长的请求 V 值越高不应驱逐。实现时记录每个请求的 prefill 结束时间戳。剩余生成长度Remaining如果请求的输出 token 数已接近max_token限制剩余计算量很少。即便驱逐后重建额外浪费有限。但这个值往往未知——可以用已生成 token 数 / 输入 token 数作为近似比值越高的请求越接近结束。KV Cache 大小Size驱逐大的能释放更多空间但重建成本也更高。单独使用这个指标会导致长 context 歧视——每次都是长请求被驱逐永远无法完成。请求优先级Priority如果服务区分在线/离线请求或用户等级高优请求的 V 值应加权保护。实际落地时不同场景的权重差异很大。在线对话场景偏重存活时间因为用户等待耐心有限离线批处理场景偏重 Cache 大小因为追求吞吐最大化。三、两级淘汰策略的 Go 实现以下代码展示一个基于阈值触发 价值函数的两级驱逐调度器package kvcache import ( context math sort sync time ) // CacheEntry 表示一个活跃请求的 KV Cache 元信息 type CacheEntry struct { RequestID string Priority int // 请求优先级1最高5最低 CacheSizeMB float64 // 当前 KV Cache 占用MB StartTime time.Time // prefill 开始时间 InputTokens int // 输入 token 数 OutputTokens int // 已生成 token 数 MaxTokens int // max_tokens 限制 cachePtr uintptr // GPU 显存中的 cache 地址不直接操作 } // EvictionScheduler 驱逐调度器 type EvictionScheduler struct { mu sync.RWMutex entries map[string]*CacheEntry totalMemGB float64 // GPU 显存总量 highWatermark float64 // 高水位线如 0.85 lowWatermark float64 // 低水位线如 0.70 currentUsage float64 // 当前显存使用量GB // 驱逐统计 evictionCount int64 lastEviction time.Time } // NewEvictionScheduler 创建驱逐调度器 func NewEvictionScheduler(totalMemGB, highWater, lowWater float64) *EvictionScheduler { return EvictionScheduler{ entries: make(map[string]*CacheEntry), totalMemGB: totalMemGB, highWatermark: highWater, // 默认 0.85 lowWatermark: lowWater, // 默认 0.70释放到低水位线为止 } } // Register 注册一个新的 KV Cache 条目 func (s *EvictionScheduler) Register(entry *CacheEntry) { s.mu.Lock() defer s.mu.Unlock() s.entries[entry.RequestID] entry } // Unregister 请求完成后注销 func (s *EvictionScheduler) Unregister(requestID string) { s.mu.Lock() defer s.mu.Unlock() delete(s.entries, requestID) } // CheckAndEvict 检查显存使用率必要时触发驱逐 // 返回被驱逐的 requestID 列表调用方负责释放 GPU 显存 func (s *EvictionScheduler) CheckAndEvict(ctx context.Context, currentUsageGB float64) []string { s.mu.Lock() defer s.mu.Unlock() s.currentUsage currentUsageGB usageRatio : currentUsageGB / s.totalMemGB // 未触及高水位线不操作 if usageRatio s.highWatermark { return nil } var evicted []string // 重复驱逐直到显存降至低水位线以下或无可驱逐项 for usageRatio s.lowWatermark len(s.entries) 0 { victim : s.selectVictim() if victim nil { break } evicted append(evicted, victim.RequestID) delete(s.entries, victim.RequestID) // 预估释放的显存量 currentUsageGB - victim.CacheSizeMB / 1024.0 usageRatio currentUsageGB / s.totalMemGB s.evictionCount s.lastEviction time.Now() } return evicted } // selectVictim 按价值函数选择驱逐目标 // 内部排序复杂度 O(n log n) func (s *EvictionScheduler) selectVictim() *CacheEntry { if len(s.entries) 0 { return nil } now : time.Now() candidates : make([]struct { entry *CacheEntry value float64 }, 0, len(s.entries)) for _, entry : range s.entries { // 价值函数 V w_age * age w_remain * remain_ratio w_size * size_reward - w_pri * priority // age 越大越保护存活时间长 重建成本高 ageMs : now.Sub(entry.StartTime).Milliseconds() ageScore : math.Log1p(float64(ageMs)) * 0.3 // 对数衰减 // remain 越少越可以驱逐快完成了重建代价低 remainTokens : entry.MaxTokens - entry.OutputTokens if remainTokens 0 { remainTokens 1 } remainRatio : float64(entry.OutputTokens) / float64(entry.InputTokensentry.OutputTokens1) remainScore : remainRatio * 0.25 // size 越大越倾向驱逐释放空间多 sizeScore : math.Log1p(entry.CacheSizeMB) * 0.25 // priority 越高越保护 priScore : float64(6-entry.Priority) * 0.2 // priority 1→1.0, 5→0.2 value : ageScore remainScore sizeScore - priScore candidates append(candidates, struct { entry *CacheEntry value float64 }{entry, value}) } // 选择价值最低的作为 victim sort.Slice(candidates, func(i, j int) bool { return candidates[i].value candidates[j].value }) return candidates[0].entry } // Stats 返回驱逐统计 func (s *EvictionScheduler) Stats() (count int64, last time.Time, activeEntries int) { s.mu.RLock() defer s.mu.RUnlock() return s.evictionCount, s.lastEviction, len(s.entries) }实现中的关键决策对数衰减函数age 和 size 都经过log1p压缩。原始值相差 100 倍时评分仅差 2~3 倍。避免单一维度主导驱逐决策。水位线双阈值高水位触发驱逐、低水位停止驱逐。避免驱逐一个后显存刚过危险线又被立即触发。原子统计evictionCount 通过 int64 mutex 保护可暴露给 Prometheus 监控。四、策略偏差、重建成本与公平性价值函数永远不完美。日志衰减压缩了绝对数值但可能导致驱逐小但快完成的请求优于驱逐大但刚开始的请求。在线问答场景可能不适合——用户正在等待一个短答案却被驱逐了。重建成本估计的偏差。StartTime表示 prefill 开始时间但 prefill 本身消耗的时间包含在 age 中。一个超长 prefill3.5 秒的请求 age 大、不容易被驱逐——但它恰恰是显存危机的罪魁祸首。修正方案是在 age 计算中扣除 prefill 耗时或单独计权重。多轮驱逐的公平性。同一请求可能在多轮驱逐中每次都被选中因为它的 V 值始终最低。这导致马太效应——某些请求可能永远无法完成。一个简单的缓解措施是引入驱逐次数计数被驱逐过的请求在下一轮获得保护加成。禁用场景仅承载 1~2 个并发请求的独占式推理无需驱逐显存充裕的场景如 H100 × 8 集群、小 batch请求长度分布极窄驱逐任何请求的成本差异极小。五、总结KV Cache 驱逐是显存紧张时的必要手段但策略优劣直接影响吞吐和用户体验。基于多维价值函数的驱逐优于简单的 FIFO 或 Largest-First。核心是平衡四个维度存活时间、剩余长度、Cache 大小和请求优先级。生产落地的优先级先接入 GPU 显存监控nvidia-smi或 DCGM 采集设置合理的水位线高 0.85、低 0.70再部署驱逐调度器权重初始化为年龄 0.3、剩余比 0.25、大小 0.25、优先级 0.2最后根据驱逐统计驱逐频率、二次驱逐率迭代调参。