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

summary

The gist

The gist: smaller, prefix-homogeneous batches can achieve higher decode throughput than larger, heterogeneous batches due to better spatial and temporal locality during KV cache accesses.

In short

The work investigates how batch size affects performance when grouping requests that share common initial text prefixes. Findings show that smaller batches with high prefix homogeneity yield much higher throughput because they maximize spatial and temporal locality in the KV cache accesses. The proposed Feather scheduler uses reinforcement learning to dynamically balance batch size and prefix sharing for optimal efficiency.

Key concepts

Prefix Homogeneity: This refers to a situation where all requests within a single batch share the same initial sequence of tokens (a common prefix). When this occurs, the GPU's memory accesses for the key-value cache become highly localized in time and space, which significantly speeds up decoding.
KV Cache Locality: The Key-Value (KV) cache stores previously computed attention results in fast GPU memory. Prefix homogeneity ensures that when processing a batch, the system repeatedly accesses the same parts of this cache concurrently. This reuse reduces slow fetches from main DRAM, leading to faster overall generation.
Feather Scheduler: Feather is an intelligent scheduling system that uses reinforcement learning to decide the best batch size and prefix sharing strategy at runtime. It employs a Chunked Hash Tree for fast prefix detection, allowing it to create batches that are as homogeneous as possible without incurring high CPU overhead.
Chunked Hash Tree (CHT): The CHT is a lightweight data structure designed to quickly find which requests in the current batch share the longest common prefix. It organizes information based on how many prefix chunks each request is missing, allowing the system to identify the best candidates for forming a homogeneous group very fast.

Terminology used across episodes

Reinforcement Learning (RL) — explained in 25 episodes
Sample Complexity — explained in 18 episodes
Federated Learning — explained in 17 episodes
Retrieval-Augmented Generation (RAG) — explained in 16 episodes
Interpretability — explained in 15 episodes
Quantum Fisher Information (QFI) — explained in 15 episodes

This episode discusses

The paper

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

Saksham Rathi, Preeti Indian Institute of Technology Bombay, Mythili Vutukuru Indian Institute of Technology Bombay

Indian Institute of Technology Bombay

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.

Jane: Today's paper: "Requests of a Feather Must Flock Together".

Tom: The gist: smaller, prefix-homogeneous batches can achieve higher decode throughput than larger, heterogeneous batches due to better spatial and temporal locality during KV cache accesses.

Jane: First, who's behind it and why it matters.

Title and authors: Tom: So, to recap, we're looking at this paper called "Requests of a Feather Must Flock Together: Batch Size vs. Prefix Homogeneity in LLM Inference." The authors are Saksham Rathi, Preeti Indian Institute of Technology Bombay Mumbai, and Mythili Vutukuru from IIT Bombay >

Jane: They’re essentially investigating how batching multiple requests for text generation affects the speed, and they focus on this idea that when requests share a common prefix, smaller batches can actually outperform larger ones that have different prefixes >

Tom: That's the main finding of the paper. It challenges the idea that just making a batch bigger is always better for decode throughput because it focuses on how much those requests overlap in what they're asking for >

Lu: It seems they’ve pointed out a gap in current scheduling strategies, noting that existing methods often focus only on maximizing batch size during the decode phase rather than considering this prefix sharing aspect >

Meng: That makes sense from an engineering viewpoint. If we just keep adding requests hoping the batch gets bigger, we might end up fetching a lot of redundant data because those requests are all asking for totally different starting points >

Lalam: And what they found is that the performance gains increase with how long that shared prefix is across all the requests in a batch, which really supports their argument >

The paper's summary: Tom: So, summarizing what they found in "Requests of a Feather Must Flock Together: Batch Size vs. Prefix Homogeneity in LLM Inference," they show that for prefix-sharing workloads, smaller batches with high homogeneity can achieve higher decode throughput >

Jane: They found that when a prefix is shared by all requests, there is significant spatial and temporal locality in the memory accesses of the KV cache located in the GPU DRAM >

Tom: This locality leads to two things: higher memory bandwidth utilization and a reduced total amount of data fetched from DRAM >

Lu: The paper also found that prefix homogeneity has an even stronger impact on performance than just maximizing batch size during decode phase, which is what they observed through controlled experiments >

Meng: So, the key takeaway here is that you can get better speed by being smart about how you group the requests so they share as much of their initial input as possible >

Lalam: And interestingly, even when there isn't much prefix sharing at all, the adaptive batching policy they developed can still match the performance of existing schedulers like vLLM's FCFS scheduler >

The paper's improvements: Tom: Now let's look at what they proposed to fix these issues. They introduced a new approach called Feather, which is a prefix-aware scheduler built using reinforcement learning to learn the best trade-off between batch size and prefix homogeneity >

Jane: It seems they designed this system with two main parts: a Chunked Hash Tree, or CHT, for really fast prefix detection and an RL-based policy to decide when to add another request to a batch >

Lu: The CHT is this lightweight data structure that uses a min-heap ordered by per-request missing count—that’s how it finds the best candidate in O of log W time >

Meng: That sounds efficient. It replaces what they call DFS-based approaches, which they say have huge CPU overhead for long sequences, with something much faster >

Lalam: And the RL policy is what lets it adapt dynamically; it learns to stop adding requests at the point that maximizes throughput based on those observed batch characteristics >

Conclusion: Tom: So, to wrap up this discussion on "Requests of a Feather Must Flock Together: Batch Size vs. Prefix Homogeneity in LLM Inference," the main message is that we should look beyond just making batches huge and start paying attention to how requests share prefixes >

Jane: They showed that for prefix-shared workloads, smaller, homogeneous batches are actually better because of the way they reuse data in memory >

Tom: The results show up as a two to ten times higher end-to-end throughput compared to existing schedulers on some benchmarks when they leverage this locality properly >

Lu: The implication is that we need schedulers that actively manage prefix sharing, not just batch size constraints, which opens up new avenues for how we organize inference pipelines >

Meng: And from a practical standpoint, the overhead of their prefix detection is very low—they report it can be reduced by up to one thousand times compared to older DFS methods >

Lalam: It's neat that this paper also highlights how adaptive batch size management, driven by reinforcement learning, can handle situations where there isn't much prefix sharing either, matching other schedulers in those cases >

More episodes

← Home