wehub-resource-sync
612 lines
25 KiB
Markdown
612 lines
25 KiB
Markdown
<!-- WEHUB_ZH_README -->
|
||
> [!NOTE]
|
||
> 本文档由 WeHub 基于上游 README 翻译整理，属于社区翻译，非官方中文文档。
|
||
> [English](./README.en.md) · [原始项目](https://github.com/Mininglamp-AI/cider) · [上游 README](https://github.com/Mininglamp-AI/cider/blob/HEAD/README.md)
|
||
> 原作者、版权与许可证归属以原始项目及本仓库 LICENSE 文件为准。
|
||
|
||
# cider
|
||
|
||
Cider 基于 MLX 为 macOS 开发。它提供：
|
||
|
||
1. MLX 所不具备的**在线激活量化**（W8A8/W4A8），配有自定义 INT8 TensorOps 内核，作为 MLX 自定义原语（custom primitives）构建，支持完整的惰性求值（lazy evaluation）。
|
||
2. 受 [FlashInfer](https://github.com/flashinfer-ai/flashinfer), 启发的**优化 SDPA 解码内核**，采用连续块 GQA 调度与寄存器分块（register tiling），以实现更快的自回归解码。
|
||
|
||
它还包含服务端扩展以及对 mlx_vlm 的非侵入式兼容性补丁（已在 mlx_vlm 0.4.3 上验证），包括修复与 RoPE 位置处理及分块预填充（chunked prefill）相关的 Qwen3-VL 多图推理问题。
|
||
|
||
## 条件编译（M4 / M5）
|
||
|
||
Cider 采用**条件编译**：INT8 TensorOps C++ 扩展仅在 Apple M5+ 上构建。
|
||
|
||
| Chip | `pip install -e .` behavior | `import cider` behavior |
|
||
|------|----------------------------|------------------------|
|
||
| **M5+** | 完整构建（CMake + Metal kernels） | 全部功能可用 |
|
||
| **M4 及以下** | 跳过 C++ 构建，安装纯 Python 包 | `is_available()` → False，`convert_model()` 为警告空操作（warning no-op） |
|
||
|
||
**通过环境变量覆盖：**
|
||
```bash
|
||
CIDER_FORCE_BUILD=1 pip install -e . # Force build (e.g., CI)
|
||
CIDER_FORCE_BUILD=0 pip install -e . # Force skip
|
||
```
|
||
|
||
## 模式
|
||
|
||
| Mode | Weights | Activations | Compute Path | Status |
|
||
|------|---------|-------------|--------------|--------|
|
||
| **W8A8** | INT8 symmetric | INT8 per-token | TensorOps matmul2d | ✅ 已实现 |
|
||
| **W4A8** | INT4 packed (uint8) | INT8 per-token | Unpack→TensorOps | ✅ 已实现 |
|
||
| W4A16 | — | — | MLX built-in | 基线 |
|
||
| W8A16 | — | — | MLX built-in | 基线 |
|
||
|
||
**W4A16 和 W8A16 已由 MLX 原生支持** — 本 SDK 提供 MLX 未实现的 W8A8 与 W4A8 模式。
|
||
|
||
## 优化 SDPA（解码注意力）
|
||
|
||
Cider 包含面向自回归解码（`Q_seq=1`）的优化**缩放点积注意力**（Scaled Dot-Product Attention）内核，受 [FlashInfer](https://github.com/flashinfer-ai/flashinfer). 启发。关键技术：
|
||
|
||
- **连续块 GQA 调度（Contiguous Chunk GQA Scheduling）**：不再让每个 Q 头独立派发（读取 KV `H_q` 次），而是由单个线程组处理共享同一 KV 组的所有 Q 头。这可减少最多 `gqa_factor`× 的 DRAM 带宽。
|
||
- **寄存器分块（Register Tiling，TILE=4）**：每个线程在寄存器中保存 4 个 Q 头的部分结果 — 无需共享内存，无需屏障同步 — 最大化分组查询解码的 ALU 吞吐。
|
||
- **两遍 FlashDecoding（2-Pass FlashDecoding）**：对于长序列（≥1024 tokens），两遍归约将 KV 序列划分到各线程组，以提升占用率。
|
||
- **AutoTune**：针对（架构、GQA 比率、头数、序列长度）自动扫描块大小，并缓存至 `~/.cider_sdpa_tune.json`。
|
||
|
||
### 用法
|
||
|
||
```python
|
||
import cider
|
||
|
||
# One-line patch — all downstream frameworks (mlx_lm, mlx_vlm) benefit automatically
|
||
cider.patch_sdpa()
|
||
|
||
# Optional: run AutoTune once to find optimal block sizes for your chip
|
||
cider.autotune_sdpa()
|
||
|
||
# Undo patch if needed
|
||
cider.unpatch_sdpa()
|
||
```
|
||
|
||
该补丁**完全透明**：预填充（`Q_seq > 1`）与带掩码注意力会自动回退到 MLX 原生实现。仅解码步骤（`Q_seq=1`，无掩码）会走 Cider 的优化内核。**无需修改** mlx_lm 或 mlx_vlm 代码。
|
||
|
||
### SDPA 微基准（Apple M5 Pro，D=128）
|
||
|
||
相对 MLX `mx.fast.scaled_dot_product_attention` 的加速比：
|
||
|
||
| Config | N=1K | N=2K | N=4K | N=8K | N=16K | N=32K |
|
||
|--------|------|------|------|------|-------|-------|
|
||
| MHA (H_q=32, H_kv=32) | 1.03x | 1.01x | 1.02x | 1.00x | 1.00x | 0.99x |
|
||
| GQA4 (H_q=32, H_kv=8) | 1.04x | 1.06x | 1.06x | 1.12x | 1.16x | 1.21x |
|
||
| GQA8 (H_q=32, H_kv=4) | 1.04x | 1.07x | 1.13x | 1.21x | 1.34x | **1.57x** |
|
||
|
||
加速范围：**0.99×–1.57×**。正确性：18/18 PASS（相对 MLX 最大差异 ≤ 0.000122，fp16）。14 WIN / 4 TIE / 0 LOSS — 零回退。
|
||
|
||
### 启用 SDPA 的端到端解码速度（Qwen3-VL-4B，Apple M5 Pro）
|
||
|
||
使用 `benchmarks/bench_e2e_wxa16.py` 在真实多图 VLM 轨迹上测试（1–3 张截图，1334–3455 prompt tokens）：
|
||
|
||
| Config | Decode tok/s | vs Baseline |
|
||
|--------|:-----------:|:-----------:|
|
||
| FP16 | 32.0 | — |
|
||
| FP16 + Cider SDPA | 32.6 | +1.9% |
|
||
| W8A16 | 54.0 | — |
|
||
| W8A16 + Cider SDPA | **58.0** | **+7.4%** |
|
||
| W8A8 per-group | 51.0 | — |
|
||
| W8A8 per-group + Cider SDPA | 52.3 | +2.5% |
|
||
| W8A8 per-channel | 49.1 | — |
|
||
| W8A8 per-channel + Cider SDPA | **52.6** | **+7.1%** |
|
||
|
||
Cider SDPA 可叠加在 W8A8 预填充收益之上。收益随 GQA 比率与 KV cache 长度而放大 — GQA 比率更高（例如 GQA8 的 Qwen3-32B）且上下文更长的模型获益更大。
|
||
|
||
MLX 的量化是**仅权重量化（weight-only）**：QuantizedLinear 将权重反量化为 FP16 并使用 FP16 GEMM。虽然 MLX 的 Steel NAX 模板足够通用，可实例化为 INT8 类型（且能达到相同的原始 matmul 吞吐 — [参见我们的透明基准测试](benchmarks/mlx_native/cider_vs_mlx_int8.md)），但 MLX 并未提供实际 W8A8 推理所需的量化/反量化流水线。Cider 通过融合的 quantize-matmul-dequant 原语填补这一缺口。
|
||
|
||
本 SDK 为受支持的推理路径实现在线 INT8 激活量化，以及基于 INT8 TensorOps 的计算。
|
||
|
||
### W8A8 量化粒度
|
||
|
||
| Granularity | Description | Speed | Precision |
|
||
|-------------|-------------|-------|-----------|
|
||
| **Per-channel** | 每个输出通道一个 scale | 最快（预填充 1.8x） | 略低 |
|
||
| **Per-group (gs=128)** | 每 128 个元素一个 scale | 较快（预填充 1.5x） | 精度保留适中 |
|
||
| **Per-group (gs=64)** | 每 64 个元素一个 scale | 中等（预填充 1.3x） | 精度更高 |
|
||
|
||
## 性能（Apple M5 Pro）
|
||
|
||
### 单算子延迟
|
||
|
||
Shape [N=10240, K=2560]
|
||
| M | PC(ms) | PG(ms) | w8a16 | w4a16 | PC/w8 | PC/w4 | PG/w8 | PG/w4|
|
||
|-----|------|------|-----|-----|-----|------|-----|----|
|
||
| 1 | 0.27ms | 0.26ms | 0.26ms | 0.18ms | 0.96x | 0.67x | 0.99x | 0.69x |
|
||
|128 | 0.34ms | 0.39ms | 0.49ms | 0.44ms | 1.43x | 1.28x | 1.26x | 1.13x |
|
||
|1024 | 1.23ms | 1.52ms | 2.24ms | 2.04ms | 1.82x | 1.66x | 1.47x | 1.34x|
|
||
|4096 | 4.41ms | 5.65ms | 8.12ms | 7.72ms | 1.84x | 1.75x | 1.44x | 1.37x |
|
||
|8192 | 8.71ms | 11.40ms | 16.23ms | 15.09ms | 1.86x | 1.73x | 1.42x | 1.32x|
|
||
|
||
|
||
Shape [N=2560, K=10240]
|
||
| M | PC(ms) | PG(ms) | w8a16 | w4a16 | PC/w8 | PC/w4 | PG/w8 | PG/w4 |
|
||
|--------|------|--------|-------| ---|--------|------|-------------|------------------|
|
||
| 1 | 0.25ms | 0.26ms | 0.26ms | 0.20ms | 1.03x | 0.78x | 0.98x | 0.75x |
|
||
|128 | 0.39ms | 0.41ms | 0.55ms | 0.46ms | 1.43x | 1.19x | 1.35x | 1.12x |
|
||
| 1024 | 1.31ms | 1.65ms | 2.35ms | 2.14ms | 1.80x | 1.64x | 1.43x | 1.30x |
|
||
| 4096 | 5.37ms | 6.79ms | 8.54ms | 8.04ms | 1.59x | 1.50x | 1.26x | 1.18x |
|
||
| 8192 | 10.97ms | 12.94ms | 17.28ms | 16.23ms | 1.58x | 1.48x | 1.34x | 1.25x |
|
||
|
||
### 端到端 VLM
|
||
|
||
**Qwen3-VL-2B**
|
||
|
||
| Prompt Tokens | FP16 Prefill (tok/s) | W8A16 Prefill (tok/s) | **W8A8 PC Prefill (tok/s)** | FP16 Decode (tok/s) | W8A16 Decode (tok/s) | **W8A8 PC Decode (tok/s)** |
|
||
|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
|
||
| 1334 | 3010 | 2065 | **3242** | 70 | 107 | **104** |
|
||
| 2393 | 2868 | 1847 | **2983** | 69 | 97 | **100** |
|
||
| 3455 | 2777 | 1741 | **2796** | 66 | 90 | **95** |
|
||
|
||
**Qwen3-VL-4B**
|
||
|
||
Prompt Tokens | FP16 Prefill (tok/s) | W8A16 Prefill (tok/s) | **W8A8 PC Prefill (tok/s)** | FP16 Decode (tok/s) | W8A16 Decode (tok/s) | **W8A8 PC Decode (tok/s)** |
|
||
|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
|
||
| 1334 | 1884 | 1786 | **2186** | 32 | **56** | 54 |
|
||
| 2393 | 1815 | 1700 | **2028** | 31 | **55** | 52 |
|
||
| 3455 | 1755 | 1603 | **1881** | 30 | **52** | 49 |
|
||
|
||
|
||
### LLM 量化：精度与速度对比
|
||
|
||
|
||
<table>
|
||
<thead>
|
||
<tr>
|
||
<th>Models</th>
|
||
<th>Quantization Configuration</th>
|
||
<th>wikitext2 PPL（↓）</th>
|
||
<th>Prefill Time (s)（↓）</th>
|
||
<th>Peak Memory (GB)（↓）</th>
|
||
</tr>
|
||
</thead>
|
||
<tbody>
|
||
<tr>
|
||
<td rowspan="5"><b>Qwen3-8B</b></td>
|
||
<td>FP16</td>
|
||
<td>9.726</td>
|
||
<td>179.9</td>
|
||
<td>18.93</td>
|
||
</tr>
|
||
<tr>
|
||
<td>W8A16 (mlx RTN)</td>
|
||
<td>9.707</td>
|
||
<td>221.3</td>
|
||
<td>12.07</td>
|
||
</tr>
|
||
<tr>
|
||
<td>W8A8 (per-channel)</td>
|
||
<td>9.756</td>
|
||
<td><b>123.5</b></td>
|
||
<td><b>11.32</b></td>
|
||
</tr>
|
||
<tr>
|
||
<td>W8A8 (per-group gs=64)</td>
|
||
<td>9.744</td>
|
||
<td>179.1</td>
|
||
<td>11.83</td>
|
||
</tr>
|
||
<tr>
|
||
<td>W8A8 (per-group gs=128)</td>
|
||
<td>9.727</td>
|
||
<td>165.8</td>
|
||
<td>11.61</td>
|
||
</tr>
|
||
</tbody>
|
||
<tr style="border-top: 1px solid #333;">
|
||
<td colspan="5" style="padding: 0; height: 3px;"></td>
|
||
</tr>
|
||
<tbody>
|
||
<tr>
|
||
<td rowspan="5"><b>Llama3-8B</b></td>
|
||
<td>FP16</td>
|
||
<td>6.138</td>
|
||
<td>175.8</td>
|
||
<td>18.32</td>
|
||
</tr>
|
||
<tr>
|
||
<td>W8A16 (mlx RTN)</td>
|
||
<td>6.147</td>
|
||
<td>236.9</td>
|
||
<td>11.46</td>
|
||
</tr>
|
||
<tr>
|
||
<td>W8A8 (per-channel)</td>
|
||
<td>6.271</td>
|
||
<td><b>123.3</b></td>
|
||
<td><b>10.69</b></td>
|
||
</tr>
|
||
<tr>
|
||
<td>W8A8 (per-group, gs=64)</td>
|
||
<td>6.269</td>
|
||
<td>178.7</td>
|
||
<td>11.19</td>
|
||
</tr>
|
||
<tr>
|
||
<td>W8A8 (per-group, gs=128)</td>
|
||
<td>6.270</td>
|
||
<td>155.7</td>
|
||
<td>10.98</td>
|
||
</tr>
|
||
</tbody>
|
||
</table>
|
||
|
||
## 说明
|
||
需要注意的是，cider 假定权重是 <span style="color: red;">**量化友好（quantization-friendly）**</span> 的。这意味着你需要确保模型已经经过某种训练后量化校准方法，例如 GTPQ、SmoothQuant 或 Quarot 等，以处理异常值（outliers），或者模型本身是通过 QAT（量化感知训练）训练的。当 cider 转换模型（例如使用 w8a8）并计算量化 scale 时，默认采用最简单的 **min-max** 方法，同时也支持第 99 百分位方法（clip_percentile=99.9）。如果遇到乱码输出，很大概率是模型受到了 <span style="color: red;">**异常值（outliers）**</span> 的影响。因此，某些预处理步骤是 **必要的**。
|
||
|
||
|
||
## 环境要求
|
||
|
||
|
||
- **Apple M5+** 用于 INT8 TensorOps（M4 及以下：以纯 Python 方式安装，`is_available()` 返回 False）
|
||
- Python 3.12+
|
||
- MLX >= 0.31
|
||
- nanobind >= 2.12（仅在 M5+ 上进行 C++ 编译时需要）
|
||
- CMake >= 3.27（仅在 M5+ 上进行 C++ 编译时需要）
|
||
|
||
## 安装
|
||
|
||
```bash
|
||
pip install -e .
|
||
```
|
||
|
||
在 M5+ 上，这会运行 CMake 编译 C++ 扩展，然后安装 Python 包。
|
||
在 M4 及以下设备上，仅安装 Python 包（不会出现编译错误）。
|
||
|
||
## 快速开始
|
||
|
||
### 一行命令完成模型转换（推荐）
|
||
|
||
```python
|
||
from cider import convert_model, is_available
|
||
|
||
model, proc = load("path/to/model")
|
||
|
||
if is_available():
|
||
convert_model(model)
|
||
# CiderLinear auto-detects:
|
||
# seq_len > 1 → W8A8 INT8 TensorOps (faster prefill)
|
||
# seq_len == 1 → INT8 MV kernel (near-native decode speed)
|
||
else:
|
||
pass # Falls back to standard MLX inference on M4
|
||
```
|
||
|
||
**重要提示**
|
||
在量化视觉语言模型（Vision-Language Models，VLMs）时，视觉 Transformer（ViT）通常不会被替换。直接使用 convert_model 也会量化视觉模型的线性层，这通常会导致精度下降。对于 VLMs，我们建议调用 convert_model(model.language_model)，仅对语言模型应用 GPTQ、SmoothQuant 和 AWQ 等现有量化方法。
|
||
|
||
已在选定的 MLX Transformer 模型上测试，包括 Qwen3、Qwen3-VL 和 Llama3 系列。其他架构可能需要适配。
|
||
|
||
|
||
### 层级别 API
|
||
|
||
```python
|
||
import numpy as np
|
||
import mlx.core as mx
|
||
from cider import W8A8Linear, W4A8Linear, is_available
|
||
|
||
assert is_available(), "Requires Apple M5+"
|
||
|
||
# Prepare weight
|
||
W = np.random.randn(4096, 4096).astype(np.float16)
|
||
|
||
# W8A8 linear (per-channel)
|
||
from cider.ops import quantize_weight_int8
|
||
w_int8, scale = quantize_weight_int8(W)
|
||
layer = W8A8Linear(
|
||
w_int8=mx.array(w_int8), scale_w=mx.array(scale),
|
||
group_size=0, in_features=4096, out_features=4096
|
||
)
|
||
x = mx.random.normal((32, 4096)).astype(mx.float16)
|
||
y = layer(x) # lazy — builds MLX graph
|
||
mx.eval(y) # GPU executes
|
||
|
||
# W4A8 linear (half the weight memory)
|
||
layer4 = W4A8Linear.from_weights(W)
|
||
y4 = layer4(x)
|
||
mx.eval(y4)
|
||
```
|
||
|
||
## 底层 API
|
||
|
||
```python
|
||
from cider import perchannel_linear, w4a8_linear, quantize_weight_int8, pack_weight_int4
|
||
|
||
# Quantize weights (numpy, offline)
|
||
w_int8, scale = quantize_weight_int8(W_np)
|
||
packed_w4, scale4 = pack_weight_int4(W_np)
|
||
|
||
# Primitive calls (return lazy mx.array)
|
||
y = perchannel_linear(x, mx.array(w_int8), mx.array(scale))
|
||
y4 = w4a8_linear(x, mx.array(packed_w4), mx.array(scale4))
|
||
```
|
||
|
||
## 项目结构
|
||
|
||
```
|
||
cider/
|
||
├── cider/ # Python package
|
||
│ ├── __init__.py # Public API (conditional on is_available)
|
||
│ ├── ops.py # Primitive wrappers + quantize helpers
|
||
│ ├── nn.py # CiderLinear, W4A8Linear (nn.Module)
|
||
│ ├── convert.py # convert_model() high-level API
|
||
│ ├── attention/ # Optimized SDPA
|
||
│ │ ├── __init__.py
|
||
│ │ └── sdpa.py # patch_sdpa / unpatch_sdpa / autotune_sdpa
|
||
│ └── kernels/ # Metal shaders (bundled)
|
||
│ ├── w8a8_matmul.metal # W8A8 GEMM (prefill, M>1)
|
||
│ ├── w8a8_int8_mv.metal # W8A8 per-channel MV (decode, M=1)
|
||
│ ├── w8a8_quantize.metal # Per-token activation quantization
|
||
│ ├── w4a8_matmul.metal # W4A8 GEMM (prefill)
|
||
│ ├── pergroup_int8_gemm.metal # Per-group GEMM (prefill)
|
||
│ ├── pergroup_int8_mv.metal # Per-group MV (decode)
|
||
│ ├── cider_sdpa_vector.h # SDPA kernel header (1-pass & 2-pass)
|
||
│ └── cider_sdpa_vector.metal # SDPA kernel template instantiations
|
||
├── csrc/ # C++ MLX primitives (nanobind, M5+ only)
|
||
│ ├── include/
|
||
│ │ ├── w8a8_primitive.h
|
||
│ │ ├── w4a8_primitive.h
|
||
│ │ ├── pergroup_primitive.h
|
||
│ │ └── sdpa_primitive.h
|
||
│ └── src/
|
||
│ ├── w8a8_primitive.mm
|
||
│ ├── w4a8_primitive.mm
|
||
│ ├── pergroup_primitive.mm
|
||
│ ├── sdpa_primitive.mm # SDPA 1-pass & 2-pass primitives
|
||
│ └── prim_bindings.cpp
|
||
├── benchmarks/
|
||
│ ├── bench_e2e_wxa16.py # End-to-end VLM benchmark (W8A8 + SDPA)
|
||
│ ├── bench_sdpa.py # SDPA micro-benchmark (correctness + speed vs MLX)
|
||
│ ├── bench_full.py # Isolated kernel latency (per-channel/per-group vs MLX)
|
||
│ ├── test_bitexact.py # Numerical correctness verification
|
||
│ └── mlx_native/ # MLX native INT8 comparison
|
||
├── tutorial/
|
||
│ ├── how_to_write_efficient_int_gemm_m5_en.md
|
||
│ └── how_to_write_efficient_int_gemm_m5_zh.md
|
||
├── tools/
|
||
│ ├── eval_ppl_all.py # Unified PPL eval (FP16/W8A16/per-channel/per-group)
|
||
│ ├── convert_compressed_tensors_to_mlx.py
|
||
│ └── smoothquant.py # SmoothQuant calibration
|
||
├── examples/
|
||
│ └── basic_usage.py
|
||
├── vlm_service/ # OpenAI-style VLM inference server
|
||
│ ├── server.py # FastAPI server (streaming + non-streaming)
|
||
│ ├── core_infer.py # HMInference engine (singleton)
|
||
│ ├── custom_qwen3vl.py # Custom Qwen3-VL generation loop
|
||
│ ├── config.py # Config loader
|
||
│ ├── bench_client.py # Server benchmark client
|
||
│ └── client.py # API client example
|
||
├── config/
|
||
│ └── config.yaml # Server & model configuration
|
||
├── experimental/ # ANE+GPU hybrid tensor parallelism (M4)
|
||
│ ├── split_linear.py # SplitLinear + ANEBridge + patch_model()
|
||
│ ├── bench.py # End-to-end benchmark
|
||
│ ├── libane_bridge_v6.m # ANE private API bridge (Obj-C source)
|
||
│ └── README.md
|
||
├── CMakeLists.txt
|
||
├── pyproject.toml
|
||
├── setup.py # Conditional build (M5+: full, M4: pure-Python)
|
||
└── README.md
|
||
```
|
||
|
||
## VLM 推理服务
|
||
|
||
`vlm_service/` 提供了一个开箱即用的 **OpenAI 风格** VLM 推理服务器，支持 W8A8 加速。
|
||
|
||
### 快速开始
|
||
|
||
1. **配置** `config/config.yaml`：
|
||
|
||
```yaml
|
||
model_name_or_path: /path/to/your/model # MLX VLM model (e.g., Qwen3-VL-2B W8A16)
|
||
sampling:
|
||
max_new_tokens: 1024
|
||
temperature: 1.0
|
||
top_p: 1.0
|
||
server:
|
||
host: 0.0.0.0
|
||
port: 8341
|
||
ttl: 1800
|
||
w8a8:
|
||
mode: 'off' # 'auto' | 'on' | 'off'
|
||
```
|
||
|
||
- `auto`：若硬件支持则启用 W8A8，否则回退到默认设置
|
||
- `on`：强制启用 W8A8（不支持时报错）。"选择 'on' 时，表示你的模型需要执行在线激活量化（online activation quantization）。在这种情况下，Cider 本身**不**保证量化精度，你需要自行应用某些量化算法，例如 SmoothQuant、QuaRot、GPTQ 或 QAT，以确保激活量化后精度不会明显下降。此选项只是让你在模型采用 W8A8 时能够利用硬件的计算优势，而不仅仅是模拟量化。"
|
||
- `off`：禁用 W8A8，使用标准 MLX 推理
|
||
|
||
2. **启动服务器**：
|
||
|
||
```bash
|
||
cd vlm_service
|
||
python server.py --config ../config/config.yaml
|
||
```
|
||
|
||
3. **发送请求**（OpenAI 风格 API）：
|
||
|
||
```bash
|
||
# Text-only
|
||
curl http://localhost:8341/v1/chat/completions \
|
||
-H "Content-Type: application/json" \
|
||
-d '{
|
||
"model": "vlm",
|
||
"messages": [{"role": "user", "content": "Hello"}],
|
||
"stream": false
|
||
}'
|
||
|
||
# With image (base64)
|
||
curl http://localhost:8341/v1/chat/completions \
|
||
-H "Content-Type: application/json" \
|
||
-d '{
|
||
"model": "vlm",
|
||
"messages": [{"role": "user", "content": [
|
||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,..."} },
|
||
{"type": "text", "text": "What is in this image?"}
|
||
]}],
|
||
"stream": true
|
||
}'
|
||
```
|
||
|
||
### API 端点
|
||
|
||
| Endpoint | Method | Description |
|
||
|----------|--------|-------------|
|
||
| `/v1/chat/completions` | POST | 对话补全（流式 / 非流式） |
|
||
| `/v1/models` | GET | 列出可用模型 |
|
||
| `/health` | GET | 健康检查 |
|
||
| `/v1/queue` | GET | 请求队列状态 |
|
||
|
||
### 服务中 W8A8 的工作方式
|
||
|
||
当 `w8a8.mode` 为 `auto` 或 `on` 时，服务器在启动时会调用 `cider.convert_model()`，将所有 Linear 层替换为 `CiderLinear`。推理过程中：
|
||
|
||
- **Prefill**（处理输入 token，seq_len > 1）：使用 W8A8 INT8 GEMM 以加速计算
|
||
- **Decode**（逐 token 生成，seq_len == 1）：使用 INT8 MV kernel（接近原生速度）
|
||
|
||
无需修改代码——切换会根据输入序列长度自动进行。
|
||
|
||
## 架构
|
||
|
||
### MLX 自定义原语
|
||
|
||
W8A8Linear 与 W4A8Linear 均实现为 `mlx::core::Primitive` 的子类。这意味着：
|
||
|
||
1. **惰性求值（lazy evaluation）**：`y = layer(x)` 构建图节点，而非立即计算
|
||
2. **图组合**：多个原语调用组合为单个 MLX 图
|
||
3. **流调度**：MLX 的调度器处理 GPU 派发顺序
|
||
|
||
### Metal Kernel 流水线
|
||
|
||
每个原语都会派发 Metal 计算 kernel：
|
||
|
||
**Prefill（M > 1）：**
|
||
1. **quantize_per_token**：FP16 激活 → INT8 + 逐 token 缩放
|
||
2. **matmul_fused_dequant**：INT8 × INT8 → INT32 → FP16（融合缩放反量化）
|
||
|
||
**Decode（M = 1）：**
|
||
- **int8_mv**：直接 INT8 矩阵-向量乘积，权重反量化即时进行（无需激活量化）
|
||
|
||
对于 W4A8，GEMM 步骤在 fragment load 中包含内联 INT4→INT8 解包。
|
||
|
||
### TensorOps matmul2d
|
||
|
||
INT8 GEMM 使用 Apple 的 `mpp::tensor_ops::matmul2d(16, 32, 16)`——通过 Metal 4 的 `cooperative_tensor` API，在 M5+ 上可用的硬件加速 INT8×INT8→INT32 矩阵乘法。这与 MLX 的 NAX 模板可用的硬件指令相同。Cider 的 kernel 在存储阶段加入融合反量化（INT32 × scales → FP16），避免额外的设备内存往返。详见 [kernel comparison](benchmarks/mlx_native/cider_vs_mlx_int8.md)。
|
||
|
||
### Tile 配置
|
||
|
||
| Config | BM | BN | BK | SK | Threads | Use When |
|
||
|--------|----|----|----|----|---------|----------|
|
||
| Large | 128 | 128 | 512 | 32 | 512 | M > 64 |
|
||
| Small | 32 | 128 | 512 | 32 | 128 | M ≤ 64 |
|
||
|
||
根据 M 自动选择。包含 L2 cache swizzle 派发。
|
||
|
||
## ANE+GPU 异构张量并行（实验性）
|
||
|
||
我们发现，在 Mac 上进行推理时，仅利用了 GPU 与 CPU 两个硬件计算单元，而 Mac 上的 ANE（Apple Neural Engine）计算单元处于空闲状态。我们将此视为潜在的优化机会。受 [maderix/ANE](https://github.com/maderix/ANE), 启发，我们开展了 ANE+GPU 混合推理模式的实验工作。目前我们将该方案应用于张量并行计算。在 M4 芯片上，在仅同步的前向推理期间（MLX 原生使用称为惰性求值的技术以降低同步开销；在端到端测试中，混合推理目前尚未显示出优势，主要因为我们尚未基于 MLX 的惰性求值实现该方案——这仍是后续工作），与纯 GPU 推理在同步流水线（synchronize pipeline）下相比，我们观察到约 **3%~16%** 的性能提升。这仍属探索性工作，端到端收益目前受限于缺少与惰性求值兼容的实现。
|
||
|
||
|
||
在 LLM prefill 阶段，GPU 的矩阵单元已满载——但 **Apple Neural Engine 完全空闲**。ANE Split 利用这一点，沿输出通道拆分每个线性层的 GEMM：
|
||
|
||
- **ANE** 计算约 65% 的输出通道（FP32，通过逆向工程的私有 `_ANEClient` API）
|
||
- **GPU** 计算剩余约 35%（FP16，标准 MLX matmul）
|
||
- 两者**并发**运行，结果拼接
|
||
|
||
这是一种**异构张量并行**——不是数据并行，也不是流水线并行——利用同一 SoC 上的两个不同计算单元。
|
||
|
||
### 性能（Apple M4，Qwen3-VL-2B Prefill）
|
||
|
||
| seq | W8A16 GPU | SplitLinear | Speedup vs W8A16 |
|
||
|-----|----------|-----------|-------------|
|
||
| 512 | 639.9 ms | **615.9 ms** | **1.039×** |
|
||
| 1024 | 1348.6 ms | **1156.9 ms** | **1.17×** |
|
||
|
||
在测试的基准案例中，余弦相似度接近 1.0，top-1 token 一致率为 100%。
|
||
|
||
|
||
### 关键设计选择
|
||
|
||
- **仅 Prefill**：Decode 回退到原始 GPU linear（零开销）
|
||
- **共享输入准备**：Q/K/V 与 Gate/Up 投影通过 `_InputGroup` 共享单次输入转置+numpy 拷贝
|
||
- **自动路由**：Down 投影（IC > 2×OC）在 ANE 效率低下时保持仅 GPU
|
||
- **短序列绕过**：序列 < 192 token 时跳过分拆（开销大于收益）
|
||
|
||
完整文档、用法、构建说明与限制见 [`experimental/README.md`](experimental/README.md)。
|
||
|
||
> **注意：** ANE Split 已在 M4 上测试。M5 引入了 ANE 架构变更，可能破坏私有 API 桥接——尚未在 M5 上验证。
|
||
|
||
## 量化
|
||
|
||
| Component | Scheme | Granularity |
|
||
|-----------|--------|-------------|
|
||
| W8A8 weights | Symmetric INT8 | Per-channel or per-group (gs=64/128) |
|
||
| W4A8 weights | Symmetric INT4 (zp=8) | Per-column |
|
||
| Activations | Symmetric INT8 | Per-token |
|
||
| Accumulation | INT32 | — |
|
||
| Output dequant | `C_fp16 = C_int32 * s_act * s_weight` | Per-element |
|
||
|
||
## 限制
|
||
|
||
- **M=1 单算子**：逐通道 MV kernel 在独立 decode 调用时慢于 MLX W4A16。逐组 MV kernel 在端到端基准中与 MLX W8A16 decode 速度相差在 5% 以内。
|
||
- **INT8 TensorOps 仅支持 Apple M5+**：M4 及以下可安装但 `is_available()` 返回 False。
|
||
- **W4A8 慢于 W8A8**：INT4→INT8 解包 ALU 开销（Metal 4 matmul2d 无原生 INT4 操作数）。
|
||
|
||
## 工具
|
||
|
||
### 统一 PPL 评估
|
||
|
||
```bash
|
||
# Run all 5 configurations in one script
|
||
python tools/eval_ppl_all.py --num-samples 50
|
||
|
||
# Evaluates: FP16, W8A16 (MLX native), W8A8 per-channel, per-group(gs=64), per-group(gs=128)
|
||
# Outputs comparison table at the end
|
||
```
|
||
|
||
## 路线图
|
||
|
||
- [x] 一行式模型转换 API（`convert_model`，自动 prefill/decode）
|
||
- [x] 自动 dtype 处理（float16 / bfloat16）
|
||
- [x] 逐通道与逐组 W8A8 量化
|
||
- [x] 专用 decode MV kernel（匹配原生 MLX 速度）
|
||
- [x] 条件编译（M4 优雅回退）
|
||
- [x] mlx_vlm 与 mlx_lm 集成示例
|
||
- [x] 优化的 SDPA decode kernel（受 FlashInfer 启发，支持 GQA）
|
||
- [x] 一行式 `patch_sdpa()`，透明加速 decode
|
||
- [x] SDPA AutoTune（按架构块大小扫描）
|
||
- [ ] ANE 原语惰性求值
|
||
- [ ] 集成剪枝功能
|
||
- [ ] KVCache 量化
|
||
|
||
## 作者
|
||
|
||
Multimodal Team, Mininglamp Technology
|
||
|
||
如有 bug 报告、功能请求或使用问题，请在本仓库提交 issue。
|
||
|
||
|
||
|
||
## 引用
|
||
|
||
若您认为本工作有价值，请引用：
|
||
|
||
```bibtex
|
||
@software{wang2026cider,
|
||
author = {Multimodal Team, Mininglamp Technology},
|
||
title = {Cider: Exploiting Unused INT8 TensorOps for Faster LLM Prefill on Apple Silicon},
|
||
year = {2026},
|
||
howpublished = {https://github.com/Mininglamp-AI/cider}
|
||
}
|
||
```
|
||
|
||
## 许可证
|
||
|
||
MIT
|
||
|
||
## 致谢
|
||
|
||
- [MLX](https://github.com/ml-explore/mlx) by Apple — primitive API, NAXFrag kernel architecture
|
||
- [FlashInfer](https://github.com/flashinfer-ai/flashinfer) — inspired contiguous chunk GQA scheduling and register tiling for SDPA decode kernel
|
||
- Metal 4 MetalPerformancePrimitives for INT8 TensorOps
|
||
- [maderix/ANE](https://github.com/maderix/ANE) — inspired and informed our ANE+GPU tensor-parallel implementation
|