We're excited to announce the availability of Paged Attention and Prefix Caching in MAX Serve, bringing state-of-the-art LLM inference optimizations. These features are available in MAX nightly and the MAX Serve nightly Docker image.

Try them now

To proceed, please make sure to install the magic CLI

Or update it via

Now install the max-pipelines package with a single command

Serve with optimizations enabled

‍
Check out what’s available with

These features are available in Modular's officially supported models, leveraging the highly optimized MAX Graph APIs.

Why do Paged Attention and Prefix Caching matter?

Multi-Head Attention (MHA) is a core building block of modern LLMs, but it can be computationally intensive during inference. MHA's computational complexity scales quadratically with sequence length O(n²) and linearly with batch size, making it particularly demanding for long sequences or large batches. KV Cache optimizes this by storing previously computed Key and Value projections, avoiding redundant computations during autoregressive generation. However, traditional KV caching faces memory management challenges with long sequences.

PagedAttention and Prefix Caching address these challenges.

Paged Attention: Memory-efficient KV Cache management

Paged attention, introduced by vLLM, revolutionizes how we handle attention computation in LLMs with:

• Block-based memory management:
  • Organizes KV cache into fixed-size memory blocks (pages)
  • Each block typically contains 16 or 32 tokens
  • Enables efficient memory allocation and deallocation
• Key benefits:
  • Continuous memory guarantee: No memory fragmentation
  • Dynamic sequence management: Efficiently handles variable-length sequences
  • Memory pooling: Shares memory across multiple requests
  • GPU memory savings: Up to 40% reduction in memory usage

Learn more about paged attention vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention‍

Prefix Caching: Optimizing similar prompts

‍Prefix caching, introduced by SGLang, provides powerful optimization for structured LLM programs:

• Core concept:
  • Identifies and caches common prefix patterns in text prompts
  • Leverages program structure for optimal cache reuse
  • Implements intelligent cache management in the prefix trees
• Key advantages:
  • Smart prefix detection: Automatically identifies reusable prompt segments
  • Program-aware caching: Optimizes for common patterns in LLM applications
  • Throughput improvement: Up to 3x speedup for structured workflows
  • Resource optimization: Efficient memory usage through structured sharing

Learn more prefix caching SGLang: Efficient Execution of Structured Language Model Programs

‍What’s next?

These improvements optimize GPU memory by up to 40% and throughput up to 3x. Here are a few resources to get you started:

• ‍Get started with MAX
• Explore MAX Serve and MAX Container
• Check out the tutorial on how to deploy Llama 3 on GPU with MAX Serve
• Check out the relevant concept page
• Join our Discord and our Modular forum

We're excited to see what you'll build with MAX! Share your projects and experiences with us using #ModularAI on social media.