MoE KV Cache

In a standard transformer, the KV cache is a contiguous block per layer. In MoE, each expert maintains its own independent KV cache for the tokens it processes. This leads to memory fragmentation, as the cache is split across multiple, potentially non-contiguous buffers. The system must manage numerous small, variable-sized cache blocks instead of a few large ones, increasing memory allocation overhead and complicating memory-efficient strategies like PagedAttention which must now track pages across many experts.

The size of the KV cache for a given request is not fixed; it depends on the router's decisions for each token at each MoE layer. Two identical-length sequences can generate different total cache sizes if their tokens are routed to different numbers of experts. This dynamic allocation makes precise memory pre-allocation and scheduling difficult, leading to either memory over-provisioning (waste) or capacity overflow (dropped tokens/performance cliffs).

During autoregressive decoding, the KV cache must be retrieved for previous tokens to compute attention. In expert parallelism, where experts are distributed across devices, this requires an All-to-All communication pattern for both the cached KV states and the new token's queries. This cross-device data movement for cache access becomes a major latency and bandwidth bottleneck, often dominating the step time compared to the actual expert computation.

Due to sparse activation, the KV cache for a specific expert may contain states from only a small subset of the batch's sequences. This results in low cache density per expert buffer. Standard cache eviction policies (e.g., least recently used) become less effective. Deciding which cached sequences to evict when an expert's buffer is full is complex, as evicting a sequence from one expert may not free up meaningful space in others, leading to suboptimal global memory usage.

Attention performance relies on memory locality—quickly accessing the KV states for related tokens. In MoE, tokens that attend to each other may have their KV states cached in different experts on different devices. This poor locality increases the latency of attention computation. The system must either gather distant cache blocks (expensive) or use approximations, potentially impacting model quality. The router's decisions directly dictate this costly access pattern.

Continuous batching dynamically adds/removes requests from a running batch to maximize GPU utilization. For MoE, this is severely complicated by the KV cache. Adding a new request requires allocating cache space across all experts it might use, which is unknown upfront. Removing a request requires finding and freeing its potentially scattered cache blocks across all experts and devices. This makes efficient dynamic batching for MoE a significant systems engineering challenge.

The systematic optimization of storing and retrieving Key (K) and Value (V) matrices in transformer attention layers to avoid recomputation for previously processed tokens. In standard dense models, this is a single, contiguous cache. For Mixture of Experts (MoE), complexity arises because each expert may process a unique subset of the sequence, potentially requiring separate cache states per expert if routing is not synchronized across layers.

• Core Purpose: Drastically reduces latency for autoregressive generation.
• MoE Challenge: Naive implementation can lead to expert-specific caches, multiplying memory footprint.

A model parallelism strategy where different expert networks are placed on separate devices (e.g., GPUs). Tokens are routed to their assigned expert's device, processed, and results are gathered. This directly impacts KV Cache management, as the cache for a sequence may be distributed across devices based on expert assignment.

• Communication Pattern: Relies heavily on All-to-All operations to shuffle tokens.
• Cache Implication: Requires careful design to either maintain a unified cache view or manage distributed cache slices to avoid excessive device-to-device transfers during attention.

Also known as iteration-level batching or incremental batching, this is a scheduling technique for inference servers that dynamically adds new requests to and removes completed ones from a running batch. For MoE models, this is complex because the computational graph changes per request based on routing decisions.

• Efficiency Goal: Maximizes GPU utilization by eliminating padding and idle time.
• MoE KV Cache Integration: The system must manage non-contiguous cache allocations for requests with different routing histories within the same physical batch.

An algorithm, central to systems like vLLM, that manages the KV cache using concepts from operating system virtual memory. It allows non-contiguous physical memory blocks to be allocated for the logical KV cache of different sequences. This is exceptionally useful for MoE where cache requirements are irregular.

• Key Innovation: Eliminates external fragmentation, allowing near-100% memory utilization.
• MoE Application: Enables efficient caching for sequences where activated experts vary, by mapping variable-sized 'pages' of cache per expert-sequence combination.

The fundamental property of conditional computation models like Mixture of Experts, where for a given input, only a small, dynamically selected subset of the model's total parameters is used. This contrasts with dense models that activate all parameters.

• Computational Benefit: Enables massive model scale (e.g., trillion parameters) with manageable FLOPs per token.
• Cache Consequence: The KV Cache only needs to be updated for the parts of the model (experts) that were actually activated for a given token, though managing this sparsity in cache storage is a key systems challenge.