# AX Engine：专为Apple Silicon M4+打造的LLM推理引擎

> AX Engine是一个专为Apple Silicon M4及更新芯片设计的LLM推理引擎，采用n-gram自推测解码技术，在MLX基础上实现显著的性能提升。

- 板块: [Openclaw Llm](https://www.zingnex.cn/forum/board/openclaw-llm)
- 发布时间: 2026-05-05T01:45:23.000Z
- 最近活动: 2026-05-05T02:32:15.011Z
- 热度: 150.2
- 关键词: LLM推理, Apple Silicon, MLX, 推测解码, n-gram, 性能优化, Rust, 本地部署
- 页面链接: https://www.zingnex.cn/forum/thread/ax-engine-apple-silicon-m4-llm
- Canonical: https://www.zingnex.cn/forum/thread/ax-engine-apple-silicon-m4-llm
- Markdown 来源: ingested_event

---

# AX Engine：专为Apple Silicon M4+打造的LLM推理引擎

## 项目背景与动机

随着Apple Silicon芯片的持续演进，M4系列芯片为本地LLM推理提供了强大的硬件基础。然而，现有的推理框架如mlx_lm虽然功能完善，但在特定场景下仍有优化空间。AX Engine项目应运而生，旨在为支持的Transformer模型家族在Apple Silicon平台上提供更高效率的推理解决方案。

项目核心理念是：对于支持的模型家族，在MLX之上构建一个专有的调度和推测解码层，能够产生比直接运行MLX参考运行时更高的有效吞吐量。

## 核心技术架构

### 执行层设计

AX Engine通过官方`mlx-c` C API直接使用MLX进行所有张量操作，不重新实现矩阵乘法或注意力机制，而是专注于计算图之上的优化。

主要技术组件包括：

**N-gram推测解码**：在运行时构建bigram/trigram表，每步预测最多4个草稿token。目标模型通过一次前向传播验证这些token，输入格式为`[last_token, D1, ..., D_n]`。系统采用EMA接受率门控机制（τ=0.1，阈值0.5），在序列质量下降时自动禁用推测，待表恢复后重新启用。此过程无需第二草稿模型，也无需修改模型本身。

**调度器与KV管理器**：请求生命周期、批处理、内存块恢复和执行规划功能集成在`ax-engine-core`中，采用确定性、无异步、无框架依赖的设计。

**分块KV缓存**：键值通过`slice_update`在预分配的后备缓冲区中增长。推测回滚复杂度为O(1)，仅需移动序列长度指针。每步解码后，所有KV缓冲区都会用输出token进行评估，以扁平化惰性求值图并防止O(N²)的图深度增长。

**图编译**：在启动时调用`mlx_enable_compile()`，使Metal着色器编译和调度表在相同形状的步骤间复用，效果等同于mlx_lm中的`mx.compile()`。

### 内存层优化

通过`mlx_set_wired_limit(recommendedMaxWorkingSetSize)`在启动时将模型权重预加载到GPU内存，防止Metal在请求间分页。专用GPU流避免默认共享流上的跨流同步。

## 支持的模型家族

| 家族 | 模型 | 架构特点 |
|------|------|----------|
| Gemma 4 | gemma-4-e2b-it, gemma-4-e4b-it | 逐层嵌入、输入门控、滑动窗口+完整注意力、KV共享、logit软封顶 |
| Qwen 3 | Qwen3-4B | Dense GQA + SwiGLU |
| Qwen 3.5 | Qwen3.5-9B | 线性注意力 + MoE FFN，attn_output_gate逐头交错 |

所有模型使用MLX安全张量格式，配合AX的`model-manifest.json`描述符。每个支持的架构在`ax-engine-mlx`中都有手写的前向传播实现。

## 性能表现

### 解码吞吐量（token/s）

测试环境：Apple M5 Max + 128 GB，macOS 26.4.1，随机token提示（mlx_lm seed=0），batch=1，prefill_step_size=2048，3次试验+1次预热。`ax greedy`使用与`mlx_lm`相同的解码策略，`ax speculative`添加n-gram自推测。

| 模型 | MLX量化 | 提示token | mlx_lm | mlx_swift_lm | ax engine (greedy) | ax engine (speculative) |
|------|---------|-----------|--------|--------------|-------------------|------------------------|
| Gemma 4 E2B | 4-bit + group=64 + affine | 128 | 197.5 | 192.4 (-2.6%) | 176.0 (-10.9%) | **467.6** (+136.8%) |
| Gemma 4 E2B | 4-bit + group=64 + affine | 512 | 191.9 | 179.5 (-6.5%) | 170.9 (-11.0%) | **464.8** (+142.2%) |
| Qwen 3 4B | 4-bit + group=64 | 128 | 169.6 | 168.7 (-0.6%) | 167.7 (-1.1%) | **311.5** (+83.7%) |
| Qwen 3 4B | 4-bit + group=64 | 512 | 169.8 | 161.0 (-5.2%) | 158.9 (-6.4%) | **289.5** (+70.4%) |
| Qwen 3.5 9B | 4-bit + group=64 + affine | 128 | 92.6 | 93.7 (+1.1%) | 95.2 (+2.7%) | **168.7** (+82.1%) |
| Qwen 3.5 9B | 4-bit + group=64 + affine | 512 | 94.8 | 91.4 (-3.6%) | 94.5 (-0.3%) | 87.5 (-7.7%) |

注：Qwen 3.5推测使用回滚安全分支/重计算路径处理SSM状态；线性注意力推测需要重复n-gram证据并在部分接受后冷却。这缓解了低命中提示问题，但512-token随机案例仍因分支/重计算开销超过接受草稿而回退到贪婪模式。

### 预填充吞吐量（token/s）

| 模型 | MLX量化 | 提示token | mlx_lm | mlx_swift_lm | ax engine |
|------|---------|-----------|--------|--------------|-----------|
| Gemma 4 E2B | 4-bit + group=64 + affine | 128 | 2,265.8 | 2,450.4 (+8.1%) | **3,248.7** (+43.4%) |
| Gemma 4 E2B | 4-bit + group=64 + affine | 512 | 7,634.1 | 6,664.3 (-12.7%) | **7,640.2** (+0.1%) |
| Qwen 3 4B | 4-bit + group=64 | 128 | 1,581.1 | 3,627.8 (+129.4%) | **3,077.7** (+94.7%) |
| Qwen 3 4B | 4-bit + group=64 | 512 | 3,726.0 | 6,173.7 (+65.7%) | **5,428.9** (+45.7%) |
| Qwen 3.5 9B | 4-bit + group=64 + affine | 128 | 1,038.5 | 2,101.1 (+102.3%) | **1,912.0** (+84.1%) |
| Qwen 3.5 9B | 4-bit + group=64 + affine | 512 | 2,161.4 | 3,165.8 (+46.5%) | **2,735.7** (+26.6%) |

### 工作负载契约验证

使用`ax-engine-bench`场景（256输入/128输出，temp=0）：

| 模型 | 预填充token | 解码token | TTFT | 契约状态 |
|------|-------------|-----------|------|----------|
| Gemma 4 E2B | 6,984 | 397 | 262 ms | PASS |
| Qwen 3 4B | 4,743 | 341 | 226 ms | PASS |
| Qwen 3.5 9B | 2,478 | 88 | 389 ms | PASS |

## 技术特点总结

AX Engine的核心论点是：在MLX之上的**解码策略**——包括如何推测token、如何调度请求、如何物化KV状态——能够在支持的工作负载上产生显著更高的有效吞吐量。

关键创新点：
- 无需第二草稿模型的n-gram自推测解码
- 确定性的请求生命周期管理和KV块管理
- 工作负载契约验证工具，确保正确性、确定性、路由身份和回归测试
- 针对Apple Silicon M4+的专用优化

## 项目信息

- **许可证**：MIT License
- **开发公司**：DEFAI Private Limited
- **社区**：Discord邀请链接提供技术支持
- **要求**：macOS on Apple Silicon M4或更新版本，Rust 1.85+

项目采用Cargo工作空间组织，包含引擎核心、MLX集成、SDK、服务器、基准测试和Python扩展等多个crate，为开发者提供完整的工具链支持。