Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Attention Sinks › Streaming Inference Pipeline

Michael BrenndoerferUpdated July 4, 20255 min read

Language AI HandbookAll chapters

Part 18 of 67Long Context

Chapter 130Attention Sinks › Streaming Inference Pipeline

Reading section 9 of 16 · Streaming Inference Pipeline · 5 minSection 9 of 16 · 5 min

Section 9 of 16 · 5 min read

Attention Sinks › Streaming Inference PipelineOne section · Sections in this chapter9 / 16 ↓

Sections in this chapter16

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Let's put everything together into a complete streaming inference pipeline:

In[17]:

Code

class StreamingInferenceEngine:
    """
    Complete streaming inference engine with attention sink preservation.
    """

    def __init__(self, model, tokenizer, num_sink_tokens=4, window_size=2044):
        self.model = model
        self.tokenizer = tokenizer
        self.num_sink_tokens = num_sink_tokens
        self.window_size = window_size

        # Extract model configuration
        config = model.config
        self.num_layers = config.num_hidden_layers
        self.num_heads = config.num_attention_heads
        self.head_dim = config.hidden_size // config.num_attention_heads

        # Initialize cache
        self.cache = None
        self.total_tokens_generated = 0

    def reset(self):
        """Reset the cache for a new generation session."""
        self.cache = None
        self.total_tokens_generated = 0

    def _init_cache(self, batch_size, device, dtype):
        """Initialize the streaming cache."""
        self.cache = StreamingLLMCache(
            num_sink_tokens=self.num_sink_tokens,
            window_size=self.window_size,
            num_layers=self.num_layers,
            num_heads=self.num_heads,
            head_dim=self.head_dim,
        )

    def generate_token(self, input_ids, past_key_values=None):
        """
        Generate a single token using the streaming cache.

        In a real implementation, this would integrate with the model's
        forward pass to use the streaming cache instead of standard KV cache.
        """
        # This is a simplified demonstration
        # Real implementation would modify the model's attention layers
        pass

Let's demonstrate the memory savings achieved by StreamingLLM compared to a full KV cache:

In[18]:

Code

# Simulate a long generation scenario
total_tokens = 10000
sink_tokens = 4
window_size = 2044
max_cache_size = sink_tokens + window_size

# Calculate memory usage (assuming 768-dim embeddings, 12 layers, float32)
# KV cache stores keys and values (2x) for each layer
full_cache_memory = total_tokens * 768 * 12 * 2 * 4 / (1024**2)  # MB
streaming_cache_memory = max_cache_size * 768 * 12 * 2 * 4 / (1024**2)  # MB
memory_saved = full_cache_memory - streaming_cache_memory
savings_percent = (1 - streaming_cache_memory / full_cache_memory) * 100

Out[19]:

Console

Streaming Inference Demonstration
==================================================

Generating 10,000 tokens with:
  - 4 sink tokens (always preserved)
  - 2,044 window tokens (sliding)
  - Maximum cache size: 2,048 tokens

Memory usage comparison:
  Full KV cache: 703.1 MB
  StreamingLLM cache: 144.0 MB
  Memory saved: 559.1 MB (79.5%)

The streaming approach enables generation of arbitrarily long sequences while using fixed memory. In this example, generating 10,000 tokens with a full cache would require over 700 MB, but StreamingLLM uses only about 145 MB, an 80% reduction.

Out[20]:

Visualization

Line plot showing linear memory growth for full cache versus flat horizontal line for StreamingLLM cache. — Memory usage comparison between full KV cache and StreamingLLM as sequence length increases. Full cache grows linearly (eventually exhausting GPU memory), while StreamingLLM stays constant regardless of generation length. At 100K tokens, full caching would require 7 GB while StreamingLLM still uses only 145 MB.

The visualization makes the scaling difference clear. Full KV caching follows a diagonal line that quickly exceeds typical GPU memory limits. At 100,000 tokens, a full cache would require nearly 7 GB. StreamingLLM's flat line at the bottom shows constant memory usage, enabling generation of arbitrarily long sequences without hitting memory limits.

To put the numbers in context: for a GPT-2-medium sized model with 768-dimensional hidden states and 12 layers, each cached token consumes 768 dimensions $\times$ 12 layers $\times$ 2 (keys and values) $\times$ 4 bytes = roughly 72 KB per token. A StreamingLLM cache of 4 sinks + 2044 window tokens uses about 147 MB total, regardless of how many tokens have been generated. A full cache for a 100,000-token sequence would require over 7 GB. For larger models like the 7B-parameter class with 4096 dimensions and 32 layers, each cached token uses about 1 MB, so StreamingLLM's 2048-token cache costs about 2 GB while a 100,000-token full cache would cost roughly 100 GB. At that scale, full caching is simply not feasible on any single GPU, so streaming approaches become necessary rather than merely convenient.