Towards Data Science

Prefill Once, Fan Out: KV Snapshot Sharing for Multi-Agent LLM Pipelines

8.5内容质量

TL;DR · AI 摘要

通过共享KV缓存,多代理LLM推理效率提升近两倍,减少冗余计算。

核心要点

  • 共享KV缓存可使两代理流水线端到端加速48.69%(约1.95倍)。
  • 第二个代理的激活延迟降低98.09%(约52倍)。
  • 使用llama_state_get_data和llama_state_seq_set_data实现KV缓存序列化与恢复。

结构提纲

按章节快速跳转。

  1. 介绍多代理LLM推理中冗余计算的问题及SwarmKV解决方案。

  2. 多个代理读取相同文档时,传统方法会导致重复的prefill计算。

  3. 通过共享KV缓存,实现一次prefill,多分支复用。

  4. 使用llama_state_get_data和llama_state_seq_set_data进行KV缓存序列化与恢复。

  5. GTX 1080上实现端到端加速48.69%,激活延迟降低98.09%。

  6. 减少8,685 ms冗余计算,无需新算法,仅需系统工程优化。

思维导图

用一张图看清主题之间的关系。

查看大纲文本(无障碍 / 无 JS 友好)
  • KV缓存共享优化多代理LLM推理
    • 问题
      • 冗余prefill计算
      • GPU重复执行相同计算
    • 解决方案
      • 共享KV缓存
      • 序列化与恢复
    • 性能提升
      • 端到端加速48.69%
      • 激活延迟降低98.09%

金句 / Highlights

值得收藏与分享的关键句。

#LLM#KV缓存#多代理推理#系统工程
打开原文

一次预填充,分发出去:用于多智能体LLM流水线的KV快照共享 | Towards Data Science

智能体AI

一次预填充,分发出去:用于多智能体LLM流水线的KV快照共享

一种系统工程方法,用于KV缓存的重复使用,证明了一次计算共享提示比N次并行预填充更高效。

Anubhab Banerjee

2026年6月9日

28分钟阅读

分享

一次幽默但真实的SwarmKV之旅——KV快照分发、复制-分叉主机缓冲区,以及如何通过稍微“苛刻”地对待llama.cpp,使一个双智能体分析流水线的效率提高了约1.95倍(并且第二个分支的激活延迟提高了约52倍)。

这是“生产级智能体推理”系列的第一部分。每一部分都会从智能体LLM流水线中移除一种冗余的工作。第一部分(本文)消除了冗余的预填充。第二部分将解决冗余等待的问题——如何通过时间切片,使50个微智能体共享一个GPU。第三部分通过自定义的CUDA Top-K内核,使RAG检索在GPU上进行。第四部分则在智能体之间传递状态,使下一个智能体永远不会遇到冷启动问题。

关键要点

问题:当多个智能体阅读同一份文档时,一个默认的服务堆栈会使每个智能体重新运行完全相同的预填充。这种冗余的密集注意力传递纯粹是浪费。

解决方案:运行一次预填充,将KV缓存序列化到主机缓冲区,按分支进行memcpy,然后在解码前恢复它。“计算一次,分发出去。”

结果:在一台七年前的GTX 1080上,一个双智能体流水线的端到端速度提高了48.69%(约1.95倍),第二个智能体的激活延迟减少了98.09%(约52倍),消除了8,685毫秒的冗余计算。

关键点:这不是一种新的算法。这属于系统工程,而且这与5G基站自LTE以来每80毫秒所做的“广播共享状态”决定是一样的。

TL;DR:标准的LLM服务会使得每个分析智能体重新进行相同的共享文档的预填充。你的GPU忠实地重新执行数十亿次冗余的前缀预填充乘法。相同的字节。相同的权重。相同的量化。所有这些只是为了重新计算一个四秒前就已经完成计算的状态。SwarmKV只运行一次预填充,通过llama_state_get_data将生成的KV状态序列化到主机缓冲区,将该缓冲区复制到每个分支的分配中,并让每个分支在解码前通过llama_state_seq_set_data恢复快照。是的,这是一个真正的往返过程——序列化、复制、恢复,但由于冗余预填充计算呈二次方增长,而KV状态传输呈线性增长,因此将数据通过受限的Pascal内存总线传输仍然比从头开始重新计算注意力矩阵便宜得多。这一点在七年前的GTX 1080上的结果中得到了体现:双智能体流水线的端到端速度提高了48.69%,分支2的激活延迟减少了98.09%(约52倍),消除了8,685毫秒的冗余密集计算,没有新的Transformer技巧。只是系统工程的立场:“计算一次,分发出去”胜过“计算N次,希望没人注意到”。

GitHub仓库:https://github.com/AnubhabBanerjee/swarmkv

(在我们开始之前,快速坦白一下:我来自5G/6G RAN工程背景。事实证明,将共享计算分发给许多下游消费者,与自LTE以来每80毫秒基站广播SIB1时所做的操作惊人地相似。下面有一个专门的部分(第8部分)来讨论这一点,这也是我写这篇文章的原因之一。)

架构思维模型 —— 阅读时请保持这个模型在脑海中。文档 → PrefillNode → llama_state_get_data → host KV 缓冲区 → 每个分支的 memcpy → llama_state_seq_set_data → AnalyticalNode 解码(RoPE 在 prefix_seq_len 处继续)所有内容都只是对这条线中某一部分的评论。

SwarmKV 架构概述

1. 一个坦白:你第二个代理的“工作”大部分是重复执行

如果你曾经通过 vanilla llama.cpp 将两个分析代理指向同一份文档,这里就是真正发生的事情(带有一点刻意的夸张):

你:“请提供这份 3,500 个 token 规格的概述,并单独列出其许可义务。” llama.cpp(代理 1):“当然。加载模型。预填充文档。解码答案。” GPU 在密集注意力上花费了 4,346 毫秒。 llama.cpp(代理 1):“完成。这里是一个 6 个 token 的答案。” 你:“很好。现在是代理 2。” llama.cpp(代理 2):“当然。加载模型。预填充文档 ——” 你:“你刚刚才做过这个。” llama.cpp(代理 2):“我是一个独立的 llama_context。我对代理 1 没有记忆。我对任何事情都没有记忆。我是一个美丽、无状态的新生命。” 🫡

GPU 再次在 bit-for-bit 完全相同的注意力计算上花费了 4,339 毫秒。

你的 GPU 热传感器:得到了很好的锻炼;你的 AWS 账单:发展出幽默感;你第二个代理的 TTFT:在它回答一个 4 个 token 的问题之前,需要 4.3 秒。

这就是笑话。这就是每一个“代理式”流水线的肮脏秘密,这些流水线从共享文档中分支出去。每个分支都从一张白纸开始,重新构建前一个分支刚刚完成构建的 KV 缓存。文档越长,这个代价越严重。在 Pascal GPU 上处理 3,500 个 token 的文档时,第二个代理的感知延迟的大部分并不是答案本身,而是再次阅读文档。

SwarmKV 就是你决定第二次阅读是可选的,而不是让每个代理反复构建相同的 KV 缓存,宁愿写 1,500 行 C++ 的结果。

现在想象一下,这个仓库中的玩具演示是关于两个代理对摘要/许可检查的。它实际构建的工作负载是 N 个专门的评估器对一份密集的技术文档。设想一个 AI 专利和现有技术的流水线:一个 50,000 个 token 的技术规格作为根,五十个并发分支评估新颖性、映射权利要求、检索现有技术、检查实施自由度、评估伦理合规性,并将其翻译成特定司法管辖区的语言。在默认服务堆栈上,这个流水线的基础成本是五十次完整的预填充。SwarmKV 的成本是一次预填充加上五十次 memcpy。这种不对称性是故意设计的,也是这个仓库存在的全部原因。我曾单独写过关于检测 AI 在发明报告中的内容 —— 这是该部分的基础设施部分。发明者笔记本的问题正是 SwarmKV 被构建的原因。

2. 为什么预填充存在?(一分钟速成课程)

如果你已经了解,可以跳过。对于其他人,这里是简短的版本。

一个自回归的 LLM 在两个阶段中处理请求。预填充是密集的处理阶段,它将每个提示 token 通过每个 transformer 层一次,并填充每个层的 key/value(KV)缓存。然后解码阶段逐个 token 运行,关注预填充的 KV 缓存,并逐步扩展它。

预填充成本大致与提示长度呈线性增长。相比之下,解码每个标记的成本较低。在使用 Qwen2.5-7B Q4_K_M 的 Pascal 类 GTX 1080 上,预填充一个约 3,500 个标记的文档大约需要 4.3 秒;解码一个简短的分支提示则需要数百毫秒,因为其主要受设置时间影响,而不是算术运算。预填充和解码之间的时间差异正是 SwarmKV 所利用的杠杆。

主流的服务堆栈(vLLM、TGI、SGLang、llama.cpp 自带的服务器)将每个请求视为独立的上下文。其中一些支持前缀缓存,但通常是请求作用域或会话作用域,而不是图作用域。它们的设计目标是最大化许多独立用户提示的吞吐量,而不是在一个从单一共享文档发散出多个分支的分析流水线中共享状态。对于这种 DAG 形状的工作负载——一个根节点,多个叶子节点,共享相同的数据——我尝试过的所有公共堆栈都要求我为每个叶子节点单独支付一次根节点的费用。

SwarmKV 作为显式的编排层,通过 C++ 实现,绕过运行时抽象,确保指针生命周期的确定性,并提高硬件级别的 memcpy 效率。

3. “只需快照 KV”的灵光一现(以及为什么它听起来简单却更难实现)

这个思路很简单:

  • 在共享文档上使用序列 ID kSwarmkvPrefixSeqId 运行一次预填充。
  • 通过 llama_state_get_data 将生成的 KV 状态序列化到主机缓冲区中。
  • 对于每个下游分支,将该缓冲区复制到每个分支的分配空间中。
  • 启动一个新的 llama_context,调用 llama_state_seq_set_data 安装快照,然后使用从 prefix_seq_len 继续的 RoPE 位置解码分支提示。

这就是“计算一次,发散出去”的范式。唯一的原因是,要实现它需要超过 30 行的 llama.cpp 补丁,因为有三个繁琐的边缘情况会立即破坏这种简单的方法。这个概念非常简单,应该是一个轻松的周末项目,但底层硬件和系统现实使它成为一个巨大的工程挑战。

问题 A:KV 的大小有多大?

一个简单的答案是:n_layers × n_head_kv × n_ctx × head_dim × dtype × 2。不过,这个手动推导出的数字每次量化格式变化、每次 GQA 比例变化,或者每次引擎添加新状态字段时都会发生变化。唯一诚实的数字是引擎在当前构建下告诉你的数字。

因此,MemoryPool 启动一个临时的 llama_context 仅仅是为了询问:

cpp
size_t MemoryPool::get_required_kv_size(uint32_t n_ctx) {
    // 从库的默认值开始,确保我们不关心的字段保持合理。
    llama_context_params params = llama_context_default_params();
    params.n_ctx = n_ctx;

    // 仅为了查询序列化状态占用空间而构造一个临时的上下文。
    llama_context * ctx = llama_init_from_model(model_ref, params);
    if (!ctx) {
        throw std::runtime_error("MemoryPool::get_required_kv_size: llama_init_from_model failed.");
    }

    // 向 llama.cpp 询问这个上下文的完整状态 blob 所占用的字节数。
    const size_t sz = llama_state_get_size(ctx);
    llama_free(ctx);

    // 如果引擎报告为零,则回退到一个较小的非零分配,以确保测试仍然可以执行注册,而无需假装我们了解张量的精确布局。
    if (sz == 0) {
        return size_t{1} << 20;
    }
    return sz;
}

这里的逻辑很简单:礼貌地向引擎询问,而不是盲目相信 PDF 或数学公式。只需启动一个上下文,询问所需大小并精确分配相应资源 —— 一个简单却非常成功的方案。

问题 B:llama.cpp 对并发解码的兼容性要求较高

在本项目中所使用的固定上游 llama.cpp 版本和 GPU 配置下,单个 GPU 上从多个线程进行并发解码并不总是安全可靠。具体行为取决于后端、版本和图调度器 —— 在更新的版本或使用独立流时,可能表现得更好 —— 但在我们的设置中,失败模式通常有以下几种:(a)崩溃,(b)KV 被损坏,或(c)在你 Google ggml 是否已经有了线程本地 arena 时,程序会卡上十分钟。剧透:在固定上游版本中,答案是否定的。

稳健的解决方案是在边界处对 llama API 接口进行序列化处理:

code
namespace swarmkv {

// llama.cpp 的 CUDA 路径在没有外部序列化的情况下,不能保证在单个 GPU 上从多个线程进行并发解码的安全性。
// 所有节点的 execute() 函数体必须在 llama_init / llama_decode / llama_free / 状态 I/O 周围持有一个互斥锁。
inline std::mutex & llama_api_mutex() {
    static std::mutex m;
    return m;
}

} // namespace swarmkv

struct LlamaGuard {
    std::lock_guard<std::mutex> lock;
    LlamaGuard() : lock(swarmkv::llama_api_mutex()) {}
};

一个简单的 20 行头文件定义了整个并发策略。每个节点的 execute() 函数体在 llama_init_from_model / llama_decode / llama_state_seq_set_data / llama_free 周围都持有这个互斥锁。DAG 级别的并发是真实存在的(如 future、依赖、扇出);GPU 计算在全局锁下交错进行。挑剔的人会正确指出,与假设的并发解码上游相比,这确实影响了性能。先记住这一点 —— 这正是本系列第二部分要解决的瓶颈。

问题 C:没有稳定的外部 KV 绑定 API

从美学角度来看,最理想的实现方式是分配一个连续的 KV 缓冲区,直接将其附加到新的上下文中,从而完全跳过 memcpy。上游的 llama.cpp 提供了 llama_memory_t 和图解码路径,但本仓库中固定使用的公共头文件并没有提供一个稳定、导出的 llama_kv_cache_bind 风格符号。

因此,SwarmKV 采取了次优的方案:它保留了调用点,诚实地命名它,并在 llama_state_set_data 之上实现这条路径。

code
void KVHandoff::bind_contiguous_cache(llama_context * ctx, ggml_backend_buffer_t cache) {
    // 验证参数,以便在启动和 CI 烟雾测试期间快速发现误用。
    // 空上下文无法进行解码;空缓存句柄是配置错误。
    if (!ctx || !cache) {
        throw std::invalid_argument("KVHandoff::bind_contiguous_cache: null context or buffer.");
    }

    // 明确标记这两个参数在当前版本中故意未使用。
    // 这可以防止在启用严格警告标志时出现 -Wunused-parameter 警告。
    (void) ctx;
    (void) cache;

    // 此处不发出任何稳定的绑定调用;请参见上方文件级注释。
    // 当上游添加了支持的附加 API 时,仅在此函数中实现它。
}

我知道,我知道。这是一个什么也不做的函数。它有完整的参数验证,文档字符串长度是函数体的两倍,并且在调用图中有一个稳定的位置。它耐心地等待着上游让它真正完成它的工作。我这辈子写过更诚实的代码,只是记不清是什么时候了。

这也是细心的读者可能会提出疑问的地方:“如果 bind_contiguous_cache 是一个无操作(no-op),那么 MemoryPool 缓冲区又有什么用呢?”这是一个非常好的问题。MemoryPool 缓冲区是暂存区域 —— PrefillNode 将其 llama_state_get_data 数据块写入的规范缓冲区,同时也是每个分支通过 memcpy 读取的来源。Decode 本身使用的是上下文内部管理的 KV。Pool 缓冲区 = 主机端的分发临时空间;上下文 KV = 引擎自己的内容。两个内存区域,一个快照,没有魔法。

4. 五步流水线(真正酷的部分)

code
步骤 0: 验证文档 + max_branch + 128 ≤ n_ctx        (context_budget.h, 快速失败)
步骤 1: 构建 DAG;DFS 检查循环            (Orchestrator)
步骤 2: 启动 std::async 工作线程;等待 future 的结果      (Orchestrator)
步骤 3: 预填充一次,将 KV 序列化到主机缓冲区      (PrefillNode + MemoryPool)
步骤 4: memcpy 快照 → 分支缓冲区 → 解码       (AnalyticalNode + KVHandoff)

让我们通过实际代码逐个分析这些步骤。这些代码片段特意保持简短,而完整的文件非常小,值得一读。

步骤 0 — 快速失败的上下文预算

以下三行代码可以避免你收到自己过去发来的凌晨 3 点的 Slack 消息:

code
const int32_t required = prefix_tokens + max_branch + generation_headroom;
    if (required > limit) {
        throw std::runtime_error(
            "Context budget exceeded: prefix_tokens=" + std::to_string(prefix_tokens) +
            " max_branch_tokens=" + std::to_string(max_branch) +
            " headroom=" + std::to_string(generation_headroom) +
            " required=" + std::to_string(required) + " n_ctx=" + std::to_string(limit));
    }

在任何上下文被构建、任何 MemoryPool 缓冲区被分配或任何 GPU 内存被访问之前,都会运行这段代码。如果你请求 SwarmKV 将 4,000 个 token 预填充到一个 n_ctx=4096 的上下文中,并且有两个分支和 128 个 token 的解码预留空间,它会告诉你数学上无法实现,并进入睡眠状态。你对未来的自己最友善的事情,就是在分配第一个字节之前就拒绝这些不可能的配置。

步骤 1 — DAG 循环检测

Orchestrator 对依赖关系邻接表执行标准的 3 色 DFS:

code
// dfs lambda 遍历邻接表,并在发现回边表示循环时抛出异常。
    auto dfs = [&](auto self, const std::string & u) -> void {
        // 标记节点 u 当前在递归栈中(正在访问)。
        state[u] = 1;
        // 探索所有从 u 到下游节点 v 的出边依赖关系。
        for (const auto & v : adj[u]) {
            // 如果 v 正在访问,说明发现了循环 u -> v,必须中止流水线设置。
            if (state[v] == 1) {
                // 抛出包含边名称的异常,使图配置错误易于诊断。
                throw std::runtime_error("Dependency cycle detected: " + u + " -> " + v);
            }
            // 仅当 v 尚未被完全处理时才递归。
            if (state[v] == 0) {
                // 从子节点 v 继续 DFS。
                self(self, v);
            }
        }

我知道这听起来很无聊,但相信我,这是必要的。这相当于在跑步前检查你的鞋带。一旦跳过这一步,你的流水线将花费其短暂的生命周期等待自己。错误信息包含导致问题的边,因此你可以轻松找到拼写错误,而无需使用 grep。

code
worker_tasks.push_back(std::async(
            std::launch::async,
            [this, name, state, dependencies, &completion_promises, &completion_futures]() {
                // 读取此节点的水印要求一次,用于依赖项控制决策。
                const int32_t req = nodes.at(name)->required_prefix_tokens();
                // 根据 V2 水印规则等待每个上游依赖项。
                for (const auto & dep_name : dependencies) {
                    // 解析上游节点指针以检测预填充提供者。
                    ExecutionNode * dep = nodes.at(dep_name).get();
                    // 如果上游是预填充提供者,并且此分支使用水印控制,则等待水印。
                    if (dep->is_prefill_provider() && req >= 0) {
                        // 阻止直到 PipelineState 水印 >= 所需前缀标记(推测性开始)。
                        state->wait_for_watermark(req);
                    } else {
                        // 否则保留 V1 行为:等待上游节点线程完成。
                        completion_futures.at(dep_name).wait();
                    }
                }
                // 使用协调器默认的 n_ctx 预算构建 llama_context_params。
                llama_context_params params = llama_context_default_params();
                // 将 n_ctx 提升到 SwarmKV 默认的管道上下文,用于多 k 标记文档。
                params.n_ctx = kSwarmkvDefaultPipelineCtx;
                // 将模型/池/名称捆绑到 OrchestratorContext 中,用于节点 execute()。
                OrchestratorContext ctx = {
                    this->memory_pool->get_model(),
                    params,
                    this->memory_pool,
                    name.c_str(),
                };
                // 运行节点逻辑并履行承诺,以便依赖项可以继续。
                try {
                    // 分发到 PrefillNode 或 AnalyticalNode 实现。
                    nodes.at(name)->execute(state, &ctx);
                    // 向 shared_future 等待者发出成功完成的信号。
                    completion_promises.at(name).set_value();
                } catch (...) {
                    if (req > 0) {
                        state->signal_milestone_consumed(req);
                    }
                    try {
                        completion_promises.at(name).set_exception(std::current_exception());
                    } catch (...) {
                    }
                    throw;
                }
            }));

每个节点一个 std::promise,使用 std::shared_future 以便多个下游分支可以等待同一个上游完成,而无需传递 future。失败路径始终设置异常,因此依赖项不会无限等待。我们已经调试了其他选择,而且哦,我们真的不喜欢那样!

注意这个循环中没有的内容:任何关于预填充、KV 或分支的逻辑。协调器并不知道 PrefillNode 是什么。它知道名称、边和承诺。节点特定的工作在 execute() 中进行,并且完全通过 ExecutionNode 虚拟接口实现多态。对于子节点来说,只有一项职责,一点都不令人感到压力!

第 3 步 — 预填充一次,导出 KV

PrefillNode 按照以下顺序执行四个操作:

  • 从 examples/base_doc.txt 读取文档文本。
  • 对其进行标记化(使用 resize-on-negative-return llama 习语)。
  • 以 llama_n_batch(lctx) 为块大小,对序列通道 kSwarmkvPrefixSeqId 上的标记进行分块解码,确保绝对 RoPE 位置与绝对标记索引匹配:
code
// 绝对 RoPE 位置等于完整文档标记流中的索引。
batch.pos[i] = cur + i;
// 每个标记恰好属于一个序列 ID 列表。
batch.n_seq_id[i] = 1;
// 将所有文档标记绑定到共享前缀序列通道常量。
batch.seq_id[i][0] = kSwarmkvPrefixSeqId;
// 除了此处所有标记保留零值外,预填充期间禁用 logits。
batch.logits[i] = 0;
  1. 将前缀序列的 KV 导出到标准主机缓冲区,并为分支打上水印。
code
KVHandoff::bind_contiguous_cache(lctx, state->materialized_branch_buffer);
// 标记 prefill_complete,这样使用 kSwarmkvWaitForPrefillComplete 的分支可以继续。
state->mark_prefill_complete();

这就是整篇文章的要点,用两行话概括。这个仓库中的其他内容 —— 协调器、LlamaGuard、预算检查、文档中提到的 no-op —— 都是为了支持这两行,并将它们的输出通过一次 memcpy 传递给分支,而无需额外的往返。

第 4 步 —— 在 LlamaGuard 下进行分支解码

  1. 为每个分支分配一个与前缀相同 n_ctx 大小的缓冲区:
code
// 为 n_ctx 分配一个分支缓冲区,以便后续解码在相同的 blob 策略下有空间。
branch_buf = ctx->memory_pool->allocate_branch_cache(static_cast<uint32_t>(ctx->ctx_params.n_ctx));
  1. 将标准快照复制到分支缓冲区,也即著名的“复制-分支”操作:
code
// 完整预填充路径从标准暂存区复制到分支分配中。
KVHandoff::materialize_branch_cache(
    state->materialized_branch_buffer,
    branch_buf,
    fork_kv_bytes);

…点击后,变成:

code
std::memcpy(dst_ptr, src_ptr, ncopy);

就是这样。这就是“在存储层进行复制-分支”,字面上就是 memcpy。它是最基本的操作。你所读到的所有关于前缀共享的高级内容 —— RadixAttention 的引用计数、分页注意力的块表间接寻址 —— 都是基于这个相同的思路:不要重新计算,复制字节。

  1. 启动一个新的 llama_context 并恢复快照:
code
// 仅恢复前缀序列通道,使分支解码在序列 0 上保持隔离。
const size_t n = llama_state_seq_set_data(
    lctx,
    static_cast<const uint8_t *>(base),
    fork_kv_bytes,
    kSwarmkvPrefixSeqId);
    // 验证 llama 消耗的字节数恰好等于我们复制到分支缓冲区的字节数。
if (n != fork_kv_bytes) {
    // 在抛出异常之前释放上下文,以避免在失败路径上泄漏 VRAM。
    llama_free(lctx);
    // 抛出一个清晰的错误信息,以便操作员可以快速调试大小不匹配的问题。
    throw std::runtime_error("AnalyticalNode: llama_state_seq_set_data size mismatch.");
}
  1. 为短分支提示构建一个 llama_batch,其中 RoPE 位置从前缀结束的位置继续:
code
for (int i = 0; i < batch.n_tokens; ++i) {
    // 将第 i 个分支的 token id 复制到 batch 的对应位置。
    batch.token[i] = tokens[static_cast<size_t>(i)];
    // 将分支 token 置于 forked 前缀位置之后,以确保 RoPE 正确。
    batch.pos[i] = static_cast<llama_pos>(fork_prefix_len) + static_cast<llama_pos>(i);
    // 每个 token 仅参与一个序列 id 列表条目。
    batch.n_seq_id[i] = 1;
    // 将所有分支 token 绑定到共享前缀序列车道常量。
    batch.seq_id[i][0] = kSwarmkvPrefixSeqId;
    // 禁用所有非该分支最后 token 的 logits。
    batch.logits[i] = 0;
}

这是每个人第一次都会搞错的地方。如果你忘记了偏移量,把分支位置从零开始,旋转嵌入会悄悄出错,模型会从一个前缀从未训练过的的位置进行解码。症状是自信地产生连贯的胡言乱语。欢迎来到最糟糕的调试地狱;请在离开时留下小费。

  1. 最后,一次 llama_decode。只需将诊断字符串写入 PipelineState::node_outputs,记录 timings_ms[name],释放 batch 和上下文。

每个分支有三行业务逻辑。解码时有两个上下文在运行。每次一个 memcpy。一个全局锁。其余的都是管道工作。

5. 证据(即数字)

现在是时候将其与基线进行对比,看看是否值得经历这么多麻烦。所有数字都来自 examples/example-run-results/。

在有人开始扔石头之前,先快速说明一下方法论:以下所有比较都运行相同的模型(Qwen2.5-7B-Instruct-Q4_K_M.gguf),相同的文档(一个确定性的 3,501-token 合成文档,通过重复 “The quick brown fox jumps over the lazy dog.” 直到达到 token 目标生成 — examples/base_doc.txt),相同的 GPU(GTX 1080,8 GiB,Pascal sm_61),相同的 n_ctx=4096,相同的 dtype。基线是两个顺序的 llama_context 实例,每个都预填充整个文档,然后解码其分支提示。SwarmKV = 一次 PrefillNode + 两次 AnalyticalNode 分支在快照上。工作负载类型:以预填充为主导的文档分析(RAG 风格),而不是自回归聊天。进行了三次连续的试验,每次试验之间有 GPU 空闲等待;最佳结果由 2·TTFT_pct + E2E_pct 选择。

在表格之前,先定义一个指标,因为它很重要:我们使用“分支-2 激活延迟(TTFT 代理)”——而不是教科书式的服务文献中“请求到达 → 第一个输出 token”TTFT。我们指的是第二个分支在分支特定工作上花费的时间:在共享预填充分摊到所有分支之后的激活延迟。在分发管道中,下游消费者感知到的成本正好是这个数字,因为上游预填充是为整个管道一次性支付的。这个指标的基线值是第二个 llama_context 被迫在回答之前重复执行的冗余文档预填充;SwarmKV 的值是分支 + 恢复 + 短提示解码。

标题:GTX 1080,Qwen2.5-7B Q4_K_M,3,501-token 文档,两个分支

指标

基线(HF 风格)

SwarmKV

差异

端到端墙钟时间

10,275 毫秒

5,272 毫秒

−48.69 %

(~

1.95×

)

分支-2 激活延迟(TTFT 代理)

4,339 毫秒

83 毫秒

−98.09 %

52.3×

基线 Agent-1 预填充

4,346 毫秒

基线 Agent-2 预填充

SwarmKV 每个分支解码(平均)

77 毫秒

冗余预填充消除

8,685 毫秒

每个分支的 ~83 毫秒去哪了

整篇文章的论点建立在一个不等式之上:每个分支的恢复 + 解码操作,远远比重复的文档预填充便宜得多。Harness 在聚合层面上直接测量了这一点——每个分支的墙钟时间(从分配、复制、恢复到解码,端到端)根据我们查看的分支不同,大约在 71–83 毫秒之间,而重复预填充的成本约为 4,339 毫秒。在聚合层面上,大约 52 倍的比率是使整篇文章中其他一切内容得以运作的关键。

如需更多结果和数字,我建议直接查看示例运行报告。

6. “好的,但这与 vLLM / 前缀缓存 / SGLang RadixAttention 有什么不同?”

这是一个非常合理的问题,值得直接回答,因为推理基础设施领域有许多重叠的原始技术,而 HPC 读者会在第一个评论中提出这个问题。

  • vLLM / 连续批处理 / 分页注意力。优化用于多租户解码时服务:在不同解码步骤上有许多并发请求,流式负载下在它们之间调度下一个标记。主要技术:分页注意力。工作单位:一组独立用户提示的流式火 hose。
  • TGI / vLLM 前缀缓存。如果你的共享前缀是请求作用域或会话作用域,效果非常好。它并未设计成将 KV 快照作为第一类对象暴露出来,以便传递给同一进程中运行不同下游任务的另一个 llama_context。
  • SGLang RadixAttention。在服务运行时内部的树状前缀共享——最接近的亲戚,但它是一个服务器,而不是单进程的编排原语。
  • llama.cpp 自己的状态保存/恢复。每个上下文都存在。SwarmKV 是管道级别的粘合剂:一个 DAG,一个由引擎本身确定大小的主机缓冲区竞技场,一个 memcpy 扇出,一个 LlamaGuard 策略,以及一个记录在案的 no-op,耐心等待上游绑定 API。

7. 那么……我该如何实际尝试它?

嗯,我已经在文章开头发布了 GitHub 链接。如果你已经读到这里,再努力一次,滚动回顶部。

相关文件位于 examples/example-run-results/ 目录下:best_run.json、all_trials.csv、plots/*.png,以及一份叙述性的 final_result.docx,它会逐步介绍方法和限制。

要求:Linux、CUDA 工具包、NVIDIA GPU(Pascal 或更新版本;消费级或数据中心级均可)、一个适合你显存的 GGUF 模型,以及耐心阅读一次 CMake 文件的耐心。

8. 剧情反转——这其实只是 SIB 广播披着变压器的外衣

我应该在此时坦白:我并不是一个“GPU 专家”。我的背景是电信行业——5G NR,一只脚坚定地踏入 6G 研究——我开始关注 LLM 推理基础设施,是因为这个代码库中的每一个问题都奇怪地熟悉。

一个简短的解码器,供没有3GPP背景的读者理解:在5G网络中,基站不会单独向每部手机发送网络配置信息,而是通过共享信道一次广播一组小的系统信息块(SIB1、SIB2等),范围内的每部手机都会读取相同的广播,而每个用户的数据则通过专用信道在该共享上下文中传输。下表中的缩写词——MIB(主信息块,每部手机读取的第一项信息)、PBCH和PDSCH(共享广播和下行数据信道)、HARQ(接收端的“保留已解码内容,只重传缺失部分”的重传协议)、RNTI(用于区分一部手机流量与其他手机流量的临时ID)——只是用于标识共享信息(仅计算一次)与每个用户唯一信息之间的信道和标识符的名称。这种区分就是整个类比的核心。

请将以下内容并排对比,告诉我这些是否是不同的问题:

5G NR小区广播(在gNB)

SwarmKV(在GPU)

每个SS突发在PBCH上有一个MIB

一次在MemoryPool中共享文档并进行标记化

在PDSCH上重复SIBs(SIB1、SIB2等)

在MemoryPool中序列化KV快照

小区内每个驻留的UE读取相同的SI

每个分析分支读取相同的快照

为单播用户数据的UE提供专用PDSCH

每个分支

code
llama_context

解码分支提示

每个UE的RNTI区分单播流

每个分支缓冲区 + 序列ID区分分支状态

HARQ软缓冲区在重传中保留

KV快照在分支中保留

跳过广播 → 每个UE强制单播SI → 空中接口崩溃

跳过快照 → 每个分支重新预填充文档 → GPU崩溃

给两个截然不同的受众的简短说明

对于正在阅读本文的HPC和CUDA优先的朋友们:我知道。KV重用并不是一个新概念。vLLM有前缀缓存,SGLang有RadixAttention,llama.cpp本身也暴露了状态保存和恢复。SwarmKV的贡献不是原始技术;而是单进程编排的形状——一个微型C++ DAG运行时,将“预填充一次,分发N个分支”作为一等操作,适用于一个8 GiB的消费者GPU,并带有研究人员实际需要的安全机制(LlamaGuard、swarmkv_validate_context_budget、文档中提到的bind no-op),以便在周二推出一个演示。请放下你们的叉子。

对于我的电信朋友们:如果“KV缓存”在十分钟后听起来像是一种外语,你们不是落后了,而是领先了。二十年来,我们的世界是FPGAs、ASICs和PRBs。我们优化的是频谱,而不是硅。然后AI-RAN、NWDAF、NVIDIA Aerial、AI-RAN联盟和3GPP Rel-20研究项目都在大约十八个月内相继发生,现在电信职业的下一个十年要求在频谱世界和GPU世界之间双语精通。直觉可以清晰地转换。从第一个CRS试点开始,你们就已经将共享计算分发给许多消费者了。同样的动物,只是换了一个动物园。

9. 诚实的注意事项(因为评论即将来临)

如果你来这里是为了找出项目中的问题——恭喜你,项目找到了它的第一个读者。从final_result.docx的限制部分和源代码中的内联注释中:

  • KV staging 是在主机端进行的。MemoryPool 从 CPU 设备( ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU) )分配 ggml_backend_buffer_t。分支解码仍然在 GPU 上运行;只有快照传输是通过 llama_state_get_data → memcpy → llama_state_seq_set_data 在主机端进行的。一个设备感知的 materialize 功能已经在路线图上,但目前被同样的上游 KV 绑定 API 阻挡,该 API 也是 bind_contiguous_cache 正在等待的。
  • 共享解码互斥锁(基于固定上游版本)。LlamaGuard 对来自工作线程的每个 llama_* 调用进行序列化。在本项目所使用的 llama.cpp 版本和 GPU 配置下,单个 GPU 上多个线程并发解码并不总是安全的 —— 具体行为取决于后端、版本和图调度,但在我们的设置中,保守的选择是使用全局锁。DAG 级别的并发是真实的,但每个请求的 GPU 计算仍然是顺序的。这是 V1 版本中最大的性能限制,而本系列的第二部分正是从这里开始。
  • SwarmKV_Prefill_Ms 报告为 0。在 PrefillNode 内部如何使用 OrchestratorContext::node_name 的已知监控缺口。预填充确实运行了(你可以在 End_To_End_Ms 和由此推导出的有效共享预填充成本中看到其成本),只是没有正确地记录到 timings_ms 中。有效共享预填充的计算方式为 SwarmKV_End_To_End_Ms − max(SwarmKV_AgentA_Ms, SwarmKV_AgentB_Ms) ≈ 5,189 ms。这是一个报告错误,而非正确性错误。已记录。
  • 合成文档。基准测试通过重复 “The quick brown fox jumps over the lazy dog. ” 直到达到目标 token 数量,构建了一个确定的 3,501-token 文档。这将性能信号与内容影响隔离,并确保每次试验在比特级别上可复现。真实文档的每次试验绝对时间会更嘈杂;但结构比例不会改变。
  • 单 GPU 类别。报告中的所有数字都来自一个 Pascal 类 GTX 1080。更新的 GPU(如 Ada、Hopper)预填充速度更快 —— 绝对毫秒数会减少,但完整预填充成本与短解码成本之间的结构比例(SwarmKV 所利用的)不会改变。
  • bind_contiguous_cache 是一个已记录的无操作函数。是的,目前仍然如此。直到上游提供一个稳定的外部 KV 附加 API,该函数会验证其参数,将其转换为 void 类型,然后返回。

但不用担心,列表中的所有内容都在路线图上。不过,这些内容不会改变主要结果。将其写下来的目的,是让你不必去挖掘这些信息 —— 一旦基准测试博客文章隐藏了它的限制条件,它的数据就不再值得信赖了。

10. V1 的上限(以及第二部分的设置)

SwarmKV 证明了你可以停止重新预填充。但如果你重新阅读了注意事项 #2,你已经注意到了下一个上限:GPU 计算本身仍然是串行的。

在实际的墙上时钟时间中,DAG 级别的并发是真实的 —— 分支确实是真实的 std::async 工作者,具有真实的依赖关系控制。但每个分支的 llama_decode 都在 LlamaGuard 内运行,这是一个全局的单个互斥锁。因此,尽管调度扩展了,GPU 的工作却只能按顺序进行。两个分支轮流执行,五十个分支轮流执行五十次。GPU 从未真正共享过;它是一次又一次手动时间复用的,没有公平性保证,也没有办法衡量谁在等待谁。

这在两个代理的演示中是可行的。但一旦运行 SwarmKV 实际设计用于的工作负载:50 个专用微代理竞争一个 GPU,这种方案就无法维持。在这样的规模下,你不再关心“是否避免了重新预填充”,而是开始关注手动实现的互斥锁无法回答的问题:

  • 当 50 个代理同时请求 GPU 时,谁先获得使用权,如何确保公平?
  • 每个代理在共享一个 GPU 卡时,其 p50、p95 和 p99 的延迟是多少?
  • 竞争引入了多少抖动,吞吐量在什么情况下会下降?
  • 我们如何有目的地划分 GPU 计算周期,而不是偶然地?

这是本系列文章的第二部分:为并发代理群集进行 GPU 时间切片。玩具代理在 Python 中按顺序运行。生产代理则在裸金属上并发运行,而当多个微代理共享一个 NVIDIA GPU 时,管理 VRAM 和计算资源本身就是一门专门的学科。第二部分构建了一个 Kubernetes 级别的时间切片分析器,它可以动态分配计算周期,并在通过 Kubernetes Device Plugin 与 CUDA 时间切片共享 GPU 的代理推理工作负载中,测量 p50/p95/p99 延迟、抖动和吞吐量代理指标。SwarmKV 中的全局互斥锁正是被这种可衡量的机制所取代的东西。

(对好奇的读者:还有一个独立的、正交的 V1 限制,值得在未来 SwarmKV V2 的文章中讨论——当前的流水线在任何分支开始之前,必须等待整个预填充完成,即使某个分支只需要前 500 个上下文标记。允许分支在所需前缀切片被实现的瞬间立即开始,确实是一个真正的优势,但这是另一个故事,另一个基准。这不是第二部分的内容。第二部分是关于如何在多个代理之间共享 GPU;而这个预填充流式处理的想法则是后续内容。)

我们第二部分再见。

免责声明:本文中的插图(头图、架构图、电信与 SwarmKV 对比图、GPU 时间切片图)是使用 AI(Claude Opus 4.8)生成的。这些图是示意性的,而非照片,图像中可见的任何标签都是风格化的,而非权威性的——请参阅文章正文和代码本身,以获取精确的功能名称、指标值和架构细节。

作者

查看 Anubhab Banerjee 的所有文章

人工智能

,

深入解析

深度学习

Kv 缓存

多代理系统

分享本文

  • 在 Facebook 上分享
  • 在 LinkedIn 上分享
  • 在 X 上分享

Towards Data Science 是一个社区出版物。提交你的见解,以触达全球读者,并通过 TDS 作者支付计划获得报酬。

将 href 更新为你的实际提交 URL

为 TDS 撰写文章

✦ 结束 CTA ✦