动态记忆压缩

vLLM 实战：突破上下文限制，打造不会遗忘的 RAG

前言

在构建基于大语言模型（LLM）的 RAG（检索增强生成）应用时，开发者常常面临两个核心痛点：

1. 上下文爆满：检索到的文档 + 历史对话超过了模型的最大窗口限制。
2. 记忆遗忘：为了节省上下文使用滑动窗口，导致模型“忘记”了早期的重要信息。

很多开发者寄希望于推理框架 vLLM 的 PagedAttention 技术能自动解决这些问题。但事实真的如此吗？如果上下文真的爆满了，我们该如何在保留记忆的同时优化显存？

本文将深入剖析 PagedAttention 的真实作用，并提供一套基于 Query 改写 + 动态记忆压缩 的完整架构方案

---

一、PagedAttention 会自动扩展上下文吗？

结论：PagedAttention 是自动执行的，但它不能突破模型本身的长度限制。

1. 它是自动的吗？

是的。 只要你使用 vLLM 作为推理后端，PagedAttention 默认就会启用。它是 vLLM 内核的一部分，用于管理 KV Cache（键值缓存）。你不需要在代码中额外调用它。

2. 它的作用是什么？

PagedAttention 的核心价值是 显存管理优化，而非 上下文长度扩展。

• 解决碎片化：它像操作系统的虚拟内存一样，将 KV Cache 分块存储，消除了显存碎片。
• 提高并发：它允许 vLLM 在相同的显存下支持更大的 Batch Size，提升吞吐量。
• 跑满上限：它能帮你更稳定地跑满模型支持的最大上下文长度（例如模型原生支持 32k，它能帮你更高效地用满这 32k），但它不能让一个 8k 的模型支持 100k 的输入。

3. 为什么还会爆满？

当 历史对话 Tokens + 检索文档 Tokens > 模型最大上下文窗口 时，无论是否开启 PagedAttention，vLLM 都会根据配置（--max-model-len）进行截断或报错。因此，上下文管理必须在应用层（Application Layer）解决，而不是依赖推理层。

---

二、滑动窗口与记忆遗忘

为了解决上下文爆满，最直观的方法是滑动窗口（Sliding Window）：只保留最近的 N 轮对话，丢弃旧的。

但这会带来严重问题：

• 关键信息丢失：用户在第 1 轮说“我叫张三”，在第 10 轮问“我叫什么”，滑动窗口会导致模型回答“我不知道”。
• 逻辑断层：早期的项目背景、约束条件被遗忘，导致后续回答不一致。

我们需要一种既能节省 Token，又能保留关键记忆的机制。

---

三、解决方案：基于“记忆检索 + 动态压缩”的架构

“如果检索出来的历史记录太多，一样会撑爆怎么办？”

解决这个问题的核心原则是：检索的是“线索”，注入的是“结论”，而不是“原文”。

1. 架构设计图

graph TD
    A[用户新 Query] --> B(Query 改写/意图识别)
    B --> C{向量检索历史会话}
    C --> D[召回 Top-K 个历史片段]
    D --> E[重排序 Re-rank]
    E --> F{是否超过 Token 预算？}
    F -- 是 --> G[LLM 二次压缩/合并摘要]
    F -- 否 --> H[直接注入]
    G --> H
    H --> I[组装 Prompt -> vLLM 推理]

    J[历史会话结束] --> K[异步任务：总结会话 + 提取事实]
    K --> L[存入向量库 + 结构化数据库]

2. 核心流程详解

第一步：异步预处理（关键！）

不要等到用户提问时才去处理历史记录。 在每次会话结束（或每隔几轮）时，后台异步调用一个小模型对历史对话进行处理：

1. 会话摘要（Summary）：将 10 轮对话压缩成 100 字的摘要。
2. 事实提取（Facts）：提取关键信息（如：用户名字、偏好、项目预算）存入结构化数据库（如 SQLite/Redis）。

• 优势：检索时，你检索的是“摘要”和“事实”，体积比原文小 90%。

第二步：Query 改写（意图明确化）

用户问：“那个项目多少钱？”

• 改写前：向量检索可能匹配不到“项目”或“钱”。
• 改写后：调用 LLM 将 Query 改写为“查询之前对话中关于‘项目预算’、‘费用’、‘金额’相关的历史记录”。
• 目的：提高历史记忆检索的命中率。

第三步：检索与“预算控制”（解决爆满担忧）

这是防止上下文爆满的核心防线。

1. 设置 Token 预算（Token Budget）：
   • 假设模型总窗口 32k。
   • 预留 10k 给新文档 RAG。
   • 预留 2k 给 System Prompt。
   • 记忆模块只允许占用 5k Tokens。
2. 检索策略：
   • 先检索“事实库”（结构化数据）：体积几乎为 0，但信息密度极高。
   • 再检索“摘要库”（向量数据）：限制只取 Top-3 最相关的会话摘要。
3. 动态合并（Map-Reduce）：
   • 如果 Top-3 摘要加起来还是超过了 5k 预算？
   • 触发二次压缩：把这 3 个摘要再丢给 LLM，指令：“请将以下三段历史记忆合并为一段 200 字的总结，保留与‘预算’相关的数字。”
   • 结果：无论历史多长，注入 Prompt 的记忆部分永远被控制在 5k 以内。

---

四、代码实现示例

以下是一个简化的 Python MemoryManager 类，展示了如何实现预算控制和动态压缩。

import tiktoken

class MemoryManager:
    def __init__(self, max_memory_tokens=4000):
        self.max_memory_tokens = max_memory_tokens
        self.vector_store = ... # 存会话摘要的向量库
        self.fact_db = ...      # 存结构化事实的数据库
        self.encoder = tiktoken.get_encoding("cl100k_base")

    def get_relevant_memory(self, current_query):
        # 1. Query 改写 (调用小模型)
        expanded_query = self.rewrite_query(current_query)

        # 2. 检索结构化事实 (几乎不占 Token)
        # 例如：{'name': 'Alice', 'budget': '50w'}
        facts = self.fact_db.search(expanded_query) 

        # 3. 检索会话摘要 (向量检索)
        # 例如：["2023-10-01 讨论了项目需求...", "2023-10-05 确认了预算..."]
        summaries = self.vector_store.search(expanded_query, top_k=5)

        # 4. 计算 Token 数
        current_tokens = self.count_tokens(facts + summaries)

        # 5. 如果爆满，进行二次压缩 (关键步骤)
        if current_tokens > self.max_memory_tokens:
            # 调用 LLM 进行压缩，目标长度 = 预算 - 事实占用
            target_length = self.max_memory_tokens - self.count_tokens(facts)
            compressed_memory = self.llm_compress(
                text=summaries, 
                max_length=target_length
            )
            summaries = compressed_memory

        return facts + summaries

    def count_tokens(self, text_list):
        return len(self.encoder.encode(" ".join(text_list)))

    def rewrite_query(self, query):
        # 使用 Prompt 将 "那个项目" 改写为 "Alpha 项目预算"
        pass

    def llm_compress(self, text, max_length):
        # 调用 vLLM 或其他模型进行摘要压缩
        prompt = f"请总结以下记忆，保留关键数字和实体，限制在{max_length}tokens 内：\n" + "\n".join(text)
        return call_llm(prompt)

---

五、vLLM 配置与优化建议

为了配合上述应用层架构，vLLM 的启动参数也需要相应调整：

1. 最大化上下文窗口：

   python -m vllm.entrypoints.api_server \
   --model <model_path> \
   --max-model-len 32768 \ # 设置为模型支持的最大值 
   --gpu-memory-utilization 0.95 # 尽可能利用显存`  

2. 应用层截断优先： 不要依赖 vLLM 的自动截断。在 Python 代码中先计算 Token 数，确保发送给 vLLM 的 Prompt 长度小于 --max-model-len。这样你可以控制截断什么（例如优先截断不相关的 RAG 文档，保留记忆摘要）。
3. 异步总结任务： 使用 Celery 或 Redis Queue，在用户会话空闲时异步调用 vLLM 对历史对话进行总结，避免阻塞主对话流程。

---

六、总结

构建长上下文 RAG 系统，不能仅依赖推理框架的优化。

1. PagedAttention 是 vLLM 的自动特性，它优化显存效率，但不突破模型长度限制。
2. 滑动窗口 会导致记忆遗忘，不可单独使用。
3. 最佳实践 是结合 Query 改写、历史记忆检索 与 动态压缩。
   • 存：异步存储摘要和事实，而非原始对话。
   • 取：根据意图检索相关记忆。
   • 控：设置严格的 Token 预算，超限则二次压缩。

通过这套架构，你可以在有限的上下文窗口内，实现“无限”的记忆能力，既解决了显存爆满问题，又让模型记住了用户的所有关键信息。