wehub-resource-sync
45 lines
1.9 KiB
Python
45 lines
1.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Basic usage example for Cider.
|
|
|
|
Demonstrates W8A8 and W4A8 quantized linear layers.
|
|
All operations return lazy mx.array — evaluated on mx.eval().
|
|
"""
|
|
import numpy as np
|
|
import mlx.core as mx
|
|
from cider import W8A8Linear, W4A8Linear, is_available
|
|
|
|
# Check hardware
|
|
if not is_available():
|
|
print("INT8 TensorOps not available (requires Apple M5+)")
|
|
exit(1)
|
|
print("INT8 TensorOps: available")
|
|
|
|
# ── Prepare weights ──────────────────────────────────────────────
|
|
K, N = 4096, 4096
|
|
W_fp16 = np.random.randn(K, N).astype(np.float16)
|
|
|
|
# ── W8A8 ─────────────────────────────────────────────────────────
|
|
w8a8 = W8A8Linear.from_weights(W_fp16)
|
|
print(f"W8A8: [{w8a8.input_dims}, {w8a8.output_dims}]")
|
|
|
|
x = mx.random.normal((32, K)).astype(mx.float16)
|
|
y = w8a8(x) # lazy — not computed yet
|
|
mx.eval(y) # now the GPU runs
|
|
print(f"W8A8 output: {y.shape}, dtype={y.dtype}")
|
|
|
|
# ── W4A8 ─────────────────────────────────────────────────────────
|
|
w4a8 = W4A8Linear.from_weights(W_fp16)
|
|
print(f"W4A8: [{w4a8.input_dims}, {w4a8.output_dims}]")
|
|
print(f"W4A8 weight storage: {w4a8.packed_weight.nbytes / 1024:.0f} KB "
|
|
f"(vs W8A8: {w8a8.weight.nbytes / 1024:.0f} KB)")
|
|
|
|
y4 = w4a8(x)
|
|
mx.eval(y4)
|
|
print(f"W4A8 output: {y4.shape}, dtype={y4.dtype}")
|
|
|
|
# ── Lazy composition ─────────────────────────────────────────────
|
|
# Multiple ops compose into a single MLX graph, evaluated together
|
|
z = w8a8(x) + w4a8(x) # builds graph, no GPU work yet
|
|
mx.eval(z) # single evaluation
|
|
print(f"Composed output: {z.shape}")
|