Attention Sink

An attention sink refers to the initial tokens of a sequence that receive disproportionately high attention scores in streaming language models. This phenomenon provides a stable anchoring point for the attention mechanism, enabling efficient infinite-length generation without catastrophic forgetting of the context. It is a critical component for models operating beyond their trained context window, as it prevents attention scores from becoming unstable or vanishing when processing very long sequences.

The need for an attention sink arises from the softmax operation in the attention mechanism. Softmax requires all attention scores for a given query to sum to 1. In very long sequences, as new tokens are generated, the attention probability mass must be distributed across an ever-growing number of prior tokens. Without a stable sink, the probability assigned to any single relevant token from the distant past becomes infinitesimally small, leading to effective context forgetting. The initial tokens act as a 'dump' for excess probability mass, stabilizing the distribution.

The StreamingLLM framework operationalizes the attention sink concept. It modifies the standard attention computation to always retain the first few tokens (e.g., 4) of the sequence in the KV Cache, alongside a sliding window of the most recent tokens.

• Fixed Sink Tokens: The initial tokens are permanently cached.
• Sliding Recent Tokens: A window of the N most recent tokens is kept.
• Eviction Policy: Tokens outside the sink and the sliding window are evicted from the KV cache. This design allows a model trained on a finite window (e.g., 4K tokens) to generalize to sequences of infinite length during inference without retraining.

Attention sinks are particularly effective with certain positional encoding schemes. Models using absolute positional encodings (like in the original Transformer) struggle with out-of-window positions. However, models with relative positional encodings, such as Rotary Positional Embedding (RoPE), can more naturally handle the sink. RoPE's relative nature allows the model to compute attention between a current token and the sink tokens based on their relative distance, even as absolute positions grow far beyond the training range. The sink's fixed starting position provides a consistent reference point for these relative calculations.

Leveraging an attention sink involves key engineering trade-offs:

• Memory Efficiency: Enables long contexts with a fixed, small KV cache size (sink + sliding window), preventing linear memory growth.
• Contextual Limitation: The model's effective 'working memory' is constrained to the sliding window. Information outside the window (and not in the sink) is permanently lost.
• Sink Token Selection: The content of the initial sink tokens matters. Typically, they are the start-of-sequence token and immediate follow-ons. Their semantic neutrality helps them function as a pure structural anchor.
• Throughput: By bounding cache size, it maintains high inference throughput and enables efficient continuous batching in serving systems like vLLM.

Attention sink management is often implemented alongside advanced KV cache memory management systems. PagedAttention, used in vLLM, is highly complementary. It organizes the KV cache into non-contiguous, fixed-size blocks.

• The attention sink's fixed tokens can be allocated to permanent, pinned pages.
• The sliding window's tokens can be efficiently managed within dynamically allocated pages.
• This combination allows for zero internal fragmentation and efficient sharing of cache blocks across sequences in a batch, maximizing GPU memory utilization for long-context, high-throughput streaming inference.

The KV cache is the foundational memory buffer that stores computed key and value tensors from previous tokens during autoregressive generation. It is the primary data structure optimized by attention sink techniques. Without a KV cache, each new token generation would require recomputing attention over the entire prior context, making inference impractically slow.

• Purpose: Eliminates redundant computation for past tokens.
• Impact: Directly determines the maximum context window a model can handle based on available GPU memory.

Sliding window attention is an attention mechanism where a token can only attend to a fixed number of preceding tokens within a local window. It is often combined with attention sinks in streaming scenarios.

• Function: Naturally bounds KV cache growth to the window size, enabling linear-time complexity.
• Synergy with Sinks: The window manages recent context, while the attention sink tokens (like the first few tokens) provide a global, stable reference point outside the window, preventing catastrophic forgetting.

Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) are attention variants designed to reduce the memory footprint of the KV cache.

• MQA: All query heads share a single key head and value head. This drastically shrinks the cache size.
• GQA: A tunable middle ground where groups of query heads share a single key/value head.
• Relation to Sinks: Reducing per-token cache size via MQA/GQA allows for longer contexts or more sink tokens to be retained within the same memory budget, enhancing streaming capabilities.

Transformer inference occurs in two distinct phases that interact with the KV cache and attention sinks.

• Prefill Phase: The initial, parallel processing of the entire input prompt. This phase computes and populates the initial KV cache, including the foundational attention sink tokens.
• Decode Phase: The token-by-token autoregressive generation. In this phase, the model relies almost entirely on the cached keys and values, reading the sink and recent tokens to generate the next output efficiently. Optimizing the memory access pattern here is critical for low latency.