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