论文

只需一次关注:使用有状态 Transformer 进行高效流式推理

Attention Once Is All You Need: Efficient Streaming Inference with Stateful Transformers

模型推理KV Cache

摘要

传统的 Transformer 推理引擎是请求驱动的,每次查询都会付出 O(n) 的预填充成本。在流式工作负载中,数据不断到达并且查询探测不断增长的上下文,这种成本是令人望而却步的。我们引入了一种以有状态会话为中心的数据驱动的计算模型:随着新数据的到来,持久的 KV 缓存会逐渐前进,因此预填充将从关键路径中移出,查询延迟变为 O(|q|),与累积的上下文大小无关。在此基础上,Flash 查询回收数据到达之间的空闲 GPU 周期,以预先评估已注册的问题并在用户询问之前返回缓存的答案,这种模式在无状态引擎中结构上是不可能的,因为它们会丢弃请求之间的中间状态。具有单元预算准入和前缀感知分组预填充功能的多租户连续批处理调度程序可以让数十个有状态会话在单个 GPU 上共存,同时保留完整的二次自注意力。在流市场数据基准测试中,参考实现比传统推理引擎(vLLM、SGLang、TensorRT-LLM、llama.cpp)实现了高达 5.9 倍的加速,随着累积上下文的增长,保持查询延迟恒定。