大模型推理上下文管理与切换机制：KV Cache、显存与调度实战

做了几年大模型推理服务的工作，我越来越觉得"上下文"这个词被用得太随意了。业务同学说上下文是聊天记录，算法同学说上下文是 KV Cache，运维同学说上下文是显存里的那一堆数据。同一个词，三种理解，等线上出问题的时候，谁也说不清模型执行上下文到底该管什么、该怎么切。

这篇文章我把"模型执行上下文管理与切换机制"这件事，从系统工程师的视角完整讲一遍。既包括执行上下文的组成拆解，也包括多用户、多模型场景下上下文切换的机制设计和落地经验。内容主要面向正在做大模型服务部署、推理加速、模型网关的同学，也适合打算把推理服务从"单路跑通"推到"多路稳定"的团队参考。全文不贴框架源码，只讲原理和实操决策，你拿回去对着自己的服务排查，一定能用上。

1. 先搞清楚：模型执行上下文到底由什么组成

1.1 业务上下文和执行上下文，是两个不同维度

刚开始做服务化的时候，我犯过一个典型错误：把对话历史、用户 session 数据当成唯一需要管理的"上下文"。后来线上出现了一个诡异问题：用户发同样的话，请求时好时坏，有的推理结果语义对不上，有的首 token 延迟明显变高。排了很久才发现，问题出在推理服务的运行时状态上——同一份业务对话，落到不同的执行上下文里，表现完全不同。

所谓"模型执行上下文"，从系统的角度看，至少包含四层信息：

• 模型权重与计算图状态：已经加载的权重、算子图、算子编译缓存、CUDA kernel 的编译产物等；
• KV Cache：自回归解码时保存的中间 Key/Value 张量，这是大模型推理里最占资源、最影响切换成本的一层；
• 运行时资源状态：显存分配器状态、GPU stream 状态、CUDA context、TensorRT engine 的 profile、CUDA graph 捕获结果等；
• 业务会话状态：prompt 历史、采样参数、多轮截断策略、工具调用记录等。

业务上下文只是最上面那一层。真正影响性能、稳定性和响应速度的，是下面那三层。我在做推理网关调度时最深的体会是：如果你只按"业务会话"维度管理上下文，那你在资源维度上的切换开销几乎不可控。

提示：排查上下文类问题时，建议先区分"业务上下文错误"和"执行上下文错误"。前者表现为生成内容答非所问、记忆混乱；后者表现为性能劣化、CUDA 报错、显存异常。两类问题修起来完全是两条路。

1.2 KV Cache 才是切换成本的大头

很多人不理解为什么上下文切换这么贵，觉得无非是把对话历史再拼一遍。问题就出在 KV Cache 上。

自回归模型生成每个 token 时，都需要把当前 token 与之前所有 token 的关系算一遍。如果不缓存，每次生成都得重新处理整个输入序列，显存省了但计算量翻了几倍。所以主流推理框架都会把历史 token 经过 Attention 计算后的 Key 和 Value 向量缓存下来，这就是 KV Cache。

KV Cache 的占用有多大？可以算一笔账。假设一个 7B 参数的模型，32 层 Transformer，32 个 attention head，每个 head 的维度是 128，KV Cache 用 FP16 存储。那么每增加一个 token，缓存大小大约是：

2（K 和 V 两组）× 32（层数）× 32（head 数）× 128（head 维度）× 2 字节（FP16）= 512KB

也就是说，一个用户多轮对话累积到 2048 token，光 KV Cache 就要吃掉约 1GB 显存。如果服务里同时挂着 16 路这样的会话，光缓存就要 16GB。这个数据一摆出来，你就明白为什么上下文管理的第一优先级不是"存对话记录"，而是"管理 KV Cache 的存储、换入和换出"。

权重和计算图是"静态上下文"，整个服务启动时加载一次，切换模型才动它。KV Cache 是"动态上下文"，每个请求、每轮对话都在变，真正需要频繁切换的其实就是它。

2. 哪些场景逼着你做上下文切换

2.1 多用户、多会话复用单卡

最典型的场景是单卡部署一个模型，同时服务多个用户会话。模型权重只有一份，但每个用户会话都有独立的 KV Cache 和采样状态。显存不够时，就只能把一部分会话的上下文切出去（保存到 CPU 内存），把另一部分会话的上下文切进来（恢复到显存）。

这个过程和操作系统的进程切换非常像。进程切换要保存寄存器、页表、内存映射，模型执行上下文切换要保存 KV Cache、分配器状态、CUDA 环境。只不过模型服务的切换粒度更大，一次切换可能就是几百 MB 甚至上 GB 的数据搬运。

我曾经在一张 24GB 的卡上推理一个 7B 模型，权重 FP16 大约占 14GB，满打满算只够放 10 路 2048 token 左右的会话缓存。业务方要求同时支持 30 路会话，靠的就是上下文切换：活跃会话留在显存，不活跃的会话把 KV Cache 落盘到内存或 SSD，等用户再次发言时再换回来。

2.2 多模型、多 LoRA 之间的切换

另一种常见场景是同一个服务里挂着多个模型，或者同一基座模型挂了多个 LoRA 适配器。

全量模型切换很贵，因为权重本身就是十几个 GB 的搬运，一般不会做成高频操作，更多是配合流控在服务低峰期做。LoRA 切换则便宜得多。LoRA 参数通常只有几十到几百 MB，切换时只需要替换适配层权重、重算缓存前缀的 KV Cache、重建采样参数，整个流程能控制在百毫秒级别。

我在做多租户推理网关时，一个基座模型下面挂了 20 多个 LoRA 适配器，每个租户请求来了先按租户 ID 查路由表，再决定当前 GPU 上需要切到哪个适配器。如果频繁切换导致性能下降太明显，还可以把热门适配器做常驻，冷门适配器按需加载，这本质上是把上下文管理从"单请求级别"提到了"租户级别"。

2.3 抢占式调度与资源回收

第三种场景是你在做推理服务集群，多个推理实例共享 GPU 资源。一个实例内存吃紧时，需要抢占另一个实例的显存；或者一个长任务占着显存不动，后面排队的短任务一直等不到资源。

这时候就需要一套上下文切换策略：被抢占的任务先把 KV Cache 固化保存，让出显存；等资源充裕时再恢复执行。这套机制如果做得好，长任务可以"暂停"，短任务能"插队"，集群的整体吞吐会明显提升。

做设计之前先把场景列清楚，因为不同场景下"切换"这个词的含义完全不同，后面每一层的技术选型都会跟着变。

3. 核心机制拆解：上下文保存、恢复与切换流程

3.1 KV Cache 的保存：不只是"把张量拷出来"

先说最基础的保存。很多第一次做的人以为 KV Cache 保存就是把 tensor 序列化写个文件。实际没那么简单，里面的坑至少有三个。

第一，KV Cache 在显存里往往不是连续块。推理框架为了提升显存利用率，普遍用类似分页的机制管理缓存，一个会话的 KV Cache 可能分散在多个物理块里。保存时你手上拿到的可能是一堆"块描述符"，需要先知道每个块的大小、位置和逻辑顺序，才能正确重组。直接拿原始指针去拷贝，拷出来的数据顺序是乱的。

第二，KV Cache 依赖模型的配置信息。保存时必须带上模型版本、层数、head 数、head 维度、dtype、注意力计算方式等元数据。否则恢复时别说对不上号，连张量形状都校验不过。实际操作中我建议在序列化文件头里固定写一串 magic number 加版本号，恢复时先校验再分配。

第三，保存动作本身要尽量不阻塞推理线程。KV Cache 拷贝是内存搬运，动辄几百 MB，如果直接在推理线程上执行，GPU 那几百微秒的空闲就会被放大成几十毫秒的延迟毛刺。我这边通常的做法是：先通过 cudaMemcpyAsync 把数据从显存拷到固定的 staging buffer，然后放到后台线程里做 CPU 内存和磁盘之间的搬运，推理线程只负责把"脏页"标记出来。

python复制def save_context(model, session_id, kv_blocks, meta):
    # 1. 从显存池中收集该 session 的 KV 物理块
    blocks = memory_pool.collect_blocks(session_id)
    # 2. 把分散块映射为逻辑连续的缓冲区
    staged = staging_buffer.acquire(len(blocks) * block_size)
    stream.copy_async(blocks, staged, direction="D2H")
    # 3. 带上元数据，异步写入 CPU 内存或 SSD
    payload = build_context_payload(staged, meta)
    background_writer.submit(session_id, payload)
    # 4. 释放显存块，归还内存池
    memory_pool.release_blocks(blocks)
    return ContextHandle(session_id, meta.version)

这个伪代码里的 memory_pool 和 staging_buffer 才是关键，前者负责显存复用，后者负责让拷贝和推理重叠。没有这两个组件，你的保存逻辑性能天花板会很低。

3.2 KV Cache 的恢复：先校验，再分配，最后重建依赖

恢复阶段比保存更容易出问题，因为你要面对的是一个"不干净"的环境。我踩过最大的坑是版本不一致：模型微调过一次，注意力头数没变但层数变了，旧会话的 KV Cache 直接塞进来，推理结果就全乱了，而且不是报错，是那种很隐蔽的"越聊越差"。

所以在恢复流程里，第一件事必须是校验 meta。我把模型版本、tokenizer 版本、dtype、max_seq_len 全部写进上下文头，恢复时逐项比对，不一致就拒绝恢复并降级为重新处理历史对话。这个降级逻辑一定要提前设计好，否则用户那边体验就直接断裂。

第二件事是显存分配。恢复前要确认当前的显存池里有没有足够大的连续可分配空间。如果不够，要么先切换走另一个不活跃会话，要么放弃恢复走降级路径。分配拿到之后，先把 KV Cache 数据搬回显存，再重建依赖关系——比如把显存块注册到 Context 管理器、绑定到新的 CUDA stream、重新设置 attention mask 等。

我给一个经验参数：对于 7B 模型、2048 token 的上下文，在 PCIe 4.0 环境下从 CPU 内存恢复到显存，耗时大约 30 到 80 毫秒；如果从 SSD 恢复，要在此基础上加 20 到 50 毫秒的磁盘读取时间。这个数字决定了你调度策略的阈值设计——如果用户平均离开时间只有几十毫秒，来回切换不如直接让会话占着显存。

python复制def restore_context(handle, model, device):
    meta = handle.read_meta()
    assert meta.version == model.version, "context version mismatch"
    blocks = memory_pool.allocate(handle.blocks_needed)
    stream.copy_async(handle.payload, blocks, direction="H2D")
    model.attach_kv_cache(handle.session_id, blocks, meta)
    rebuild_dependencies(handle.session_id, model)
    return ContextSession(handle.session_id, blocks)

注意：恢复 KV Cache 后，务必重新构建与该 context 关联的 CUDA stream 和 attention mask。我见过有人只把数据搬回显存，忘了重建 mask，结果生成出的内容前半段正常、后半段混乱，极难排查。

3.3 切换决策：什么时候该切、先切谁

机制层解决"怎么切"，策略层解决"什么时候切、切谁"。我实践下来，上下文切换决策主要看四个信号：

• 显存水位：达到预设阈值（比如总量的 85%）就触发清理；
• 会话活跃度：按最后活跃时间排序，最久没动的先切换出去；
• 会话长度：KV Cache 越长的会话切换成本越高，不能简单按时间排序；
• 任务优先级：实时交互会话优先保留，离线批处理任务可以随时切。

最简单的策略是 LRU：维护一个按最后活跃时间排序的队列，显存不够就从队尾开始切。但纯 LRU 有个问题：长会话的切换成本很高，切一个 8K token 的会话比切一个 256 token 的会话贵几十倍。所以我在生产环境用的是带权重的策略——综合最后活跃时间和切换成本，优先切换"性价比高"的会话（不怎么活跃、缓存又小）。

text复制evict_candidate = min(
    sessions,
    key=lambda s: s.idle_time / max(1.0, s.kv_cache_size)
)

如果同一个用户反复进出，还可以加一个"预热列表"：把最近切换出去的会话句柄保留在 CPU 内存里，用户再次请求时不需要重新读盘，直接走内存恢复，速度能快一截。

4. 实操中的性能优化与参数配置

4.1 显存预算：权重、KV 池、暂存区，三项分开规划

做上下文管理之前，第一步是给显存做预算。我会把显存分成三个逻辑区域：权重区、KV 缓存池、切换暂存区。权重区放模型权重和 CUDA graph 捕获结果，KV 缓存池是所有会话共享的动态缓存区，暂存区是切换时用来做 D2H/H2D 中转的固定缓冲。

以 24GB 显存、7B FP16 模型为例，我一般这么分：

这个比例不是死的。如果业务要求并发会话数更高，可以把 KV 池加大、权重区压缩到量化版本；如果切换非常频繁，暂存区就得按最大单会话 KV 大小来预留，否则大会话切换时会临时申请不到内存。我遇到过暂存区过小导致切换失败的情况，后来直接把暂存区大小设置为"最大上下文长度对应的 KV 上限"。

4.2 切换预加载：在用户开口之前就把上下文准备好

上下文的加载是异步的，这给了我们一个天然的优化窗口：预测用户的下一步动作，提前把可能用到的上下文恢复到显存，让切换延迟从用户感知中消失。

实现上有两种做法。一种是根据用户活跃时间做定时预加载：用户上一次会话结束后的空闲期间，后台线程就把他的上下文从 SSD 搬回 CPU 内存，等用户真正发起请求时只剩一次 D2H 拷贝。另一种是在网关上做意图预测：比如用户在某个页面停留超过一定时间，就提前把该租户的 LoRA 和 KV 状态准备好。

实测下来，预加载能把首 token 延迟从 800ms 以上压到 200ms 以内。代价是 CPU 内存占用会明显上升，因为你要把所有"可能马上用到"的 KV 都放在内存里。所以这个优化要和业务方对齐：哪些用户属于高优先级，值得预加载，哪些用户让他冷启动就行。

4.3 网络与硬件对切换性能的影响

切换性能不只是软件问题。KV Cache 的搬运本质是数据移动，硬件决定了速度上限。我列几个实测参考值，大家按自己的硬件情况对照：

• CPU 内存到显存（PCIe 4.0 x16）：理论带宽约 32GB/s，实测 1GB KV 大约 50~80ms；
• SSD 到 CPU 内存（NVMe）：1GB 大约 100~150ms；
• 显存到显存（同卡内 D2D）：1GB 大约 10ms 以内，所以如果有多卡，优先考虑在卡间搬运而不是走 CPU。

如果集群里有多张卡，还可以做一个优化：把高频会话的 KV 副本存在另一张卡的空闲显存里，切换时走 peer-to-peer 拷贝，速度比 D2H + H2D 快一个量级。这个方案对显存碎片要求高，一般单机多卡且业务量大的团队才需要上。

5. 常见问题与排查技巧实录

5.1 切换后首 token 延迟突然拉高

这是最常被误判的问题。我见过有人把首 token 延迟上升归因于网络或模型本身，折腾半天，最后发现是 CUDA graph 失效了。

推理框架为了降低 kernel 启动开销，通常会把一整套生成流程捕获成 CUDA graph。问题是，CUDA graph 里的显存地址是固定的，一旦上下文切换导致 KV Cache 换了物理块，旧 graph 就不能直接复用了。如果框架没有做好 graph 的按上下文重建，就会退化成逐 kernel 启动的模式，首 token 延迟直接上涨 30% 到 50%。

排查思路：先看首 token 延迟是否有周期性、是否都在切换事件之后；再查推理日志里有没有 graph recapture 的记录；最后看 CUDA graph 是否按 session 维度管理。解决办法是给每个上下文状态维护独立的 graph 版本，只在 KV 物理块变化时才重新捕获。

5.2 恢复上下文后生成内容语义漂移

比性能问题更隐蔽的是语义问题。KV Cache 恢复之后生成结果不对，通常有几个原因：

• 模型版本不一致，KV Cache 是旧版模型的，新模型读不了旧缓存；
• tokenizer 版本不一致，同样的历史对话切出来的 token 不同，KV 就对不上；
• 恢复时 attention mask 或 position id 没重建对，导致位置编码错位。

我的建议是：在上下文 payload 里把所有跟"token 切分"有关的版本信息全部记录下来，包括模型版本、tokenizer 版本、采样配置。恢复时任何一项对不上，就走降级路径——直接用历史对话重新跑一遍 prefill，而不是硬恢复 KV Cache。虽然慢一点，但至少正确。

注意：上下文恢复宁可慢，不可错。一旦把错误的 KV Cache 恢复进去，不只是当前回答错，整个会话后续的生成都会错，而且用户很难分辨是哪一步开始错的。

5.3 显存碎片化和恢复失败

上下文频繁切换最容易带来的副作用就是显存碎片化。KV Cache 反复分配、释放，不连续的小块越来越多，最终导致某个大上下文恢复时找不到连续空间。

处理方案有三个层次。第一层，在 KV 池内部做"块级"管理，不要求上下文占连续空间，而是用块表映射逻辑位置和物理块，这其实就是分页思想的落地。第二层，定期做显存整理，把活跃上下文的 KV 块搬到大连续区域，但整理本身也有成本，建议在低峰期做。第三层，切换到"切换出最贵上下文"的策略，宁可切走一个大的空出空间，也不要拆散多个小块。

我个人的经验是：碎片化问题要从设计层面解决，不能靠运维清理。KV 池一定要用固定大小的块作为分配单元，哪怕偶尔浪费一点内存，也比把池子搞成"砖头堆"强得多。

5.4 问题速查表

6. 写在最后的一点个人体会

做上下文管理最忌讳的一件事，就是把它当成"存储系统"来做。我早期犯的错就是把上下文管理做成 KV 存储：序列化、落盘、恢复，每一步都"优雅"，唯独忘了这套机制服务的是 GPU 上的推理过程。真正有效的设计，一定是从显存池、CUDA stream、调度策略这些"系统层"入手，让上下文的保存和恢复与推理本身重叠起来，而不是变成推理链路外的一个串行环节。

另一个体会是，上下文切换的调度策略一定要做成可观测的。我后来给线上服务加了一组指标：切换次数、切换耗时分布、恢复失败率、KV 池碎片率。有了这些数据，调参数才不是拍脑袋。比如你发现"切换耗时 90 分位远超均值"，那大概率是某个大上下文会话在拖累整体，这时候就应该给大会话单独走一条不落盘的常驻路径，而不是继续在池子里切来切去。

最后分享一个小技巧：切换机制上线前，一定要做一次"断电恢复演练"。模拟 GPU 进程被杀、上下文全部丢失的场景，验证降级逻辑能不能自动把会话重建起来。我在实际测试中发现，很多团队切来切去都很顺，一旦异常退出就全乱套，原因就是降级路径从来没被真正触发过。把降级路径当第一公民对待，而不是当兜底补丁，这才是上下文管理系统的成熟标志。