长文档OCR的KV Cache破局：常数级显存优化实践

长文档OCR的KV Cache破局：常数级显存优化实践

在长文档OCR（光学字符识别）与文档理解领域，随着大模型能力的跃升，我们对于“处理多长文本”的期望也在不断攀升。从早期的单页识别，到几十页的连续解析，再到如今动辄数千页的超长序列处理，业务场景对模型的吞吐能力与上下文窗口提出了严苛要求。然而，让模型一口气读完数百万字并非易事，其核心拦路虎便是推理过程中的显存开销。

一、长序列OCR推理的显存瓶颈：KV Cache为何线性膨胀？

在基于Transformer架构的视觉语言大模型中，模型在自回归解码时，需要将历史token的键值对缓存保留在显存中，以避免重复计算，这就是业界熟知的KV Cache。对于文档OCR场景，这一机制面临双重压力：

1. 视觉Token的沉重负担：高分辨率文档图像经过视觉编码器后，会切分成大量的视觉Patch，每个Patch都会转化为Token序列。在标准解码流程中，若将所有页面的视觉缓存同时保留在显存中，其占用空间会随页数呈线性甚至指数级增长。
2. 峰值显存的致命一击：以百页级单序列测试为例，若采用标准全量缓存策略，峰值KV Cache占用可轻易飙升至上百GB量级。这不仅远超常规推理服务器单卡显存容量，甚至连整机内存都难以承受。一旦文档长度突破物理显存上限，推理过程便会因OOM（Out of Memory）而崩溃。

这就意味着，如果不从底层机制上进行革新，单纯堆砌算力根本无法实现“页数没上限”的连续文档解析。

二、破局之道：借鉴“抄书”智慧的持续解码架构

要打破显存随页数增长的魔咒，核心思路是改变“全量保留”的缓存策略。近期，百度文心Unlimited OCR 2开源了其在这一领域的探索，提出了一套将KV Cache压至常数级的创新架构。其核心设计思路，其实巧妙借鉴了人类抄书的智慧：人在抄写长卷时，不需要记住前文的所有细节，只需盯着当前正在抄写的页面，并偶尔回溯目录或关键页即可。

该架构主要由两大技术支柱构成：稀疏参考滑动窗口注意力机制（SR-SWA）与持续解码机制。

三、核心技术拆解：SR-SWA如何实现“少算”？

传统注意力机制中，当前token需要与历史所有token进行注意力计算。在长文档场景下，这不仅计算量巨大，且大量历史信息对当前页的识别并无实质帮助。SR-SWA通过以下策略实现计算的稀疏化：

1. 按需选取与相关性过滤：SR-SWA打破了全局注意力的刚性连接。模型在解码当前页面内容时，会动态评估历史页面缓存的相关性。对于无关的视觉信息或已充分转写的文本特征，直接切断注意力连接，从而大幅减少无谓的矩阵运算。
2. 滑动窗口的局部聚焦：引入滑动窗口机制，确保模型在处理跨页上下文（如跨页表格、段落延续）时，能保持对相邻页面的高精度感知。窗口内的信息参与稠密计算，窗口外的信息则通过稀疏参考机制按需调用。
3. 计算复杂度降维：通过这种稀疏化设计，模型无需在每一步解码时都对全量历史进行打分。计算开销从O(N²)级别的全局注意力，降至与当前文档活跃区域相关的近常数级复杂度，彻底摆脱了“文档越长、算得越慢”的窘境。

四、核心技术拆解：分块持续解码如何实现“少存”？

解决了“算得慢”的问题，还要解决“存不下”的问题。持续解码机制配合SR-SWA，重构了显存的管理方式。

1. 常数级缓存控制：标准解码将所有页面的视觉缓存常驻显存，而持续解码机制则采用“阅后即焚”与“按需驻留”相结合的策略。除了当前解码块及滑动窗口内的必要上下文外，已处理完成的页面KV Cache会被及时释放或沉淀。
2. 单步开销恒定：得益于这种动态显存管理，无论模型已经处理了100页还是3000页，单步解码时的显存占用始终维持在一个极低且稳定的水平。技术报告显示，该机制可将峰值KV Cache占用从标准策略的极高GB量级，压缩至不足2GB，降幅达96%。
3. 工程化实践示意：在实际部署中，这种机制通常需要配合高效的外存换入换出策略。其简化逻辑如下：

# 伪代码示意：持续解码的缓存管理逻辑
def continuous_decode(document_pages):
    active_cache = init_constant_size_cache() # 初始化常数级大小缓存
    for page in document_pages:
        visual_tokens = encode(page)
        # SR-SWA：仅选取相关历史信息进入活跃缓存
        ref_cache = sparse_select(active_cache, visual_tokens)
        # 解码并更新当前页特征
        decoded_text = decode(visual_tokens, ref_cache)
        # 持续解码机制：淘汰非必要历史，维持缓存恒定
        active_cache = evict_and_update(active_cache, decoded_text)
    return full_text

五、架构价值与行业启示

这种将KV Cache压至常数级的架构，为长序列视觉理解带来了显著的工程价值：

1. 突破长度物理限制：显存占用的解耦，使得模型能够真正处理理论上无上限的文档长度，为古籍数字化、超长网文解析、大型工程图纸审查等极端场景提供了可行方案。
2. 降低推理硬件门槛：原本需要多卡分布式推理的超长任务，现在有望在单张消费级或入门级专业显卡上流畅运行，大幅降低了企业的算力成本。
3. 与云原生基础设施的协同：在实际企业级部署中，这种低显存占用模型可完美适配标准的容器平台与弹性计算实例。结合对象存储读取海量文档图片，利用消息队列调度解析任务，再通过日志服务监控推理状态，能够以极低的资源消耗构建高吞吐的文档处理流水线。

结语

长文档OCR从“页级”迈向“万页级”的跨越，不仅是模型识别精度的提升，更是底层推理架构的一次硬核突围。通过稀疏参考滑动窗口注意力与持续解码机制的有机结合，将线性增长的KV Cache成功压成常数，百度文心Unlimited OCR 2的开源为整个开发者社区提供了一个极具参考价值的范式。这证明了在追求无限上下文窗口的道路上，除了暴力扩容算力，精巧的算法设计与显存调度策略同样能绽放出决定性的力量。