oMLX Bonsai 自定义内核:Qwen 系列快速 QMV 路径完整解析

oMLX Bonsai 自定义内核:Qwen 系列快速 QMV 路径完整解析

【免费下载链接】omlxLLM inference server with continuous batching & SSD caching for Apple Silicon — managed from the macOS menu bar项目地址: https://gitcode.com/GitHub_Trending/om/omlx

oMLX是一款专为 Apple Silicon 打造的 LLM 推理服务器,而 Bonsai 自定义内核 正是它让 1-bit / 2-bit 极致量化模型(如 Qwen 系列 Bonsai 版)高速解码的秘密武器。本文将用通俗的方式拆解这条"快速 QMV 路径"是什么、为什么快、以及它如何在 oMLX 中自动生效。

一、先搞懂三个名词:Qwen、Bonsai、QMV

🔍为什么这三个词经常一起出现?

在 Apple Silicon 上,单 token 解码是典型的内存带宽瓶颈任务:权重必须从统一内存(UMD/DRAM)读出来算完即弃。Qwen 这类大模型的解码速度,几乎直接取决于"每秒能从内存里读出多少 GB 权重"。

oMLX 的 Bonsai 内核 就是为此重写了 Metal 层算子,把每次 QMV 的内存流量压到最低。

二、三条核心加速路径:一条比一条省

oMLX 按比特数和批大小 M自动路由到不同的 Metal 内核,逻辑全部在 fast.py 中,调度规则如下:

1️⃣ qmv_fast:1-bit 单 token 解码的"最快一公里"

• 专为1-bit 亲和量化(affine)权重设计,从 Bonsai MLX 分支移植而来
• 权重以 uint32 打包存储(一个 uint32 装 32 个 1-bit 值),解码时逐字节移位解包 + 点积,全程不落地浮点权重
• 对应源码:csrc/bonsai_quantized.metal

💡 一个容易被忽略的细节:标准 MLX 的 metallib没有 bits=1 的 affine 反量化核,oMLX 在 bonsai_qmv.py 里额外打补丁,补齐了 1-bit 权重在 prefill 阶段(先显式反量化为 float16 再做矩阵乘)的缺口。

2️⃣ qmv_wide:小批量解码的"权重复用"魔法

当多个请求并发解码(M = 2~5)时,每行单独算 QMV 会把同一份权重反复读出 M 次。qmv_wide内核让一次读入的权重同时服务 M 个向量:

• 触发条件:M ≥ 3 且硬件为 gen-15+(M3 及以后)
• 官方基准实测:Qwen 系 gate/up/down_proj 在 M=5 时,带宽从71 GB/s 提升到 104 GB/s(约 1.3~1.5 倍)
• 还有一个更狠的sym 对称变体:当检测到bias = -scale × ratio(恒等式 I-B)时,内核直接跳过 bias 张量的 DRAM 读取,再省一笔流量

3️⃣ t5 三值格式:约 1.585 bit/权重的"压缩之王"

这是 oMLX 为 Qwen3.5 等模型准备的三进制(取值 {0,1,2})打包格式:

• 用 uint8 存储,每权组 13/26 字节,平均约1.585 bit/权重
• 权重常驻内存时,比 2-bit uint32 格式再省约 23% RAM(见 bonsai_t5_load.py 文档字符串)
• 解码走bonsai_t5_qmv/bonsai_t5_qmv_wide;长序列 prefill(M > 16)则切换到bonsai_t5_qmm—— 一个"解包 + simdgroup 矩阵乘"融合核,每个 t5 字节只读一次
• 仓库还提供 tools/repack_ternary_t5.py 用于把现有模型权重重打包成 t5 格式

三、如何自动生效?用户几乎"零操作"

🪄 这是 Bonsai 内核设计上最值得点赞的地方:你不需要改任何模型代码。

模型加载完成后,model_loading.py 会依次调用两个补丁:

1. apply_bonsai_qmv_patch()(定义于 patches/bonsai_qmv.py) 全局替换mlx.nn.QuantizedLinear.__call__。此后凡是 1/2-bit affine 层进入解码区间(M ≤ 5),就自动被路由到 Bonsai 快速 QMV 内核;prefill 阶段则原样交还标准 MLX。
2. apply_bonsai_t5_load_patch()(定义于 patches/bonsai_t5_load.py) 让Module.load_weights接受 t5 格式的 uint8 权重,并拦截mx.quantized_matmul的直接调用(mlx_vlm 等库会绕过QuantizedLinear直接调它)。

🛟优雅降级机制:这些 Metal 内核以 C++/nanobind 扩展(csrc/)形式编译(构建入口见 setup.py 第 54 行)。若本地未编译该扩展,fast.py会自动探测 ABI 并静默回退到标准mx.quantized_matmul——慢一些,但结果正确、绝不报错。

另一个亮点是spec_decode_verify(csrc/spec_decode.metal):把投机解码中"比对草稿 token → 统计接受长度 → 生成修正 token"三步融合成单个 Metal 内核一次完成,进一步压缩每 token 的往返开销。

四、动手验证:跑一跑 Bonsai 解码微基准

想亲眼看看这些内核跑出了多少 GB/s?仓库自带微基准脚本,按 Qwen3.6-27B 的真实投影层形状(q_proj、gate_proj、down_proj 等 7 种尺寸)测 M ∈ {1,2,3,4,5}、bits ∈ {1,2}、group_size ∈ {64,128}:

python benchmarks/bonsai_decode_bench.py --M 1,2,3,4,5 --bits 1,2 --csv

脚本会输出每种内核变体的实测 DRAM 带宽(GB/s)与延迟(µs),带宽按"权重 + scales + biases + 激活"的实际字节流精确核算——sym 变体的 biases 流量会如实记为 0,方便你验证"省掉 bias 读取"这条优化是否真的命中。

正确性方面有 tests/test_bonsai_qmv.py 与 tests/test_bonsai_t5_load.py 两组测试兜底,ABI 匹配检查逻辑则见 tests/test_custom_kernel_abi_probe.py。

五、关键文件速查表

小结:为什么值得了解 Bonsai QMV 路径

✅极致省内存:1-bit 到 1.585-bit 三值格式,大模型装进 Mac 的统一内存 ✅不牺牲解码速度:qmv_wide 权重复用让并发小批量解码达到 1.3~1.5 倍带宽 ✅零配置生效:加载模型即自动打补丁,缺失原生扩展时优雅降级 ✅可验证:内置按 Qwen 真实形状设计的微基准,带宽数据透明可复现

如果你正在 Apple Silicon 上跑 Qwen 的 Bonsai 量化模型,oMLX 的这套自定义内核就是"小模型体积、大模型速度"背后最核心的工程答案。

【免费下载链接】omlxLLM inference server with continuous batching & SSD caching for Apple Silicon — managed from the macOS menu bar项目地址: https://gitcode.com/GitHub_Trending/om/omlx