Krea 2图像模型的Apple Silicon原生C++移植实践

我想做一个普通的 macOS 应用来运行 Krea 2，不依赖 Python 安装，也不需要后台运行的服务器。最终做出了 Krealize，但对我来说真正有意思的部分是把模型塞进原生应用，并让它在内存量差异巨大的 Mac 上都能良好运行。

我从 MFLUX 入手。MFLUX 是一个简洁的 Python 实现，利用 Apple 的 MLX 框架运行图像模型。它写得非常直接，所以是理解整个 Krea 2 流水线的良好参考，不会被大型框架绕晕。

它对 Krea 2 的实现涵盖了 Qwen 文本编码器、文本融合、28 层图像变换器、调度器以及 Qwen-Image VAE。更重要的是，在移植过程中我可以拿它的输出作为已知正确的基准进行比对。

MFLUX 流水线的移植

这是一个真正的 C++ 移植，不是把 Python 嵌入到应用中。我把 MLX 操作翻译成了 MLX C++ API，最初尽量保持与 MFLUX 相近的操作顺序。这让早期版本有点无聊，但更容易验证。

原生引擎有四个主要部分：分词器、权重加载器、模型实现，以及一个供 SwiftUI 使用的细小 Objective-C++ 桥接层。

分词器最后花的工作比我预想的多。MFLUX 可以直接用 Python 里的 Hugging Face 分词器。我不想带上这个依赖，所以 C++ 这边自己读取 tokenizer.json。它加载词表、BPE 合并表和追加的 token，然后执行 Qwen 的字节编码和预分词。这里一个微小的差异就会导致不同的 prompt embedding，所以哪怕是看起来很无聊的字符串解析代码也可能改变最终图像。

权重方面，引擎直接用 MLX 打开 MFLUX 格式的 safetensors 分片。量化线性层查找 weight、scale 和 bias 三个 tensor，然后调用 quantized_matmul。普通层则用常规的 matmul。保持 MFLUX 的 tensor 名称很有用，因为这样就不需要为应用再搞一套转换格式。

模型代码遵循与 Python 版本大致相同的路径。首先 prompt 经过 Qwen。Krea 2 从多个文本层取出 hidden states，融合后与图像 token 合并。然后图像变换器对每个去噪步骤跑 28 个 block。最后 VAE 把 latent 还原为像素。

Swift 通过一个 Objective-C++ 类来调用这一切，该类拥有一个 krea2::Engine。保持引擎存活很重要。这意味着模型权重和缓存的 prompt embedding 可以在多次生成之间复用，而不用每次都重新加载。桥接层还会在每个入口点安装相同的序列化 MLX CPU 和 GPU 流，因为后续的生成请求可能在不同的 dispatch worker 上到达。

此时它确实能跑了，但还不足以交给别人。MLX 是惰性求值的，所以一行简短的 C++ 代码可以构建一个非常大的计算图而不实际执行。这本来是好事，直到计算图持有多块临时缓冲区，而 Mac 内存耗尽。

真正让它变快的优化

第一个优化实际上是个内存修复。早期引擎使用一个固定的内存上限。我加入了一个低内存模式，可以在 prompt 编码后卸载文本编码器、在 VAE 解码前卸载变换器，并在各阶段之间清除 MLX 缓存。

这招奏效了，但让用户通过复选框来理解模型常驻内存感觉不对。我把它替换成了针对 16、24、32、48、64 和 96+ GB Mac 的自动硬件配置文件。

在 16 GB Mac 上，每个变换器 block 都逐一求值后才进入下一个。在 24 GB 上每两个 block 一批，在 32 GB 上每四个 block 一批。VAE 也按解码器 block 边界分批。当某阶段结束后，文本、变换器和 VAE 的权重都会被卸载。

在 48 GB 或更大的 Mac 上，引擎可以保持各组件常驻，并提交完整的 28 层变换器计算图。MLX 的 wired-memory 上限仍然限制在 Apple 建议工作集的 90%，另外留出独立保留区给 macOS。中间内存大小使用下一档配置文件，而不是寄希望于最好情况。

这种分批策略只改变数组的求值和释放时机。不改变 seed、精度、尺寸、调度器或注意力计算。

LoRA 还有另一个可以避免的开销。简单版本在每个受影响的层上为每个 LoRA 多跑两次矩阵乘法：

y = xW + (xA1^T)B1^T + (xA2^T)B2^T + ...

同 dtype 的 adapter 可以沿着它们的 rank 维度拼接。把 A 矩阵垂直拼接、把缩放后的 B 矩阵水平拼接之后，相同的求和变成：

y = xW + (x[A1; A2; ...]^T)[s1B1, s2B2, ...]^T

于是 N 个 LoRA 从 2N 次 matmul 启动变成了每对 dtype 组合只需两次启动。五个普通 adapter 现在的启动次数和一个一样。如果只是某个 LoRA 强度变了，我直接在已融合的 B 矩阵中重新缩放对应的切片，而不是重新打开每个文件并重建一切。这也避免了在短时间内同时保留两套完整的融合 adapter 集合。

注意力是另一个大量临时分配的地方。图像变换器有 48 个 query head 和 12 个 key/value head。我的第一个 C++ 版本把 K 和 V 各重复了四次来匹配 head 数量。MLX 的 scaled dot product attention 原生支持 grouped-query attention，所以我删掉了那些拷贝。这避免了每个图像注意力层都 materialize 四份 K 和 V。

我没有在 Qwen 文本编码器里做同样的修改。它的 grouped-query 路径产生了与 MFLUX 参考输出微小的数值差异，而短期蒸馏生成就可以让这些差异变得可见。文本路径仍然扩展 K 和 V，而体量大得多的图像路径则使用原生 grouped-query attention。稍微有点不整洁，但保留了有用的优化，同时不改变 1024px 的参考输出。

然后我寻找在去噪步骤之间重复的工作。Prompt fusion 不依赖 timestep。空间旋转位置表也不依赖。第一个版本在每次变换器调用内部都重建了这两者。现在它们每个生成任务只求值一次，然后在每个步骤中复用。

在 48 GB 或更大的 Mac 上，我还把完整的变换器传给 mx::compile。这让 MLX 可以融合部分计算图开销。我没有在较小配置上启用它，因为那些配置需要在 block 之间有明确的求值点来控制峰值内存。一套配置不是对所有 Mac 都最优。

Krea 2 Turbo 是一个八步蒸馏模型，而变换器调用是每一步中最昂贵的部分。当前的调度器在两个选定的中后期步骤上复用了之前的预测。这把八个调度器步骤减少到六次真正的去噪器求值。这个捷径只在经过测试的精确组合下启用：八步、ER-SDE 调度器、guidance 1.0。它不是一个通用的"跳过某些步骤"的开关。

Fast 模式还做了一件更简单的事。它以请求宽度和高度的 90% 来渲染，尺寸取 16 的倍数，然后对结果进行缩放。与内存分批不同，这是一个真正的速度对细节的权衡，所以作为独立模式保留。

在做更多调优之前我加了一个命令行基准测试。它可以测量冷启动、准备后运行和热运行，并记录 prompt 时间、每次去噪器求值、VAE 时间、MLX 峰值内存、进程占用空间、热状态以及输出的 SHA-256。固定的 seed 和输出哈希值可以清楚地看出某个"优化"之所以快只是因为悄悄改变了图像。

当前版本在配备 48 GB 内存的 M5 Max MacBook Pro 上，fast 模式生成 1024 x 1024 图像大约需要 24 秒。在配备 24 GB 内存的 M5 MacBook Air 上关闭 fast 模式大约需要 2 分 25 秒。这是两台机器，不是通用基准，但差距很好地说明了内存配置文件存在的意义。

仍然还有很多可以融合或编译的工作，但我在这些方面相当保守。对于这类模型，一个节省几秒的改动如果引入了内存峰值或者让输出慢慢偏离参考值，那就是无用的。

Krealize 可在 krealize.app 获取。我很希望收到来自其他 Apple 芯片 Mac 的基准测试结果，尤其是 16 GB 和 24 GB 机型。