PagedAttention MLX：为Apple Silicon打造的高吞吐量LLM推理引擎

paged-attention-mlx将vLLM的PagedAttention技术移植到Apple Silicon平台，通过分页内存管理和Metal加速内核，在macOS上实现内存高效的LLM推理，支持128K上下文长度。

大语言模型（LLM）推理的性能瓶颈往往不在计算，而在内存管理。传统的自回归生成需要维护KV缓存（Key-Value Cache）来存储历史token的键值表示，以便在生成新token时计算注意力。这个缓存的大小随序列长度线性增长，很快就会成为内存消耗的主力。

更棘手的是内存分配方式。传统实现通常为每个序列预分配连续的内存块，这导致了两个严重问题：

内存碎片化：不同长度的序列需要不同大小的缓存，频繁分配释放产生大量内存碎片。实际可用内存远低于理论值，系统无法充分利用硬件资源。

内存浪费：为了简化实现，系统往往按最大可能长度预分配缓存，而实际使用可能只有一小部分。在长上下文场景下，这种浪费尤为严重。

vLLM项目提出的PagedAttention技术借鉴了操作系统的虚拟内存管理思想：将KV缓存分割成固定大小的"块（Block）"，按需分配，非连续存储。这类似于操作系统将物理内存分页，通过页表映射到虚拟地址空间，彻底解决了连续内存分配的限制。

helgklaizar开发的paged-attention-mlx项目将这一革命性技术带到了Apple Silicon平台，并与Apple的MLX深度学习框架深度集成。以下是其主要技术特性：

分页内存管理：通过BlockManager将KV缓存划分为固定大小的块，按需动态分配。这消除了内存碎片化，使得16GB内存的设备可以支持128K上下文长度，而传统方法仅能支持8K。

Metal加速内核：针对Apple Silicon的GPU架构，项目实现了自定义的C++/Metal内核。通过精细的内存访问模式和并行计算优化，解码阶段的注意力计算延迟降至0.8ms，比传统实现快33%。

动态调度能力：支持请求的抢占（preemption）、缓存的交换（swapping）和变长序列的批处理。在资源紧张时，系统可以将不活跃的缓存换出到CPU内存或SSD，保证高优先级请求的响应。

即插即用集成：与mlx-lm模型完全兼容，现有代码几乎无需修改即可享受PagedAttention带来的性能提升。开发者只需几行代码即可启用分页注意力。

项目提供了令人印象深刻的性能对比数据（基于16GB内存的Apple Silicon设备）：

指标	朴素注意力	PagedAttention	提升幅度
最大上下文长度	8K	128K	16倍
内存碎片化率	25-40%	<1%	40倍以上
解码延迟	1.2ms	0.8ms	33%

这些数据表明，PagedAttention不仅大幅扩展了上下文处理能力，还显著提升了内存使用效率和推理速度。对于在Mac上运行大模型的用户来说，这是一个巨大的福音。

paged-attention-mlx的架构清晰地分离了调度、内存管理和计算三个层面：

请求管理器（Request Manager）：接收用户请求，管理请求队列和优先级。

调度器V2（SchedulerV2）：实现智能的请求调度策略，决定何时处理哪个请求，何时进行抢占和交换。

块管理器（BlockManager）：核心的分页内存管理组件，维护虚拟块到物理块的映射，处理块的分配和回收。

Metal内核：高性能的注意力计算实现，支持从物理KV缓存池高效gather数据并执行注意力计算。

物理KV缓存池：实际的内存存储区域，支持向CPU/SSD的换出操作。

这种分层设计使得每个组件职责清晰，便于独立优化和扩展。

项目的使用非常简单。首先克隆仓库并安装：

git clone https://github.com/helgklaizar/mlx-paged-attention
cd mlx-paged-attention
pip install -e .

然后构建Metal扩展：

mkdir build && cd build
cmake ../src
make

在Python代码中使用：

import mlx.core as mx
from mlx_paged_attention import PagedAttentionLayer, BlockManager

# 初始化块管理器
manager = BlockManager(num_blocks=1024, num_heads=32, head_dim=128)

# 设置注意力层
layer = PagedAttentionLayer(num_heads=32, head_dim=128, manager=manager)

# 推理
seq_id = 0
x = mx.random.normal((1, 4096))
output = layer(x, seq_id=seq_id)

这种简洁的API设计使得现有项目可以无缝迁移到PagedAttention。

将PagedAttention移植到MLX平台并非易事，团队面临了若干技术挑战：

内存对齐：Apple Silicon GPU对内存访问模式有严格要求。项目通过精心设计的块布局和访问模式，确保Metal内核能够高效执行。

统一内存优化：Apple Silicon的统一内存架构（CPU和GPU共享内存）既是优势也是挑战。项目充分利用这一特性，实现了零拷贝的数据传输和灵活的换出策略。

动态形状处理：不同序列长度和批量大小的组合产生大量动态形状。项目通过预分配策略和动态重排，在保证灵活性的同时维持高性能。

项目规划了清晰的演进路线：

里程碑1：MVP版本，提供基于Python的BlockManager实现

里程碑2：自定义C++/Metal PagedAttention内核，达到生产级性能

里程碑3：支持Multi-Head Latent Attention（DeepSeek-V3架构），进一步提升效率

里程碑4：集成Unified Memory Swapping（Flash-Dec），支持超大规模上下文

这种渐进式的开发策略既保证了早期可用性，又为长期愿景保留了空间。