arXiv:2605.06046cs.LG2026-05

通过智能调度提升大模型推理速度,尤其在共享前缀场景下效果显著。

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

论文配图:Requests of a Feather Must Flock Together: Batch Size vs. Prefix Homogeneity in LLM Inference
图 1 · 摘自论文原文
  • 用强化学习动态平衡批处理大小与前缀一致性,优化推理效率。
  • 相比现有方法,端到端吞吐提升2到10倍,且无前缀共享时性能不降。
  • 提出轻量级哈希树结构,快速检测共享前缀,降低CPU开销。

大语言模型自回归生成过程受内存限制,因需保留所有历史标记的键值缓存(KV cache)。现有工作通过批量处理多个请求来提升解码效率,但通常在显存约束下最大化批大小。本文发现,在具有共享前缀的负载中,较小且前缀一致的批次比大而异构的批次能实现更高解码吞吐,因其在访问KV缓存时具备更好的空间和时间局部性。然而,当前先进推理引擎中的前缀感知调度器仅关注减少缓存占用,未在更小的同质批次上停止组批。此外,现有调度器依赖基数树遍历检测共享前缀,导致大量CPU开销,常与GPU执行时间相当。本文提出Feather,一种基于强化学习的前缀感知调度器,可学习批大小与前缀同质性的最优权衡。我们还引入轻量级数据结构——分块哈希树(CHT),实现快速前缀检测与高效请求选择,避免昂贵的树遍历。将Feather集成至vLLM和SGLang后,评估显示其相比现有调度器实现2–10×更高的端到端吞吐,且在缺乏前缀共享的场景下性能不低于现状。该优势源于减少总KV缓存访问次数,超越了目标相同的前缀感知注意力核函数。

原文摘要 · Abstract (English)

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.

大模型推理调度优化前缀共享强化学习

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。