wehub-resource-sync
109 lines
3.5 KiB
Python
109 lines
3.5 KiB
Python
#!/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()
|