Requests of a Feather Must Flock Together: Batch Size vs. Prefix Homogeneity in LLM Inference
作者: Saksham Rathi, Preeti, Mythili Vutukuru
分类: cs.LG
发布日期: 2026-05-07
备注: 22 pages, 36 figures
💡 一句话要点
Feather：通过强化学习优化LLM推理中批大小与前缀同质性的调度器
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 LLM推理 前缀共享 强化学习 调度器 KV缓存 分块哈希树
📋 核心要点
- 现有LLM推理优化方法主要通过增大批大小来提高吞吐量，但忽略了前缀共享工作负载下前缀同质性对KV缓存访问效率的影响。
- Feather提出了一种基于强化学习的前缀感知调度器，旨在学习批大小和前缀同质性之间的最佳平衡，从而优化KV缓存访问。
- 实验结果表明，Feather在vLLM和SGLang中实现了2-10倍的端到端吞吐量提升，尤其是在具有前缀共享的工作负载下。
📝 摘要（中文）
大型语言模型中的自回归token生成是内存受限的，因为它需要“关注”所有先前token的键和值张量（KV缓存）。先前的工作旨在通过将多个请求批量处理在一起，并在GPU内存约束下最大化批大小来提高此解码过程的效率。我们工作的关键观察是，对于前缀共享的工作负载，较小的、前缀同质的批次（其中所有请求共享一个公共前缀）由于KV缓存访问期间更好的空间和时间局部性，可以实现比更大的、异构批次更高的解码吞吐量。然而，最先进的推理引擎中的前缀感知调度器仅在批次内最大化前缀重用，以减少KV缓存内存占用，但不会在可能表现更好的较小同质批次处停止批次形成。此外，我们表明现有调度器中的共享前缀检测依赖于基数树遍历，从而导致大量的CPU开销，这通常与GPU执行时间相当。本文提出了Feather，一种前缀感知调度器，它使用强化学习（RL）来学习批大小和前缀同质性之间的最佳权衡。我们还引入了分块哈希树（CHT），一种轻量级数据结构，可以实现快速前缀检测和RL调度器的有效请求选择，避免了昂贵的树遍历。我们将Feather集成到vLLM和SGLang中，我们的评估表明，与现有调度器相比，Feather实现了2-10倍更高的端到端吞吐量，而在工作负载没有足够的前缀共享时，其性能并不比现状差。Feather通过减少KV缓存访问的总数来实现这些收益，超过了具有相同目标的前缀感知注意力内核的性能。
🔬 方法详解
问题定义：现有LLM推理引擎在处理具有前缀共享的请求时，通常只关注最大化批大小以减少KV缓存的内存占用，而忽略了前缀同质性对KV缓存访问效率的影响。现有调度器中的前缀检测依赖于基数树遍历，导致显著的CPU开销，降低了整体推理效率。
核心思路：Feather的核心思路是利用强化学习来动态地权衡批大小和前缀同质性。通过学习不同批次配置下的性能表现，Feather能够选择最优的批次大小，从而在保证高吞吐量的同时，最大化KV缓存的访问效率。同时，引入轻量级数据结构CHT加速前缀检测。
技术框架：Feather的整体框架包含三个主要模块：请求队列、前缀检测模块和强化学习调度器。请求队列接收来自用户的请求，前缀检测模块使用CHT快速检测请求之间的共享前缀。强化学习调度器根据前缀信息和系统状态，决定如何将请求组成批次，并将其发送到GPU进行推理。
关键创新：Feather的关键创新在于：1) 使用强化学习来动态优化批大小和前缀同质性之间的平衡；2) 引入了分块哈希树（CHT）这一轻量级数据结构，用于快速高效地检测请求之间的共享前缀，避免了传统基数树遍历带来的CPU开销。
关键设计：强化学习调度器使用策略梯度算法进行训练，奖励函数综合考虑了吞吐量、延迟和KV缓存利用率。CHT将前缀分成多个块，并使用哈希表存储块之间的关系，从而实现快速的前缀匹配。具体来说，CHT的构建和查询复杂度均为O(n)，其中n为前缀的长度。
🖼️ 关键图片
📊 实验亮点
Feather在vLLM和SGLang上的实验结果表明，与现有调度器相比，Feather实现了2-10倍的端到端吞吐量提升。即使在没有足够前缀共享的工作负载下，Feather的性能也与现有调度器相当。通过减少KV缓存访问的总数，Feather甚至超越了具有相同目标的前缀感知注意力内核的性能。
🎯 应用场景
Feather可以应用于各种需要高效LLM推理的场景，例如聊天机器人、代码生成、文本摘要等。通过优化批处理策略，Feather能够显著提高推理吞吐量，降低延迟，从而提升用户体验。此外，Feather还可以应用于云端LLM服务，降低运营成本，提高资源利用率。
📄 摘要（原文）
Auto-regressive token generation in large language models is memory-bound because it requires "attending to" key and value tensors (KV cache) of all previous tokens. Prior work aims to improve the efficiency of this decode process by batching multiple requests together, and maximizing batch size subject to GPU memory constraints. The key observation of our work is that with prefix-sharing workloads, smaller, prefix-homogeneous batches -- where all requests share a common prefix -- can achieve higher decode throughput than larger, heterogeneous batches, due to better spatial and temporal locality during KV cache accesses. However, prefix-aware schedulers in state-of-the-art inference engines maximize prefix reuse within a batch only to reduce KV cache memory footprint, but do not stop batch formation at smaller homogeneous batches that could have performed better. Further, we show that shared prefix detection in existing schedulers relies on radix-tree traversals, incurring substantial CPU overhead that is often comparable to GPU execution time. This paper presents Feather, a prefix-aware scheduler that uses reinforcement learning (RL) to learn the optimal tradeoff between batch size and prefix homogeneity. We also introduce Chunked Hash Tree (CHT), a lightweight data structure that enables fast prefix detection and efficient request selection for the RL scheduler, avoiding expensive tree traversals. We integrate Feather into vLLM and SGLang, and our evaluation shows that Feather achieves 2--10$\times$ higher end-to-end throughput as compared to existing schedulers, while doing no worse than the status quo when the workload does not have enough prefix sharing. Feather achieves these gains by reducing the total number of KV cache accesses, surpassing the performance of prefix-aware attention kernels that have the same goal.