chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,347 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
End-to-end benchmark on REAL trajectory data:
|
||||
1. GPU FP16 baseline
|
||||
2. GPU W8A16 (native quantized_matmul)
|
||||
3. GPU W8A8_pg (INT8 TensorOps via cider, pergroup)
|
||||
4. GPU W8A8_pc (INT8 TensorOps via cider, perchannel)
|
||||
5-8. Above 4 configs + Cider SDPA patch (decode acceleration)
|
||||
|
||||
Uses replay_prompt.build_prompt_at_step() to build real prompts with real screenshots.
|
||||
Compares both accuracy (action output) and speed (prefill + decode).
|
||||
"""
|
||||
import sys, os, time, gc, re, json, io, base64
|
||||
import numpy as np
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from PIL import Image
|
||||
from pathlib import Path
|
||||
ROOT_DIR = str(Path(__file__).parent.parent)
|
||||
sys.path.insert(0, ROOT_DIR)
|
||||
sys.path.insert(0, os.path.join(ROOT_DIR, "vlm_service"))
|
||||
|
||||
from session_data.replay_prompt import build_prompt_at_step
|
||||
from vlm_service.custom_qwen3vl import custom_stream_generate
|
||||
from mlx_vlm.utils import load as vlm_load
|
||||
home_dir = os.path.expanduser("~")
|
||||
FP16_MODEL = os.path.join(home_dir, 'Downloads/checkpoint-50349')
|
||||
W8A16_MODEL = os.path.join(home_dir, 'Downloads/checkpoint-50349-mlx_w8a16')
|
||||
SESSION_DIR = os.path.join(ROOT_DIR, "session_data")
|
||||
STEPS = [0, 1, 2] # step 0: 1img, step 1: 2imgs, step 2: 3imgs
|
||||
MAX_TOKENS = 200
|
||||
PREFILL_STEP = 8192
|
||||
N_SPEED_RUNS = 2 # per-step speed runs
|
||||
|
||||
|
||||
def load_step(session_dir, step):
|
||||
"""Load prompt + images for a given step."""
|
||||
data = build_prompt_at_step(session_dir, step)
|
||||
pil_images = [Image.open(io.BytesIO(base64.b64decode(b))) for b in data['images']]
|
||||
return data, pil_images
|
||||
|
||||
|
||||
def build_messages(data):
|
||||
"""Build chat messages from replay_prompt data, with <image> placeholders."""
|
||||
prompt_text = data['prompt']
|
||||
n_images = prompt_text.count('<image>')
|
||||
parts = prompt_text.split('<image>')
|
||||
content = []
|
||||
for i, part in enumerate(parts):
|
||||
if part:
|
||||
content.append({"type": "text", "text": part})
|
||||
if i < n_images:
|
||||
content.append({"type": "image"})
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": data['system_prompt']},
|
||||
{"role": "user", "content": content},
|
||||
]
|
||||
return messages
|
||||
|
||||
|
||||
def run_generate(model, processor, messages, pil_images):
|
||||
"""Run one generation, return (text, timing_dict)."""
|
||||
prompt = processor.tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
mx.reset_peak_memory()
|
||||
|
||||
text = ""
|
||||
prefill_time = 0
|
||||
decode_tps = 0
|
||||
prompt_tokens = 0
|
||||
gen_tokens = 0
|
||||
|
||||
try:
|
||||
for resp in custom_stream_generate(
|
||||
model, processor,
|
||||
prompt=prompt,
|
||||
image=pil_images,
|
||||
max_tokens=MAX_TOKENS,
|
||||
temperature=0.0,
|
||||
prefill_step_size=PREFILL_STEP,
|
||||
verbose=False,
|
||||
):
|
||||
text += resp.text
|
||||
prefill_time = resp.prompt_tokens / resp.prompt_tps if resp.prompt_tps > 0 else 0
|
||||
decode_tps = resp.generation_tps
|
||||
prompt_tokens = resp.prompt_tokens
|
||||
gen_tokens = resp.generation_tokens
|
||||
except UnicodeDecodeError:
|
||||
# mlx_vlm detokenizer sometimes fails on partial utf-8 sequences
|
||||
pass
|
||||
|
||||
total_time = prefill_time + (gen_tokens / decode_tps if decode_tps > 0 else 0)
|
||||
|
||||
return text, {
|
||||
'prefill_ms': prefill_time * 1000,
|
||||
'prefill_tps': prompt_tokens / prefill_time if prefill_time > 0 else 0,
|
||||
'decode_tps': decode_tps,
|
||||
'prompt_tokens': prompt_tokens,
|
||||
'gen_tokens': gen_tokens,
|
||||
'total_ms': total_time * 1000,
|
||||
}
|
||||
|
||||
|
||||
def extract_action(text):
|
||||
"""Extract action from model output."""
|
||||
m = re.search(r'<action>(.*?)</action>', text, re.DOTALL)
|
||||
return m.group(1).strip() if m else None
|
||||
|
||||
|
||||
def extract_coords(action_str):
|
||||
"""Extract (x, y) from action string."""
|
||||
if not action_str:
|
||||
return None
|
||||
nums = re.findall(r'\d+', action_str)
|
||||
if len(nums) >= 2:
|
||||
return (int(nums[0]), int(nums[1]))
|
||||
return None
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 80)
|
||||
print(" E2E Benchmark: FP16 / W8A16 / W8A8 × MLX SDPA / Cider SDPA")
|
||||
print("=" * 80)
|
||||
print(f" Session: {SESSION_DIR}")
|
||||
print(f" Steps: {STEPS}")
|
||||
print(f" Max tokens: {MAX_TOKENS}")
|
||||
print(f" Speed runs per step: {N_SPEED_RUNS}")
|
||||
|
||||
# Load all steps
|
||||
print("\nLoading trajectory data...")
|
||||
step_data = {}
|
||||
for s in STEPS:
|
||||
data, images = load_step(SESSION_DIR, s)
|
||||
msgs = build_messages(data)
|
||||
step_data[s] = (data, images, msgs)
|
||||
print(f" Step {s}: {len(images)} images")
|
||||
|
||||
# Import cider for SDPA patching
|
||||
import cider
|
||||
|
||||
# configs: (key, model_path, use_cider_quant, use_cider_sdpa, label)
|
||||
configs = [
|
||||
# Without Cider SDPA (MLX default attention)
|
||||
('fp16', FP16_MODEL, False, False, "FP16"),
|
||||
('fp16+sdpa', FP16_MODEL, False, True, "FP16+CiderSDPA"),
|
||||
('w8a16', W8A16_MODEL, False, False, "W8A16"),
|
||||
('w8a16+sdpa', W8A16_MODEL, False, True, "W8A16+CiderSDPA"),
|
||||
('w8a8_pg', W8A16_MODEL, True, False, "W8A8_pg"),
|
||||
('w8a8_pg+sdpa', W8A16_MODEL, True, True, "W8A8_pg+CiderSDPA"),
|
||||
('w8a8_pc', FP16_MODEL, True, False, "W8A8_pc"),
|
||||
('w8a8_pc+sdpa', FP16_MODEL, True, True, "W8A8_pc+CiderSDPA"),
|
||||
]
|
||||
|
||||
all_results = {}
|
||||
|
||||
for cfg_key, model_path, use_cider_quant, use_cider_sdpa, label in configs:
|
||||
print(f"\n{'='*80}")
|
||||
print(f" {label}")
|
||||
print(f"{'='*80}")
|
||||
|
||||
# Load model
|
||||
print(f" Loading {model_path}...")
|
||||
model, proc = vlm_load(model_path)
|
||||
|
||||
# Apply cider quantization if needed
|
||||
if use_cider_quant:
|
||||
from cider import convert_model
|
||||
stats = convert_model(model)
|
||||
print(f" [cider quant] {stats}")
|
||||
|
||||
# Apply cider SDPA if needed
|
||||
if use_cider_sdpa:
|
||||
cider.patch_sdpa(verbose=True)
|
||||
else:
|
||||
cider.unpatch_sdpa(verbose=False)
|
||||
|
||||
cfg_results = {}
|
||||
|
||||
for s in STEPS:
|
||||
data, images, msgs = step_data[s]
|
||||
print(f"\n Step {s} ({len(images)} imgs):")
|
||||
|
||||
# Warmup
|
||||
text, timing = run_generate(model, proc, msgs, images)
|
||||
print(f" Warmup: {timing['prompt_tokens']} prompt tok, "
|
||||
f"prefill {timing['prefill_ms']:.0f}ms "
|
||||
f"({timing['prefill_tps']:.0f} tok/s), "
|
||||
f"decode {timing['decode_tps']:.1f} tok/s")
|
||||
|
||||
# Accuracy run (use warmup result)
|
||||
action = extract_action(text)
|
||||
coords = extract_coords(action)
|
||||
|
||||
# Speed runs
|
||||
timings = []
|
||||
for r in range(N_SPEED_RUNS):
|
||||
_, t = run_generate(model, proc, msgs, images)
|
||||
timings.append(t)
|
||||
print(f" Run {r+1}: prefill {t['prefill_ms']:.0f}ms "
|
||||
f"({t['prefill_tps']:.0f} tok/s) | "
|
||||
f"decode {t['decode_tps']:.1f} tok/s | "
|
||||
f"total {t['total_ms']:.0f}ms")
|
||||
|
||||
cfg_results[s] = {
|
||||
'text': text,
|
||||
'action': action,
|
||||
'coords': coords,
|
||||
'timings': timings,
|
||||
'prompt_tokens': timings[0]['prompt_tokens'],
|
||||
}
|
||||
|
||||
print(f" Action: {action}")
|
||||
|
||||
all_results[cfg_key] = cfg_results
|
||||
|
||||
# Free model before loading next
|
||||
del model, proc
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
# Ensure SDPA is unpatched after benchmark
|
||||
cider.unpatch_sdpa(verbose=False)
|
||||
|
||||
# ── Speed Summary ──
|
||||
print(f"\n{'='*80}")
|
||||
print(f" SPEED SUMMARY (median of {N_SPEED_RUNS} runs)")
|
||||
print(f"{'='*80}")
|
||||
|
||||
cfg_keys = [c[0] for c in configs]
|
||||
cfg_labels = {c[0]: c[4] for c in configs}
|
||||
|
||||
# Table: per-step results
|
||||
print(f"\n {'Config':<22s}", end="")
|
||||
for s in STEPS:
|
||||
print(f" | Step{s} prefill decode total", end="")
|
||||
print()
|
||||
print(f" {'─'*22}", end="")
|
||||
for _ in STEPS:
|
||||
print(f" | {'─'*30}", end="")
|
||||
print()
|
||||
|
||||
for k in cfg_keys:
|
||||
line = f" {cfg_labels[k]:<22s}"
|
||||
for s in STEPS:
|
||||
ts = all_results[k][s]['timings']
|
||||
pf = float(np.median([t['prefill_tps'] for t in ts]))
|
||||
dc = float(np.median([t['decode_tps'] for t in ts]))
|
||||
tt = float(np.median([t['total_ms'] for t in ts]))
|
||||
line += f" | {pf:>7.0f}t/s {dc:>5.1f}t/s {tt:>5.0f}ms"
|
||||
print(line)
|
||||
|
||||
# ── Decode speed comparison (SDPA impact) ──
|
||||
print(f"\n{'='*80}")
|
||||
print(f" DECODE SPEED: MLX SDPA vs Cider SDPA")
|
||||
print(f"{'='*80}")
|
||||
|
||||
# Group by base config
|
||||
base_configs = ['fp16', 'w8a16', 'w8a8_pg', 'w8a8_pc']
|
||||
for base in base_configs:
|
||||
sdpa_key = base + '+sdpa'
|
||||
if base not in all_results or sdpa_key not in all_results:
|
||||
continue
|
||||
|
||||
print(f"\n {cfg_labels[base]}:")
|
||||
for s in STEPS:
|
||||
ts_base = all_results[base][s]['timings']
|
||||
ts_sdpa = all_results[sdpa_key][s]['timings']
|
||||
dc_base = float(np.median([t['decode_tps'] for t in ts_base]))
|
||||
dc_sdpa = float(np.median([t['decode_tps'] for t in ts_sdpa]))
|
||||
speedup = dc_sdpa / dc_base if dc_base > 0 else 0
|
||||
pf_base = float(np.median([t['prefill_tps'] for t in ts_base]))
|
||||
pf_sdpa = float(np.median([t['prefill_tps'] for t in ts_sdpa]))
|
||||
print(f" Step {s}: decode {dc_base:.1f} → {dc_sdpa:.1f} tok/s "
|
||||
f"({speedup:.2f}x) | prefill {pf_base:.0f} → {pf_sdpa:.0f} tok/s")
|
||||
|
||||
# ── Accuracy comparison ──
|
||||
print(f"\n{'='*80}")
|
||||
print(f" ACCURACY: SDPA patch should not change outputs")
|
||||
print(f"{'='*80}")
|
||||
|
||||
for base in base_configs:
|
||||
sdpa_key = base + '+sdpa'
|
||||
if base not in all_results or sdpa_key not in all_results:
|
||||
continue
|
||||
mismatches = 0
|
||||
for s in STEPS:
|
||||
a1 = all_results[base][s]['action']
|
||||
a2 = all_results[sdpa_key][s]['action']
|
||||
if a1 != a2:
|
||||
mismatches += 1
|
||||
c1 = all_results[base][s]['coords']
|
||||
c2 = all_results[sdpa_key][s]['coords']
|
||||
if c1 and c2:
|
||||
diff = (abs(c1[0]-c2[0]), abs(c1[1]-c2[1]))
|
||||
close = all(d <= 5 for d in diff)
|
||||
tag = f"{'≈' if close else '≠'} diff=({diff[0]},{diff[1]})"
|
||||
else:
|
||||
tag = "≠"
|
||||
print(f" {cfg_labels[base]} Step {s}: {tag}")
|
||||
print(f" base: {a1}")
|
||||
print(f" sdpa: {a2}")
|
||||
if mismatches == 0:
|
||||
print(f" {cfg_labels[base]}: all {len(STEPS)} steps IDENTICAL ✅")
|
||||
|
||||
# ── Overall Summary ──
|
||||
print(f"\n{'='*80}")
|
||||
print(f" OVERALL DECODE tok/s (median across all steps)")
|
||||
print(f"{'='*80}")
|
||||
for k in cfg_keys:
|
||||
all_decode = []
|
||||
for s in STEPS:
|
||||
ts = all_results[k][s]['timings']
|
||||
all_decode.append(float(np.median([t['decode_tps'] for t in ts])))
|
||||
med = np.median(all_decode)
|
||||
print(f" {cfg_labels[k]:<22s}: {med:.1f} tok/s")
|
||||
|
||||
# Save
|
||||
out_path = '/tmp/e2e_wxa16_sdpa_results.json'
|
||||
save_data = {}
|
||||
for k in cfg_keys:
|
||||
save_data[k] = {}
|
||||
for s in STEPS:
|
||||
r = all_results[k][s]
|
||||
save_data[k][str(s)] = {
|
||||
'action': r['action'],
|
||||
'coords': r['coords'],
|
||||
'prompt_tokens': r['prompt_tokens'],
|
||||
'timings_median': {
|
||||
'prefill_ms': float(np.median([t['prefill_ms'] for t in r['timings']])),
|
||||
'prefill_tps': float(np.median([t['prefill_tps'] for t in r['timings']])),
|
||||
'decode_tps': float(np.median([t['decode_tps'] for t in r['timings']])),
|
||||
'total_ms': float(np.median([t['total_ms'] for t in r['timings']])),
|
||||
},
|
||||
}
|
||||
with open(out_path, 'w') as f:
|
||||
json.dump(save_data, f, indent=2)
|
||||
print(f"\n Results saved to {out_path}")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,106 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Full kernel benchmark: Cider per-channel & per-group vs MLX w8a16 & w4a16.
|
||||
|
||||
Tests M = {1, 128, 1024, 4096, 8192} across Qwen3-2B/7B shapes.
|
||||
"""
|
||||
import sys, os, time
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
import numpy as np
|
||||
import mlx.core as mx
|
||||
from cider import perchannel_linear, pergroup_linear, quantize_weight_int8
|
||||
|
||||
WARMUP = 5
|
||||
REPEAT = 20
|
||||
|
||||
# Qwen3-VL shapes: [N, K]
|
||||
SHAPES_NK = [
|
||||
(3584, 3584), # qkv square
|
||||
(18944, 3584), # up-proj
|
||||
(3584, 18944), # down-proj
|
||||
(2560, 2560), # 2B square
|
||||
(10240, 2560), # 2B up
|
||||
(2560, 10240), # 2B down
|
||||
]
|
||||
|
||||
M_VALUES = [1, 128, 1024, 4096, 8192]
|
||||
GROUP_SIZE = 128
|
||||
|
||||
|
||||
def timed(fn, warmup=WARMUP, repeat=REPEAT):
|
||||
for _ in range(warmup):
|
||||
y = fn(); mx.eval(y)
|
||||
times = []
|
||||
for _ in range(repeat):
|
||||
t0 = time.perf_counter()
|
||||
y = fn(); mx.eval(y)
|
||||
times.append(time.perf_counter() - t0)
|
||||
return np.median(times) * 1000
|
||||
|
||||
|
||||
def main():
|
||||
print(f"Full Kernel Benchmark (warmup={WARMUP}, repeat={REPEAT}, gs={GROUP_SIZE})")
|
||||
print()
|
||||
|
||||
for N, K in SHAPES_NK:
|
||||
print(f"=== Shape [N={N}, K={K}] ===")
|
||||
print(f"{'M':>6s} | {'PC(ms)':>8s} {'PG(ms)':>8s} {'w8a16':>8s} {'w4a16':>8s} | {'PC/w8':>7s} {'PC/w4':>7s} {'PG/w8':>7s} {'PG/w4':>7s}")
|
||||
print("-" * 95)
|
||||
|
||||
# Prepare per-channel weights
|
||||
np.random.seed(42)
|
||||
w_fp = np.random.randn(N, K).astype(np.float32)
|
||||
# Per-channel: scale per row
|
||||
absmax = np.abs(w_fp).max(axis=1, keepdims=True).clip(min=1e-8)
|
||||
scale_pc = (absmax / 127.0).astype(np.float32)
|
||||
w_int8_pc = np.clip(np.round(w_fp / scale_pc), -127, 127).astype(np.int8)
|
||||
w_int8_pc_mx = mx.array(w_int8_pc)
|
||||
scale_pc_mx = mx.array(scale_pc.squeeze())
|
||||
|
||||
# Per-group weights (symmetric, gs=GROUP_SIZE)
|
||||
ng = (K + GROUP_SIZE - 1) // GROUP_SIZE
|
||||
w_pg = w_fp.copy()
|
||||
scale_pg = np.zeros((N, ng), dtype=np.float32)
|
||||
w_int8_pg = np.zeros((N, K), dtype=np.int8)
|
||||
for g in range(ng):
|
||||
k0 = g * GROUP_SIZE
|
||||
k1 = min(k0 + GROUP_SIZE, K)
|
||||
blk = w_pg[:, k0:k1]
|
||||
amax = np.abs(blk).max(axis=1, keepdims=True).clip(min=1e-8)
|
||||
s = (amax / 127.0).astype(np.float32)
|
||||
scale_pg[:, g] = s.squeeze()
|
||||
w_int8_pg[:, k0:k1] = np.clip(np.round(blk / s), -127, 127).astype(np.int8)
|
||||
w_int8_pg_mx = mx.array(w_int8_pg)
|
||||
scale_pg_mx = mx.array(scale_pg)
|
||||
|
||||
# MLX w8a16 (bits=8, gs=64)
|
||||
w_mx = mx.array(w_fp.astype(np.float16))
|
||||
w8_q, w8_s, w8_b = mx.quantize(w_mx, bits=8, group_size=64)
|
||||
|
||||
# MLX w4a16 (bits=4, gs=64)
|
||||
w4_q, w4_s, w4_b = mx.quantize(w_mx, bits=4, group_size=64)
|
||||
|
||||
mx.eval(w_int8_pc_mx, scale_pc_mx, w_int8_pg_mx, scale_pg_mx,
|
||||
w8_q, w8_s, w8_b, w4_q, w4_s, w4_b)
|
||||
|
||||
for M in M_VALUES:
|
||||
x = mx.random.normal((M, K)).astype(mx.float16)
|
||||
mx.eval(x)
|
||||
|
||||
t_pc = timed(lambda: perchannel_linear(x, w_int8_pc_mx, scale_pc_mx))
|
||||
t_pg = timed(lambda: pergroup_linear(x, w_int8_pg_mx, scale_pg_mx, GROUP_SIZE))
|
||||
t_w8 = timed(lambda: mx.quantized_matmul(x, w8_q, w8_s, w8_b, bits=8, group_size=64))
|
||||
t_w4 = timed(lambda: mx.quantized_matmul(x, w4_q, w4_s, w4_b, bits=4, group_size=64))
|
||||
|
||||
r_pc_w8 = t_w8 / t_pc if t_pc > 0 else 0
|
||||
r_pc_w4 = t_w4 / t_pc if t_pc > 0 else 0
|
||||
r_pg_w8 = t_w8 / t_pg if t_pg > 0 else 0
|
||||
r_pg_w4 = t_w4 / t_pg if t_pg > 0 else 0
|
||||
|
||||
print(f"{M:>6d} | {t_pc:>6.2f}ms {t_pg:>6.2f}ms {t_w8:>6.2f}ms {t_w4:>6.2f}ms | {r_pc_w8:>5.2f}x {r_pc_w4:>5.2f}x {r_pg_w8:>5.2f}x {r_pg_w4:>5.2f}x")
|
||||
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,155 @@
|
||||
"""benchmarks/bench_sdpa.py — Cider SDPA vs MLX native: correctness + performance.
|
||||
|
||||
Usage:
|
||||
python benchmarks/bench_sdpa.py # Default configs
|
||||
python benchmarks/bench_sdpa.py --autotune # Run AutoTune first
|
||||
python benchmarks/bench_sdpa.py --iters 500 # More iterations for stable results
|
||||
"""
|
||||
import argparse
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
|
||||
sys.path.insert(0, ".")
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
|
||||
def bench_one(fn, q, k, v, scale, warmup, iters):
|
||||
"""Benchmark a single SDPA function. Returns p10 latency in microseconds."""
|
||||
for _ in range(warmup):
|
||||
r = fn(q, k, v, scale=scale)
|
||||
mx.eval(r)
|
||||
|
||||
times = []
|
||||
for _ in range(iters):
|
||||
t0 = time.perf_counter()
|
||||
r = fn(q, k, v, scale=scale)
|
||||
mx.eval(r)
|
||||
times.append((time.perf_counter() - t0) * 1e6)
|
||||
|
||||
times.sort()
|
||||
p10 = times[max(0, len(times) // 10)]
|
||||
median = times[len(times) // 2]
|
||||
return p10, median
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Cider SDPA benchmark")
|
||||
parser.add_argument("--warmup", type=int, default=50)
|
||||
parser.add_argument("--iters", type=int, default=200)
|
||||
parser.add_argument("--autotune", action="store_true", help="Run AutoTune before benchmark")
|
||||
parser.add_argument("--D", type=int, default=128, help="Head dimension")
|
||||
args = parser.parse_args()
|
||||
|
||||
from cider.attention.sdpa import scaled_dot_product_attention, _mlx_sdpa
|
||||
|
||||
# GPU info
|
||||
info = mx.metal.device_info() if hasattr(mx.metal, 'device_info') else mx.device_info()
|
||||
arch = info.get("architecture", "unknown")
|
||||
print(f"GPU: {arch}")
|
||||
print(f"MLX: {mx.__version__}")
|
||||
print(f"Config: warmup={args.warmup}, iters={args.iters}, D={args.D}")
|
||||
|
||||
if args.autotune:
|
||||
import cider
|
||||
print("\n[AutoTune]")
|
||||
cider.autotune_sdpa(D=args.D, iters=args.iters // 4)
|
||||
print()
|
||||
|
||||
configs = [
|
||||
# (B, H_q, H_kv, N, label)
|
||||
(1, 32, 32, 1024, "MHA N=1K"),
|
||||
(1, 32, 32, 2048, "MHA N=2K"),
|
||||
(1, 32, 32, 4096, "MHA N=4K"),
|
||||
(1, 32, 32, 8192, "MHA N=8K"),
|
||||
(1, 32, 32, 16384, "MHA N=16K"),
|
||||
(1, 32, 32, 32768, "MHA N=32K"),
|
||||
]
|
||||
for H_q in [32, 64, 96, 128]:
|
||||
for gqa_factor in [2, 4, 8, 16]:
|
||||
H_kv = H_q // gqa_factor
|
||||
for N in [1024, 2048, 4096, 8192, 16384, 32768]:
|
||||
label = f"GQA{H_q}_{gqa_factor} N={N // 1024}K"
|
||||
configs.append((1, H_q, H_kv, N, label))
|
||||
|
||||
D = args.D
|
||||
|
||||
# ── Correctness ─────────────────────────────────────────────
|
||||
print("\n" + "=" * 72)
|
||||
print("CORRECTNESS (max |cider - mlx|)")
|
||||
print("=" * 72)
|
||||
print(f"{'Config':<16} {'MaxDiff':>12} {'Status':>8}")
|
||||
print("-" * 40)
|
||||
|
||||
mx.random.seed(42)
|
||||
all_correct = True
|
||||
for B, Hq, Hkv, N, label in configs:
|
||||
q = mx.random.normal((B, Hq, 1, D)).astype(mx.float16)
|
||||
k = mx.random.normal((B, Hkv, N, D)).astype(mx.float16)
|
||||
v = mx.random.normal((B, Hkv, N, D)).astype(mx.float16)
|
||||
scale = 1.0 / math.sqrt(D)
|
||||
|
||||
ref = _mlx_sdpa(q, k, v, scale=scale)
|
||||
out = scaled_dot_product_attention(q, k, v, scale=scale)
|
||||
mx.eval(ref, out)
|
||||
|
||||
diff = mx.abs(ref - out).max().item()
|
||||
ok = diff < 0.01
|
||||
if not ok:
|
||||
all_correct = False
|
||||
status = "PASS" if ok else "FAIL"
|
||||
print(f"{label:<16} {diff:>12.6f} {status:>8}")
|
||||
|
||||
print()
|
||||
if all_correct:
|
||||
print(">>> ALL CORRECT <<<")
|
||||
else:
|
||||
print(">>> SOME FAILED <<<")
|
||||
sys.exit(1)
|
||||
|
||||
# ── Performance ─────────────────────────────────────────────
|
||||
print("\n" + "=" * 72)
|
||||
print(f"PERFORMANCE (p10 latency, {args.iters} iters)")
|
||||
print("=" * 72)
|
||||
print(f"{'Config':<16} {'MLX(us)':>10} {'Cider(us)':>10} {'Speedup':>10} {'Status':>8}")
|
||||
print("-" * 58)
|
||||
|
||||
wins = ties = losses = 0
|
||||
|
||||
mx.random.seed(42)
|
||||
for B, Hq, Hkv, N, label in configs:
|
||||
q = mx.random.normal((B, Hq, 1, D)).astype(mx.float16)
|
||||
k = mx.random.normal((B, Hkv, N, D)).astype(mx.float16)
|
||||
v = mx.random.normal((B, Hkv, N, D)).astype(mx.float16)
|
||||
scale = 1.0 / math.sqrt(D)
|
||||
mx.eval(q, k, v)
|
||||
|
||||
mlx_p10, _ = bench_one(_mlx_sdpa, q, k, v, scale, args.warmup, args.iters)
|
||||
cider_p10, _ = bench_one(scaled_dot_product_attention, q, k, v, scale, args.warmup, args.iters)
|
||||
|
||||
ratio = mlx_p10 / cider_p10
|
||||
if ratio >= 1.02:
|
||||
status = "WIN"
|
||||
wins += 1
|
||||
elif ratio >= 0.98:
|
||||
status = "TIE"
|
||||
ties += 1
|
||||
else:
|
||||
status = "LOSS"
|
||||
losses += 1
|
||||
|
||||
print(f"{label:<16} {mlx_p10:>10.1f} {cider_p10:>10.1f} {ratio:>9.2f}x {status:>8}")
|
||||
|
||||
print("-" * 58)
|
||||
print(f"Summary: {wins} wins / {ties} ties / {losses} losses")
|
||||
print()
|
||||
|
||||
if losses == 0:
|
||||
print(">>> NO REGRESSIONS <<<")
|
||||
else:
|
||||
print(f">>> {losses} REGRESSION(S) DETECTED <<<")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,353 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Benchmark: Cider W8A8 INT8 matmul vs MLX-style NAX INT8 matmul.
|
||||
|
||||
Both use the same mpp::tensor_ops::matmul2d<int8,int8,int32> hardware instruction.
|
||||
The difference is in data loading/storing patterns.
|
||||
|
||||
We write a standalone INT8 GEMM kernel in MLX Steel NAX style,
|
||||
compiled via mx.fast.metal_kernel, and compare against cider's kernel.
|
||||
"""
|
||||
import sys, os, time
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
import numpy as np
|
||||
import mlx.core as mx
|
||||
from cider import perchannel_linear, quantize_weight_int8
|
||||
from cider.ops import int8_matmul_int32
|
||||
|
||||
# =============================================================
|
||||
# MLX-style INT8 NAX GEMM kernel (standalone)
|
||||
# Same tile config as cider: BM=128, BN=128, BK=512, WM=4, WN=4
|
||||
# Uses matmul2d<int8,int8,int32> - identical HW instruction to cider
|
||||
# =============================================================
|
||||
|
||||
MLX_INT8_HEADER = r"""
|
||||
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
|
||||
#include <metal_simdgroup>
|
||||
|
||||
using namespace metal;
|
||||
|
||||
constant constexpr int kElemsPerFrag = 8;
|
||||
constant constexpr int kElemCols = 4;
|
||||
constant constexpr int kElemRowsJump = 8;
|
||||
|
||||
METAL_FUNC short2 nax_get_coord(ushort lid) {
|
||||
short qid = short(lid >> 2);
|
||||
short fm = ((qid & 4) | ((short(lid) >> 1) & 3));
|
||||
short fn = ((qid & 2) | (short(lid) & 1)) * 4;
|
||||
return short2(fn, fm);
|
||||
}
|
||||
|
||||
METAL_FUNC void load_frag_int8(thread int8_t *dst,
|
||||
const device int8_t *base,
|
||||
int stride, short2 sc,
|
||||
short row_off, short col_off) {
|
||||
const device int8_t *ptr = base + (sc.y + row_off) * stride + sc.x + col_off;
|
||||
for (short i = 0; i < 2; i++) {
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
dst[i * kElemCols + j] = ptr[i * kElemRowsJump * stride + j];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
METAL_FUNC void store_frag_int32(const thread int32_t *src,
|
||||
device int32_t *base,
|
||||
int stride, short2 sc,
|
||||
short row_off, short col_off,
|
||||
uint M, uint N,
|
||||
uint m_base, uint n_base) {
|
||||
for (short i = 0; i < 2; i++) {
|
||||
uint r = m_base + row_off + sc.y + i * kElemRowsJump;
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
uint c = n_base + col_off + sc.x + j;
|
||||
if (r < M && c < N) {
|
||||
base[r * N + c] = src[i * kElemCols + j];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
MLX_INT8_SOURCE = r"""
|
||||
uint M_val = M[0], N_val = N[0], K_val = K[0];
|
||||
uint tiles_m_val = tilesm[0], tiles_n_val = tilesn[0];
|
||||
|
||||
constexpr int BM = 128, BN = 128, BK = 512;
|
||||
constexpr int WM = 4, WN = 4;
|
||||
constexpr int SM = BM / WM;
|
||||
constexpr int SN = BN / WN;
|
||||
constexpr short TM = SM / 16;
|
||||
constexpr short TN = SN / 16;
|
||||
constexpr short TK = 32 / 16;
|
||||
constexpr short SK = 32;
|
||||
|
||||
uint2 tgid = uint2(threadgroup_position_in_grid.x, threadgroup_position_in_grid.y);
|
||||
uint sgid = simdgroup_index_in_threadgroup;
|
||||
uint lid = thread_index_in_simdgroup;
|
||||
|
||||
if (tgid.x >= tiles_n_val || tgid.y >= tiles_m_val) return;
|
||||
|
||||
short2 sc = nax_get_coord(ushort(lid));
|
||||
uint sg_row = sgid / WN;
|
||||
uint sg_col = sgid % WN;
|
||||
uint m_base = tgid.y * BM + sg_row * SM;
|
||||
uint n_base = tgid.x * BN + sg_col * SN;
|
||||
|
||||
const device int8_t *sg_A = A + m_base * K_val;
|
||||
const device int8_t *sg_B = B + n_base;
|
||||
|
||||
constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor(
|
||||
16, 32, 16, false, false, true,
|
||||
mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate);
|
||||
mpp::tensor_ops::matmul2d<desc, metal::execution_simdgroup> gemm_op;
|
||||
|
||||
auto ct_a = gemm_op.get_left_input_cooperative_tensor<int8_t, int8_t, int32_t>();
|
||||
auto ct_b = gemm_op.get_right_input_cooperative_tensor<int8_t, int8_t, int32_t>();
|
||||
auto ct_c = gemm_op.get_destination_cooperative_tensor<decltype(ct_a), decltype(ct_b), int32_t>();
|
||||
|
||||
int32_t c_frags[TM * TN][kElemsPerFrag];
|
||||
for (int f = 0; f < TM * TN; f++)
|
||||
for (int i = 0; i < kElemsPerFrag; i++)
|
||||
c_frags[f][i] = 0;
|
||||
|
||||
int gemm_k_iters = int(K_val) / BK;
|
||||
for (int kk0 = 0; kk0 < gemm_k_iters; kk0++) {
|
||||
threadgroup_barrier(mem_flags::mem_none);
|
||||
for (int kk1 = 0; kk1 < BK; kk1 += SK) {
|
||||
int8_t a_frags[TM][TK][kElemsPerFrag];
|
||||
int8_t b_frags[TK][TN][kElemsPerFrag];
|
||||
volatile int compiler_barrier;
|
||||
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
load_frag_int8(a_frags[mm][kk], sg_A + kk1, int(K_val), sc, short(mm*16), short(kk*16));
|
||||
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
load_frag_int8(b_frags[kk][nn], sg_B + kk1 * N_val, int(N_val), sc, short(kk*16), short(nn*16));
|
||||
|
||||
for (short mm = 0; mm < TM; mm++) {
|
||||
for (short nn = 0; nn < TN; nn += 2) {
|
||||
for (short kk = 0; kk < TK; kk++) {
|
||||
for (short i = 0; i < kElemsPerFrag; i++) ct_a[i] = a_frags[mm][kk][i];
|
||||
for (short i = 0; i < kElemsPerFrag; i++) {
|
||||
ct_b[i] = b_frags[kk][nn][i];
|
||||
ct_b[kElemsPerFrag + i] = b_frags[kk][nn+1][i];
|
||||
}
|
||||
short c0 = mm*TN+nn, c1 = c0+1;
|
||||
for (short i = 0; i < kElemsPerFrag; i++) {
|
||||
ct_c[i] = c_frags[c0][i];
|
||||
ct_c[kElemsPerFrag+i] = c_frags[c1][i];
|
||||
}
|
||||
gemm_op.run(ct_a, ct_b, ct_c);
|
||||
for (short i = 0; i < kElemsPerFrag; i++) {
|
||||
c_frags[c0][i] = ct_c[i];
|
||||
c_frags[c1][i] = ct_c[kElemsPerFrag+i];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
(void)compiler_barrier;
|
||||
}
|
||||
sg_A += BK;
|
||||
sg_B += BK * N_val;
|
||||
}
|
||||
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
store_frag_int32(c_frags[mm*TN+nn], out, int(N_val), sc, short(mm*16), short(nn*16), M_val, N_val, m_base, n_base);
|
||||
"""
|
||||
|
||||
WARMUP = 20
|
||||
REPEAT = 50
|
||||
|
||||
SHAPES = [
|
||||
(16, 4096, 4096),
|
||||
(64, 4096, 4096),
|
||||
(128, 4096, 4096),
|
||||
(256, 4096, 4096),
|
||||
(512, 4096, 4096),
|
||||
(1024, 4096, 4096),
|
||||
(2048, 4096, 4096),
|
||||
# MLP shapes
|
||||
(128, 3584, 18944),
|
||||
(512, 3584, 18944),
|
||||
(128, 18944, 3584),
|
||||
(512, 18944, 3584),
|
||||
]
|
||||
|
||||
|
||||
def bench_cider_w8a8(x, w_int8, scale_w):
|
||||
"""Cider full pipeline: quantize_act + int8_matmul + fused_dequant -> FP16"""
|
||||
for _ in range(WARMUP):
|
||||
y = perchannel_linear(x, w_int8, scale_w)
|
||||
mx.eval(y)
|
||||
times = []
|
||||
for _ in range(REPEAT):
|
||||
t0 = time.perf_counter()
|
||||
y = perchannel_linear(x, w_int8, scale_w)
|
||||
mx.eval(y)
|
||||
times.append(time.perf_counter() - t0)
|
||||
return np.median(times) * 1000
|
||||
|
||||
|
||||
def bench_cider_int8_raw(A_int8, B_int8):
|
||||
"""Cider raw INT8 matmul: int8 x int8 -> int32 (no quantize, no dequant)"""
|
||||
for _ in range(WARMUP):
|
||||
y = int8_matmul_int32(A_int8, B_int8)
|
||||
mx.eval(y)
|
||||
times = []
|
||||
for _ in range(REPEAT):
|
||||
t0 = time.perf_counter()
|
||||
y = int8_matmul_int32(A_int8, B_int8)
|
||||
mx.eval(y)
|
||||
times.append(time.perf_counter() - t0)
|
||||
return np.median(times) * 1000
|
||||
|
||||
|
||||
def bench_mlx_int8(A_int8, B_int8, M, N, K):
|
||||
"""MLX-style standalone INT8 matmul: int8 x int8 -> int32 (raw GEMM only)"""
|
||||
BM, BN = 128, 128
|
||||
tiles_m = (M + BM - 1) // BM
|
||||
tiles_n = (N + BN - 1) // BN
|
||||
|
||||
M_buf = mx.array([M], dtype=mx.uint32)
|
||||
N_buf = mx.array([N], dtype=mx.uint32)
|
||||
K_buf = mx.array([K], dtype=mx.uint32)
|
||||
tm_buf = mx.array([tiles_m], dtype=mx.uint32)
|
||||
tn_buf = mx.array([tiles_n], dtype=mx.uint32)
|
||||
|
||||
kernel = mx.fast.metal_kernel(
|
||||
name='mlx_style_int8_gemm',
|
||||
input_names=['A', 'B', 'M', 'N', 'K', 'tilesm', 'tilesn'],
|
||||
output_names=['out'],
|
||||
source=MLX_INT8_SOURCE,
|
||||
header=MLX_INT8_HEADER,
|
||||
)
|
||||
|
||||
# grid = total threads, NOT threadgroup count
|
||||
# threadgroup=(512,1,1), so grid_x = tiles_n * 512, grid_y = tiles_m * 1
|
||||
grid_x = tiles_n * 512
|
||||
grid_y = tiles_m
|
||||
|
||||
for _ in range(WARMUP):
|
||||
result = kernel(
|
||||
inputs=[A_int8, B_int8, M_buf, N_buf, K_buf, tm_buf, tn_buf],
|
||||
output_shapes=[(M, N)],
|
||||
output_dtypes=[mx.int32],
|
||||
grid=(grid_x, grid_y, 1),
|
||||
threadgroup=(512, 1, 1),
|
||||
)
|
||||
mx.eval(result[0])
|
||||
|
||||
times = []
|
||||
for _ in range(REPEAT):
|
||||
t0 = time.perf_counter()
|
||||
result = kernel(
|
||||
inputs=[A_int8, B_int8, M_buf, N_buf, K_buf, tm_buf, tn_buf],
|
||||
output_shapes=[(M, N)],
|
||||
output_dtypes=[mx.int32],
|
||||
grid=(grid_x, grid_y, 1),
|
||||
threadgroup=(512, 1, 1),
|
||||
)
|
||||
mx.eval(result[0])
|
||||
times.append(time.perf_counter() - t0)
|
||||
return np.median(times) * 1000
|
||||
|
||||
|
||||
def verify_correctness(K, N):
|
||||
"""Quick correctness check: MLX-style int8 matmul vs numpy."""
|
||||
M = 16
|
||||
np.random.seed(123)
|
||||
a_np = np.random.randint(-5, 5, (M, K), dtype=np.int8)
|
||||
b_np = np.random.randint(-5, 5, (K, N), dtype=np.int8)
|
||||
ref = a_np.astype(np.int32) @ b_np.astype(np.int32)
|
||||
|
||||
A_int8 = mx.array(a_np)
|
||||
B_int8 = mx.array(b_np)
|
||||
|
||||
BM, BN = 128, 128
|
||||
tiles_m = (M + BM - 1) // BM
|
||||
tiles_n = (N + BN - 1) // BN
|
||||
|
||||
kernel = mx.fast.metal_kernel(
|
||||
name='mlx_style_int8_gemm',
|
||||
input_names=['A', 'B', 'M', 'N', 'K', 'tilesm', 'tilesn'],
|
||||
output_names=['out'],
|
||||
source=MLX_INT8_SOURCE,
|
||||
header=MLX_INT8_HEADER,
|
||||
)
|
||||
|
||||
result = kernel(
|
||||
inputs=[A_int8, B_int8,
|
||||
mx.array([M], dtype=mx.uint32),
|
||||
mx.array([N], dtype=mx.uint32),
|
||||
mx.array([K], dtype=mx.uint32),
|
||||
mx.array([tiles_m], dtype=mx.uint32),
|
||||
mx.array([tiles_n], dtype=mx.uint32)],
|
||||
output_shapes=[(M, N)],
|
||||
output_dtypes=[mx.int32],
|
||||
grid=(tiles_n * 512, tiles_m, 1),
|
||||
threadgroup=(512, 1, 1),
|
||||
)
|
||||
mx.eval(result[0])
|
||||
got = np.array(result[0])
|
||||
|
||||
max_diff = np.max(np.abs(ref - got))
|
||||
match = np.array_equal(ref, got)
|
||||
print(f"Correctness check ({M}x{K} @ {K}x{N}): max_diff={max_diff}, exact_match={match}")
|
||||
if not match:
|
||||
print(f" ref[0,:8] = {ref[0,:8]}")
|
||||
print(f" got[0,:8] = {got[0,:8]}")
|
||||
return match
|
||||
|
||||
|
||||
def main():
|
||||
print("=" * 70)
|
||||
print("Cider W8A8 (full pipeline) vs MLX-style INT8 GEMM (raw matmul)")
|
||||
print("=" * 70)
|
||||
print(f"Warmup={WARMUP}, Repeat={REPEAT}")
|
||||
print()
|
||||
print("NOTE: Cider includes quantize_act + INT8 matmul + fused dequant")
|
||||
print(" MLX-style is raw INT8 matmul only (no quantize/dequant)")
|
||||
print(" So MLX-style SHOULD be faster (does less work)")
|
||||
print()
|
||||
|
||||
# Verify correctness first
|
||||
print("--- Correctness Check ---")
|
||||
ok = verify_correctness(4096, 4096)
|
||||
if not ok:
|
||||
print("WARNING: correctness check failed! Results may be invalid.")
|
||||
print()
|
||||
|
||||
print(f"{'M':>5s} {'K':>6s} {'N':>6s} | {'Cider-full':>10s} | {'Cider-raw':>10s} | {'MLX-INT8':>9s} | {'full/raw':>8s} | {'raw/MLX':>7s}")
|
||||
print("-" * 80)
|
||||
|
||||
prev_kn = None
|
||||
for M, K, N in SHAPES:
|
||||
if (K, N) != prev_kn:
|
||||
np.random.seed(42)
|
||||
w_fp32 = np.random.randn(K, N).astype(np.float32)
|
||||
w_int8_np, scale_w_np = quantize_weight_int8(w_fp32)
|
||||
w_int8 = mx.array(w_int8_np)
|
||||
scale_w = mx.array(scale_w_np)
|
||||
B_int8 = mx.array(w_int8_np)
|
||||
mx.eval(w_int8, scale_w, B_int8)
|
||||
prev_kn = (K, N)
|
||||
|
||||
x_fp16 = mx.random.normal((M, K)).astype(mx.float16)
|
||||
A_int8 = mx.array(np.random.randint(-127, 127, (M, K), dtype=np.int8))
|
||||
mx.eval(x_fp16, A_int8)
|
||||
|
||||
t_full = bench_cider_w8a8(x_fp16, w_int8, scale_w)
|
||||
t_raw = bench_cider_int8_raw(A_int8, B_int8)
|
||||
t_mlx = bench_mlx_int8(A_int8, B_int8, M, N, K)
|
||||
|
||||
r_full_raw = t_full / t_raw if t_raw > 0 else 0
|
||||
r_raw_mlx = t_raw / t_mlx if t_mlx > 0 else 0
|
||||
print(f"{M:>5d} {K:>6d} {N:>6d} | {t_full:>8.2f}ms | {t_raw:>8.2f}ms | {t_mlx:>7.2f}ms | {r_full_raw:>6.2f}x | {r_raw_mlx:>5.2f}x")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,117 @@
|
||||
# Cider W8A8 Kernel vs MLX INT8 GEMM: A Transparent Comparison
|
||||
|
||||
|
||||
MLX's Steel NAX GEMM templates appear generic enough at the shader level to support `<int8_t, int8_t, int32_t>` instantiations. In our standalone test, an MLX-style INT8×INT8→INT32 kernel can be built on the same `mpp::tensor_ops::matmul2d` TensorOps instruction used by Cider. At the raw matmul level, both kernels achieve **comparable throughput** (within 2–5% across all tested sizes).
|
||||
|
||||
Cider's contribution is **not** a novel matmul kernel. It is an end-to-end W8A8 inference pipeline for Apple Silicon, combining online activation quantization, INT8 TensorOps matmul, and dequantization in a usable MLX-integrated path.
|
||||
|
||||
**Apple M5 introduced cooperative_tensor INT8 TensorOps (matmul2d 16×32×16), delivering 2× the TOPS of FP16 — but no ML framework on macOS uses them.** In the MLX quantization path we examined (v0.31), quantized weights are dequantized for computation and the corresponding matmul path remains FP16-based rather than end-to-end INT8. This means prefill (the compute-bound phase of LLM inference) gets zero benefit from INT8 hardware. W8A8 is the only quantization scheme that accelerates both compute and memory: activations and weights are both INT8, so the matmul runs on the faster INT8 datapath. Cider provides the missing pipeline — online activation quantization, INT8 TensorOps matmul, and fused dequantization — to unlock this dormant hardware capability. In our current experiments on Qwen3-VL-2B (M5 Pro), Cider achieves 1.15–1.21× prefill speedup over a W8A16 baseline, while preserving comparable language-model perplexity in our Llama-3-8B check (PPL Δ < 0.01).
|
||||
|
||||
|
||||
## Background: What MLX Does Today
|
||||
|
||||
In the MLX v0.31 code paths we examined, we did not find a public end-to-end INT8 TensorOps inference path.
|
||||
|
||||
More specifically:
|
||||
|
||||
- **Standard matmul** (`steel_gemm_fused_nax.metal`): Instantiated for `float16`, `bfloat16`, `float32` only. No `int8` instantiation exists.
|
||||
- **Quantized matmul** (`quantized_nax.h`): Dequantizes packed int{2,3,4,5,6,8} weights → FP16, then uses **FP16 TensorOps**. This is weight-only quantization — activations stay in FP16.
|
||||
|
||||
This does **not** mean INT8 TensorOps are impossible in MLX. At the shader level, the underlying GEMM template structure appears flexible enough to support INT8 instantiations. The missing piece is a complete and exposed W8A8 execution pipeline.
|
||||
|
||||
|
||||
## The Benchmark
|
||||
|
||||
We wrote a standalone INT8 GEMM kernel in MLX's NAX style, compiled via `mx.fast.metal_kernel`, using the same:
|
||||
- `matmul2d<int8_t, int8_t, int32_t>` TensorOps instruction
|
||||
- BM=128, BN=128, BK=512 tile configuration
|
||||
- NAXFrag cooperative tensor layout
|
||||
- Device memory direct load (no threadgroup staging)
|
||||
|
||||
This is **not** a hypothetical comparison — the MLX-style kernel compiles, runs, and produces bit-exact results (verified: `max_diff=0`).
|
||||
|
||||
### Three-Way Comparison (Apple M5 Pro, Warmup=30, Repeat=100, 3 runs best-of)
|
||||
|
||||
**Square shapes (K=4096, N=4096)**:
|
||||
|
||||
| M | Cider Full Pipeline | Cider Raw INT8 | MLX-style INT8 | Full/Raw | Raw/MLX |
|
||||
|---|-------------------|---------------|----------------|----------|---------|
|
||||
| 16 | 0.24ms | 0.21ms | 0.25ms | 1.12x | 0.84x |
|
||||
| 64 | 0.25ms | 0.24ms | 0.26ms | 1.04x | 0.95x |
|
||||
| 128 | 0.28ms | 0.27ms | 0.26ms | 1.03x | 1.02x |
|
||||
| 256 | 0.40ms | 0.40ms | 0.37ms | 1.01x | 1.08x |
|
||||
| 512 | 0.59ms | 0.54ms | 0.52ms | 1.09x | 1.04x |
|
||||
| 1024 | 0.94ms | 0.90ms | 0.86ms | 1.05x | 1.05x |
|
||||
| 2048 | 1.68ms | 1.55ms | 1.49ms | 1.08x | 1.04x |
|
||||
|
||||
**MLP shapes (Llama-3 8B dimensions)**:
|
||||
|
||||
| M | K | N | Cider Full | Cider Raw | MLX-style | Full/Raw | Raw/MLX |
|
||||
|---|---|---|-----------|-----------|-----------|----------|---------|
|
||||
| 128 | 3584 | 18944 | 0.73ms | 0.74ms | 0.70ms | 0.99x | 1.06x |
|
||||
| 512 | 3584 | 18944 | 1.66ms | 1.65ms | 1.55ms | 1.00x | 1.07x |
|
||||
| 128 | 18944 | 3584 | 0.83ms | 0.78ms | 0.75ms | 1.06x | 1.05x |
|
||||
| 512 | 18944 | 3584 | 1.87ms | 1.81ms | 1.68ms | 1.04x | 1.07x |
|
||||
|
||||
Where:
|
||||
- **Cider Full Pipeline** = `perchannel_linear()`: FP16→INT8 quantize + INT8 matmul + fused INT32→FP16 dequant
|
||||
- **Cider Raw INT8** = `int8_matmul_int32()`: Pure INT8×INT8→INT32 (no quantize, no dequant)
|
||||
- **MLX-style INT8** = Standalone NAX kernel via `mx.fast.metal_kernel`: Same pure INT8→INT32
|
||||
|
||||
## What the Numbers Suggest
|
||||
|
||||
### 1. Raw INT8 TensorOps throughput is comparable
|
||||
|
||||
Across the tested shapes, Cider Raw INT8 and the standalone MLX-style INT8 kernel deliver similar throughput. This is consistent with the fact that both rely on the same underlying `matmul2d` TensorOps instruction with similar tile configurations.
|
||||
|
||||
|
||||
|
||||
### 2. The main cost of a usable W8A8 path is pipeline integration
|
||||
|
||||
Compared with raw INT8 matmul, the full Cider path additionally includes:
|
||||
- activation quantization,
|
||||
- scale handling,
|
||||
- and dequantization back to the desired output dtype.
|
||||
|
||||
These steps introduce overhead, but they are also what make W8A8 usable in practice.
|
||||
|
||||
### 3. Cider's main value is exposing an end-to-end W8A8 path
|
||||
Our results suggest that the key contribution of Cider is not a claim of unique raw GEMM throughput, but the availability of an integrated W8A8 execution path on Apple Silicon.
|
||||
|
||||
|
||||
## Why Cider Uses a Custom Kernel Path
|
||||
|
||||
If a standalone MLX-style INT8 kernel can achieve similar raw throughput, why does Cider still implement its own kernel path?
|
||||
|
||||
|
||||
### 1. To support dequantization as part of the usable W8A8 path
|
||||
|
||||
A practical W8A8 inference path needs more than raw INT32 accumulation. It also needs output scaling and dtype conversion in a form that fits the surrounding runtime.
|
||||
|
||||
|
||||
### 2. To keep quantization and INT8 matmul in one integrated execution path
|
||||
|
||||
Cider packages activation quantization and INT8 matmul into a custom MLX primitive, which is more practical than treating them as disconnected experimental kernels.
|
||||
|
||||
|
||||
### 3. To bypass the current lack of a public end-to-end INT8 route in standard MLX dispatch
|
||||
|
||||
Even if INT8 kernels are possible at the shader level, they are not currently exposed as a standard end-to-end W8A8 inference path in the MLX stack we examined.
|
||||
|
||||
|
||||
## Reproducing
|
||||
|
||||
```bash
|
||||
cd /path/to/cider
|
||||
python benchmarks/mlx_native/bench_cider_vs_mlx_int8.py
|
||||
```
|
||||
|
||||
Requires: Apple M5+, MLX ≥ 0.31, Cider installed.
|
||||
|
||||
The MLX-style kernel is compiled at runtime via `mx.fast.metal_kernel` with Metal 4 + MPP headers. No Xcode installation required.
|
||||
|
||||
## Notes
|
||||
|
||||
- All statements in this document are scoped to the MLX version and code paths we examined during this comparison.
|
||||
- The standalone MLX-style INT8 kernel is a raw-kernel benchmark, not a full MLX-native end-to-end W8A8 inference pipeline.
|
||||
- Cider Full Pipeline results include activation quantization and output dequantization overhead, while raw INT8 kernel results do not.
|
||||
@@ -0,0 +1,83 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Bit-exact correctness test for INT8×INT8→INT32 TensorOps kernel.
|
||||
|
||||
Pure integer matmul must produce EXACT results — no floating-point
|
||||
tolerance, no cosine similarity. Every element must match bit-for-bit.
|
||||
|
||||
Usage:
|
||||
cd /path/to/cider
|
||||
python tests/test_bitexact.py
|
||||
"""
|
||||
import sys, os
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
import numpy as np
|
||||
import mlx.core as mx
|
||||
from cider import int8_matmul_int32
|
||||
|
||||
# ── Test shapes ─────────────────────────────────────────────────
|
||||
SHAPES = [
|
||||
# (M, K, N)
|
||||
(1, 2, 2), # minimal
|
||||
(2, 4, 2), # tiny
|
||||
(1, 128, 128), # single row
|
||||
(1, 4096, 4096), # decode-size
|
||||
(4, 2048, 2048), # small batch
|
||||
(16, 4096, 4096), # medium
|
||||
(32, 4096, 4096), # tile boundary (small→large)
|
||||
(64, 4096, 4096), # large tile
|
||||
(128, 4096, 4096), # full large tile
|
||||
(16, 4096, 8192), # wide N
|
||||
(16, 8192, 4096), # wide K
|
||||
(64, 2048, 4096), # mixed
|
||||
(7, 513, 1025), # non-aligned (prime-ish)
|
||||
(33, 4097, 4095), # non-aligned near 4096
|
||||
]
|
||||
|
||||
|
||||
def test_bitexact():
|
||||
print("=" * 70)
|
||||
print("INT8×INT8→INT32 Bit-Exact Tests")
|
||||
print("=" * 70)
|
||||
passed, failed = 0, 0
|
||||
|
||||
for M, K, N in SHAPES:
|
||||
np.random.seed(42)
|
||||
a_np = np.random.randint(-128, 128, (M, K), dtype=np.int8)
|
||||
b_np = np.random.randint(-128, 128, (N, K), dtype=np.int8)
|
||||
|
||||
# GPU kernel
|
||||
c_gpu = int8_matmul_int32(mx.array(a_np), mx.array(b_np))
|
||||
mx.eval(c_gpu)
|
||||
c_gpu_np = np.array(c_gpu)
|
||||
|
||||
# Numpy reference (int64 to avoid overflow in accumulation)
|
||||
c_ref = (a_np.astype(np.int64) @ b_np.astype(np.int64).T).astype(np.int32)
|
||||
|
||||
# Bit-exact check
|
||||
match = np.array_equal(c_gpu_np, c_ref)
|
||||
if match:
|
||||
passed += 1
|
||||
tag = "PASS"
|
||||
else:
|
||||
failed += 1
|
||||
diff = (c_gpu_np != c_ref)
|
||||
n_diff = diff.sum()
|
||||
# Show first mismatch
|
||||
idx = np.argwhere(diff)[0]
|
||||
tag = f"FAIL ({n_diff}/{M*N} mismatches, first@{tuple(idx)}: gpu={c_gpu_np[tuple(idx)]} ref={c_ref[tuple(idx)]})"
|
||||
|
||||
print(f" M={M:>4d} K={K:>4d} N={N:>4d} [{tag}]")
|
||||
|
||||
print(f"\nBit-exact: {passed}/{passed+failed} passed")
|
||||
return failed == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ok = test_bitexact()
|
||||
print("\n" + "=" * 70)
|
||||
if ok:
|
||||
print("ALL BIT-EXACT TESTS PASSED ✓")
|
||||
else:
|
||||
print("SOME TESTS FAILED")
|
||||
sys.exit(1)
|
||||
Reference in New Issue
Block a user