MemStitch 为 vLLM 实现零拷贝上下文桥接，首词延迟提速 25 倍

背景

在多智能体（Multi-Agent）协作工作流中，多个智能体常常需要顺序处理同一段长文本上下文。例如，Agent A（法律审计）先读取一份200页的合同并运行合规分析，此过程中会填充GPU的KV缓存；随后Agent B（财务合规）需要读取同一份200页合同并审计财务负债。在标准推理引擎（如vLLM）下，Agent B被迫重复执行昂贵的预填充（prefill）阶段，导致GPU激活重复计算，并产生高首令牌延迟（TTFT）。这种冗余不仅浪费计算资源，还显著拖慢整体协作效率。

核心内容

MemStitch（项目名称为Context-Stitcher）通过内存级别的缓存桥接来解决上述问题。其核心机制包括以下三个部分：

• 上下文拓扑哈希（Context Topological Hashing）：将提示（prompt）按物理块大小进行分段，并将每一段映射为加密指纹（基于Merkle链）。这样系统可以快速识别不同智能体使用的相同上下文前缀。

• 零拷贝块拼接（Zero-Copy Block Stitching）：对于匹配的前缀，直接绕过预填充阶段。通过将Agent B的逻辑注意力表（logical attention table）映射到Agent A缓存块的物理GPU内存地址，实现无需数据复制的上下文共享。

• 零信任安全网关（Zero-Trust Secure Gate）：实施边界控制列表（boundary control lists），防止未授权的智能体会话访问共享物理块，确保安全性。

基准测试

针对一个共享的200页文档，在连续执行两个智能体的场景下，MemStitch与标准vLLM冷启动预填充的对比结果如下：

• TTFT预填充延迟（Agent B响应时间）——越低越好
  基线（vLLM冷启动）：1200 ms
  MemStitch：48 ms（预填充速度提升25.0倍）

• GPU物理缓存块分配（总显存）——越低越好
  基线（vLLM冷启动）：53个块（无共享）
  MemStitch：30个块（节省43.4%显存）

使用方式

安装依赖后运行 python run.py，网关将在 http://localhost:8000 上激活。提供两种接入方式：

Python SDK装饰器

通过 StitcherMesh 和 @stitch_agent 装饰器链接上下文内存：

from context_stitcher import StitcherMesh, stitch_agent

mesh = StitcherMesh(backend="vllm", model="meta-llama/Llama-3.1-8B-Instruct")
mesh.sg.add_policy("agent_a", "agent_b")

@stitch_agent(mesh)
def agent_a():
    prompt = "[长文档上下文...]\n分析知识产权条款。"
    res = mesh.generate(prompt=prompt, fingerprint="legal_doc_v1")
    return "legal_doc_v1"

@stitch_agent(mesh)
def agent_b(context_fingerprint):
    prompt = "[长文档上下文...]\n评估财务合规风险。"
    res = mesh.generate(prompt=prompt, fingerprint=context_fingerprint)
    print(f"Agent B响应: {res['generated_text']}")
    print(f"节省时间: {res['prefill_time_saved_ms']}ms")

OpenAI兼容REST API

将LLM客户端的base URL指向MemStitch代理网关：

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="not-needed")
response = client.chat.completions.create(
    model="context-stitcher-sim",
    messages=[{"role": "user", "content": "[长文档上下文...]\n评估合规风险。"}],
    extra_body={"agent_id": "AgentB", "session_id": "session_legal_audit"}
)

或通过curl：

curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "context-stitcher-sim",
    "messages": [{"role": "user", "content": "[长文档上下文...]\n评估合规风险。"}],
    "agent_id": "AgentB",
    "session_id": "session_legal_audit"
  }'

动态访问授权

支持运行时查看、添加或撤销智能体间的访问规则：

curl -X GET http://localhost:8000/policies
curl -X POST http://localhost:8000/policy \
  -H "Content-Type: application/json" \
  -d '{"owner_agent": "AgentA", "allowed_reader": "AgentB"}'
curl -X DELETE http://localhost:8000/policy \
  -H "Content-Type: application/json" \
  -d '{"owner_agent": "AgentA", "allowed_reader": "AgentB"}'

此外，MemStitch自带一个响应式开发者控制台（访问 http://localhost:8000），可实时监控物理缓存块状态（空闲、私有分配、共享/拼接页、安全警报等）。

关键要点

• 零拷贝上下文桥接：通过直接映射GPU内存地址，避免重复预填充，而非复制数据。
• 上下文拓扑哈希：利用Merkle链对提示分段进行指纹识别，快速匹配相同前缀。
• 显著性能提升：在200页文档的连续智能体场景下，TTFT延迟降低25倍，显存占用减少43.4%。
• 零信任安全模型：通过边界控制列表（BCL）隔离智能体会话，防止未授权访问共享缓存。
• 灵活整合方式：支持Python装饰器（StitcherMesh）和OpenAI兼容REST API两类接入方式。
• 动态策略管理：运行时可通过REST接口添加、查询和删除智能体间的缓存共享授权。
• 实时可观测性：提供网页控制台监控缓存块状态和安全事件。

意义与影响

MemStitch的出现直接解决了多智能体串行推理场景中普遍存在的冗余预填充问题。在LLM应用中，多轮文档分析、长上下文审查、代码仓库审计等协作工作流日益普遍，而标准推理引擎的高延迟和显存浪费已成为瓶颈。MemStitch通过内存级缓存共享，使得后续智能体几乎零开销地继承先前智能体的计算成果，这将在以下方面产生深远影响：

• 降低推理成本：减少GPU计算和显存占用，尤其适用于长文档、大规模智能体编排场景。
• 加速复杂工作流：在RAG、多步Agent、合规审查等应用中，大幅缩短端到端响应时间。
• 安全与可管理性：零信任网关确保了多租户环境下的数据隔离，避免敏感上下文泄露。
• 易于集成：既提供Python原生SDK，又兼容OpenAI协议，可无缝替换现有推理后端。

MemStitch将推理引擎的优化从单请求扩展到多智能体协作层面，是LLM系统基础设施中重要且实用的进步。