chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,94 @@
|
||||
# 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
|
||||
@@ -0,0 +1,108 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
bench_w8a16_splitlinear.py — W8A16 vs SplitLinear prefill benchmark.
|
||||
|
||||
Usage:
|
||||
python3 bench_w8a16_splitlinear.py [seq_len] # default 512
|
||||
|
||||
Requires:
|
||||
- mlx_vlm with Qwen3-VL-2B-Instruct-8bit model
|
||||
- split_linear.py (V4) + libane_bridge_v6.dylib in same directory
|
||||
"""
|
||||
import sys, time, os
|
||||
import numpy as np
|
||||
import mlx.core as mx
|
||||
from mlx_vlm.utils import load as vlm_load
|
||||
from mlx_vlm.models.cache import KVCache
|
||||
|
||||
# ─── Config ───
|
||||
SEQ = int(sys.argv[1]) if len(sys.argv) > 1 else 512
|
||||
W8A16_MODEL = '~/Downloads/sft_baseline_v2_w8a16'
|
||||
N_WARMUP = 3
|
||||
N_BENCH = 8
|
||||
|
||||
def bench_forward(lang, ids, pos, n_layers, reps):
|
||||
ts = []
|
||||
for _ in range(reps):
|
||||
cache = [KVCache() for _ in range(n_layers)]
|
||||
t0 = time.perf_counter()
|
||||
mx.eval(lang(ids, cache=cache, position_ids=pos).logits)
|
||||
ts.append((time.perf_counter() - t0) * 1000)
|
||||
return ts
|
||||
|
||||
def main():
|
||||
print(f"\n{'='*60}")
|
||||
print(f" W8A16 vs SplitLinear Prefill Benchmark (seq={SEQ})")
|
||||
print(f"{'='*60}")
|
||||
print(f" Model: {W8A16_MODEL}")
|
||||
print(f" Warmup: {N_WARMUP}, Bench: {N_BENCH}\n")
|
||||
|
||||
# Load model
|
||||
print("[1/4] Loading W8A16 model...")
|
||||
model, _ = vlm_load(W8A16_MODEL)
|
||||
lang = model.language_model
|
||||
N = lang.args.num_hidden_layers
|
||||
print(f" {N} layers loaded\n")
|
||||
|
||||
# Prepare inputs
|
||||
ids = mx.ones((1, SEQ), dtype=mx.int32)
|
||||
pos = mx.broadcast_to(
|
||||
mx.arange(SEQ).reshape(1, SEQ)[None, :, :],
|
||||
(3, 1, SEQ)
|
||||
)
|
||||
mx.eval(ids, pos)
|
||||
|
||||
# ─── W8A16 Baseline ───
|
||||
print("[2/4] W8A16 GPU baseline")
|
||||
bench_forward(lang, ids, pos, N, N_WARMUP) # warmup
|
||||
ts_w8 = bench_forward(lang, ids, pos, N, N_BENCH)
|
||||
for i, t in enumerate(ts_w8):
|
||||
print(f" Run {i+1}: {t:.1f}ms")
|
||||
med_w8 = float(np.median(ts_w8))
|
||||
print(f" Median: {med_w8:.1f}ms\n")
|
||||
|
||||
# ─── Reference logits ───
|
||||
print("[3/4] Computing reference logits for accuracy check...")
|
||||
c_ref = [KVCache() for _ in range(N)]
|
||||
ref = np.array(lang(ids, cache=c_ref, position_ids=pos).logits.astype(mx.float32))
|
||||
|
||||
# ─── SplitLinear ───
|
||||
print("[4/4] Patch with SplitLinear + benchmark")
|
||||
from split_linear import patch_model, SplitLinear
|
||||
bridge = patch_model(model, SEQ)
|
||||
SplitLinear.set_prefill(True)
|
||||
|
||||
# Warmup
|
||||
bench_forward(lang, ids, pos, N, N_WARMUP)
|
||||
|
||||
# Accuracy
|
||||
c_hyb = [KVCache() for _ in range(N)]
|
||||
hyb = np.array(lang(ids, cache=c_hyb, position_ids=pos).logits.astype(mx.float32))
|
||||
cos = float(np.dot(ref.flatten(), hyb.flatten()) /
|
||||
(np.linalg.norm(ref.flatten()) * np.linalg.norm(hyb.flatten()) + 1e-12))
|
||||
top1 = float((ref.argmax(-1) == hyb.argmax(-1)).mean() * 100)
|
||||
print(f" Accuracy: cos={cos:.6f}, top1={top1:.1f}%")
|
||||
|
||||
# Benchmark
|
||||
ts_sp = bench_forward(lang, ids, pos, N, N_BENCH)
|
||||
for i, t in enumerate(ts_sp):
|
||||
print(f" Run {i+1}: {t:.1f}ms")
|
||||
med_sp = float(np.median(ts_sp))
|
||||
|
||||
# ─── Summary ───
|
||||
speedup = med_w8 / med_sp
|
||||
delta = med_sp - med_w8
|
||||
print(f"\n{'='*60}")
|
||||
print(f" W8A16 GPU: {med_w8:.1f}ms")
|
||||
print(f" SplitLinear: {med_sp:.1f}ms ({speedup:.3f}x)")
|
||||
print(f" Delta: {delta:+.1f}ms")
|
||||
print(f" Accuracy: cos={cos:.6f} top1={top1:.1f}%")
|
||||
print(f"{'='*60}")
|
||||
|
||||
if speedup >= 1.0:
|
||||
print(f" ✅ SplitLinear faster by {(speedup-1)*100:.1f}%")
|
||||
else:
|
||||
print(f" ⚠️ SplitLinear slower by {(1-speedup)*100:.1f}%")
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,675 @@
|
||||
// libane_bridge_v6.m — ANE bridge with manual retain (no ARC), IOSurface
|
||||
// isolation
|
||||
// Compile: clang -shared -o libane_bridge_v6.dylib libane_bridge_v6.m \
|
||||
// -framework Foundation -framework IOSurface -framework Accelerate -lobjc -O2
|
||||
// -fno-objc-arc
|
||||
#import <Accelerate/Accelerate.h>
|
||||
#import <Foundation/Foundation.h>
|
||||
#import <IOSurface/IOSurface.h>
|
||||
#import <dispatch/dispatch.h>
|
||||
#import <dlfcn.h>
|
||||
#import <objc/message.h>
|
||||
#import <objc/runtime.h>
|
||||
|
||||
static Class g_D, g_I, g_AR, g_AIO;
|
||||
static int g_init = 0;
|
||||
static dispatch_queue_t g_ane_queue;
|
||||
static dispatch_group_t g_ane_group;
|
||||
|
||||
typedef struct {
|
||||
id model; // retained manually
|
||||
int ic, oc, seq;
|
||||
int ioInIdx, ioOutIdx;
|
||||
} ANEModel;
|
||||
|
||||
#define MAX_MODELS 256
|
||||
#define MAX_SURFACES 32
|
||||
static ANEModel g_models[MAX_MODELS];
|
||||
static id g_requests[MAX_MODELS]; // retained manually
|
||||
static int g_count = 0;
|
||||
|
||||
typedef struct {
|
||||
IOSurfaceRef surface;
|
||||
id aneObj; // retained manually
|
||||
size_t allocSize;
|
||||
} IOSEntry;
|
||||
static IOSEntry g_surfaces[MAX_SURFACES];
|
||||
static int g_surfCount = 0;
|
||||
|
||||
static int find_or_create_surface(size_t needed, int exclude_idx) {
|
||||
needed = (needed + 0xFFFF) & ~(size_t)0xFFFF;
|
||||
if (needed < 65536) {
|
||||
needed = 65536;
|
||||
}
|
||||
for (int i = 0; i < g_surfCount; i++) {
|
||||
if (g_surfaces[i].allocSize == needed && i != exclude_idx) {
|
||||
return i;
|
||||
}
|
||||
}
|
||||
if (g_surfCount >= MAX_SURFACES) {
|
||||
return -1;
|
||||
}
|
||||
IOSurfaceRef s = IOSurfaceCreate((__bridge CFDictionaryRef) @{
|
||||
(id)kIOSurfaceWidth : @(needed),
|
||||
(id)kIOSurfaceHeight : @1,
|
||||
(id)kIOSurfaceBytesPerElement : @1,
|
||||
(id)kIOSurfaceBytesPerRow : @(needed),
|
||||
(id)kIOSurfaceAllocSize : @(needed),
|
||||
(id)kIOSurfacePixelFormat : @0
|
||||
});
|
||||
if (!s) {
|
||||
return -1;
|
||||
}
|
||||
id obj = ((id (*)(Class, SEL, IOSurfaceRef))objc_msgSend)(
|
||||
g_AIO, @selector(objectWithIOSurface:), s);
|
||||
[obj retain];
|
||||
int idx = g_surfCount++;
|
||||
g_surfaces[idx] = (IOSEntry){s, obj, needed};
|
||||
return idx;
|
||||
}
|
||||
|
||||
static NSData *build_blob(const float *w, int oc, int ic) {
|
||||
size_t n = (size_t)oc * ic, wsize = n * 2, total = 128 + wsize;
|
||||
uint8_t *buf = (uint8_t *)calloc(total, 1);
|
||||
buf[0] = 0x01;
|
||||
buf[4] = 0x02;
|
||||
uint8_t *chunk = buf + 64;
|
||||
chunk[0] = 0xEF;
|
||||
chunk[1] = 0xBE;
|
||||
chunk[2] = 0xAD;
|
||||
chunk[3] = 0xDE;
|
||||
chunk[4] = 0x01;
|
||||
*(uint32_t *)(chunk + 8) = (uint32_t)wsize;
|
||||
*(uint32_t *)(chunk + 16) = 128;
|
||||
_Float16 *fp16 = (_Float16 *)(buf + 128);
|
||||
for (size_t i = 0; i < n; i++) {
|
||||
fp16[i] = (_Float16)w[i];
|
||||
}
|
||||
return [NSData dataWithBytesNoCopy:buf length:total freeWhenDone:YES];
|
||||
}
|
||||
|
||||
static NSString *gen_conv_mil(int ic, int oc, int seq) {
|
||||
return [NSString
|
||||
stringWithFormat:
|
||||
@"program(1.3)\n[buildInfo = dict<string, "
|
||||
@"string>({{\"coremlc-component-MIL\", \"3510.2.1\"}, "
|
||||
"{\"coremlc-version\", \"3505.4.1\"}, "
|
||||
"{\"coremltools-component-milinternal\", \"\"}, "
|
||||
"{\"coremltools-version\", \"9.0\"}})]\n{\n"
|
||||
" func main<ios18>(tensor<fp32, [1, %d, 1, %d]> x) {\n"
|
||||
" string to_fp16 = const()[name = string(\"to_fp16\"), val = "
|
||||
"string(\"fp16\")];\n"
|
||||
" tensor<fp16, [1, %d, 1, %d]> x16 = cast(dtype = to_fp16, x "
|
||||
"= x)[name = string(\"cast_in\")];\n"
|
||||
" tensor<fp16, [%d, %d, 1, 1]> W = const()[name = "
|
||||
"string(\"W\"), val = tensor<fp16, [%d, %d, 1, 1]>"
|
||||
"(BLOBFILE(path = string(\"@model_path/weights/weight.bin\"), "
|
||||
"offset = uint64(64)))];\n"
|
||||
" string c_pad_type = const()[name = string(\"c_pad_type\"), "
|
||||
"val = string(\"valid\")];\n"
|
||||
" tensor<int32, [2]> c_strides = const()[name = "
|
||||
"string(\"c_strides\"), val = tensor<int32, [2]>([1, 1])];\n"
|
||||
" tensor<int32, [4]> c_pad = const()[name = "
|
||||
"string(\"c_pad\"), val = tensor<int32, [4]>([0, 0, 0, 0])];\n"
|
||||
" tensor<int32, [2]> c_dilations = const()[name = "
|
||||
"string(\"c_dilations\"), val = tensor<int32, [2]>([1, 1])];\n"
|
||||
" int32 c_groups = const()[name = string(\"c_groups\"), val "
|
||||
"= int32(1)];\n"
|
||||
" tensor<fp16, [1, %d, 1, %d]> y16 = conv(dilations = "
|
||||
"c_dilations, groups = c_groups, pad = c_pad, "
|
||||
"pad_type = c_pad_type, strides = c_strides, weight = W, x = "
|
||||
"x16)[name = string(\"conv\")];\n"
|
||||
" string to_fp32 = const()[name = string(\"to_fp32\"), val = "
|
||||
"string(\"fp32\")];\n"
|
||||
" tensor<fp32, [1, %d, 1, %d]> y = cast(dtype = to_fp32, x = "
|
||||
"y16)[name = string(\"cast_out\")];\n"
|
||||
" } -> (y);\n}\n",
|
||||
ic, seq, ic, seq, oc, ic, oc, ic, oc, seq, oc, seq];
|
||||
}
|
||||
|
||||
// INT8 MIL text: uses constexpr_blockwise_shift_scale for INT8 weight
|
||||
// quantization
|
||||
static NSString *gen_conv_mil_int8(int ic, int oc, int seq) {
|
||||
return [NSString
|
||||
stringWithFormat:
|
||||
@"program(1.3)\n[buildInfo = dict<string, "
|
||||
@"string>({{\"coremlc-component-MIL\", \"3510.2.1\"}, "
|
||||
"{\"coremlc-version\", \"3505.4.1\"}, "
|
||||
"{\"coremltools-component-milinternal\", \"\"}, "
|
||||
"{\"coremltools-version\", \"9.0\"}})]\
|
||||
{\n"
|
||||
" func main<ios18>(tensor<fp32, [1, %d, 1, %d]> x) {\n"
|
||||
" string to_fp16 = const()[name = string(\"to_fp16\"), val = "
|
||||
"string(\"fp16\")];\n"
|
||||
" tensor<fp16, [1, %d, 1, %d]> x16 = cast(dtype = to_fp16, x "
|
||||
"= x)[name = string(\"cast_in\")];\n"
|
||||
" tensor<int8, [%d, %d, 1, 1]> W_data = const()[name = "
|
||||
"string(\"W_data\"), "
|
||||
"val = tensor<int8, [%d, %d, 1, 1]>(BLOBFILE(path = "
|
||||
"string(\"@model_path/weights/weight_data.bin\"), offset = "
|
||||
"uint64(64)))];\n"
|
||||
" tensor<fp16, [%d, 1, 1, 1]> W_scale = const()[name = "
|
||||
"string(\"W_scale\"), "
|
||||
"val = tensor<fp16, [%d, 1, 1, 1]>(BLOBFILE(path = "
|
||||
"string(\"@model_path/weights/weight_scale.bin\"), offset = "
|
||||
"uint64(64)))];\n"
|
||||
" tensor<fp16, [%d, 1, 1, 1]> W_offset = const()[name = "
|
||||
"string(\"W_offset\"), "
|
||||
"val = tensor<fp16, [%d, 1, 1, 1]>(BLOBFILE(path = "
|
||||
"string(\"@model_path/weights/weight_offset.bin\"), offset = "
|
||||
"uint64(64)))];\n"
|
||||
" tensor<fp16, [%d, %d, 1, 1]> W = "
|
||||
"constexpr_blockwise_shift_scale("
|
||||
"data = W_data, scale = W_scale, offset = W_offset)[name = "
|
||||
"string(\"dequant\")];\n"
|
||||
" string c_pad_type = const()[name = string(\"c_pad_type\"), "
|
||||
"val = string(\"valid\")];\n"
|
||||
" tensor<int32, [2]> c_strides = const()[name = "
|
||||
"string(\"c_strides\"), val = tensor<int32, [2]>([1, 1])];\n"
|
||||
" tensor<int32, [4]> c_pad = const()[name = "
|
||||
"string(\"c_pad\"), val = tensor<int32, [4]>([0, 0, 0, 0])];\n"
|
||||
" tensor<int32, [2]> c_dilations = const()[name = "
|
||||
"string(\"c_dilations\"), val = tensor<int32, [2]>([1, 1])];\n"
|
||||
" int32 c_groups = const()[name = string(\"c_groups\"), val "
|
||||
"= int32(1)];\n"
|
||||
" tensor<fp16, [1, %d, 1, %d]> y16 = conv(dilations = "
|
||||
"c_dilations, groups = c_groups, pad = c_pad, "
|
||||
"pad_type = c_pad_type, strides = c_strides, weight = W, x = "
|
||||
"x16)[name = string(\"conv\")];\n"
|
||||
" string to_fp32 = const()[name = string(\"to_fp32\"), val = "
|
||||
"string(\"fp32\")];\n"
|
||||
" tensor<fp32, [1, %d, 1, %d]> y = cast(dtype = to_fp32, x = "
|
||||
"y16)[name = string(\"cast_out\")];\n"
|
||||
" } -> (y);\n}\n",
|
||||
ic, seq, ic, seq, oc, ic, oc, ic, // W_data shape
|
||||
oc, oc, // W_scale shape
|
||||
oc, oc, // W_offset shape
|
||||
oc, ic, // W dequantized shape
|
||||
oc, seq, // y16 shape
|
||||
oc, seq]; // y shape
|
||||
}
|
||||
|
||||
// Build INT8 weight blob: header(64) + magic(64) + data
|
||||
static NSData *build_blob_int8(const int8_t *data, int count) {
|
||||
size_t wsize = (size_t)count, total = 128 + wsize;
|
||||
uint8_t *buf = (uint8_t *)calloc(total, 1);
|
||||
buf[0] = 0x01;
|
||||
buf[4] = 0x02;
|
||||
uint8_t *chunk = buf + 64;
|
||||
chunk[0] = 0xEF;
|
||||
chunk[1] = 0xBE;
|
||||
chunk[2] = 0xAD;
|
||||
chunk[3] = 0xDE;
|
||||
chunk[4] = 0x01;
|
||||
*(uint32_t *)(chunk + 8) = (uint32_t)wsize;
|
||||
*(uint32_t *)(chunk + 16) = 128;
|
||||
memcpy(buf + 128, data, wsize);
|
||||
return [NSData dataWithBytesNoCopy:buf length:total freeWhenDone:YES];
|
||||
}
|
||||
|
||||
static NSData *build_blob_fp16(const _Float16 *data, int count) {
|
||||
size_t wsize = (size_t)count * 2, total = 128 + wsize;
|
||||
uint8_t *buf = (uint8_t *)calloc(total, 1);
|
||||
buf[0] = 0x01;
|
||||
buf[4] = 0x02;
|
||||
uint8_t *chunk = buf + 64;
|
||||
chunk[0] = 0xEF;
|
||||
chunk[1] = 0xBE;
|
||||
chunk[2] = 0xAD;
|
||||
chunk[3] = 0xDE;
|
||||
chunk[4] = 0x01;
|
||||
*(uint32_t *)(chunk + 8) = (uint32_t)wsize;
|
||||
*(uint32_t *)(chunk + 16) = 128;
|
||||
memcpy(buf + 128, data, wsize);
|
||||
return [NSData dataWithBytesNoCopy:buf length:total freeWhenDone:YES];
|
||||
}
|
||||
|
||||
static int setup_model(id mdl, int ic, int oc, int seq) {
|
||||
size_t inB = (size_t)ic * seq * 4;
|
||||
size_t outB = (size_t)oc * seq * 4;
|
||||
int ioInIdx = find_or_create_surface(inB, -1);
|
||||
int ioOutIdx = find_or_create_surface(outB, ioInIdx); // exclude input surface
|
||||
if (ioInIdx < 0 || ioOutIdx < 0) {
|
||||
return -1;
|
||||
}
|
||||
id wI = g_surfaces[ioInIdx].aneObj;
|
||||
id wO = g_surfaces[ioOutIdx].aneObj;
|
||||
id req = ((id (*)(Class, SEL, id, id, id, id, id, id, id))objc_msgSend)(
|
||||
g_AR,
|
||||
@selector(requestWithInputs:inputIndices:outputs:outputIndices:
|
||||
weightsBuffer:perfStats:procedureIndex:),
|
||||
@[ wI ], @[ @0 ], @[ wO ], @[ @0 ], nil, nil, @0);
|
||||
[req retain];
|
||||
[mdl retain];
|
||||
|
||||
int handle = g_count;
|
||||
g_models[handle] = (ANEModel){mdl, ic, oc, seq, ioInIdx, ioOutIdx};
|
||||
g_requests[handle] = req;
|
||||
g_count++;
|
||||
return handle;
|
||||
}
|
||||
|
||||
// ===== Public C API =====
|
||||
|
||||
int ane_init(void) {
|
||||
if (g_init) {
|
||||
return 0;
|
||||
}
|
||||
dlopen("/System/Library/PrivateFrameworks/AppleNeuralEngine.framework/"
|
||||
"AppleNeuralEngine",
|
||||
RTLD_NOW);
|
||||
g_D = NSClassFromString(@"_ANEInMemoryModelDescriptor");
|
||||
g_I = NSClassFromString(@"_ANEInMemoryModel");
|
||||
g_AR = NSClassFromString(@"_ANERequest");
|
||||
g_AIO = NSClassFromString(@"_ANEIOSurfaceObject");
|
||||
if (!g_D || !g_I || !g_AR || !g_AIO)
|
||||
return -1;
|
||||
g_ane_queue = dispatch_queue_create("ane.eval", DISPATCH_QUEUE_SERIAL);
|
||||
g_ane_group = dispatch_group_create();
|
||||
g_init = 1;
|
||||
return 0;
|
||||
}
|
||||
|
||||
int ane_load_model(int ic, int oc, int seq, const float *w) {
|
||||
if (!g_init || g_count >= MAX_MODELS) {
|
||||
return -1;
|
||||
}
|
||||
NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init];
|
||||
|
||||
NSData *mil =
|
||||
[gen_conv_mil(ic, oc, seq) dataUsingEncoding:NSUTF8StringEncoding];
|
||||
NSData *blob = build_blob(w, oc, ic);
|
||||
NSDictionary *wd =
|
||||
@{@"@model_path/weights/weight.bin" : @{@"offset" : @0, @"data" : blob}};
|
||||
NSError *e = nil;
|
||||
|
||||
id desc = ((id (*)(Class, SEL, id, id, id))objc_msgSend)(
|
||||
g_D, @selector(modelWithMILText:weights:optionsPlist:), mil, wd, nil);
|
||||
if (!desc) {
|
||||
[pool drain];
|
||||
return -1;
|
||||
}
|
||||
id mdl = ((id (*)(Class, SEL, id))objc_msgSend)(
|
||||
g_I, @selector(inMemoryModelWithDescriptor:), desc);
|
||||
if (!mdl) {
|
||||
[pool drain];
|
||||
return -1;
|
||||
}
|
||||
|
||||
id hx = ((id (*)(id, SEL))objc_msgSend)(mdl, @selector(hexStringIdentifier));
|
||||
NSString *td = [NSTemporaryDirectory() stringByAppendingPathComponent:hx];
|
||||
[[NSFileManager defaultManager]
|
||||
createDirectoryAtPath:[td stringByAppendingPathComponent:@"weights"]
|
||||
withIntermediateDirectories:YES
|
||||
attributes:nil
|
||||
error:nil];
|
||||
[mil writeToFile:[td stringByAppendingPathComponent:@"model.mil"]
|
||||
atomically:YES];
|
||||
[blob writeToFile:[td stringByAppendingPathComponent:@"weights/weight.bin"]
|
||||
atomically:YES];
|
||||
|
||||
BOOL ok;
|
||||
ok = ((BOOL (*)(id, SEL, unsigned int, id, NSError **))objc_msgSend)(
|
||||
mdl, @selector(compileWithQoS:options:error:), 21, @{}, &e);
|
||||
if (!ok) {
|
||||
NSLog(@"compile failed %d->%d: %@", ic, oc, e);
|
||||
[pool drain];
|
||||
return -1;
|
||||
}
|
||||
ok = ((BOOL (*)(id, SEL, unsigned int, id, NSError **))objc_msgSend)(
|
||||
mdl, @selector(loadWithQoS:options:error:), 21, @{}, &e);
|
||||
if (!ok) {
|
||||
NSLog(@"load failed %d->%d: %@", ic, oc, e);
|
||||
[pool drain];
|
||||
return -1;
|
||||
}
|
||||
|
||||
int result = setup_model(mdl, ic, oc, seq);
|
||||
[pool drain];
|
||||
return result;
|
||||
}
|
||||
|
||||
// Load INT8 quantized conv model onto ANE
|
||||
// w_int8: [oc, ic] int8 row-major, scale: [oc] fp32, offset: [oc] fp32
|
||||
int ane_load_model_int8(int ic, int oc, int seq, const int8_t *w_int8,
|
||||
const float *scale_fp32, const float *offset_fp32) {
|
||||
if (!g_init || g_count >= MAX_MODELS) {
|
||||
return -1;
|
||||
}
|
||||
NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init];
|
||||
|
||||
NSData *mil =
|
||||
[gen_conv_mil_int8(ic, oc, seq) dataUsingEncoding:NSUTF8StringEncoding];
|
||||
|
||||
// Build weight blobs: INT8 data, FP16 scale, FP16 offset
|
||||
NSData *blob_data = build_blob_int8(w_int8, oc * ic);
|
||||
|
||||
_Float16 *scale_fp16 = (_Float16 *)calloc(oc, sizeof(_Float16));
|
||||
_Float16 *offset_fp16 = (_Float16 *)calloc(oc, sizeof(_Float16));
|
||||
for (int i = 0; i < oc; i++) {
|
||||
scale_fp16[i] = (_Float16)scale_fp32[i];
|
||||
offset_fp16[i] = (_Float16)offset_fp32[i];
|
||||
}
|
||||
NSData *blob_scale = build_blob_fp16(scale_fp16, oc);
|
||||
NSData *blob_offset = build_blob_fp16(offset_fp16, oc);
|
||||
free(scale_fp16);
|
||||
free(offset_fp16);
|
||||
|
||||
NSDictionary *wd = @{
|
||||
@"@model_path/weights/weight_data.bin" :
|
||||
@{@"offset" : @0, @"data" : blob_data},
|
||||
@"@model_path/weights/weight_scale.bin" :
|
||||
@{@"offset" : @0, @"data" : blob_scale},
|
||||
@"@model_path/weights/weight_offset.bin" :
|
||||
@{@"offset" : @0, @"data" : blob_offset}
|
||||
};
|
||||
NSError *e = nil;
|
||||
|
||||
id desc = ((id (*)(Class, SEL, id, id, id))objc_msgSend)(
|
||||
g_D, @selector(modelWithMILText:weights:optionsPlist:), mil, wd, nil);
|
||||
if (!desc) {
|
||||
NSLog(@"INT8 desc failed %d->%d", ic, oc);
|
||||
[pool drain];
|
||||
return -1;
|
||||
}
|
||||
id mdl = ((id (*)(Class, SEL, id))objc_msgSend)(
|
||||
g_I, @selector(inMemoryModelWithDescriptor:), desc);
|
||||
if (!mdl) {
|
||||
NSLog(@"INT8 model failed %d->%d", ic, oc);
|
||||
[pool drain];
|
||||
return -1;
|
||||
}
|
||||
|
||||
id hx = ((id (*)(id, SEL))objc_msgSend)(mdl, @selector(hexStringIdentifier));
|
||||
NSString *td = [NSTemporaryDirectory() stringByAppendingPathComponent:hx];
|
||||
[[NSFileManager defaultManager]
|
||||
createDirectoryAtPath:[td stringByAppendingPathComponent:@"weights"]
|
||||
withIntermediateDirectories:YES
|
||||
attributes:nil
|
||||
error:nil];
|
||||
[mil writeToFile:[td stringByAppendingPathComponent:@"model.mil"]
|
||||
atomically:YES];
|
||||
[blob_data
|
||||
writeToFile:[td stringByAppendingPathComponent:@"weights/weight_data.bin"]
|
||||
atomically:YES];
|
||||
[blob_scale writeToFile:[td stringByAppendingPathComponent:
|
||||
@"weights/weight_scale.bin"]
|
||||
atomically:YES];
|
||||
[blob_offset writeToFile:[td stringByAppendingPathComponent:
|
||||
@"weights/weight_offset.bin"]
|
||||
atomically:YES];
|
||||
|
||||
BOOL ok;
|
||||
ok = ((BOOL (*)(id, SEL, unsigned int, id, NSError **))objc_msgSend)(
|
||||
mdl, @selector(compileWithQoS:options:error:), 21, @{}, &e);
|
||||
if (!ok) {
|
||||
NSLog(@"INT8 compile failed %d->%d: %@", ic, oc, e);
|
||||
[pool drain];
|
||||
return -1;
|
||||
}
|
||||
ok = ((BOOL (*)(id, SEL, unsigned int, id, NSError **))objc_msgSend)(
|
||||
mdl, @selector(loadWithQoS:options:error:), 21, @{}, &e);
|
||||
if (!ok) {
|
||||
NSLog(@"INT8 load failed %d->%d: %@", ic, oc, e);
|
||||
[pool drain];
|
||||
return -1;
|
||||
}
|
||||
|
||||
int result = setup_model(mdl, ic, oc, seq);
|
||||
[pool drain];
|
||||
return result;
|
||||
}
|
||||
|
||||
int ane_eval(int handle) {
|
||||
if (handle < 0 || handle >= g_count) {
|
||||
return -1;
|
||||
}
|
||||
NSError *e = nil;
|
||||
return ((BOOL (*)(id, SEL, unsigned int, id, id, NSError **))objc_msgSend)(
|
||||
g_models[handle].model,
|
||||
@selector(evaluateWithQoS:options:request:error:), 21, @{},
|
||||
g_requests[handle], &e)
|
||||
? 0
|
||||
: -2;
|
||||
}
|
||||
|
||||
int ane_run(int handle, const float *input, float *output) {
|
||||
if (handle < 0 || handle >= g_count) {
|
||||
return -1;
|
||||
}
|
||||
ANEModel *m = &g_models[handle];
|
||||
IOSurfaceRef ioI = g_surfaces[m->ioInIdx].surface;
|
||||
IOSurfaceRef ioO = g_surfaces[m->ioOutIdx].surface;
|
||||
IOSurfaceLock(ioI, 0, NULL);
|
||||
memcpy(IOSurfaceGetBaseAddress(ioI), input,
|
||||
(size_t)m->ic * m->seq * sizeof(float));
|
||||
IOSurfaceUnlock(ioI, 0, NULL);
|
||||
int rc = ane_eval(handle);
|
||||
if (rc != 0) {
|
||||
return rc;
|
||||
}
|
||||
IOSurfaceLock(ioO, kIOSurfaceLockReadOnly, NULL);
|
||||
memcpy(output, IOSurfaceGetBaseAddress(ioO),
|
||||
(size_t)m->oc * m->seq * sizeof(float));
|
||||
IOSurfaceUnlock(ioO, kIOSurfaceLockReadOnly, NULL);
|
||||
return 0;
|
||||
}
|
||||
|
||||
int ane_run_async(int handle, const float *input) {
|
||||
if (handle < 0 || handle >= g_count) {
|
||||
return -1;
|
||||
}
|
||||
ANEModel *m = &g_models[handle];
|
||||
IOSurfaceRef ioI = g_surfaces[m->ioInIdx].surface;
|
||||
IOSurfaceLock(ioI, 0, NULL);
|
||||
memcpy(IOSurfaceGetBaseAddress(ioI), input,
|
||||
(size_t)m->ic * m->seq * sizeof(float));
|
||||
IOSurfaceUnlock(ioI, 0, NULL);
|
||||
dispatch_group_async(g_ane_group, g_ane_queue, ^{
|
||||
NSError *e = nil;
|
||||
((BOOL (*)(id, SEL, unsigned int, id, id, NSError **))objc_msgSend)(
|
||||
g_models[handle].model,
|
||||
@selector(evaluateWithQoS:options:request:error:), 21, @{},
|
||||
g_requests[handle], &e);
|
||||
});
|
||||
return 0;
|
||||
}
|
||||
|
||||
int ane_wait_read(int handle, float *output) {
|
||||
if (handle < 0 || handle >= g_count) {
|
||||
return -1;
|
||||
}
|
||||
dispatch_group_wait(g_ane_group, DISPATCH_TIME_FOREVER);
|
||||
ANEModel *m = &g_models[handle];
|
||||
IOSurfaceRef ioO = g_surfaces[m->ioOutIdx].surface;
|
||||
IOSurfaceLock(ioO, kIOSurfaceLockReadOnly, NULL);
|
||||
memcpy(output, IOSurfaceGetBaseAddress(ioO),
|
||||
(size_t)m->oc * m->seq * sizeof(float));
|
||||
IOSurfaceUnlock(ioO, kIOSurfaceLockReadOnly, NULL);
|
||||
return 0;
|
||||
}
|
||||
|
||||
void ane_wait(void) { dispatch_group_wait(g_ane_group, DISPATCH_TIME_FOREVER); }
|
||||
int ane_model_count(void) { return g_count; }
|
||||
int ane_surface_count(void) { return g_surfCount; }
|
||||
|
||||
void *ane_get_input_ptr(int handle) {
|
||||
if (handle < 0 || handle >= g_count) {
|
||||
return NULL;
|
||||
}
|
||||
return IOSurfaceGetBaseAddress(g_surfaces[g_models[handle].ioInIdx].surface);
|
||||
}
|
||||
void *ane_get_output_ptr(int handle) {
|
||||
if (handle < 0 || handle >= g_count)
|
||||
return NULL;
|
||||
return IOSurfaceGetBaseAddress(g_surfaces[g_models[handle].ioOutIdx].surface);
|
||||
}
|
||||
int ane_get_ic(int h) { return (h >= 0 && h < g_count) ? g_models[h].ic : -1; }
|
||||
int ane_get_oc(int h) { return (h >= 0 && h < g_count) ? g_models[h].oc : -1; }
|
||||
int ane_get_seq(int h) {
|
||||
return (h >= 0 && h < g_count) ? g_models[h].seq : -1;
|
||||
}
|
||||
|
||||
// Write input directly into IOSurface (caller provides [IC, SEQ] C-contiguous
|
||||
// data)
|
||||
void ane_write_input(int handle, const float *data) {
|
||||
if (handle < 0 || handle >= g_count) {
|
||||
return;
|
||||
}
|
||||
ANEModel *m = &g_models[handle];
|
||||
IOSurfaceRef ioI = g_surfaces[m->ioInIdx].surface;
|
||||
IOSurfaceLock(ioI, 0, NULL);
|
||||
memcpy(IOSurfaceGetBaseAddress(ioI), data,
|
||||
(size_t)m->ic * m->seq * sizeof(float));
|
||||
IOSurfaceUnlock(ioI, 0, NULL);
|
||||
}
|
||||
|
||||
// Read output directly from IOSurface into [OC, SEQ] C-contiguous buffer
|
||||
void ane_read_output(int handle, float *data) {
|
||||
if (handle < 0 || handle >= g_count) {
|
||||
return;
|
||||
}
|
||||
ANEModel *m = &g_models[handle];
|
||||
IOSurfaceRef ioO = g_surfaces[m->ioOutIdx].surface;
|
||||
IOSurfaceLock(ioO, kIOSurfaceLockReadOnly, NULL);
|
||||
memcpy(data, IOSurfaceGetBaseAddress(ioO),
|
||||
(size_t)m->oc * m->seq * sizeof(float));
|
||||
IOSurfaceUnlock(ioO, kIOSurfaceLockReadOnly, NULL);
|
||||
}
|
||||
|
||||
// Tile transpose: src[rows][cols] → dst[cols][rows], cache-friendly 16×16
|
||||
// blocks
|
||||
static void tile_transpose(const float *__restrict__ src,
|
||||
float *__restrict__ dst, int rows, int cols) {
|
||||
const int T = 16;
|
||||
for (int r = 0; r < rows; r += T) {
|
||||
int rend = r + T < rows ? r + T : rows;
|
||||
for (int c = 0; c < cols; c += T) {
|
||||
int cend = c + T < cols ? c + T : cols;
|
||||
for (int ri = r; ri < rend; ri++)
|
||||
for (int ci = c; ci < cend; ci++)
|
||||
dst[ci * rows + ri] = src[ri * cols + ci];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ane_run_T: input [SEQ, IC] row-major, output [SEQ, OC] row-major
|
||||
// Tile-transpose input to [IC, SEQ] for IOSurface, eval, transpose output back
|
||||
int ane_run_T(int handle, const float *input_seq_ic, float *output_seq_oc) {
|
||||
if (handle < 0 || handle >= g_count) {
|
||||
return -1;
|
||||
}
|
||||
ANEModel *m = &g_models[handle];
|
||||
IOSurfaceRef ioI = g_surfaces[m->ioInIdx].surface;
|
||||
IOSurfaceRef ioO = g_surfaces[m->ioOutIdx].surface;
|
||||
IOSurfaceLock(ioI, 0, NULL);
|
||||
tile_transpose(input_seq_ic, (float *)IOSurfaceGetBaseAddress(ioI), m->seq,
|
||||
m->ic);
|
||||
IOSurfaceUnlock(ioI, 0, NULL);
|
||||
NSError *e = nil;
|
||||
BOOL ok = ((BOOL (*)(id, SEL, unsigned int, id, id, NSError **))objc_msgSend)(
|
||||
m->model, @selector(evaluateWithQoS:options:request:error:), 21, @{},
|
||||
g_requests[handle], &e);
|
||||
if (!ok) {
|
||||
return -2;
|
||||
}
|
||||
IOSurfaceLock(ioO, kIOSurfaceLockReadOnly, NULL);
|
||||
tile_transpose((const float *)IOSurfaceGetBaseAddress(ioO), output_seq_oc,
|
||||
m->oc, m->seq);
|
||||
IOSurfaceUnlock(ioO, kIOSurfaceLockReadOnly, NULL);
|
||||
return 0;
|
||||
}
|
||||
|
||||
// ane_eval_direct: write input, eval, read output — single call, minimal
|
||||
// overhead
|
||||
int ane_eval_direct(int handle, const float *input, float *output) {
|
||||
if (handle < 0 || handle >= g_count) {
|
||||
return -1;
|
||||
}
|
||||
ANEModel *m = &g_models[handle];
|
||||
IOSurfaceRef ioI = g_surfaces[m->ioInIdx].surface;
|
||||
IOSurfaceRef ioO = g_surfaces[m->ioOutIdx].surface;
|
||||
// Write input
|
||||
IOSurfaceLock(ioI, 0, NULL);
|
||||
memcpy(IOSurfaceGetBaseAddress(ioI), input,
|
||||
(size_t)m->ic * m->seq * sizeof(float));
|
||||
IOSurfaceUnlock(ioI, 0, NULL);
|
||||
// Eval
|
||||
NSError *e = nil;
|
||||
BOOL ok = ((BOOL (*)(id, SEL, unsigned int, id, id, NSError **))objc_msgSend)(
|
||||
m->model, @selector(evaluateWithQoS:options:request:error:), 21, @{},
|
||||
g_requests[handle], &e);
|
||||
if (!ok) {
|
||||
return -2;
|
||||
}
|
||||
// Read output
|
||||
IOSurfaceLock(ioO, kIOSurfaceLockReadOnly, NULL);
|
||||
memcpy(output, IOSurfaceGetBaseAddress(ioO),
|
||||
(size_t)m->oc * m->seq * sizeof(float));
|
||||
IOSurfaceUnlock(ioO, kIOSurfaceLockReadOnly, NULL);
|
||||
return 0;
|
||||
}
|
||||
|
||||
// ane_run_rowmajor: input [L, IC] row-major, output [L, OC] row-major
|
||||
// Uses vDSP_mtrans for fast transpose (Apple Accelerate)
|
||||
// When L==seq: pure vDSP path (fastest). L<seq: vDSP + scatter.
|
||||
int ane_run_rowmajor(int handle, const float *input_rm, int L,
|
||||
float *output_rm) {
|
||||
if (handle < 0 || handle >= g_count) {
|
||||
return -1;
|
||||
}
|
||||
ANEModel *m = &g_models[handle];
|
||||
if (L > m->seq) {
|
||||
return -3;
|
||||
}
|
||||
IOSurfaceRef ioI = g_surfaces[m->ioInIdx].surface;
|
||||
IOSurfaceRef ioO = g_surfaces[m->ioOutIdx].surface;
|
||||
|
||||
// === Transpose input [L, IC] -> IOSurface [IC, seq] ===
|
||||
IOSurfaceLock(ioI, 0, NULL);
|
||||
float *ioBuf = (float *)IOSurfaceGetBaseAddress(ioI);
|
||||
if (L == m->seq) {
|
||||
// Fast path: vDSP_mtrans [L][IC] -> [IC][L=seq]
|
||||
vDSP_mtrans(input_rm, 1, ioBuf, 1, m->ic, L);
|
||||
} else {
|
||||
// L < seq: transpose to temp, then scatter with zero-pad
|
||||
// For each channel ch: ioBuf[ch*seq+0..L-1] = transposed, [L..seq-1] = 0
|
||||
// Use vDSP_mtrans to a contiguous temp, then scatter
|
||||
float *tmp = (float *)malloc((size_t)m->ic * L * sizeof(float));
|
||||
vDSP_mtrans(input_rm, 1, tmp, 1, m->ic, L);
|
||||
for (int ch = 0; ch < m->ic; ch++) {
|
||||
memcpy(&ioBuf[ch * m->seq], &tmp[ch * L], L * sizeof(float));
|
||||
if (L < m->seq) {
|
||||
memset(&ioBuf[ch * m->seq + L], 0, (m->seq - L) * sizeof(float));
|
||||
}
|
||||
}
|
||||
free(tmp);
|
||||
}
|
||||
IOSurfaceUnlock(ioI, 0, NULL);
|
||||
|
||||
// === ANE eval ===
|
||||
NSError *e = nil;
|
||||
BOOL ok = ((BOOL (*)(id, SEL, unsigned int, id, id, NSError **))objc_msgSend)(
|
||||
m->model, @selector(evaluateWithQoS:options:request:error:), 21, @{},
|
||||
g_requests[handle], &e);
|
||||
if (!ok) {
|
||||
return -2;
|
||||
}
|
||||
// === Transpose output IOSurface [OC, seq] -> [L, OC] ===
|
||||
IOSurfaceLock(ioO, kIOSurfaceLockReadOnly, NULL);
|
||||
const float *outBuf = (const float *)IOSurfaceGetBaseAddress(ioO);
|
||||
if (L == m->seq) {
|
||||
// Fast path: vDSP_mtrans [OC][seq] -> [seq][OC]
|
||||
vDSP_mtrans(outBuf, 1, output_rm, 1, L, m->oc);
|
||||
} else {
|
||||
// Gather first L elements from each of OC rows, then transpose
|
||||
float *tmp = (float *)malloc((size_t)m->oc * L * sizeof(float));
|
||||
for (int ch = 0; ch < m->oc; ch++) {
|
||||
memcpy(&tmp[ch * L], &outBuf[ch * m->seq], L * sizeof(float));
|
||||
}
|
||||
vDSP_mtrans(tmp, 1, output_rm, 1, L, m->oc);
|
||||
free(tmp);
|
||||
}
|
||||
IOSurfaceUnlock(ioO, kIOSurfaceLockReadOnly, NULL);
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,467 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
SplitLinear: Drop-in replacement for nn.Linear / nn.QuantizedLinear.
|
||||
Splits GEMM along output channels — ANE does ~65%, GPU does ~35%, concurrent.
|
||||
|
||||
Key optimization: same-input projections (Q/K/V, Gate/Up) share a single
|
||||
input preparation via _InputGroup, eliminating redundant transpose+eval+numpy.
|
||||
|
||||
Usage:
|
||||
from split_linear import patch_model
|
||||
bridge = patch_model(model, seq=512)
|
||||
|
||||
API:
|
||||
patch_model(model, seq) → high-level one-liner
|
||||
SplitLinear(layer, bridge, seq) → single layer replacement
|
||||
ANEBridge() → ANE private API wrapper
|
||||
"""
|
||||
import os, ctypes
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import numpy as np
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
LIB_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
SPLIT_ALIGN = 64
|
||||
MIN_SEQ_FOR_SPLIT = 192 # Below this, split overhead > benefit
|
||||
|
||||
|
||||
# ─── ANE Bridge ───
|
||||
|
||||
class ANEBridge:
|
||||
"""Thin wrapper around ANE private API."""
|
||||
_instance = None
|
||||
|
||||
def __init__(self):
|
||||
lib = ctypes.CDLL(os.path.join(LIB_DIR, 'libane_bridge_v6.dylib'))
|
||||
lib.ane_init.restype = ctypes.c_int
|
||||
lib.ane_load_model.restype = ctypes.c_int
|
||||
lib.ane_load_model.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_int,
|
||||
ctypes.POINTER(ctypes.c_float)]
|
||||
lib.ane_run.restype = ctypes.c_int
|
||||
lib.ane_run.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_float),
|
||||
ctypes.POINTER(ctypes.c_float)]
|
||||
lib.ane_run_rowmajor.restype = ctypes.c_int
|
||||
lib.ane_run_rowmajor.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_float),
|
||||
ctypes.c_int, ctypes.POINTER(ctypes.c_float)]
|
||||
lib.ane_model_count.restype = ctypes.c_int
|
||||
lib.ane_surface_count.restype = ctypes.c_int
|
||||
assert lib.ane_init() == 0, "ANE init failed"
|
||||
self.lib = lib
|
||||
|
||||
def load(self, ic, oc, seq, w_fp32):
|
||||
w = np.ascontiguousarray(w_fp32, dtype=np.float32)
|
||||
h = self.lib.ane_load_model(ic, oc, seq,
|
||||
w.ctypes.data_as(ctypes.POINTER(ctypes.c_float)))
|
||||
assert h >= 0, f"ANE load failed: {ic}→{oc}, seq={seq}"
|
||||
return h
|
||||
|
||||
def run(self, h, inp, out):
|
||||
self.lib.ane_run(h,
|
||||
inp.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
|
||||
out.ctypes.data_as(ctypes.POINTER(ctypes.c_float)))
|
||||
|
||||
def run_rowmajor(self, h, inp_rm, L, out_rm):
|
||||
"""ANE compute with row-major I/O. Uses vDSP for transpose.
|
||||
inp_rm: [L, IC] row-major float32
|
||||
out_rm: [L, OC] row-major float32 (pre-allocated)
|
||||
"""
|
||||
self.lib.ane_run_rowmajor(h,
|
||||
inp_rm.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
|
||||
L,
|
||||
out_rm.ctypes.data_as(ctypes.POINTER(ctypes.c_float)))
|
||||
|
||||
@property
|
||||
def model_count(self):
|
||||
return self.lib.ane_model_count()
|
||||
|
||||
@classmethod
|
||||
def shared(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = cls()
|
||||
return cls._instance
|
||||
|
||||
|
||||
# Single ANE worker thread
|
||||
_ane_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix='ane')
|
||||
|
||||
|
||||
# ─── Helpers ───
|
||||
|
||||
def _extract_weight(layer):
|
||||
"""Extract FP32 weight [OC, IC] from nn.Linear or nn.QuantizedLinear."""
|
||||
if isinstance(layer, nn.QuantizedLinear):
|
||||
mx.eval(layer.weight, layer.scales)
|
||||
b = getattr(layer, 'biases', None)
|
||||
if b is not None:
|
||||
mx.eval(b)
|
||||
w = mx.dequantize(layer.weight, layer.scales, b,
|
||||
layer.group_size, layer.bits)
|
||||
else:
|
||||
w = layer.weight
|
||||
mx.eval(w)
|
||||
return np.array(w, dtype=np.float32)
|
||||
|
||||
|
||||
class _InputGroup:
|
||||
"""
|
||||
Shared input preparation for same-input projections (Q/K/V or Gate/Up).
|
||||
|
||||
Zero-copy path: eval x as FP32 row-major, share numpy view across group.
|
||||
No GPU transpose needed — ane_run_rowmajor handles transpose in C via vDSP.
|
||||
"""
|
||||
__slots__ = ('_np', '_x2d', '_L', '_seq')
|
||||
|
||||
def __init__(self, ic=0, seq=0):
|
||||
self._np = None
|
||||
self._x2d = None
|
||||
self._L = 0
|
||||
self._seq = seq
|
||||
|
||||
def prepare(self, x):
|
||||
"""Prepare input. Called by first proj in group.
|
||||
Returns (inp_np_rowmajor [L, IC], x_2d, L)."""
|
||||
x_2d = x.reshape(-1, x.shape[-1]) if x.ndim == 3 else x
|
||||
L = x_2d.shape[0]
|
||||
|
||||
# Cast to FP32 contiguous row-major, eval, then zero-copy numpy view
|
||||
x_f32 = x_2d.astype(mx.float32) if x_2d.dtype != mx.float32 else x_2d
|
||||
x_f32 = mx.contiguous(x_f32)
|
||||
mx.eval(x_f32)
|
||||
self._np = np.array(x_f32, copy=False) # [L, IC] row-major, zero-copy
|
||||
|
||||
self._L = L
|
||||
self._x2d = x_2d
|
||||
return self._np, x_2d, L
|
||||
|
||||
def get(self):
|
||||
"""Get cached input. Called by subsequent projs.
|
||||
Returns (inp_np_rowmajor [L, IC], x_2d, L)."""
|
||||
return self._np, self._x2d, self._L
|
||||
|
||||
|
||||
# ─── SplitLinear ───
|
||||
|
||||
class SplitLinear:
|
||||
"""
|
||||
Drop-in nn.Linear replacement with ANE+GPU tensor parallelism.
|
||||
|
||||
Two modes controlled externally via set_prefill(True/False):
|
||||
- prefill: ANE ~65% + GPU ~35% concurrent
|
||||
- decode: original nn.Linear on GPU → zero overhead
|
||||
"""
|
||||
_prefill_mode = False # class-level flag
|
||||
|
||||
def __init__(self, layer, bridge, seq, ane_frac=None, name="",
|
||||
input_group=None, is_first=True):
|
||||
self.name = name
|
||||
self._orig = layer # always keep original for decode
|
||||
w_np = _extract_weight(layer)
|
||||
oc, ic = w_np.shape
|
||||
self.ic = ic
|
||||
self.oc = oc
|
||||
self.seq = seq
|
||||
self._is_first = is_first
|
||||
|
||||
# Auto-detect split fraction
|
||||
if ane_frac is None:
|
||||
if seq < MIN_SEQ_FOR_SPLIT:
|
||||
ane_frac = 0.0
|
||||
elif ic > oc * 2:
|
||||
ane_frac = 0.0 # Wide→narrow: ANE inefficient (down_proj)
|
||||
elif oc < SPLIT_ALIGN * 2:
|
||||
ane_frac = 0.0
|
||||
else:
|
||||
ane_frac = 0.65
|
||||
|
||||
self.ane_frac = ane_frac
|
||||
|
||||
if ane_frac <= 0:
|
||||
self.mode = 'gpu'
|
||||
return
|
||||
|
||||
self.mode = 'split'
|
||||
self.ane = bridge
|
||||
|
||||
self.ane_oc = (int(oc * ane_frac) // SPLIT_ALIGN) * SPLIT_ALIGN
|
||||
self.gpu_oc = oc - self.ane_oc
|
||||
|
||||
if self.ane_oc < SPLIT_ALIGN or self.gpu_oc < 1:
|
||||
self.mode = 'gpu'
|
||||
return
|
||||
|
||||
self.h_ane = bridge.load(ic, self.ane_oc, seq, w_np[:self.ane_oc, :])
|
||||
self.buf_ane = np.empty((seq, self.ane_oc), dtype=np.float32) # ANE FP32 output
|
||||
self.buf_ane_f16 = np.empty((seq, self.ane_oc), dtype=np.float16) # FP16 cast buffer
|
||||
|
||||
# GPU weight for split path
|
||||
self._is_quantized = isinstance(layer, nn.QuantizedLinear)
|
||||
if self._is_quantized:
|
||||
# Re-quantize GPU portion as QuantizedLinear → native quantized_matmul
|
||||
w_gpu_fp32 = mx.array(w_np[self.ane_oc:, :]) # [gpu_oc, ic] float32
|
||||
w_q, scales, biases = mx.quantize(w_gpu_fp32,
|
||||
group_size=layer.group_size,
|
||||
bits=layer.bits)
|
||||
self._gpu_layer = nn.QuantizedLinear(
|
||||
ic, self.gpu_oc, bias=False,
|
||||
group_size=layer.group_size, bits=layer.bits)
|
||||
self._gpu_layer.weight = w_q
|
||||
self._gpu_layer.scales = scales
|
||||
self._gpu_layer.biases = biases
|
||||
mx.eval(self._gpu_layer.parameters())
|
||||
self._w_gpu = None # not used for quantized path
|
||||
else:
|
||||
self._gpu_layer = None
|
||||
self._w_gpu = None
|
||||
|
||||
if input_group is not None:
|
||||
self._grp = input_group
|
||||
else:
|
||||
self._grp = _InputGroup(ic, seq)
|
||||
self._is_first = True
|
||||
|
||||
@classmethod
|
||||
def set_prefill(cls, enabled):
|
||||
cls._prefill_mode = enabled
|
||||
|
||||
def __call__(self, x):
|
||||
# decode / gpu-only / short seq: use original nn.Linear
|
||||
if self.mode == 'gpu' or not SplitLinear._prefill_mode:
|
||||
return self._orig(x)
|
||||
|
||||
L = x.shape[-2] if x.ndim == 3 else x.shape[0]
|
||||
if L < MIN_SEQ_FOR_SPLIT:
|
||||
return self._orig(x)
|
||||
|
||||
# ── prefill split path (zero-copy + concurrent ANE/GPU) ──
|
||||
orig_shape = x.shape
|
||||
|
||||
# Input preparation (shared or fresh) — row-major [L, IC] FP32
|
||||
if self._is_first:
|
||||
inp_np, x_2d, L = self._grp.prepare(x)
|
||||
else:
|
||||
inp_np, x_2d, L = self._grp.get()
|
||||
|
||||
# ANE: launch in worker thread (row-major path, vDSP transpose in C)
|
||||
out_buf = self.buf_ane[:L] # [L, ane_oc] view
|
||||
fut = _ane_pool.submit(self.ane.run_rowmajor, self.h_ane, inp_np, L, out_buf)
|
||||
|
||||
# GPU: matmul with GPU portion weight
|
||||
if self._gpu_layer is not None:
|
||||
# Quantized: use native quantized_matmul via QuantizedLinear
|
||||
gpu_out = self._gpu_layer(x_2d)
|
||||
elif self._w_gpu is not None:
|
||||
gpu_out = x_2d @ self._w_gpu.T
|
||||
else:
|
||||
w_gpu = self._orig.weight[self.ane_oc:, :]
|
||||
gpu_out = x_2d @ w_gpu.T
|
||||
mx.eval(gpu_out) # sync GPU — enables concurrent ANE execution
|
||||
|
||||
# Wait for ANE
|
||||
fut.result()
|
||||
|
||||
# Merge: numpy FP32→FP16 cast + lazy concat
|
||||
np.copyto(self.buf_ane_f16[:L], out_buf, casting='same_kind')
|
||||
ane_out = mx.array(self.buf_ane_f16[:L]) # [L, ane_oc] FP16
|
||||
merged = mx.concatenate([ane_out, gpu_out], axis=-1)
|
||||
|
||||
if len(orig_shape) == 3:
|
||||
merged = merged.reshape(orig_shape[0], orig_shape[1], -1)
|
||||
return merged
|
||||
|
||||
def __repr__(self):
|
||||
if self.mode == 'gpu':
|
||||
return f"SplitLinear({self.name}, gpu_only, {self.ic}→{self.oc})"
|
||||
return (f"SplitLinear({self.name}, {self.ane_oc}ane+{self.gpu_oc}gpu, "
|
||||
f"{self.ic}→{self.oc}, frac={self.ane_frac:.0%})")
|
||||
|
||||
|
||||
# ─── Tree Walk ───
|
||||
|
||||
def _find_linears(module, prefix=""):
|
||||
"""Walk MLX model tree, yield (parent, key, full_name, linear)."""
|
||||
if not hasattr(module, 'children'):
|
||||
return
|
||||
for attr_name, child in module.children().items():
|
||||
full_name = f"{prefix}.{attr_name}" if prefix else attr_name
|
||||
if isinstance(child, (nn.Linear, nn.QuantizedLinear)):
|
||||
yield (module, attr_name, full_name, child)
|
||||
elif isinstance(child, nn.Module):
|
||||
yield from _find_linears(child, full_name)
|
||||
elif isinstance(child, list):
|
||||
for i, v in enumerate(child):
|
||||
fname = f"{full_name}.{i}"
|
||||
if isinstance(v, (nn.Linear, nn.QuantizedLinear)):
|
||||
yield (child, i, fname, v)
|
||||
elif hasattr(v, 'children'):
|
||||
yield from _find_linears(v, fname)
|
||||
elif isinstance(child, dict):
|
||||
for k, v in child.items():
|
||||
fname = f"{full_name}.{k}"
|
||||
if isinstance(v, (nn.Linear, nn.QuantizedLinear)):
|
||||
yield (module, k, fname, v)
|
||||
elif hasattr(v, 'children'):
|
||||
yield from _find_linears(v, fname)
|
||||
elif hasattr(child, 'children'):
|
||||
yield from _find_linears(child, full_name)
|
||||
|
||||
|
||||
def _get_input_dims(layer):
|
||||
"""Get true input dimensions for nn.Linear or nn.QuantizedLinear."""
|
||||
if isinstance(layer, nn.QuantizedLinear):
|
||||
return layer.weight.shape[-1] * 32 // layer.bits
|
||||
return layer.weight.shape[-1]
|
||||
|
||||
|
||||
def patch_model(model, seq, bridge=None, verbose=True):
|
||||
"""
|
||||
Patch all linear layers with SplitLinear.
|
||||
|
||||
Same-input projections share an InputGroup:
|
||||
- Q/K/V share one (same hidden_state input)
|
||||
- Gate/Up share one (same hidden_state after attn)
|
||||
- O and Down each get their own (unique inputs)
|
||||
|
||||
Returns: bridge instance.
|
||||
"""
|
||||
if bridge is None:
|
||||
bridge = ANEBridge.shared()
|
||||
|
||||
lang = model.language_model
|
||||
lm = lang.model
|
||||
N = len(lm.layers)
|
||||
n_split = n_gpu = 0
|
||||
|
||||
for li in range(N):
|
||||
la = lm.layers[li]
|
||||
attn = la.self_attn
|
||||
mlp = la.mlp
|
||||
|
||||
ic_attn = _get_input_dims(attn.q_proj)
|
||||
ic_mlp = _get_input_dims(mlp.gate_proj)
|
||||
|
||||
# Q/K/V share input group
|
||||
qkv_grp = _InputGroup(ic_attn, seq)
|
||||
for i, name in enumerate(('q_proj', 'k_proj', 'v_proj')):
|
||||
orig = getattr(attn, name)
|
||||
sl = SplitLinear(orig, bridge, seq,
|
||||
name=f"layer.{li}.{name}",
|
||||
input_group=qkv_grp,
|
||||
is_first=(i == 0))
|
||||
setattr(attn, name, sl)
|
||||
n_split += 1 if sl.mode == 'split' else 0
|
||||
n_gpu += 1 if sl.mode == 'gpu' else 0
|
||||
|
||||
# O proj — own input
|
||||
sl = SplitLinear(attn.o_proj, bridge, seq,
|
||||
name=f"layer.{li}.o_proj")
|
||||
attn.o_proj = sl
|
||||
n_split += 1 if sl.mode == 'split' else 0
|
||||
n_gpu += 1 if sl.mode == 'gpu' else 0
|
||||
|
||||
# Gate/Up share input group
|
||||
gu_grp = _InputGroup(ic_mlp, seq)
|
||||
for i, name in enumerate(('gate_proj', 'up_proj')):
|
||||
orig = getattr(mlp, name)
|
||||
sl = SplitLinear(orig, bridge, seq,
|
||||
name=f"layer.{li}.{name}",
|
||||
input_group=gu_grp,
|
||||
is_first=(i == 0))
|
||||
setattr(mlp, name, sl)
|
||||
n_split += 1 if sl.mode == 'split' else 0
|
||||
n_gpu += 1 if sl.mode == 'gpu' else 0
|
||||
|
||||
# Down proj — own input (auto GPU-only due to IC>OC*2)
|
||||
sl = SplitLinear(mlp.down_proj, bridge, seq,
|
||||
name=f"layer.{li}.down_proj")
|
||||
mlp.down_proj = sl
|
||||
n_split += 1 if sl.mode == 'split' else 0
|
||||
n_gpu += 1 if sl.mode == 'gpu' else 0
|
||||
|
||||
if verbose:
|
||||
print(f"[SplitLinear] {N} layers: {n_split} split, {n_gpu} gpu-only")
|
||||
print(f"[SplitLinear] ANE models: {bridge.model_count}")
|
||||
|
||||
return bridge
|
||||
|
||||
|
||||
# ─── Self-test ───
|
||||
if __name__ == '__main__':
|
||||
import sys, time
|
||||
from mlx_vlm.utils import load as vlm_load
|
||||
from mlx_vlm.models.cache import KVCache
|
||||
from mlx.utils import tree_flatten
|
||||
|
||||
SEQ = int(sys.argv[1]) if len(sys.argv) > 1 else 512
|
||||
N_W = 3; N_B = 10
|
||||
FP16_MODEL = '~/Downloads/weights/mlx/Qwen3-VL-2B-Instruct-16bit'
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f" SplitLinear Self-Test (seq={SEQ})")
|
||||
print(f"{'='*60}\n")
|
||||
|
||||
model, _ = vlm_load(FP16_MODEL)
|
||||
# Cast to true FP16
|
||||
flat = tree_flatten(model.trainable_parameters())
|
||||
fp16 = [(k, v.astype(mx.float16)) for k, v in flat]
|
||||
model.load_weights(fp16)
|
||||
mx.eval(model.parameters())
|
||||
|
||||
lang = model.language_model
|
||||
N = lang.args.num_hidden_layers
|
||||
ids = mx.ones((1, SEQ), dtype=mx.int32)
|
||||
pos = mx.broadcast_to(mx.arange(SEQ).reshape(1, SEQ)[None, :, :], (3, 1, SEQ))
|
||||
mx.eval(ids, pos)
|
||||
|
||||
# GPU baseline
|
||||
print("[1/3] GPU FP16 baseline")
|
||||
for _ in range(N_W):
|
||||
c = [KVCache() for _ in range(N)]
|
||||
mx.eval(lang(ids, cache=c, position_ids=pos).logits)
|
||||
ts = []
|
||||
for _ in range(N_B):
|
||||
c = [KVCache() for _ in range(N)]
|
||||
t0 = time.perf_counter()
|
||||
mx.eval(lang(ids, cache=c, position_ids=pos).logits)
|
||||
ts.append((time.perf_counter()-t0)*1000)
|
||||
bl = float(np.median(ts))
|
||||
print(f" {bl:.1f}ms\n")
|
||||
|
||||
# Reference logits (FP32 for cos_sim)
|
||||
c_ref = [KVCache() for _ in range(N)]
|
||||
ref = np.array(lang(ids, cache=c_ref, position_ids=pos).logits.astype(mx.float32))
|
||||
|
||||
# Patch
|
||||
print("[2/3] Patch + benchmark")
|
||||
bridge = patch_model(model, SEQ)
|
||||
SplitLinear.set_prefill(True)
|
||||
|
||||
for _ in range(N_W):
|
||||
c = [KVCache() for _ in range(N)]
|
||||
mx.eval(lang(ids, cache=c, position_ids=pos).logits)
|
||||
|
||||
# Accuracy
|
||||
c_hyb = [KVCache() for _ in range(N)]
|
||||
hyb = np.array(lang(ids, cache=c_hyb, position_ids=pos).logits.astype(mx.float32))
|
||||
cos = float(np.dot(ref.flatten(), hyb.flatten()) /
|
||||
(np.linalg.norm(ref.flatten()) * np.linalg.norm(hyb.flatten()) + 1e-12))
|
||||
top1 = float((ref.argmax(-1) == hyb.argmax(-1)).mean() * 100)
|
||||
print(f" cos={cos:.6f}, top1={top1:.1f}%")
|
||||
|
||||
# Benchmark
|
||||
ts = []
|
||||
for i in range(N_B):
|
||||
c = [KVCache() for _ in range(N)]
|
||||
t0 = time.perf_counter()
|
||||
mx.eval(lang(ids, cache=c, position_ids=pos).logits)
|
||||
t = (time.perf_counter()-t0)*1000
|
||||
ts.append(t)
|
||||
print(f" Run {i+1}: {t:.1f}ms")
|
||||
med = float(np.median(ts))
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f" GPU FP16: {bl:.1f}ms")
|
||||
print(f" SplitLinear: {med:.1f}ms ({bl/med:.3f}x)")
|
||||
print(f" cos={cos:.6f} top1={top1:.1f}%")
|
||||
print(f" delta: {med - bl:+.1f}ms")
|
||||
print(f"{'='*60}")
|
||||
Reference in New Issue
Block a user