针对 Apple Silicon 优化端侧推理
一个可提高预填充和解码吞吐量的自定义本地引擎。
Apple 芯片上的混合计算协调云端前沿智能与 Mac 上的本地模型之间的任务。云端模型处理研究和推理,而本地模型则处理 Mac 上的私有文件和应用。
为了让这种劳动分工感觉无缝,本地推理必须跟上其余任务的步伐。这需要一个能够快速处理提示词并维持高 token 生成速率的引擎。
Lily 是我们专为 Apple 芯片和 Qwen3.6-35B-A3B 构建的轻量级本地推理引擎,针对预填充(prefill)和解码(decode)进行了单独优化。独立演示版已在 GitHub 上公开提供。
简介
在 Mac 上运行 LLM 的常见方法是使用 MLX,这是 Apple 针对 Apple 芯片的开源机器学习框架。其配套库 MLX-LM 添加了使用各种语言模型加载和生成文本所需的组件。MLX 和 MLX-LM 共同为本地 LLM 推理提供了一个现成的、通用的技术栈。
Qwen3.6-35B-A3B 是一个稀疏的混合模型:它使用混合专家(MoE)路由,并将固定大小的循环状态与完整注意力机制结合起来。这些架构选择减少了所需的计算量,但也造成了不规则的工作负载。Token 会路由到不同的专家权重,且循环状态本质上具有顺序性。
MLX-LM 已经为推理阶段和常见工作负载形状选择了优化的内核,但其可复用的操作必须支持许多模型架构。专用于 Qwen 的引擎可以在模型和运行时级别进行专门化,围绕模型的固定结构来协调内核、数据移动和调度。
Lily 在单个进程中端到端地实现了这种专业化。Rust 运行时加载模型检查点并管理会话状态和生成循环,与 OpenAI 兼容的聊天补全 API 接受请求并流式传输 token,自定义 Metal kernels 执行特定于 Qwen 的操作。PyTorch 和 MLX 都不在执行路径中。
我们分别测量预填充和解码性能。预填充吞吐量反映了引擎处理提示词的速度;解码吞吐量反映了它生成输出 token 的速度。
我们在一台配备 M5 Max、40 核 GPU 和 128 GB 统一内存的单台 MacBook Pro 上对 Qwen3.6-35B-A3B 进行基准测试。在用于预填充的十个提示长度和用于解码的十个上下文长度上,范围从 256 到 128K 个 token(K = 1,024),该引擎的预填充吞吐量平均为 MLX-LM 的 1.22 倍,解码吞吐量平均为其 1.32 倍。在 4K token 提示词和 4K token 解码上下文下,自定义引擎达到每秒 5,749.9 个预填充 token 和每秒 186.6 个解码 token,而 MLX-LM 分别为 4,737.5 和 140.9。在多轮会话中,这些时间节省会随着每一次额外的模型调用不断累积。
接下来,我们将解释 Qwen 的架构如何在 Apple 芯片上创造特定于模型的优化机会。然后,我们将逐步介绍由此产生的预填充和解码更改。我们还将探讨额外优化停止带来收益的地方,最后通过与 MLX-LM 的端到端对比来做结。
在 Apple 芯片上针对 Qwen 的特定优化机会
Qwen 创建了三种不同的工作负载形状
Qwen3.6-35B-A3B 包含 350 亿参数,但每个 token 仅激活约 30 亿参数。路由器对 256 个专家子网络进行打分并选择 8 个,同时还有一个处理每个 token 的共享专家。这种稀疏 MoE 设计减少了计算量,但产生了不均衡的工作:专家接收不同数量的 token,并且每个 token 都需要来自不同专家组合的权重。
Qwen 还将 10 个全注意力层与 30 个 Gated DeltaNet 层结合在一起。这两种层类型以不同的方式保留较早的信息。
注意力层使用分组查询注意力(GQA)。Qwen 拥有 16 个查询头和两个键值(KV)头,八个查询头共享每个 KV 头。共享使 KV 缓存更小,并允许在查询头之间重复使用缓存的数据。缓存仍然会为每个 token 存储新的键和值,因此随着上下文的增长,每个解码步骤会读取更多数据。
Gated DeltaNet 将早期信息压缩为固定大小的循环状态。学习到的门控控制保留多少现有状态,而 delta 更新则结合来自当前当前 token 的信息。模型以循环方式定义这些更新,因此每个 token 都依赖于前一个 token 产生状态。然而,在预填充期间,引擎可以通过两种方式评估相同的计算。它可以直接扫描 token 同时向前传递状态,或者将更新重新组织为暴露更多矩阵操作和 token 级并行性的块。哪种方法更快取决于模型维度、工作负载和硬件。
这些结构共同构成了三种计算模式:不均衡的专家组、对不断增长的缓存的注意力机制,以及可直接或分块评估的固定大小循环。
Apple 芯片为不同的工作负载提供不同的路径
预填充会同时处理许多提示词 token 激活行。这里考虑的本地工作负载通常一次解码一个请求(batch 1)并在每步处理一个新行。这种差异改变了相同模型权重的应用方式。预填充可以在数百或数千行中重复使用每个权重块。解码在很大程度上做不到这一点,因为每个新 token 都需要再次遍历权重。
Apple 芯片将 CPU 和 GPU 置于统一内存之后,这是两者都可以访问的单个物理内存池。这允许模型在不维护单独 GPU 副本的情况下保持驻留,但这并不意味着数据移动是免费的。读取权重和中间值仍然会消耗内存带宽,而寄存器和其他片上存储更快,但要小得多。
M5 GPU 还提供了不同的计算路径。预填充的线性层使用通用矩阵-矩阵乘法(GEMM),将权重矩阵同时应用于多行。兼容的 GEMM 可以通过 Metal 4 张量操作使用每个 GPU 核心中的神经加速器。Batch-1 解码改为使用通用矩阵-向量乘法(GEMV),将相同的权重应用于一行。由于权重重复使用率很低,GEMV 主要受内存带宽限制,因此比专为具有更大数据重用的矩阵操作而设计的神经加速器更适合 GPU 的向量算术逻辑单元(ALU)。
这些执行路径并非 Lily 所独有。MLX 在相同的统一内存上运行,并根据工作负载形状选择优化的矩阵和向量内核。MLX-LM 的 Qwen 实现已经对专家工作进行了分组,使用融合的循环 Metal 内核评估 Gated DeltaNet,并使用了感知 GQA 的注意力机制。这些功能是 Apple 芯片上实现高效 Qwen 推理的共同起点。
优化策略
Lily 较窄的范围允许它围绕 Qwen 的精确架构和维度来协调这些共享的执行路径。它使用特定于阶段的 GPU 路径,映射 Qwen 的专家、循环和注意力工作负载以最小化数据移动,并根据测得的工作负载形状选择内核和布局。该策略包含三个部分:
- 使 GPU 路径匹配推理阶段。当预填充可以在多行中复用权重时,使用面向矩阵的执行;当 batch-1 解码一次处理一行时,使用面向向量的执行。
- 将 Qwen 的结构映射到 GPU 上,同时最小化数据移动。保持权重压缩直到使用,在不返回 CPU 的情况下组织路由的专家工作,在循环扫描期间将 Gated DeltaNet 状态保留在芯片上,并复用分组查询注意力共享的 KV 数据。
- 使内核适应工作负载形状。在每个阶段中,根据可用行数、行在专家间的分布、操作的维度以及当前上下文长度来选择瓦片大小、执行布局和注意力路径。
以下各节将解释这些选择。对于在 M5 Max 上的匹配消融实验中评估的优化,我们通过比较除正在研究的优化之外完全相同的引擎配置来估计它们的效果。由于这些实验将我们的引擎版本与自身进行比较,因此它们解释的是机制,而不是将最终结果与 MLX-LM 进行分解对比。
预填充:重逢权重并将路由保留在 GPU 上
预填充暴露了许多 token 行,但 Qwen 将这些行不均衡地路由给专家,并在整个序列中更新循环状态。其优化分为三组:围绕路由的行组织稀疏专家工作,将 Gated DeltaNet 扫描保留在芯片上,以及将长提示词划分为有界的块。
优化稀疏专家计算
在矩阵乘法期间进行权重反量化
Qwen3.6-35B-A3B 检查点使用分组仿射 4 位量化。每个权重存储为 4 位整数代码,而每组 64 个权重共享一个 bfloat16 比例和偏置,用于重建其值。这使得 350 亿参数模型从大约 70 GB 的 bfloat16 权重缩减为 19.4 GB 的检查点,从而使在 Mac 上保持模型驻留变得切实可行。
用于矩阵乘法的 Metal 4 张量操作消耗 bfloat16 操作数,而不是打包的 4 位表示形式。在乘法之前,GPU 必须以 bfloat16 重建权重。Lily 中优化的分组 GEMM 一次只对一小块权重瓦片执行此转换,并将其结果保存在片上线程组内存中,时间刚好够将其与路由的激活行相乘。累加使用 32 位浮点,输出以 bfloat16 写入。完整的展开权重数组从未在统一内存中创建。
在消融实验中,反量化作为单独的操作运行:它将 4 位权重扩展到统一内存中的 bfloat16 数组中,之后矩阵内核再将该数组读回。在 512 token 提示词下,将反量化移入分组 GEMM 通过消除此中间写入和读取,使端到端预填充吞吐量提升了 77.4%。
将专家路由保留在 GPU 上
分组 GEMM 需要将分配给每个专家的激活行存储在一起。在为每个 token 选择八个专家后,直方图会计算分配给每个专家的分配数。前缀扫描将这些计数转换为起始偏移量,散布步骤将行放到其专家组中,块映射则列出了分组 GEMM 必须处理的固定大小矩阵块。
对于每个提示词块,优化后的路径将整个序列保留在一个命令缓冲区(一组有序的 GPU 操作)中。消融实验则会暂停,以便 CPU 检查路由中间体并提交下一个操作。将直方图和前缀扫描保留在 GPU 上会增加两个内核,但消除了每个 MoE 层内部的 CPU-GPU 同步。
在 512 token 提示词下,启用驻留 GPU 的路由使端到端预填充性能提升了 89%。这也说明了为什么仅凭内核数量可能会产生误导:更快的路径会启动更多内核,但从不在层内等待 CPU。
使瓦片大小与专家负载相匹配
在 2K token 提示词下,将每个 token 路由到 256 个专家中的 8 个会产生 16,384 个 token-专家分配,平均每个专家有 64 个激活行。实际分布是不均衡的:一些专家接收许多行,而另一些专家接收很少的行。
分组 GEMM 将每个专家的输出划分为瓦片,它们是矩阵乘法输出的小型矩形块。每个瓦片分配给一个 GPU 线程组。在 Apple 芯片的 GPU 上,线程组包含一个或多个 simdgroup,每个 simdgroup 由 32 个以锁步(lockstep)执行指令的问题组成。
更大的瓦片(tile)可以将设置成本分摊到更多行中,并暴露更多并行工作,但当专家只接收少数几行时,大瓦片的一部分仍处于闲置状态。因此,瓦片大小和 simdgroup 数量是耦合的。
消融实验将瓦片大小固定为 16 行。与该对照组相比,在 2K token 下启用带有四个 simdgroup 的 32 行瓦片使端到端预填充性能提升了 13.2%。
将循环状态保留在芯片上
在预填充期间,每个 Gated DeltaNet 层按顺序扫描提示词,同时向前传递其循环状态。如果禁用寄存器驻留,消融实验将使用分块扫描(blockwise scan)。在 2K token 提示词下,该路径每层会移动 256 MiB(兆字节)的状态,并反复在屏障(所有参与的问题必须互相等待的同步点)处停止协作问题。
循环状态是一个矩阵。优化后的内核将每一列分配给一个 simdgroup。simdgroup 将列分配给其问题,将列加载到其寄存器中一次,并在整个扫描过程中保持状态。问题通过 simdgroup 操作而不是跨线程组共享的片上存储(即线程组内存)来交换中间结果。完成的状态仅在扫描结束后才写回。
状态及其门控使用 32 位浮点格式,因为小的舍入误差会在连续更新中累积。查询和键激活保持为 bfloat16。
在 2K token 提示词下,启用寄存器驻留扫描使端到端预填充性能提升了 5.6%。专家 GEMM 占了预填充时间的 90% 左右。顺序扫描没有暴露出足够的、可复用的矩阵工作来受益于神经加速器。
通过提示词分块限制临时内存
运行时将长提示词处理为一系列有界的块,而不是一次性将每个提示词 token 的临时数据全部保留在内存中。模型权重保留在统一内存中,而循环状态和 KV cache 则将上下文从一个块传递到下一个块。不会丢弃任何先前的上下文。
如果没有分块,临时激活数组会随着完整的提示词而增长,并与模型权重、循环状态和 KV 缓存争夺统一内存。分块一次只保持一个片段的临时值处于活动状态,然后在处理下一个片段之前释放或复用该存储。这限制了峰值工作内存,使引擎能够在不改变模型输出的情况下处理更长的提示词。
分块预填充在许多引擎中很受欢迎,对于在这些内存受限的环境中服务长多轮轨迹至关重要。注意力层的总预填充时间在提示词长度上仍然是二次方的,并且由于早期块的重复 KV 加载而带来了一些额外的开销。
解码:最小化每个 token 移动的字节数
Batch-1 解码每次处理一个新行。由于权重重复使用率很低,其吞吐量主要取决于引擎为每个 token 移动了多少字节。解码更改分为四组:优化单行权重路径、将每个步骤保留在 GPU 上、减少中间和状态流量,以及高效读取注意力缓存。
优化单行权重路径
MLX 已经将单行工作分派给专用的矩阵-向量内核。由于 Lily 不使用 MLX,自定义运行时必须提供相同的基本策略。我们的行并行 GEMV 专为单行激活而设计。simdgroup 协同处理输出,同时并行读取权重矩阵的不同部分。
将每个解码步骤保留在 GPU 上
将 token 交接保留在 GPU 上
每个解码步骤以选择下一个 token 结束;随后的步骤以该 token 作为输入开始。将选择发送到 CPU 然后再发回 GPU 会为每个 token 增加一个同步点。我们的运行时改为在两个命令缓冲区和两个驻留 GPU 的 token 槽之间交替。GPU 选择得分最高的 token 并将其 token ID 直接写入下一个解码步骤的输入槽,同时 CPU 准备后续工作。
重叠独立的 GPU 工作负载
在记录的一个 batch-1 解码步骤中,生成一个 token 启动了 795 个 GPU 内核。它们的依赖关系形成了 555 个顺序阶段,使一些内核可以自由并发运行。然而,Metal 的串行执行模式仍然按顺序运行了每一个内核。
优化后的解码路径在并发 Metal 传递中记录实际数据依赖关系。当 GPU 资源允许时,独立的内核启动可以同时运行。仅当后续工作需要早期结果时才会插入屏障。
减少中间状态和流量
独立的内核通常会产生中间结果:一个内核将临时结果写入内存,下一个内核将结果读回。优化后的解码路径融合了四个链:带有门控激活的两个专家输入投影;带有路由分数和共享专家结果的专家输出投影;注意力前的主键和查询准备;以及带有归一化的循环更新。每个融合内核将临时值保存在寄存器中,而不是通过内存发送它们。
融合还缩短了依赖图:当中间写入消失时,保护其消费者的屏障也随之消失。
高效读取注意力缓存
高效合并注意力缓存读取
注意力机制在每个解码步骤中从 KV 缓存中读取键和值。在消融实验中,相邻的 GPU 问题并不总是请求相邻的字节,这迫使内存系统服务更多的独立事务。启用合并加载会使相邻问题请求相邻字节,以便硬件可以合并它们的读取。
在 bfloat16 配置下,合并将键带宽从 33.8 GB/s 提高到 47.9 GB/s,将值带宽从 42.0 GB/s 提高到 61.8 GB/s,并在 3,840 token 的上下文下将端到端解码性能提升了 2.1%。
打包查询头以复用 KV 行
分组查询注意力机制允许八个查询头共享一个 KV 头。在消融实验中,每个查询头在单独的 simdgroup 中运行,因此所有八个查询头都会独立请求相同的缓存 KV 行。优化后的内核将四个查询头打包到一个线程组中,该线程组加载每个 KV 行一次,并在四次注意力计算中重复使用它。第二个线程组处理其余四个头。
这种技术通常被称为 GQA 打包,它执行相同的算术并生成相同的输出字节,同时将八个独立的 KV 请求减少为两个共享加载。与未打包的消融相比,它在 32K token 上下文下将端到端解码吞吐量提升了 23.8%。
在长上下文中切换注意力布局
全注意力层中的每一个解码步骤都会扫描现有的 KV 缓存。固定块布局将缓存划分为相等的片段,GPU 可以并行处理这些片段。当缓存较小时,其额外的调度并不划算,但随着上下文的增长,固定块布局能更均匀地平衡工作负载。
对于此模型,运行时将 32K token 以下的常规注意力路径保持不变,并在 32K 或更长时使用固定块路径。当每个头有 256 个值且八个查询头共享一个 KV 头时,会触发该切换;其他形状仍保留在常规路径上。消融实验禁用了此切换并始终使用常规路径。在 32K 时,启用固定块路由使端到端解码性能提升了 7.7%,在 64K 时提升了 27.4%,在 128K 时提升了 40.2%。
进一步优化的局限性
某些更改改善了孤立的操作,但并没有改善端到端推理。
投机解码(Speculative decoding)使用较小的模型来提出 token 供完整模型验证,这使 batch-1 解码变慢了 18%。验证处理了两到五行的组,对于该硬件来说这是一个低效的形状,并且这些行通常选择不同的专家,从而增加了读取的专家权重数据量。减少起草者的输出词汇量使起草者吞吐量提升了 4.7–5.1%,但并没有使完整的投机循环更快。此结果具有工作负载特定性:我们在 Blackwell 上的批量 Qwen 部署在不同条件下使用投机解码。
其他实验包括减少 GPU 启动次数、重叠整个阶段、使用更大的预填充瓦片、应用更广泛的融合、加速路由器,以及将输出投影与 token 选择结合起来。没有哪一项改善了完整的推理循环。
硬件限制的测量还表明,主要预填充和解码操作中剩下的提升空间很小。MoE GEMM 和 GEMV 达到了其访问模式下最快持续权重读取速率的 97.9% 和 90.3%。从稀疏 GEMV 中移除算术运算仅使吞吐量改变了 0.2%,这证实了权重读取而不是计算是限制资源。预填充的矩阵乘法在孤立状态下同样达到了理论矩阵限制的 93%,在测试模型内部达到了 80–86%。
端到端性能
端到端比较在两个引擎中加载相同的 4 位检查点字节,并在单个 40 核、128 GB M5 Max 上一次运行一个请求。在每一轮中,两个引擎以交替的顺序运行,以减少来自背景负载和芯片温度变化的偏差。我们与 MLX-LM 最快的直接生成路径进行比较,而不是与其服务器进行比较,因此测量重点在于模型执行,而不是服务开销。
扫描涵盖了从 256 到 128K token 的十种用于预填充的提示词长度和十种用于解码的上下文长度。随着引擎将固定的设置成本分摊到更多 token 上,预填充吞吐量最初会上升。预填充在 4K token 提示词附近达到峰值,然后下降,因为随着提示词的增长,十个全注意力层执行了更多的工作。解码在短上下文中几乎保持平坦,一旦读取不断增长的 KV 缓存变得显著,它就会下降。自定义引擎在记录的每个长度上都更快。
由于专门的执行可能会改变浮点运算的顺序,我们还对照 MLX-LM 检查了数值一致性。在教师强制(teacher-forced)比较中,两个引擎在 192 个位置的每一个位置都从相同的参考前缀预测下一个 token,从而防止早期的差异影响后期的输入。Lily 的困惑度仅高出 0.04%,并且它在 96.35% 的测试位置上选择了相同的最高排名 token。
为本地平台打造
Apple 芯片不是较小的数据中心 GPU。它是一个完整的本地推理平台,拥有自己独特的硬件和软件特性。统一内存赋予了单个节点极高的上限,使其能够容纳大量的模型和状态。M5 神经加速器负责预填充中的密集矩阵工作。向量 ALU 则处理解码中受限于带宽、低复用率的其余部分。
Qwen 提供了更多专业化机会:将专家路由和循环状态保留在 GPU 上,消除不必要的中间体,重叠独立工作,复用共享的 KV 数据,以及使内核适应工作负载形状。
通过针对模型和平台的优化,单台 Mac 就可以高效运行大型稀疏模型。未来的工作将扩大对各种模型、芯片和提供工作负载的覆盖范围,并将此处在一个配置上验证的机制转化为更通用的运行时策略。
更广泛的原则是将引擎与模型的架构以及硬件特定的计算和内存路径相匹配。随着前沿开源权重模型和硬件的演进,高性能本地推理将越来越依赖于针对这两者进行定制的引擎,而不是抽象掉它们差异的引擎。
引用
@online{perplexity2026OptimizingOnDeviceInferenceForAppleSilicon,
author = {{Perplexity} and Yibo Wu and Lequn Chen and Kevin Hu and Phi Hoang and Claire Beamer and Jeremy Yang and Denis Yarats},
title = {{Optimizing on-device inference for Apple Silicon}},
date = {2026-09-01},
url = {https://www.perplexity.ai/hub/blog/optimizing-on-device-inference-for-apple-silicon}
}