Cache-Aware Scheduling

The scheduler groups incoming inference requests based on the similarity of their initial prompts or ongoing conversation context. This grouping maximizes KV cache reuse across requests within the same batch. For example, multiple user queries about "weather in London" can share a significant portion of their cached key-value tensors, reducing the total memory footprint and the number of memory reads required per generated token. This is a higher-order optimization built on top of continuous batching, focusing on semantic rather than just temporal grouping.

This principle is borrowed from computer architecture and applied to transformer inference. By scheduling requests with overlapping context, the system ensures that the data needed for the attention computation (the KV cache) is already resident in the fast GPU memory or cache hierarchy. This minimizes costly memory-bound stalls where the processor waits for data to be fetched from DRAM. The goal is to transform random, scattered memory accesses into sequential, predictable patterns that hardware prefetchers can optimize.

The scheduler dynamically allocates portions of the physical KV cache memory to different request groups or users. Unlike static partitioning, this allows for efficient use of the total cache capacity. Policies might include:

• Least-Recently-Used (LRU) Eviction: Removing cached tensors for inactive conversation threads.
• Priority-Based Allocation: Giving more cache space to high-priority or paid-tier requests.
• Working Set Estimation: Monitoring the active context length of each request group to pre-allocate sufficient cache blocks, reducing fragmentation.

Cache-aware scheduling works synergistically with memory management systems like PagedAttention (used in vLLM). While PagedAttention eliminates internal fragmentation by managing the cache in fixed-size blocks, the scheduler reduces external fragmentation and improves block reuse. The scheduler's grouping decisions inform which KV cache "pages" are likely to be needed together, allowing the memory manager to co-locate them, further improving access locality and reducing TLB (Translation Lookaside Buffer) misses.

A core engineering trade-off exists. To form optimal groups, the scheduler may need to wait briefly for similar requests to arrive, introducing scheduling latency. The system must balance this against the throughput gains from a higher cache hit rate. Sophisticated implementations use:

• Batching Windows: A configurable time window for collecting requests.
• Cost Models: Predicting the compute/memory savings from grouping specific requests.
• Early Flushing: Sending a batch for execution even if not fully optimized to meet latency Service Level Agreements (SLAs).

The benefits manifest differently across inference stages:

• Prefill Phase: Context-similar requests can share the computationally expensive initial processing of common prompt prefixes. Some systems may even compute the KV cache for a shared prefix once and replicate it across requests.
• Decode Phase: The major benefit is realized here. As tokens are generated autoregressively, requests in the same batch read from a largely shared cache, dramatically reducing memory bandwidth pressure. This is critical for maintaining high tokens/second throughput, especially in memory-bound regimes common in large model inference.

The KV cache is the foundational memory buffer in transformer-based language models. During autoregressive decoding, it stores pre-computed key and value tensors for all previous tokens in a sequence. This eliminates the need to recompute these tensors for each new token, trading memory for a massive reduction in computational latency. The size of the KV cache grows linearly with sequence length and batch size, making its management a primary bottleneck for long-context or high-throughput inference.

PagedAttention is a memory management algorithm that organizes the KV cache into non-contiguous, fixed-size blocks or "pages." Inspired by virtual memory in operating systems, it allows:

• Efficient sharing of physical memory between different requests in a batch.
• Elimination of internal fragmentation, where allocated but unused cache space is wasted due to variable sequence lengths.
• Dynamic memory allocation, enabling new sequences to start without pre-allocating maximum context. This is the core innovation behind high-throughput engines like vLLM.

Continuous batching is a dynamic scheduling paradigm that groups incoming inference requests into a shared computational batch. Unlike static batching, it allows:

• New requests to join the batch as soon as resources are available.
• Completed sequences to exit the batch immediately, freeing their KV cache. This maximizes GPU utilization and aggregate throughput. Cache-aware scheduling operates within this framework, grouping requests with similar cache characteristics to improve locality and hit rates within the continuously batched workload.

These are attention variants designed to reduce KV cache memory pressure.

• Multi-Query Attention (MQA): All query heads share a single key head and a single value head. This drastically reduces the cache size but can impact model quality.
• Grouped-Query Attention (GQA): A tunable hybrid where groups of query heads share a single key and value head. It offers a practical trade-off, providing most of MQA's memory savings with quality closer to standard Multi-Head Attention. Models like Llama 2 and Gemma use GQA.

KV cache quantization reduces the numerical precision of the cached key and value tensors. Instead of storing them in FP16 or BF16, they are stored in formats like INT8 or FP8. This technique:

• Halves or quarters the cache memory footprint, enabling longer contexts or larger batches.
• Reduces memory bandwidth pressure, as fewer bytes are moved for each attention operation.
• Typically requires minimal calibration and has a negligible impact on output quality, making it a highly effective optimization for memory-bound inference.

Transformer inference is split into two distinct computational phases that define cache usage:

• Prefill Phase: The initial, parallel processing of the entire input prompt. This phase is compute-bound, as it calculates the initial KV cache for all prompt tokens with full attention.
• Decode Phase: The token-by-token autoregressive generation. This phase is memory-bound, as it performs small matmul operations dominated by reading the growing KV cache. Cache-aware scheduling optimizes the decode phase by improving the efficiency of these memory reads.