基于kvcache的RAG系统

本片博客介绍一下最近做的一个实验基于大模型的kvcache来做数据库的语义化检索。不过实验不成功因为该方法的目的可以在当前RAG实现的基础上以一种更简单、成本更低的方法来实现。项目地址https://github.com/hhk-png/faq-kv-cache。我们在做RAG的时候一般是对某些特定领域的知识来搭建向量数据库在用户查询信息的时候先使用embedding模型将这些信息向量化然后在向量数据库中检索跟用户输入的信息相似的数据进而得到已有的知识数据交给大模型使用。大模型有一个kvcache机制即对于已经计算过的前缀可以直接通过缓存跳过计算只计算新增的token减少整体的计算时间。具体的可以查看https://www.bilibili.com/video/BV17CPkeEEzk/?vd_source36bfa49fd2dca513136af48ec97cffbe。前提到的我做的 faq-kv-cache 正是基于kvcache机制来做的信息的语义检索。比方说用户的问题是“小明几岁了”数据库中有“小明15岁“”小红15岁“这两个数据。在进行数据的查找时就会将查询文本拼接成下面的形式以下是可以使用的数据 [id1]小明15岁 [id2]小红15岁 用户的问题 小明几岁了 帮我查找数据中与用户问题相关的数据id并输出且仅输出对应的id以,隔开。在将上面的文本发送给大模型之后一般大模型就会将[id1]作为输出然后此时就可以根据这个摘要的id去数据库中查询完整的数据拼接到对话的上下文中去回答用户的问题。在上面的文本中用户的问题前面所有的文本可以是固定的这也是数据量最大的部分每次将相同的前缀发送到大模型侧针对这一部分是可以直接缓存命中需要重新计算的只是后面用户的问题部分这样就实现了较为快速的语义化检索。项目使用的是deepseek的api从api-key扣费来看每一轮的数据前缀确实是从缓存中获取的。测试使用的数据集是https://github.com/Toyhom/Chinese-medical-dialogue-data/tree/master/Data_%E6%95%B0%E6%8D%AE/Oncology_%E8%82%BF%E7%98%A4%E7%A7%91。受限于api的并发限制在小数据量时后端的返回较为快速且比较准确。在较大数据量时大概36MB的数据每个块300w token后端的返回大概要十秒。如果将并发限制去掉TTFT应该可以控制在5秒左右。在扣费方面除了warm up时花费较高以外后面有了kvcache花费大概在每轮查询几毛钱。warm up前缀时花了20块钱左右。现在的RAG基于向量数据库可以实现语义化检索但造成检索结果不准确的缺点有切块不按语义切、切块太大或者切块太小等原因这些缺点似乎可以在建库时使用大模型来解决让大模型进行语义化的切块而不用基于kvcache的方式在整体速度上也可以更快。