PersistentKV: Page-Aware Decode Scheduling for Long-Context LLM Serving on Commodity GPUs
作者: Muhammad Ahmed
分类: cs.LG
发布日期: 2026-06-25
备注: 7 pages, 3 tables; workshop paper
💡 一句话要点
提出PersistentKV以解决长上下文LLM服务中的KV缓存问题
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 长上下文 自回归模型 键值缓存 页面感知调度 GPU优化 分组查询注意力 性能提升
📋 核心要点
- 现有的自回归LLM服务受到KV缓存移动的限制,导致GPU利用率低下,尤其是在长上下文解码时。
- PersistentKV通过设计原生块表解码注意力引擎和页面感知调度,优化了分组查询注意力的工作映射和任务调度。
- 实验结果表明,PersistentKV在多种工作负载下提高了同步墙吞吐量,尤其在长上下文步骤中表现出显著的性能提升。
📝 摘要(中文)
自回归大型语言模型(LLM)的服务越来越受到键值(KV)缓存移动的限制,而非密集矩阵乘法。现代分页注意力系统减少了KV缓存的碎片化,成熟的内核如FlashInfer提供了高度优化的原生分页解码注意力。然而,最佳的单内核实现并不总是最佳的服务调度:低活跃的长上下文解码可能会低效利用商品GPU,而混合序列长度则在多个精确长度启动和粗略填充批次之间引入了紧张关系。本文提出PersistentKV,一个原生块表解码注意力引擎和针对分组查询注意力(GQA)的页面感知调度研究。
🔬 方法详解
问题定义:本文旨在解决自回归大型语言模型服务中KV缓存移动导致的性能瓶颈,现有方法在长上下文解码时未能有效利用GPU资源。
核心思路:PersistentKV通过设计一个原生块表解码注意力引擎,结合页面感知调度策略,优化了分组查询注意力的工作映射和任务调度,以提高GPU的利用率。
技术框架:整体架构包括KV头组的工作映射、K,V块的重用、原生页面表的支持,以及一个紧凑的工作队列调度,专注于执行非空的行-KV头-序列拆分任务。
关键创新:PersistentKV的主要创新在于其页面感知调度策略和工作队列设计,使得工作分配成为影响服务系统性能的关键变量,而不仅仅是注意力计算。
关键设计:在实验中,使用了RTX 3060 GPU,采用FP16精度,页面大小为16,Hq=32,Hkv=8,d=128,并通过校准自适应策略选择合适的解码策略以优化性能。实验结果显示,针对不同工作负载的阈值和拆分计数经过固定校准后,显著提升了吞吐量。
📊 实验亮点
实验结果显示,在B8双峰、均匀和Zipf-like工作负载下,PersistentKV的同步墙吞吐量提高了1.063-1.265倍,在B1桶状追踪中提升了1.399倍。针对B4双峰边界情况,策略选择FlashInfer避免了PersistentKV的回归,展示了自适应页面感知调度的有效性。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等场景,能够有效提升长上下文处理的效率和性能。随着大型语言模型的广泛应用,PersistentKV的设计理念和技术框架将为未来的LLM服务提供重要的参考和借鉴。
📄 摘要(原文)
Autoregressive large language model (LLM) serving is increasingly limited by key-value (KV) cache movement rather than dense matrix multiplication. Modern paged-attention systems reduce KV-cache fragmentation and mature kernels such as FlashInfer provide highly optimized native-paged decode attention. However, the best single-kernel implementation is not always the best serving schedule: low-active long-context decode can under-utilize commodity GPUs, while mixed sequence lengths introduce a tension between many exact-length launches and coarse padded batches. We present PersistentKV, a native block-table decode attention engine and page-aware scheduling study for grouped-query attention (GQA). PersistentKV maps work by KV-head group, is designed to reuse K,V tiles across grouped query heads, supports native page tables, and adds a compact workqueue schedule that executes only non-empty row-KV-head-sequence-split tasks. On an RTX 3060 with FP16, page size 16, Hq=32, Hkv=8, d=128, and identical correctness tolerance against FlashInfer, a calibrated adaptive policy selects FlashInfer for small active batches, PersistentKV sequence splitting for B1 long-context steps, and PersistentKV workqueue scheduling for B8 long-context steps. With thresholds and split counts fixed on calibration traces, one held-out trace seed improves synchronized wall throughput by 1.063-1.265x on B8 bimodal, uniform, and Zipf-like workloads and by 1.399x on a B1 bucketed trace. On the B4 bimodal boundary case, the policy avoids the PersistentKV regression by selecting FlashInfer. These results identify a concrete systems niche for adaptive page-aware decode scheduling and show that work assignment, not only attention math, is a decisive serving-system variable.