wehub-resource-sync
95 lines
3.5 KiB
Markdown
95 lines
3.5 KiB
Markdown
# ANE Split — ANE+GPU Tensor Parallelism for Prefill Acceleration
|
||
*Status: experimental research prototype. Uses private ANE APIs and is not part of Cider's stable inference path.*
|
||
|
||
Split each linear layer's GEMM along output channels: **ANE** and **GPU** computes running concurrently. This explores whether the otherwise idle Apple Neural Engine can help accelerate prefill, with minimal degradation observed in our current tests.
|
||
|
||
|
||
> **Platform:** Apple M4 (tested). M5 ANE API changes may cause failures — not yet validated.
|
||
|
||
## How It Works
|
||
|
||
```
|
||
Input x ──┬── ANE (r output channels, FP32, private API) ──┐
|
||
│ ├── concat → output
|
||
└── GPU (1 - r output channels, FP16, MLX matmul) ──┘
|
||
```
|
||
|
||
1. **SplitLinear** wraps each `nn.Linear` / `nn.QuantizedLinear`
|
||
2. In **prefill mode** (seq ≥ 192): split path with ANE+GPU concurrency
|
||
3. In **decode mode**: falls back to original `nn.Linear` on GPU (zero overhead)
|
||
4. Same-input projections (Q/K/V, Gate/Up) share input preparation via `_InputGroup`
|
||
|
||
### Automatic Layer Routing
|
||
|
||
| Layer Type | Routing | Reason |
|
||
|------------|---------|--------|
|
||
| Q, K, V, O projections | ANE+GPU split | Expand: IC → OC, ANE efficient |
|
||
| Gate, Up projections | ANE+GPU split | Expand: IC → OC, ANE efficient |
|
||
| Down projection | GPU only | Narrow: IC > 2×OC, ANE inefficient |
|
||
| Short sequences (< 192) | GPU only | Split overhead > benefit |
|
||
|
||
## Performance (Apple M4, Qwen3-VL-2B)
|
||
|
||
| 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×** |
|
||
|
||
- Accuracy: cos ≈ 1.0, top-1 match = 100%
|
||
- 168 layers split (28 layers × 6 projections), 28 GPU-only (down_proj)
|
||
|
||
|
||
## Quick Start
|
||
|
||
```python
|
||
from split_linear import patch_model, SplitLinear
|
||
|
||
# Load any MLX VLM model
|
||
from mlx_vlm.utils import load as vlm_load
|
||
model, processor = vlm_load("path/to/model")
|
||
|
||
# Patch all linear layers (one-liner)
|
||
bridge = patch_model(model, seq=512)
|
||
|
||
# Enable split for prefill, disable for decode
|
||
SplitLinear.set_prefill(True) # prefill: ANE+GPU parallel
|
||
# ... run prefill ...
|
||
SplitLinear.set_prefill(False) # decode: original GPU path
|
||
```
|
||
|
||
## Benchmark
|
||
|
||
```bash
|
||
# Default: seq=512
|
||
python3 bench.py
|
||
|
||
# Custom seq length
|
||
python3 bench.py 1024
|
||
|
||
# Custom model path
|
||
MODEL_PATH=/path/to/model python3 bench.py 512
|
||
```
|
||
|
||
## Building the ANE Bridge
|
||
|
||
To build libane_bridge_v6.dylib from the source:
|
||
|
||
```bash
|
||
clang -shared -O2 -framework Foundation -framework CoreML \
|
||
-framework Accelerate -o libane_bridge_v6.dylib libane_bridge_v6.m
|
||
```
|
||
|
||
> Requires macOS with ANE private frameworks. Uses undocumented `_ANEClient` API.
|
||
|
||
## Limitations
|
||
|
||
- **M4 only** — M5 ANE internal changes may break the private API bridge
|
||
- **Fixed sequence length** — ANE models are compiled for a specific seq; re-patch needed for different lengths
|
||
- **FP32 on ANE** — ANE operates in FP32 (no INT8/FP16 GEMM); benefit comes from parallelism, not precision
|
||
- **Memory overhead** — ANE models consume additional system memory (~200MB for 2B model)
|
||
- **No decode benefit** — Decode is single-token, falls back to GPU (no split overhead)
|
||
- **No E2E benefit** - MLX natively employs lazy evaluation to reduce synchronization overhead. In end-to-end testing, our hybrid approach currently shows no advantage because we haven't yet implemented it with MLX's lazy evaluation.
|
||
## License
|
||
|
||
MIT
|