论文

羽毛请求必须聚集在一起:LLM 推理中的批量大小与前缀同质性

Requests of a Feather Must Flock Together: Batch Size vs. Prefix Homogeneity in LLM Inference

AI 基础设施推理服务

摘要

大语言模型 中的自回归词元生成受内存限制,因为它需要“处理”所有先前词元的键和值张量(KV 缓存)。之前的工作旨在通过将多个请求一起批处理并在 GPU 内存限制下最大化批处理大小来提高解码过程的效率。我们工作的主要观察结果是,通过前缀共享工作负载,较小的、前缀同质的批次(所有请求共享一个公共前缀)可以比较大的异构批次实现更高的解码吞吐量,因为 KV 缓存访问期间具有更好的空间和时间局部性。然而,最先进的推理引擎中的前缀感知调度程序最大限度地提高了批次内的前缀重用,只是为了减少 KV 缓存内存占用,但不会阻止以较小的同质批次形成批次,而这些批次本来可以表现得更好。此外,我们还表明,现有调度程序中的共享前缀检测依赖于基数树遍历,从而产生大量的 CPU 开销,通常与 GPU 执行时间相当。本文提出了 Feather,一种前缀感知调度器,它使用强化学习 (RL) 来学习批量大小和前缀同质性之间的最佳权衡。我们还引入了分块哈希树 (CHT),这是一种轻量级数据结构,可为 RL 调度程序实现快速前缀检测和高效请求选择,从而避免昂贵的树遍历。我们将 Feather 集成到 vLLM 和 SGLang 中,我们的评估表明,与现有调度程序相比,Feather 实现了 2--10$\times 更高的端到端吞吐量,同时在工作负载没有足够的前缀共享时,表现并不比现状差。 Feather 通过减少 KV 缓存访问总数来实现这些收益,超越了具有相同目标的前缀感知注意内核的性能。