Mac 上的 Megakernel:融合省下了什么,又付出了什么代价
Megakernels on Mac: What Fusion Saves and What It Costs
在 Mac 上测试 Megakernel 融合方案时,FFN-to-GDN-input 窗口耗时 521.2 µs,其中内核间空隙仅 7.7 µs,大部分时间已在运算内部。
在降维投影之后,均方根归一化需要整个隐藏向量的平方和。一个完成了四分之一的 worker 无法使用未完成的和来对其进行归一化。融合改变了部分和所在的位置以及 worker 发出完成信号的方式;剩余的四分之一仍然必须完成。
循环具有相同的约束。如果下一个状态是 St+1=F(St,xt),那么下一个 token 依赖于前一个状态。独立的状态列可以并行运行,但 token 更新保持其顺序。
持久化调度现在必须自行协调这些依赖关系。队列、完成标志和屏障会消耗资源。等待另一个阶段的 worker 也可能占用该阶段所需的空间。线程组屏障只能协调其自身的组;它无法安全地替代整个网格的同步方案。
捕获到的 FFN 到 GDN 输入窗口耗时 521.2 µs,其内核之间仅间隔 7.7 µs。大部分时间已经花在操作内部。
持久化 worker 的布置对某些投影的适配程度优于其他投影。将输入投影保留为单独的原生内核恢复了大约半毫秒,但混合方案仍然落后。改进完整原型后,它比其前身快 1.045%,比优化后的原版慢 2.701%。
worker 时间线使调度问题变得可见。在插桩运行中,160 个 worker 中有 24 个在早期降维投影票据上等待了约 205 µs,而门控/上投影工作已经完成。其他 worker 在那段等待期间处于忙碌状态。平均空闲时间约为每个 worker 85 µs,这有助于定位不平衡,而不是一整块可以全部恢复的时间。
额外的暂存也没有帮助。通过线程组存储加载降维投影权重,单槽增加了 20.714 µs,双槽前瞻增加了 22.670 µs。两种变体都输掉了全部 24 对。
寄存器压力是更广泛回归的一个可能解释。更好的对照产生了 +0.043 和 -0.044 ms/步的配对中位数,没有可重复的惩罚,也没有报告溢出。先前的寄存器归因被撤回。
Lithos 的 GDN 研究揭示了类似的区别。其优化后的融合配置耗时 599.72 µs,低于 646.25 µs。但匹配的打包、分离调度对照耗时 596.02 µs。打包和投影几何结构提供了大部分改进。
这一区别改变了下一步努力的方向。更深的 K 拆分将一个依赖调用的降维投影探针从 118.8 µs 降至 99.7 µs,但输出变化约为最后一个位置上的一个 BF16 单位。短 K 输出投影略微变差。形状、算术和所需的数值行为都会影响哪些更改可用。
实际的融合边界是那种能够消除特定读取、写入或重复变换,同时不拖慢周围工作的边界。Metal Performance Shaders、Metal Performance Primitives 或 MLX 中的现有原语可能已经可以完成这项工作。更大的融合区域需要证明其引入的额外协调是值得的。
方法与测量
带宽与内存测量
有用的起点是所选执行路径移动了多少数据,以及融合可以消除其中哪一部分的数据移动。这些实验区分了观测到的流式传输速率、源自源数据的有效载荷估计,以及真实内核之间花费的时间。
| 发现 | 结果 | 证据类型 | 解读 |
|---|---|---|---|
| M5 Max 上的最佳流式传输配置 | 563.7、570.2 和 572.5 GB/s | 实测,三次 4 GiB 探测运行 | 在此有限访问模式和几何扫描中约为 570 GB/s |
| 普通 B1 投影有效载荷 | 15.758 GB | 根据固定的模型布局和调用图计算 | 逻辑目标、草稿和上下文投影字节数;并非实测 DRAM 流量 |
| 以 570 GB/s 传输的有效载荷 | 27.65 ms | 条件计算 | 若该逻辑有效载荷以该速率通过接口的参考传输时间 |
| FFN 到 GDN 输入窗口中的内部间隙 | 521.2 µs 中的 7.7 µs | 实测分析器窗口 | 可见间隙仅占该区域的一小部分;生产者和消费者工作也占据了内核区间 |
流式探测扫描了三种访问模式、四种网格大小和两种线程组大小。上述每个数字都是其运行中的最佳配置,而不是对某一固定配置的重复采样。一次较早的 8 GiB 探测报告了 512.9 GB/s,表明访问模式和探测几何形状对结果的影响有多大。
保存的流式传输记录未记录持续预热时长或热平台期。它们支持一个观测到的带宽参考值,而非经过验证的预热最大值。一个单独的分阶段测试装置在其计时窗口之外记录了 3.02329 累计 GPU 秒的预热时间;该准备工作仅适用于该测试装置。测量说明和源记录
27.65 ms 的估计值
投影审计统计了 14.434467840 GB 的目标有效载荷和 1.323417600 GB 的草稿及上下文有效载荷。目标前馈门/上/下子集占 9.625927680 GB。受限草稿头已包含其可变段,因此只计一次。
treference=15.757885440 GB570 GB/s=27.6454 ms。
此计算以特定参考速率对特定逻辑有效载荷进行定价。缓存复用、重复读取、丢弃的草稿、上下文长度、模型格式和所选图都可能改变实际流量。严格的下界需要同时满足跨限制接口的最小流量和其持续服务速率的上界。最佳观测流式传输结果对这两者均无法为所有未来设计确立。
因此,从实测的 40 ms 周期中减去 27.65 ms 并不能确定 12.35 ms 的可移除开销。其余部分包括计算、状态和缓存流量以及依赖关系。这些成本可能重叠。融合需要移除特定的读、写或调度边界,而不是在不同范围测量的两个数字之间的差距。
GPU 预热、持续发热和缓存驻留之间的区别在 Pulsar 的运行条件测量 中有所涉及。
投影形状改变实现的速率
流式参考并不能预测每一个矩阵运算。一次八行投影扫描报告了以下最佳有效载荷速率:
| 投影 | 报告的速率 |
|---|---|
| 草稿上下文 | 241 GB/s |
| FFN 下 | 247 GB/s |
| 混合器输出 | 272 GB/s |
| 全注意力输入 | 422 GB/s |
| GDN 输入 | 483 GB/s |
| FFN 门/上 | 530 GB/s |
| 词汇头 | 554 GB/s |
这种分布使得针对特定形状的控制变得重要。更深的 K 拆分将依赖调用的 FFN-down 探针从 118.8 微秒降至 99.7 微秒,但结果变化约一个 BF16 最后一位单位。短 K 输出投影则从 36.7 微秒变为 37.3 微秒,独立调用的计时大致不变。数值差异对字节完全一致的路径很重要;收益还取决于投影形状和调用顺序。投影与控制回执
持久化原型的结果
测量涵盖三个层面:捕获的内核窗口、隔离的投影变体以及配对的 serving 对比。
调度之间的小间隙
在一个捕获的 FFN 到 GDN 窗口中,分离内核路径耗时约 521.2 微秒。其内核之间约有 7.7 微秒。
这约占该窗口的 1.5%。这并不意味着全部 7.7 微秒都能被消除,也不是完整的 serving 周期测量。它表明可见间隙是这一特定区域的一小部分。
生产者的最终工作和消费者的初始工作也在各自的内核区间内耗时。即使调度边界消失,这些工作仍然存在。
投影几何与持久化调度
不同的投影需要不同数量和形状的并行工作。一种持久化 worker 安排并不适合所有投影。
在混合实验中,将输入投影保留为原生内核恢复了约半毫秒。该混合方案仍未击败分离内核基线。
随后将改进的持久化原型与其前身及原版执行进行了对比:
| 对比 | 观察结果 |
|---|---|
| 改进的 megakernel 对比原始 megakernel | 快 1.045% |
| 改进的 megakernel 对比优化的分离内核 | 慢 2.701% |
这是 9 月 27 日在 Qwen3.8-27B 上的筛选,在 1K、5K 和 10K 输入长度下各采样三个种子。比率是九次配对比较的几何平均值;输出哈希和有效 token 数匹配。每个请求都有代表性的 GPU 预热。指标是解码输出 token 数除以引擎记录的解码墙钟时间,而非到完成答案的延迟。
调度与同步成本
插桩暴露了被调度次数掩盖的成本。在一次 worker 时间线实验中,分离链耗时 407.6 微秒,持久化内核 428.0 微秒,插桩的持久化内核 450.7 微秒。软件时钟本身增加了 5.3%,因此其 worker 计时描述的是插桩运行。
| 诊断 | 记录结果 | 测量边界 |
|---|---|---|
| 早期 down-projection 票据 | 160 个 worker 中有 24 个等待约 205 微秒,而 136 个 gate/up 任务占据一波 | 插桩的 worker 时间线;其他 worker 正在计算 |
| 平均 worker 空闲时间 | 约 每 worker 85 微秒,包括就绪等待和最后一个不完整波 | worker 占用核算,而非 85 微秒的可恢复关键路径时间 |
| 内核内 RMS 归一化 | 对于测试的多 worker 变体,每次转换约多 3.1–4.3 微秒 | 三次 24 轮调用,对照匹配的 down-kernel 控制组,随后进行单独的 RMS 归一化;Chrome 打开 |
| 精简持久化变体对比原版 | Chrome 关闭时 +4.55 微秒/转换;Chrome 打开时 +4.97 微秒/转换 | 分离的 672 轮和 432 轮直接配对队列;没有因果性的 Chrome 对比 |
这些诊断覆盖了不同的变体,无法合并为对服务回归的单一解释。它们说明了为什么移除一次调度是不够的:工作分配、就绪信号和归一化的代价可能与它们所替代的边界一样高。工作线程时钟也会扰动执行,而平均工作线程空闲时间可能与其他地方的有用工作重叠。诊断凭证与范围
线程组暂存
另外两个变体通过线程组存储加载 Q4 下投影权重。
在 24 个交替配对中,单槽版本每个投影增加约 20.714 微秒,双槽前瞻版本增加约 22.670 微秒。两者都没有赢下任何一对。
这些算子测量使用了真实权重和合成激活。两个被测实现都失败了;该结果不能推广到所有暂存设计,也不能将占用率单独确定为原因。
额外的拷贝、同步和资源使用必须由它们所实现的复用偿还。
寄存器压力仍未得到证实
megakernel 在驱动元数据中使用了更多寄存器,因此寄存器压力是一个有吸引力的解释。
后来的一项对照尝试在改变值生命周期的同时匹配额外的加载和算术。直接配对中位数为 +0.043 和 -0.044 ms/步,没有可重复的惩罚,且所检查的内核报告没有寄存器溢出。该对照并不完美,但足以撤回先前的数值归因。
保存的帧级计数器报告占用率约为 15.5% 和 18.0%,而活动核心计数器报告约为 92.0% 和 95.3%。这些测量的是不同的东西;仅凭低占用率并不能确定核心空闲。这两个观察都没有确定损失的原因。
未融合对照中的打包与几何
Lithos 发表的 GDN 研究提供了一个有用的独立示例:
| 实现 | 报告的逐层最小时间均值 |
|---|---|
| 原始实现 | 646.25 微秒 |
| 优化的自动配置 | 599.72 微秒 |
| 匹配的打包多调度对照 | 596.02 微秒 |
| 独立调优的打包对照 | 598.87 微秒 |
该研究将显著改进归因于操作数打包和投影几何。融合并未确立相对于相应打包对照的优势。这些是固定的八行层测量,而非生成吞吐量,并且与原始模型的数值等价性是在容差范围内评估的。作者的研究
选择融合边界
一个有用的候选方案会移除可识别的工作:一次中间读/写、冗余变换、重复的权重加载或测量到的调度间隙。运行时、Metal Performance Shaders (MPS)、Metal Performance Primitives (MPP) 或 MLX 中已有的原语可能已经覆盖了它。当新实现需要更好的打包或瓦片几何时,分离内核对照也需要同样的改进,以便比较能够隔离融合。
数值要求也决定了哪些变更是可用的。字节相等、有文档记录的容差、目标分布保持和任务质量需要不同的检查。一个具有明确支持的形状和失败行为的窄候选方案更容易针对该要求进行验证,包括一个打破其核心假设的案例。然后,算子测量结果会输入到图、有状态模型和服务测试中,之后实现才会扩展。
融合改变了目标传递的成本。推测解码改变了从中获得的有效进展;Pulsar 的引擎测量在完整的服务路径中覆盖了这两者。
延伸阅读
调度与正确性
在持久化内核内部,编译器或手写调度器描述了存在哪些任务、它们的输入位于何处以及何时可以运行。
概念上:
while there is work in this region:
obtain a task whose prerequisites are satisfied
compute the task's output
make the output visible to its consumers
publish completion
这个调度草图留下了未解决的同步问题。Metal 实现需要一种安全的方式来检测已完成的前置条件,并使其写入对消费者可见。
有几种合理的融合边界:
| 边界 | 可能的好处 | 新的责任 |
|---|---|---|
| 逐元素链 | 消除中间读/写 | 保持操作顺序和尾部 |
| 投影加尾声 | 在存储前重用累加器 | 保持归约和舍入 |
| 注意力或 GDN 混合器 | 重用局部统计量、状态切片和布局 | 协调依赖任务 |
| 解码器层 | 减少更多边界 | 容纳非常不同的阶段 |
| 整个解码步骤 | 将整个调度保留在 GPU 上 | 处理广泛的同步、状态和资源生命周期 |
更大的区域提供更多重用,但也将更多样化的工作纳入一个调度之下。
例如,一个层的多层感知机(MLP)可能已经拥有非常有效的矩阵内核。将其拉入持久化工作器布局可能会使该矩阵操作变慢,即使周围的调度次数下降。
间接命令缓冲区或另一种预编码重放形式解决了不同的问题。它可以减少重复的主机编码工作,同时保留多个 GPU 调度。引擎可以同时使用命令重放和混合器巨型内核。
依赖关系与同步
在下投影产生隐藏向量后,均方根(RMS)归一化需要其所有元素的平方和。完成前四分之一的工人仍必须等待剩余的四分之一,然后才能获得最终的归一化因子。
融合可以改变部分和所在的位置以及工人发出完成信号的方式。它通常无法使用未完成的求和进行归一化并保持等价。
同样的原则适用于递归。如果状态更新是:
St+1=F(St,xt),
下一个 token 的更新依赖于前一个状态。并行化独立的状态列可能会有所帮助。简单地同时启动所有 token 更新并不能保持递归。
工作队列、完成标志和软件屏障会增加协调成本。适合一种投影的 worker 安排可能无法满足另一种投影的需求。线程组屏障不是整个网格的屏障,而一个网格中常驻 worker 若为尚未调度的生产者自旋等待,可能无法取得进展。
实现必须明确每个值由哪个 worker 拥有、写入何时可见,以及超时或取消时会发生什么。在已提交的工作仍可能读取缓冲区时,不能释放这些缓冲区。下一个消费者必须看到逻辑上已提交的状态,即使该状态的物化被推迟了。
来源:Hacker News · abhishek.it