chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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