mlx-paged-attention: Bringing vLLM-level High Throughput Inference to Apple Silicon

mlx-paged-attention is a project that ports vLLM's PagedAttention technology to Apple's MLX framework, bringing vLLM-level high-throughput large language model (LLM) inference capabilities to macOS and Apple Silicon users. This is an important port of PagedAttention to non-CUDA platforms, demonstrating the technology's versatility and portability.

In LLM inference, the attention mechanism is the most resource-intensive part. Traditional implementations pre-allocate continuous GPU memory for each request's key-value (KV) cache, leading to severe memory waste due to varying sequence lengths and fragmentation. PagedAttention, from vLLM, uses a virtual memory-inspired paging mechanism: it splits KV cache into fixed-size pages and dynamically allocates them as needed, significantly improving memory efficiency and enabling more concurrent requests on the same hardware.

The mlx-paged-attention project adapts vLLM's PagedAttention to MLX. Its core features include:

1. Paged KV Cache Management: Divides KV cache into fixed-size blocks for dynamic allocation, avoiding static waste.
2. Memory Sharing & Copy-on-Write: Allows shared prefixes (e.g., system prompts) to share KV blocks, copying only when modified—ideal for batch requests with common prefixes.
3. Continuous Batching: Adds new requests immediately when old ones finish, boosting GPU utilization and throughput.

Porting to Apple Silicon faces unique challenges:

1. Unified Memory Architecture (UMA): CPU/GPU share physical memory; the project optimizes to avoid unnecessary copies and leverage zero-copy.
2. Metal Performance Shaders: Adapts to Metal's programming model and memory semantics (different from CUDA).
3. Memory Bandwidth: Optimizes access patterns, reduces fragmentation, and improves cache hit rates to maximize bandwidth use.

mlx-paged-attention delivers significant performance gains:

• Higher Concurrency: Throughput increases by 2-4x vs static allocation (depending on workload/model size).
• Lower Memory Usage: Eliminates fragmentation, enabling larger models or longer sequences on memory-limited devices.
• Stable Latency: Continuous batching reduces latency jitter from batch synchronization.

Key application scenarios:

1. Local API Services: Build high-performance LLM inference services on Mac for local apps or LAN devices (great for privacy/low latency).
2. Multi-User Concurrency: Supports more concurrent users on shared Macs.
3. Long Text Processing: Efficient memory management enables handling long documents (e.g., analysis, code review).
4. Batch Tasks: Continuous batching maximizes hardware use for bulk text processing (e.g., data annotation, content generation).

Comparison with vLLM:

• Platform: vLLM targets NVIDIA GPU/CUDA; mlx-paged-attention focuses on Apple Silicon/MLX.
• Features: mlx-paged-attention implements core PagedAttention but may lack advanced features like speculative decoding or prefix caching (as a newer project).
• Integration: vLLM runs as an independent server; mlx-paged-attention integrates tightly with MLX for easy combination with other MLX apps.

Future plans for mlx-paged-attention:

• Support more optimizations (quantization, speculative decoding).
• Deeper integration with MLX ecosystem components.
• Specialized optimizations for latest Apple Silicon chips (e.g., M3 Ultra).
• Improved APIs and documentation to lower usage barriers. This project represents the trend of LLM inference optimizations expanding to diverse hardware platforms, offering a powerful tool for Apple Silicon LLM developers.