chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,131 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Step-by-step kernel benchmark. Run each step sequentially with gc."""
|
||||
import sys, os, time, gc
|
||||
import numpy as np
|
||||
sys.path.insert(0, "~/work/cider/cider/lib")
|
||||
import mlx.core as mx
|
||||
import _cider_prim as ext
|
||||
|
||||
DEV = "~/work/cider/dev/step_kernels"
|
||||
STEPS = ["step1", "step2", "step3", "step4", "step5"]
|
||||
WARMUP = 5
|
||||
REPEAT = 50
|
||||
|
||||
# Large tile shapes (M%128==0)
|
||||
LARGE = [
|
||||
(128, 4096, 4096),
|
||||
(256, 4096, 4096),
|
||||
(128, 3584, 18944),
|
||||
(256, 3584, 18944),
|
||||
]
|
||||
|
||||
# Small tile shapes (step3+ only)
|
||||
SMALL = [
|
||||
(1, 4096, 4096),
|
||||
(16, 4096, 4096),
|
||||
(32, 4096, 4096),
|
||||
(64, 4096, 4096),
|
||||
]
|
||||
|
||||
def bench_one(A_mx, B_mx, kernel_dir):
|
||||
"""Warmup + benchmark, return median ms."""
|
||||
for _ in range(WARMUP):
|
||||
mx.eval(ext.int8_matmul_int32(A_mx, B_mx, kernel_dir))
|
||||
times = []
|
||||
for _ in range(REPEAT):
|
||||
t0 = time.perf_counter()
|
||||
mx.eval(ext.int8_matmul_int32(A_mx, B_mx, kernel_dir))
|
||||
t1 = time.perf_counter()
|
||||
times.append((t1 - t0) * 1000)
|
||||
times.sort()
|
||||
return times[len(times) // 2]
|
||||
|
||||
# === Large tile benchmark ===
|
||||
print("=" * 90)
|
||||
print("Part 1: Large Tile (M>=128)")
|
||||
print("=" * 90)
|
||||
|
||||
large_results = {s: {} for s in STEPS}
|
||||
|
||||
for M, K, N in LARGE:
|
||||
key = f"({M},{K},{N})"
|
||||
np.random.seed(42)
|
||||
A_np = np.random.randint(-127, 128, (M, K), dtype=np.int8)
|
||||
B_np = np.random.randint(-127, 128, (K, N), dtype=np.int8)
|
||||
A_mx = mx.array(A_np)
|
||||
B_mx = mx.array(B_np)
|
||||
mx.eval(A_mx, B_mx)
|
||||
del A_np, B_np
|
||||
gc.collect()
|
||||
|
||||
for step in STEPS:
|
||||
kdir = os.path.join(DEV, step)
|
||||
med = bench_one(A_mx, B_mx, kdir)
|
||||
large_results[step][key] = med
|
||||
print(f" {step} {key}: {med:.4f} ms")
|
||||
|
||||
del A_mx, B_mx
|
||||
mx.clear_cache()
|
||||
gc.collect()
|
||||
|
||||
# Print large table
|
||||
print()
|
||||
hdr = f"{'Shape':<22}"
|
||||
for s in STEPS:
|
||||
hdr += f" | {s:>10}"
|
||||
print(hdr)
|
||||
print("-" * 82)
|
||||
for M, K, N in LARGE:
|
||||
key = f"({M},{K},{N})"
|
||||
row = f"{key:<22}"
|
||||
for s in STEPS:
|
||||
v = large_results[s].get(key, -1)
|
||||
row += f" | {v:>10.3f}"
|
||||
print(row)
|
||||
|
||||
# === Small tile benchmark (step3/4/5 only) ===
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("Part 2: Small Tile (M<128, step3/4/5)")
|
||||
print("=" * 60)
|
||||
|
||||
small_steps = ["step3", "step4", "step5"]
|
||||
small_results = {s: {} for s in small_steps}
|
||||
|
||||
for M, K, N in SMALL:
|
||||
key = f"({M},{K},{N})"
|
||||
np.random.seed(42)
|
||||
A_np = np.random.randint(-127, 128, (M, K), dtype=np.int8)
|
||||
B_np = np.random.randint(-127, 128, (K, N), dtype=np.int8)
|
||||
A_mx = mx.array(A_np)
|
||||
B_mx = mx.array(B_np)
|
||||
mx.eval(A_mx, B_mx)
|
||||
del A_np, B_np
|
||||
gc.collect()
|
||||
|
||||
for step in small_steps:
|
||||
kdir = os.path.join(DEV, step)
|
||||
med = bench_one(A_mx, B_mx, kdir)
|
||||
small_results[step][key] = med
|
||||
print(f" {step} {key}: {med:.4f} ms")
|
||||
|
||||
del A_mx, B_mx
|
||||
mx.clear_cache()
|
||||
gc.collect()
|
||||
|
||||
# Print small table
|
||||
print()
|
||||
hdr = f"{'Shape':<22}"
|
||||
for s in small_steps:
|
||||
hdr += f" | {s:>10}"
|
||||
print(hdr)
|
||||
print("-" * 58)
|
||||
for M, K, N in SMALL:
|
||||
key = f"({M},{K},{N})"
|
||||
row = f"{key:<22}"
|
||||
for s in small_steps:
|
||||
v = small_results[s].get(key, -1)
|
||||
row += f" | {v:>10.3f}"
|
||||
print(row)
|
||||
|
||||
print("\nDone!")
|
||||
@@ -0,0 +1,381 @@
|
||||
// ============================================================
|
||||
// Step 1: Naive INT8 TensorOps GEMM with Threadgroup Memory Staging
|
||||
// - BM=128, BN=128, BK=128, SK=32
|
||||
// - Only large tile (WM=4, WN=4, 512 threads)
|
||||
// - Threadgroup memory staging: load A/B tiles to TG memory first
|
||||
// - Swizzle decode included (required by C++ dispatch)
|
||||
// ============================================================
|
||||
|
||||
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
constant constexpr short kElemsPerFrag = 8;
|
||||
constant constexpr short kElemCols = 4;
|
||||
constant constexpr short kElemRowsJump = 8;
|
||||
|
||||
inline 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};
|
||||
}
|
||||
|
||||
// Swizzle decode: convert swizzled grid coords to real tile coords
|
||||
inline void swizzle_decode(uint2 tgid, uint swizzle_log, uint tiles_n,
|
||||
thread uint &tid_y, thread uint &tid_x) {
|
||||
uint tile = 1u << swizzle_log;
|
||||
uint tgid_y_raw = tgid.y * tile + (tgid.x % tile);
|
||||
uint tgid_x_raw = tgid.x / tile;
|
||||
tid_y = tgid_y_raw;
|
||||
tid_x = tgid_x_raw;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void nax_frag_load_tg(thread T *dst, const threadgroup T *src, int ld,
|
||||
short2 sc, short off_m = 0, short off_n = 0) {
|
||||
src += (sc.y + off_m) * ld + (sc.x + off_n);
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++)
|
||||
dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j];
|
||||
}
|
||||
|
||||
inline void nax_frag_store_int32(const thread int32_t *src, device int32_t *dst,
|
||||
int ld, short2 sc, short off_m, short off_n,
|
||||
uint M, uint N, uint m_base, uint n_base) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
uint mi = m_base + sc.y + off_m + i * kElemRowsJump;
|
||||
uint ni = n_base + sc.x + off_n + j;
|
||||
if (mi < M && ni < N)
|
||||
dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] =
|
||||
src[i * kElemCols + j];
|
||||
}
|
||||
}
|
||||
|
||||
inline void nax_frag_store_dequant(const thread int32_t *src, device half *dst,
|
||||
int ld, short2 sc, short off_m, short off_n,
|
||||
uint M, uint N, uint m_base, uint n_base,
|
||||
const device float *scale_a,
|
||||
const device float *scale_w) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
uint mi = m_base + sc.y + off_m + i * kElemRowsJump;
|
||||
uint ni = n_base + sc.x + off_n + j;
|
||||
if (mi < M && ni < N) {
|
||||
float val = float(src[i * kElemCols + j]) * scale_a[mi] * scale_w[ni];
|
||||
dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] =
|
||||
half(val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Step 1: TG-staged GEMM, BM=128 only
|
||||
template <int BM, int BN, int BK, int SK, int WM, int WN>
|
||||
void step1_gemm_int32_impl(
|
||||
const device int8_t *A, const device int8_t *B,
|
||||
device int32_t *C, uint M, uint N, uint K,
|
||||
uint swizzle_log, uint tiles_m, uint tiles_n,
|
||||
threadgroup int8_t *tg_A,
|
||||
threadgroup int8_t *tg_B,
|
||||
uint2 tgid, uint sgid, uint lid, uint tid_in_tg) {
|
||||
|
||||
constexpr int SM = BM / WM; // 32
|
||||
constexpr int SN = BN / WN; // 32
|
||||
constexpr short TM = SM / 16; // 2
|
||||
constexpr short TN = SN / 16; // 2
|
||||
constexpr short TK = SK / 16; // 2
|
||||
constexpr int NUM_THREADS = WM * WN * 32; // 512
|
||||
|
||||
uint tid_y, tid_x;
|
||||
swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x);
|
||||
if (tid_x >= tiles_n || tid_y >= tiles_m) return;
|
||||
|
||||
short2 sc = nax_get_coord(ushort(lid));
|
||||
uint sg_row = sgid / WN;
|
||||
uint sg_col = sgid % WN;
|
||||
uint m_base = tid_y * BM + sg_row * SM;
|
||||
uint n_base = tid_x * BN + sg_col * SN;
|
||||
|
||||
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;
|
||||
|
||||
for (uint k_base = 0; k_base < K; k_base += BK) {
|
||||
// Cooperative load A[BM, BK] → threadgroup
|
||||
uint a_elems = BM * BK;
|
||||
for (uint idx = tid_in_tg; idx < a_elems; idx += NUM_THREADS) {
|
||||
uint row = idx / BK;
|
||||
uint col = idx % BK;
|
||||
uint gm = tid_y * BM + row;
|
||||
uint gk = k_base + col;
|
||||
tg_A[row * BK + col] = (gm < M && gk < K) ? A[gm * K + gk] : int8_t(0);
|
||||
}
|
||||
// Cooperative load B[BK, BN] → threadgroup
|
||||
uint b_elems = BK * BN;
|
||||
for (uint idx = tid_in_tg; idx < b_elems; idx += NUM_THREADS) {
|
||||
uint row = idx / BN;
|
||||
uint col = idx % BN;
|
||||
uint gk = k_base + row;
|
||||
uint gn = tid_x * BN + col;
|
||||
tg_B[row * BN + col] = (gk < K && gn < N) ? B[gk * N + gn] : int8_t(0);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (int kk1 = 0; kk1 < BK; kk1 += SK) {
|
||||
int8_t a_frags[TM][TK][kElemsPerFrag];
|
||||
int8_t b_frags[TK][TN][kElemsPerFrag];
|
||||
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
nax_frag_load_tg(a_frags[mm][kk],
|
||||
tg_A + (sg_row * SM) * BK + kk1,
|
||||
BK, sc, short(mm * 16), short(kk * 16));
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_load_tg(b_frags[kk][nn],
|
||||
tg_B + kk1 * BN + (sg_col * SN),
|
||||
BN, 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];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
device int32_t *D = C + m_base * N + n_base;
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_store_int32(c_frags[mm * TN + nn], D, int(N), sc,
|
||||
short(mm * 16), short(nn * 16), M, N, m_base, n_base);
|
||||
}
|
||||
|
||||
// Dequant version
|
||||
template <int BM, int BN, int BK, int SK, int WM, int WN>
|
||||
void step1_gemm_dequant_impl(
|
||||
const device int8_t *A, const device int8_t *B,
|
||||
device half *C, uint M, uint N, uint K,
|
||||
const device float *scale_a, const device float *scale_w,
|
||||
uint swizzle_log, uint tiles_m, uint tiles_n,
|
||||
threadgroup int8_t *tg_A,
|
||||
threadgroup int8_t *tg_B,
|
||||
uint2 tgid, uint sgid, uint lid, uint tid_in_tg) {
|
||||
|
||||
constexpr int SM = BM / WM;
|
||||
constexpr int SN = BN / WN;
|
||||
constexpr short TM = SM / 16;
|
||||
constexpr short TN = SN / 16;
|
||||
constexpr short TK = SK / 16;
|
||||
constexpr int NUM_THREADS = WM * WN * 32;
|
||||
|
||||
uint tid_y, tid_x;
|
||||
swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x);
|
||||
if (tid_x >= tiles_n || tid_y >= tiles_m) return;
|
||||
|
||||
short2 sc = nax_get_coord(ushort(lid));
|
||||
uint sg_row = sgid / WN;
|
||||
uint sg_col = sgid % WN;
|
||||
uint m_base = tid_y * BM + sg_row * SM;
|
||||
uint n_base = tid_x * BN + sg_col * SN;
|
||||
|
||||
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;
|
||||
|
||||
for (uint k_base = 0; k_base < K; k_base += BK) {
|
||||
uint a_elems = BM * BK;
|
||||
for (uint idx = tid_in_tg; idx < a_elems; idx += NUM_THREADS) {
|
||||
uint row = idx / BK;
|
||||
uint col = idx % BK;
|
||||
uint gm = tid_y * BM + row;
|
||||
uint gk = k_base + col;
|
||||
tg_A[row * BK + col] = (gm < M && gk < K) ? A[gm * K + gk] : int8_t(0);
|
||||
}
|
||||
uint b_elems = BK * BN;
|
||||
for (uint idx = tid_in_tg; idx < b_elems; idx += NUM_THREADS) {
|
||||
uint row = idx / BN;
|
||||
uint col = idx % BN;
|
||||
uint gk = k_base + row;
|
||||
uint gn = tid_x * BN + col;
|
||||
tg_B[row * BN + col] = (gk < K && gn < N) ? B[gk * N + gn] : int8_t(0);
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
for (int kk1 = 0; kk1 < BK; kk1 += SK) {
|
||||
int8_t a_frags[TM][TK][kElemsPerFrag];
|
||||
int8_t b_frags[TK][TN][kElemsPerFrag];
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
nax_frag_load_tg(a_frags[mm][kk],
|
||||
tg_A + (sg_row * SM) * BK + kk1,
|
||||
BK, sc, short(mm * 16), short(kk * 16));
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_load_tg(b_frags[kk][nn],
|
||||
tg_B + kk1 * BN + (sg_col * SN),
|
||||
BN, 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];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
}
|
||||
|
||||
device half *D = C + m_base * N + n_base;
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_store_dequant(c_frags[mm * TN + nn], D, int(N), sc,
|
||||
short(mm * 16), short(nn * 16), M, N, m_base, n_base,
|
||||
scale_a, scale_w);
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Kernel entry points
|
||||
// ============================================================
|
||||
|
||||
kernel void w8a8_matmul_fused_dequant(
|
||||
const device int8_t *A [[buffer(0)]],
|
||||
const device int8_t *B [[buffer(1)]],
|
||||
device half *C [[buffer(2)]],
|
||||
constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]],
|
||||
constant uint &K [[buffer(5)]],
|
||||
const device float *scale_a [[buffer(6)]],
|
||||
const device float *scale_w [[buffer(7)]],
|
||||
constant uint &swizzle_log [[buffer(8)]],
|
||||
constant uint &tiles_m [[buffer(9)]],
|
||||
constant uint &tiles_n [[buffer(10)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]],
|
||||
uint tid_in_tg [[thread_index_in_threadgroup]]) {
|
||||
threadgroup int8_t tg_A[128 * 128];
|
||||
threadgroup int8_t tg_B[128 * 128];
|
||||
step1_gemm_dequant_impl<128, 128, 128, 32, 4, 4>(
|
||||
A, B, C, M, N, K, scale_a, scale_w, swizzle_log, tiles_m, tiles_n,
|
||||
tg_A, tg_B, tgid, sgid, lid, tid_in_tg);
|
||||
}
|
||||
|
||||
kernel void w8a8_matmul_fused_dequant_small(
|
||||
const device int8_t *A [[buffer(0)]],
|
||||
const device int8_t *B [[buffer(1)]],
|
||||
device half *C [[buffer(2)]],
|
||||
constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]],
|
||||
constant uint &K [[buffer(5)]],
|
||||
const device float *scale_a [[buffer(6)]],
|
||||
const device float *scale_w [[buffer(7)]],
|
||||
constant uint &swizzle_log [[buffer(8)]],
|
||||
constant uint &tiles_m [[buffer(9)]],
|
||||
constant uint &tiles_n [[buffer(10)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]],
|
||||
uint tid_in_tg [[thread_index_in_threadgroup]]) {
|
||||
// Step 1 has no small tile — fallback to large tile
|
||||
threadgroup int8_t tg_A[128 * 128];
|
||||
threadgroup int8_t tg_B[128 * 128];
|
||||
step1_gemm_dequant_impl<128, 128, 128, 32, 4, 4>(
|
||||
A, B, C, M, N, K, scale_a, scale_w, swizzle_log, tiles_m, tiles_n,
|
||||
tg_A, tg_B, tgid, sgid, lid, tid_in_tg);
|
||||
}
|
||||
|
||||
kernel void int8_matmul_int32(
|
||||
const device int8_t *A [[buffer(0)]],
|
||||
const device int8_t *B [[buffer(1)]],
|
||||
device int32_t *C [[buffer(2)]],
|
||||
constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]],
|
||||
constant uint &K [[buffer(5)]],
|
||||
constant uint &swizzle_log [[buffer(6)]],
|
||||
constant uint &tiles_m [[buffer(7)]],
|
||||
constant uint &tiles_n [[buffer(8)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]],
|
||||
uint tid_in_tg [[thread_index_in_threadgroup]]) {
|
||||
threadgroup int8_t tg_A[128 * 128];
|
||||
threadgroup int8_t tg_B[128 * 128];
|
||||
step1_gemm_int32_impl<128, 128, 128, 32, 4, 4>(
|
||||
A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n,
|
||||
tg_A, tg_B, tgid, sgid, lid, tid_in_tg);
|
||||
}
|
||||
|
||||
kernel void int8_matmul_int32_small(
|
||||
const device int8_t *A [[buffer(0)]],
|
||||
const device int8_t *B [[buffer(1)]],
|
||||
device int32_t *C [[buffer(2)]],
|
||||
constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]],
|
||||
constant uint &K [[buffer(5)]],
|
||||
constant uint &swizzle_log [[buffer(6)]],
|
||||
constant uint &tiles_m [[buffer(7)]],
|
||||
constant uint &tiles_n [[buffer(8)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]],
|
||||
uint tid_in_tg [[thread_index_in_threadgroup]]) {
|
||||
// Step 1 has no small tile — fallback to large tile
|
||||
threadgroup int8_t tg_A[128 * 128];
|
||||
threadgroup int8_t tg_B[128 * 128];
|
||||
step1_gemm_int32_impl<128, 128, 128, 32, 4, 4>(
|
||||
A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n,
|
||||
tg_A, tg_B, tgid, sgid, lid, tid_in_tg);
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
// ============================================================
|
||||
// Per-token quantization: FP16 → INT8 + float32 scale
|
||||
// Target: Apple M5, Metal 4
|
||||
//
|
||||
// Each threadgroup handles one row (one token).
|
||||
// Threads cooperate to find absmax via simdgroup reduce,
|
||||
// then quantize in parallel.
|
||||
//
|
||||
// Host dispatch:
|
||||
// threadgroup = (min(256, ceil(K/32)*32), 1, 1)
|
||||
// grid = (M, 1, 1)
|
||||
// ============================================================
|
||||
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
kernel void
|
||||
quantize_per_token(const device half *X [[buffer(0)]], // [M, K] FP16 input
|
||||
device int8_t *A [[buffer(1)]], // [M, K] INT8 output
|
||||
device float *scale [[buffer(2)]], // [M] float32 scale
|
||||
constant uint &M [[buffer(3)]],
|
||||
constant uint &K [[buffer(4)]],
|
||||
uint gid [[threadgroup_position_in_grid]], // row index
|
||||
uint lid [[thread_index_in_threadgroup]],
|
||||
uint tg_size [[threads_per_threadgroup]]) {
|
||||
if (gid >= M) {
|
||||
return;
|
||||
}
|
||||
|
||||
const device half *row_in = X + gid * K;
|
||||
device int8_t *row_out = A + gid * K;
|
||||
|
||||
// Step 1: Find local absmax
|
||||
float local_max = 0.0f;
|
||||
for (uint i = lid; i < K; i += tg_size) {
|
||||
float v = abs(float(row_in[i]));
|
||||
local_max = max(local_max, v);
|
||||
}
|
||||
|
||||
// Step 2: Simdgroup reduce max
|
||||
float sg_max = simd_max(local_max);
|
||||
|
||||
// Step 3: Threadgroup reduce across simdgroups via shared memory
|
||||
threadgroup float sg_maxes[8]; // up to 8 simdgroups (256/32)
|
||||
threadgroup float shared_scale;
|
||||
threadgroup float shared_inv_scale;
|
||||
uint sg_id = lid / 32;
|
||||
uint sg_lid = lid % 32;
|
||||
if (sg_lid == 0) {
|
||||
sg_maxes[sg_id] = sg_max;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Final reduce (first simdgroup only)
|
||||
if (sg_id == 0) {
|
||||
float row_max = 0.0f;
|
||||
uint num_sgs = (tg_size + 31) / 32;
|
||||
if (sg_lid < num_sgs) {
|
||||
row_max = sg_maxes[sg_lid];
|
||||
}
|
||||
row_max = simd_max(row_max);
|
||||
|
||||
// Compute and broadcast scale
|
||||
float s = row_max / 255.0f;
|
||||
if (s == 0.0f) {
|
||||
s = 1.0f;
|
||||
}
|
||||
|
||||
if (sg_lid == 0) {
|
||||
shared_scale = s;
|
||||
shared_inv_scale = 1.0f / s;
|
||||
// Store scale to output
|
||||
scale[gid] = s;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Step 4: All threads read broadcasted scale
|
||||
float inv_s = shared_inv_scale;
|
||||
|
||||
// Step 5: Quantize
|
||||
for (uint i = lid; i < K; i += tg_size) {
|
||||
float v = float(row_in[i]) * inv_s;
|
||||
v = clamp(round(v), -128.0f, 127.0f);
|
||||
row_out[i] = int8_t(v);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,285 @@
|
||||
// ============================================================
|
||||
// Step 2: Direct Device Read (No Threadgroup Memory)
|
||||
// - BM=128, BN=128, BK=128, SK=32
|
||||
// - Only large tile (WM=4, WN=4, 512 threads)
|
||||
// - Direct device memory read (no TG staging)
|
||||
// - Swizzle decode included (required by C++ dispatch)
|
||||
// ============================================================
|
||||
|
||||
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
constant constexpr short kElemsPerFrag = 8;
|
||||
constant constexpr short kElemCols = 4;
|
||||
constant constexpr short kElemRowsJump = 8;
|
||||
|
||||
inline 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};
|
||||
}
|
||||
|
||||
inline void swizzle_decode(uint2 tgid, uint swizzle_log, uint tiles_n,
|
||||
thread uint &tid_y, thread uint &tid_x) {
|
||||
uint tile = 1u << swizzle_log;
|
||||
tid_y = tgid.y * tile + (tgid.x % tile);
|
||||
tid_x = tgid.x / tile;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void nax_frag_load(thread T *dst, const device T *src, int ld, short2 sc,
|
||||
short off_m = 0, short off_n = 0) {
|
||||
src += (sc.y + off_m) * ld + (sc.x + off_n);
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++)
|
||||
dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j];
|
||||
}
|
||||
|
||||
inline void nax_frag_store_int32(const thread int32_t *src, device int32_t *dst,
|
||||
int ld, short2 sc, short off_m, short off_n,
|
||||
uint M, uint N, uint m_base, uint n_base) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
uint mi = m_base + sc.y + off_m + i * kElemRowsJump;
|
||||
uint ni = n_base + sc.x + off_n + j;
|
||||
if (mi < M && ni < N)
|
||||
dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] =
|
||||
src[i * kElemCols + j];
|
||||
}
|
||||
}
|
||||
|
||||
inline void nax_frag_store_dequant(const thread int32_t *src, device half *dst,
|
||||
int ld, short2 sc, short off_m, short off_n,
|
||||
uint M, uint N, uint m_base, uint n_base,
|
||||
const device float *scale_a,
|
||||
const device float *scale_w) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
uint mi = m_base + sc.y + off_m + i * kElemRowsJump;
|
||||
uint ni = n_base + sc.x + off_n + j;
|
||||
if (mi < M && ni < N) {
|
||||
float val = float(src[i * kElemCols + j]) * scale_a[mi] * scale_w[ni];
|
||||
dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] =
|
||||
half(val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Direct device read, BK=128, large tile only
|
||||
template <int BM, int BN, int BK, int SK, int WM, int WN>
|
||||
void step2_gemm_int32_impl(
|
||||
const device int8_t *A, const device int8_t *B,
|
||||
device int32_t *C, uint M, uint N, uint K,
|
||||
uint swizzle_log, uint tiles_m, uint tiles_n,
|
||||
uint2 tgid, uint sgid, uint lid) {
|
||||
|
||||
constexpr int SM = BM / WM;
|
||||
constexpr int SN = BN / WN;
|
||||
constexpr short TM = SM / 16;
|
||||
constexpr short TN = SN / 16;
|
||||
constexpr short TK = SK / 16;
|
||||
|
||||
uint tid_y, tid_x;
|
||||
swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x);
|
||||
if (tid_x >= tiles_n || tid_y >= tiles_m) return;
|
||||
|
||||
short2 sc = nax_get_coord(ushort(lid));
|
||||
uint sg_row = sgid / WN;
|
||||
uint sg_col = sgid % WN;
|
||||
uint m_base = tid_y * BM + sg_row * SM;
|
||||
uint n_base = tid_x * BN + sg_col * SN;
|
||||
|
||||
const device int8_t *sg_A = A + m_base * K;
|
||||
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) / 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++)
|
||||
nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc,
|
||||
short(mm * 16), short(kk * 16));
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), 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;
|
||||
}
|
||||
|
||||
// Handle K remainder
|
||||
int rem_k = int(K) - gemm_k_iters * BK;
|
||||
if (rem_k > 0) {
|
||||
for (int kk1 = 0; kk1 < rem_k; kk1 += SK) {
|
||||
int actual_sk = min(SK, rem_k - kk1);
|
||||
int8_t a_frags[TM][TK][kElemsPerFrag];
|
||||
int8_t b_frags[TK][TN][kElemsPerFrag];
|
||||
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short kk = 0; kk < TK; kk++) {
|
||||
const device int8_t *ptr = sg_A + kk1 + (sc.y + mm * 16) * int(K) + (sc.x + kk * 16);
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
int ki = kk * 16 + int(sc.x) + j;
|
||||
if (ki < actual_sk) {
|
||||
int row = int(sc.y) + mm * 16 + i * kElemRowsJump;
|
||||
a_frags[mm][kk][i * kElemCols + j] = sg_A[kk1 + row * int(K) + ki];
|
||||
} else {
|
||||
a_frags[mm][kk][i * kElemCols + j] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
int ki = kk * 16 + int(sc.y) + i * kElemRowsJump;
|
||||
if (ki < actual_sk) {
|
||||
int col = nn * 16 + int(sc.x) + j;
|
||||
b_frags[kk][nn][i * kElemCols + j] = sg_B[(kk1 + ki) * int(N) + col];
|
||||
} else {
|
||||
b_frags[kk][nn][i * kElemCols + j] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
device int32_t *D = C + m_base * N + n_base;
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_store_int32(c_frags[mm * TN + nn], D, int(N), sc,
|
||||
short(mm * 16), short(nn * 16), M, N, m_base, n_base);
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Entry points
|
||||
// ============================================================
|
||||
|
||||
kernel void w8a8_matmul_fused_dequant(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device half *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
const device float *scale_a [[buffer(6)]],
|
||||
const device float *scale_w [[buffer(7)]],
|
||||
constant uint &swizzle_log [[buffer(8)]],
|
||||
constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
step2_gemm_int32_impl<128, 128, 128, 32, 4, 4>(
|
||||
A, B, (device int32_t *)C, M, N, K, swizzle_log, tiles_m, tiles_n,
|
||||
tgid, sgid, lid);
|
||||
// NOTE: for dequant we'd need a separate impl, but for benchmark
|
||||
// we only test int32 path. This entry exists for compilation only.
|
||||
}
|
||||
|
||||
kernel void w8a8_matmul_fused_dequant_small(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device half *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
const device float *scale_a [[buffer(6)]],
|
||||
const device float *scale_w [[buffer(7)]],
|
||||
constant uint &swizzle_log [[buffer(8)]],
|
||||
constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
step2_gemm_int32_impl<128, 128, 128, 32, 4, 4>(
|
||||
A, B, (device int32_t *)C, M, N, K, swizzle_log, tiles_m, tiles_n,
|
||||
tgid, sgid, lid);
|
||||
}
|
||||
|
||||
kernel void int8_matmul_int32(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
constant uint &swizzle_log [[buffer(6)]],
|
||||
constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
step2_gemm_int32_impl<128, 128, 128, 32, 4, 4>(
|
||||
A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid);
|
||||
}
|
||||
|
||||
kernel void int8_matmul_int32_small(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
constant uint &swizzle_log [[buffer(6)]],
|
||||
constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
// Step 2 has no small tile — fallback to large
|
||||
step2_gemm_int32_impl<128, 128, 128, 32, 4, 4>(
|
||||
A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid);
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
// ============================================================
|
||||
// Per-token quantization: FP16 → INT8 + float32 scale
|
||||
// Target: Apple M5, Metal 4
|
||||
//
|
||||
// Each threadgroup handles one row (one token).
|
||||
// Threads cooperate to find absmax via simdgroup reduce,
|
||||
// then quantize in parallel.
|
||||
//
|
||||
// Host dispatch:
|
||||
// threadgroup = (min(256, ceil(K/32)*32), 1, 1)
|
||||
// grid = (M, 1, 1)
|
||||
// ============================================================
|
||||
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
kernel void
|
||||
quantize_per_token(const device half *X [[buffer(0)]], // [M, K] FP16 input
|
||||
device int8_t *A [[buffer(1)]], // [M, K] INT8 output
|
||||
device float *scale [[buffer(2)]], // [M] float32 scale
|
||||
constant uint &M [[buffer(3)]],
|
||||
constant uint &K [[buffer(4)]],
|
||||
uint gid [[threadgroup_position_in_grid]], // row index
|
||||
uint lid [[thread_index_in_threadgroup]],
|
||||
uint tg_size [[threads_per_threadgroup]]) {
|
||||
if (gid >= M) {
|
||||
return;
|
||||
}
|
||||
|
||||
const device half *row_in = X + gid * K;
|
||||
device int8_t *row_out = A + gid * K;
|
||||
|
||||
// Step 1: Find local absmax
|
||||
float local_max = 0.0f;
|
||||
for (uint i = lid; i < K; i += tg_size) {
|
||||
float v = abs(float(row_in[i]));
|
||||
local_max = max(local_max, v);
|
||||
}
|
||||
|
||||
// Step 2: Simdgroup reduce max
|
||||
float sg_max = simd_max(local_max);
|
||||
|
||||
// Step 3: Threadgroup reduce across simdgroups via shared memory
|
||||
threadgroup float sg_maxes[8]; // up to 8 simdgroups (256/32)
|
||||
threadgroup float shared_scale;
|
||||
threadgroup float shared_inv_scale;
|
||||
uint sg_id = lid / 32;
|
||||
uint sg_lid = lid % 32;
|
||||
if (sg_lid == 0) {
|
||||
sg_maxes[sg_id] = sg_max;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Final reduce (first simdgroup only)
|
||||
if (sg_id == 0) {
|
||||
float row_max = 0.0f;
|
||||
uint num_sgs = (tg_size + 31) / 32;
|
||||
if (sg_lid < num_sgs) {
|
||||
row_max = sg_maxes[sg_lid];
|
||||
}
|
||||
row_max = simd_max(row_max);
|
||||
|
||||
// Compute and broadcast scale
|
||||
float s = row_max / 255.0f;
|
||||
if (s == 0.0f) {
|
||||
s = 1.0f;
|
||||
}
|
||||
|
||||
if (sg_lid == 0) {
|
||||
shared_scale = s;
|
||||
shared_inv_scale = 1.0f / s;
|
||||
// Store scale to output
|
||||
scale[gid] = s;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Step 4: All threads read broadcasted scale
|
||||
float inv_s = shared_inv_scale;
|
||||
|
||||
// Step 5: Quantize
|
||||
for (uint i = lid; i < K; i += tg_size) {
|
||||
float v = float(row_in[i]) * inv_s;
|
||||
v = clamp(round(v), -128.0f, 127.0f);
|
||||
row_out[i] = int8_t(v);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,431 @@
|
||||
// ============================================================
|
||||
// Step 3: Multi-tile Dispatch (Large BM=128 + Small BM=32)
|
||||
// - BK=128, SK=32
|
||||
// - Direct device read
|
||||
// - Swizzle decode (required by C++ dispatch)
|
||||
// - Large: BM=128, BN=128, WM=4, WN=4 (512 threads)
|
||||
// - Small: BM=32, BN=128, WM=1, WN=4 (128 threads)
|
||||
// ============================================================
|
||||
|
||||
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
constant constexpr short kElemsPerFrag = 8;
|
||||
constant constexpr short kElemCols = 4;
|
||||
constant constexpr short kElemRowsJump = 8;
|
||||
|
||||
inline 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};
|
||||
}
|
||||
|
||||
inline void swizzle_decode(uint2 tgid, uint swizzle_log, uint tiles_n,
|
||||
thread uint &tid_y, thread uint &tid_x) {
|
||||
uint tile = 1u << swizzle_log;
|
||||
tid_y = tgid.y * tile + (tgid.x % tile);
|
||||
tid_x = tgid.x / tile;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void nax_frag_load(thread T *dst, const device T *src, int ld, short2 sc,
|
||||
short off_m = 0, short off_n = 0) {
|
||||
src += (sc.y + off_m) * ld + (sc.x + off_n);
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++)
|
||||
dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j];
|
||||
}
|
||||
|
||||
inline void nax_frag_store_int32(const thread int32_t *src, device int32_t *dst,
|
||||
int ld, short2 sc, short off_m, short off_n,
|
||||
uint M, uint N, uint m_base, uint n_base) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
uint mi = m_base + sc.y + off_m + i * kElemRowsJump;
|
||||
uint ni = n_base + sc.x + off_n + j;
|
||||
if (mi < M && ni < N)
|
||||
dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] =
|
||||
src[i * kElemCols + j];
|
||||
}
|
||||
}
|
||||
|
||||
inline void nax_frag_store_dequant(const thread int32_t *src, device half *dst,
|
||||
int ld, short2 sc, short off_m, short off_n,
|
||||
uint M, uint N, uint m_base, uint n_base,
|
||||
const device float *scale_a,
|
||||
const device float *scale_w) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
uint mi = m_base + sc.y + off_m + i * kElemRowsJump;
|
||||
uint ni = n_base + sc.x + off_n + j;
|
||||
if (mi < M && ni < N) {
|
||||
float val = float(src[i * kElemCols + j]) * scale_a[mi] * scale_w[ni];
|
||||
dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] =
|
||||
half(val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Generic GEMM with swizzle decode
|
||||
template <int BM, int BN, int BK, int SK, int WM, int WN>
|
||||
void step3_gemm_int32_impl(
|
||||
const device int8_t *A, const device int8_t *B,
|
||||
device int32_t *C, uint M, uint N, uint K,
|
||||
uint swizzle_log, uint tiles_m, uint tiles_n,
|
||||
uint2 tgid, uint sgid, uint lid) {
|
||||
|
||||
constexpr int SM = BM / WM;
|
||||
constexpr int SN = BN / WN;
|
||||
constexpr short TM = SM / 16;
|
||||
constexpr short TN = SN / 16;
|
||||
constexpr short TK = SK / 16;
|
||||
|
||||
uint tid_y, tid_x;
|
||||
swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x);
|
||||
if (tid_x >= tiles_n || tid_y >= tiles_m) return;
|
||||
|
||||
short2 sc = nax_get_coord(ushort(lid));
|
||||
uint sg_row = sgid / WN;
|
||||
uint sg_col = sgid % WN;
|
||||
uint m_base = tid_y * BM + sg_row * SM;
|
||||
uint n_base = tid_x * BN + sg_col * SN;
|
||||
|
||||
const device int8_t *sg_A = A + m_base * K;
|
||||
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) / 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++)
|
||||
nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc,
|
||||
short(mm * 16), short(kk * 16));
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), 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;
|
||||
}
|
||||
|
||||
// Remainder K
|
||||
int rem_k = int(K) - gemm_k_iters * BK;
|
||||
if (rem_k > 0) {
|
||||
for (int kk1 = 0; kk1 < rem_k; kk1 += SK) {
|
||||
int actual_sk = min(SK, rem_k - kk1);
|
||||
int8_t a_frags[TM][TK][kElemsPerFrag];
|
||||
int8_t b_frags[TK][TN][kElemsPerFrag];
|
||||
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short kk = 0; kk < TK; kk++) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
int ki = kk * 16 + int(sc.x) + j;
|
||||
if (ki < actual_sk) {
|
||||
int row = int(sc.y) + mm * 16 + i * kElemRowsJump;
|
||||
a_frags[mm][kk][i * kElemCols + j] = sg_A[kk1 + row * int(K) + ki];
|
||||
} else {
|
||||
a_frags[mm][kk][i * kElemCols + j] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
int ki = kk * 16 + int(sc.y) + i * kElemRowsJump;
|
||||
if (ki < actual_sk) {
|
||||
int col = nn * 16 + int(sc.x) + j;
|
||||
b_frags[kk][nn][i * kElemCols + j] = sg_B[(kk1 + ki) * int(N) + col];
|
||||
} else {
|
||||
b_frags[kk][nn][i * kElemCols + j] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
device int32_t *D = C + m_base * N + n_base;
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_store_int32(c_frags[mm * TN + nn], D, int(N), sc,
|
||||
short(mm * 16), short(nn * 16), M, N, m_base, n_base);
|
||||
}
|
||||
|
||||
// Dequant version
|
||||
template <int BM, int BN, int BK, int SK, int WM, int WN>
|
||||
void step3_gemm_dequant_impl(
|
||||
const device int8_t *A, const device int8_t *B,
|
||||
device half *C, uint M, uint N, uint K,
|
||||
const device float *scale_a, const device float *scale_w,
|
||||
uint swizzle_log, uint tiles_m, uint tiles_n,
|
||||
uint2 tgid, uint sgid, uint lid) {
|
||||
|
||||
constexpr int SM = BM / WM;
|
||||
constexpr int SN = BN / WN;
|
||||
constexpr short TM = SM / 16;
|
||||
constexpr short TN = SN / 16;
|
||||
constexpr short TK = SK / 16;
|
||||
|
||||
uint tid_y, tid_x;
|
||||
swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x);
|
||||
if (tid_x >= tiles_n || tid_y >= tiles_m) return;
|
||||
|
||||
short2 sc = nax_get_coord(ushort(lid));
|
||||
uint sg_row = sgid / WN;
|
||||
uint sg_col = sgid % WN;
|
||||
uint m_base = tid_y * BM + sg_row * SM;
|
||||
uint n_base = tid_x * BN + sg_col * SN;
|
||||
|
||||
const device int8_t *sg_A = A + m_base * K;
|
||||
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) / 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++)
|
||||
nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc,
|
||||
short(mm * 16), short(kk * 16));
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), 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;
|
||||
}
|
||||
|
||||
int rem_k = int(K) - gemm_k_iters * BK;
|
||||
if (rem_k > 0) {
|
||||
for (int kk1 = 0; kk1 < rem_k; kk1 += SK) {
|
||||
int actual_sk = min(SK, rem_k - kk1);
|
||||
int8_t a_frags[TM][TK][kElemsPerFrag];
|
||||
int8_t b_frags[TK][TN][kElemsPerFrag];
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short kk = 0; kk < TK; kk++) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
int ki = kk * 16 + int(sc.x) + j;
|
||||
if (ki < actual_sk) {
|
||||
int row = int(sc.y) + mm * 16 + i * kElemRowsJump;
|
||||
a_frags[mm][kk][i * kElemCols + j] = sg_A[kk1 + row * int(K) + ki];
|
||||
} else {
|
||||
a_frags[mm][kk][i * kElemCols + j] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
int ki = kk * 16 + int(sc.y) + i * kElemRowsJump;
|
||||
if (ki < actual_sk) {
|
||||
int col = nn * 16 + int(sc.x) + j;
|
||||
b_frags[kk][nn][i * kElemCols + j] = sg_B[(kk1 + ki) * int(N) + col];
|
||||
} else {
|
||||
b_frags[kk][nn][i * kElemCols + j] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
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];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
device half *D = C + m_base * N + n_base;
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_store_dequant(c_frags[mm * TN + nn], D, int(N), sc,
|
||||
short(mm * 16), short(nn * 16), M, N, m_base, n_base,
|
||||
scale_a, scale_w);
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Entry points — Large tile
|
||||
// ============================================================
|
||||
kernel void w8a8_matmul_fused_dequant(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device half *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
const device float *scale_a [[buffer(6)]],
|
||||
const device float *scale_w [[buffer(7)]],
|
||||
constant uint &swizzle_log [[buffer(8)]],
|
||||
constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
step3_gemm_dequant_impl<128, 128, 128, 32, 4, 4>(
|
||||
A, B, C, M, N, K, scale_a, scale_w, swizzle_log, tiles_m, tiles_n,
|
||||
tgid, sgid, lid);
|
||||
}
|
||||
|
||||
// Small tile entry point (NEW in Step 3!)
|
||||
kernel void w8a8_matmul_fused_dequant_small(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device half *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
const device float *scale_a [[buffer(6)]],
|
||||
const device float *scale_w [[buffer(7)]],
|
||||
constant uint &swizzle_log [[buffer(8)]],
|
||||
constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
step3_gemm_dequant_impl<32, 128, 128, 32, 1, 4>(
|
||||
A, B, C, M, N, K, scale_a, scale_w, swizzle_log, tiles_m, tiles_n,
|
||||
tgid, sgid, lid);
|
||||
}
|
||||
|
||||
// INT32 entry points
|
||||
kernel void int8_matmul_int32(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
constant uint &swizzle_log [[buffer(6)]],
|
||||
constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
step3_gemm_int32_impl<128, 128, 128, 32, 4, 4>(
|
||||
A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid);
|
||||
}
|
||||
|
||||
kernel void int8_matmul_int32_small(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
constant uint &swizzle_log [[buffer(6)]],
|
||||
constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
step3_gemm_int32_impl<32, 128, 128, 32, 1, 4>(
|
||||
A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid);
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
// ============================================================
|
||||
// Per-token quantization: FP16 → INT8 + float32 scale
|
||||
// Target: Apple M5, Metal 4
|
||||
//
|
||||
// Each threadgroup handles one row (one token).
|
||||
// Threads cooperate to find absmax via simdgroup reduce,
|
||||
// then quantize in parallel.
|
||||
//
|
||||
// Host dispatch:
|
||||
// threadgroup = (min(256, ceil(K/32)*32), 1, 1)
|
||||
// grid = (M, 1, 1)
|
||||
// ============================================================
|
||||
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
kernel void
|
||||
quantize_per_token(const device half *X [[buffer(0)]], // [M, K] FP16 input
|
||||
device int8_t *A [[buffer(1)]], // [M, K] INT8 output
|
||||
device float *scale [[buffer(2)]], // [M] float32 scale
|
||||
constant uint &M [[buffer(3)]],
|
||||
constant uint &K [[buffer(4)]],
|
||||
uint gid [[threadgroup_position_in_grid]], // row index
|
||||
uint lid [[thread_index_in_threadgroup]],
|
||||
uint tg_size [[threads_per_threadgroup]]) {
|
||||
if (gid >= M) {
|
||||
return;
|
||||
}
|
||||
|
||||
const device half *row_in = X + gid * K;
|
||||
device int8_t *row_out = A + gid * K;
|
||||
|
||||
// Step 1: Find local absmax
|
||||
float local_max = 0.0f;
|
||||
for (uint i = lid; i < K; i += tg_size) {
|
||||
float v = abs(float(row_in[i]));
|
||||
local_max = max(local_max, v);
|
||||
}
|
||||
|
||||
// Step 2: Simdgroup reduce max
|
||||
float sg_max = simd_max(local_max);
|
||||
|
||||
// Step 3: Threadgroup reduce across simdgroups via shared memory
|
||||
threadgroup float sg_maxes[8]; // up to 8 simdgroups (256/32)
|
||||
threadgroup float shared_scale;
|
||||
threadgroup float shared_inv_scale;
|
||||
uint sg_id = lid / 32;
|
||||
uint sg_lid = lid % 32;
|
||||
if (sg_lid == 0) {
|
||||
sg_maxes[sg_id] = sg_max;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Final reduce (first simdgroup only)
|
||||
if (sg_id == 0) {
|
||||
float row_max = 0.0f;
|
||||
uint num_sgs = (tg_size + 31) / 32;
|
||||
if (sg_lid < num_sgs) {
|
||||
row_max = sg_maxes[sg_lid];
|
||||
}
|
||||
row_max = simd_max(row_max);
|
||||
|
||||
// Compute and broadcast scale
|
||||
float s = row_max / 255.0f;
|
||||
if (s == 0.0f) {
|
||||
s = 1.0f;
|
||||
}
|
||||
|
||||
if (sg_lid == 0) {
|
||||
shared_scale = s;
|
||||
shared_inv_scale = 1.0f / s;
|
||||
// Store scale to output
|
||||
scale[gid] = s;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Step 4: All threads read broadcasted scale
|
||||
float inv_s = shared_inv_scale;
|
||||
|
||||
// Step 5: Quantize
|
||||
for (uint i = lid; i < K; i += tg_size) {
|
||||
float v = float(row_in[i]) * inv_s;
|
||||
v = clamp(round(v), -128.0f, 127.0f);
|
||||
row_out[i] = int8_t(v);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,431 @@
|
||||
// ============================================================
|
||||
// Step 4: Deep K-loop (BK=512) Dispatch (Large BM=128 + Small BM=32)
|
||||
// - BK=512, SK=32
|
||||
// - Direct device read
|
||||
// - Swizzle decode (required by C++ dispatch)
|
||||
// - Large: BM=128, BN=128, WM=4, WN=4 (512 threads)
|
||||
// - Small: BM=32, BN=128, WM=1, WN=4 (128 threads)
|
||||
// ============================================================
|
||||
|
||||
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
constant constexpr short kElemsPerFrag = 8;
|
||||
constant constexpr short kElemCols = 4;
|
||||
constant constexpr short kElemRowsJump = 8;
|
||||
|
||||
inline 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};
|
||||
}
|
||||
|
||||
inline void swizzle_decode(uint2 tgid, uint swizzle_log, uint tiles_n,
|
||||
thread uint &tid_y, thread uint &tid_x) {
|
||||
uint tile = 1u << swizzle_log;
|
||||
tid_y = tgid.y * tile + (tgid.x % tile);
|
||||
tid_x = tgid.x / tile;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void nax_frag_load(thread T *dst, const device T *src, int ld, short2 sc,
|
||||
short off_m = 0, short off_n = 0) {
|
||||
src += (sc.y + off_m) * ld + (sc.x + off_n);
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++)
|
||||
dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j];
|
||||
}
|
||||
|
||||
inline void nax_frag_store_int32(const thread int32_t *src, device int32_t *dst,
|
||||
int ld, short2 sc, short off_m, short off_n,
|
||||
uint M, uint N, uint m_base, uint n_base) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
uint mi = m_base + sc.y + off_m + i * kElemRowsJump;
|
||||
uint ni = n_base + sc.x + off_n + j;
|
||||
if (mi < M && ni < N)
|
||||
dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] =
|
||||
src[i * kElemCols + j];
|
||||
}
|
||||
}
|
||||
|
||||
inline void nax_frag_store_dequant(const thread int32_t *src, device half *dst,
|
||||
int ld, short2 sc, short off_m, short off_n,
|
||||
uint M, uint N, uint m_base, uint n_base,
|
||||
const device float *scale_a,
|
||||
const device float *scale_w) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
uint mi = m_base + sc.y + off_m + i * kElemRowsJump;
|
||||
uint ni = n_base + sc.x + off_n + j;
|
||||
if (mi < M && ni < N) {
|
||||
float val = float(src[i * kElemCols + j]) * scale_a[mi] * scale_w[ni];
|
||||
dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] =
|
||||
half(val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Generic GEMM with swizzle decode
|
||||
template <int BM, int BN, int BK, int SK, int WM, int WN>
|
||||
void step4_gemm_int32_impl(
|
||||
const device int8_t *A, const device int8_t *B,
|
||||
device int32_t *C, uint M, uint N, uint K,
|
||||
uint swizzle_log, uint tiles_m, uint tiles_n,
|
||||
uint2 tgid, uint sgid, uint lid) {
|
||||
|
||||
constexpr int SM = BM / WM;
|
||||
constexpr int SN = BN / WN;
|
||||
constexpr short TM = SM / 16;
|
||||
constexpr short TN = SN / 16;
|
||||
constexpr short TK = SK / 16;
|
||||
|
||||
uint tid_y, tid_x;
|
||||
swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x);
|
||||
if (tid_x >= tiles_n || tid_y >= tiles_m) return;
|
||||
|
||||
short2 sc = nax_get_coord(ushort(lid));
|
||||
uint sg_row = sgid / WN;
|
||||
uint sg_col = sgid % WN;
|
||||
uint m_base = tid_y * BM + sg_row * SM;
|
||||
uint n_base = tid_x * BN + sg_col * SN;
|
||||
|
||||
const device int8_t *sg_A = A + m_base * K;
|
||||
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) / 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++)
|
||||
nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc,
|
||||
short(mm * 16), short(kk * 16));
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), 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;
|
||||
}
|
||||
|
||||
// Remainder K
|
||||
int rem_k = int(K) - gemm_k_iters * BK;
|
||||
if (rem_k > 0) {
|
||||
for (int kk1 = 0; kk1 < rem_k; kk1 += SK) {
|
||||
int actual_sk = min(SK, rem_k - kk1);
|
||||
int8_t a_frags[TM][TK][kElemsPerFrag];
|
||||
int8_t b_frags[TK][TN][kElemsPerFrag];
|
||||
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short kk = 0; kk < TK; kk++) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
int ki = kk * 16 + int(sc.x) + j;
|
||||
if (ki < actual_sk) {
|
||||
int row = int(sc.y) + mm * 16 + i * kElemRowsJump;
|
||||
a_frags[mm][kk][i * kElemCols + j] = sg_A[kk1 + row * int(K) + ki];
|
||||
} else {
|
||||
a_frags[mm][kk][i * kElemCols + j] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
int ki = kk * 16 + int(sc.y) + i * kElemRowsJump;
|
||||
if (ki < actual_sk) {
|
||||
int col = nn * 16 + int(sc.x) + j;
|
||||
b_frags[kk][nn][i * kElemCols + j] = sg_B[(kk1 + ki) * int(N) + col];
|
||||
} else {
|
||||
b_frags[kk][nn][i * kElemCols + j] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
device int32_t *D = C + m_base * N + n_base;
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_store_int32(c_frags[mm * TN + nn], D, int(N), sc,
|
||||
short(mm * 16), short(nn * 16), M, N, m_base, n_base);
|
||||
}
|
||||
|
||||
// Dequant version
|
||||
template <int BM, int BN, int BK, int SK, int WM, int WN>
|
||||
void step4_gemm_dequant_impl(
|
||||
const device int8_t *A, const device int8_t *B,
|
||||
device half *C, uint M, uint N, uint K,
|
||||
const device float *scale_a, const device float *scale_w,
|
||||
uint swizzle_log, uint tiles_m, uint tiles_n,
|
||||
uint2 tgid, uint sgid, uint lid) {
|
||||
|
||||
constexpr int SM = BM / WM;
|
||||
constexpr int SN = BN / WN;
|
||||
constexpr short TM = SM / 16;
|
||||
constexpr short TN = SN / 16;
|
||||
constexpr short TK = SK / 16;
|
||||
|
||||
uint tid_y, tid_x;
|
||||
swizzle_decode(tgid, swizzle_log, tiles_n, tid_y, tid_x);
|
||||
if (tid_x >= tiles_n || tid_y >= tiles_m) return;
|
||||
|
||||
short2 sc = nax_get_coord(ushort(lid));
|
||||
uint sg_row = sgid / WN;
|
||||
uint sg_col = sgid % WN;
|
||||
uint m_base = tid_y * BM + sg_row * SM;
|
||||
uint n_base = tid_x * BN + sg_col * SN;
|
||||
|
||||
const device int8_t *sg_A = A + m_base * K;
|
||||
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) / 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++)
|
||||
nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc,
|
||||
short(mm * 16), short(kk * 16));
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), 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;
|
||||
}
|
||||
|
||||
int rem_k = int(K) - gemm_k_iters * BK;
|
||||
if (rem_k > 0) {
|
||||
for (int kk1 = 0; kk1 < rem_k; kk1 += SK) {
|
||||
int actual_sk = min(SK, rem_k - kk1);
|
||||
int8_t a_frags[TM][TK][kElemsPerFrag];
|
||||
int8_t b_frags[TK][TN][kElemsPerFrag];
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short kk = 0; kk < TK; kk++) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
int ki = kk * 16 + int(sc.x) + j;
|
||||
if (ki < actual_sk) {
|
||||
int row = int(sc.y) + mm * 16 + i * kElemRowsJump;
|
||||
a_frags[mm][kk][i * kElemCols + j] = sg_A[kk1 + row * int(K) + ki];
|
||||
} else {
|
||||
a_frags[mm][kk][i * kElemCols + j] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
for (short kk = 0; kk < TK; kk++)
|
||||
for (short nn = 0; nn < TN; nn++) {
|
||||
for (short i = 0; i < 2; i++)
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
int ki = kk * 16 + int(sc.y) + i * kElemRowsJump;
|
||||
if (ki < actual_sk) {
|
||||
int col = nn * 16 + int(sc.x) + j;
|
||||
b_frags[kk][nn][i * kElemCols + j] = sg_B[(kk1 + ki) * int(N) + col];
|
||||
} else {
|
||||
b_frags[kk][nn][i * kElemCols + j] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
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];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
device half *D = C + m_base * N + n_base;
|
||||
for (short mm = 0; mm < TM; mm++)
|
||||
for (short nn = 0; nn < TN; nn++)
|
||||
nax_frag_store_dequant(c_frags[mm * TN + nn], D, int(N), sc,
|
||||
short(mm * 16), short(nn * 16), M, N, m_base, n_base,
|
||||
scale_a, scale_w);
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Entry points — Large tile
|
||||
// ============================================================
|
||||
kernel void w8a8_matmul_fused_dequant(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device half *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
const device float *scale_a [[buffer(6)]],
|
||||
const device float *scale_w [[buffer(7)]],
|
||||
constant uint &swizzle_log [[buffer(8)]],
|
||||
constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
step4_gemm_dequant_impl<128, 128, 512, 32, 4, 4>(
|
||||
A, B, C, M, N, K, scale_a, scale_w, swizzle_log, tiles_m, tiles_n,
|
||||
tgid, sgid, lid);
|
||||
}
|
||||
|
||||
// Small tile entry point (NEW in Step 3!)
|
||||
kernel void w8a8_matmul_fused_dequant_small(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device half *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
const device float *scale_a [[buffer(6)]],
|
||||
const device float *scale_w [[buffer(7)]],
|
||||
constant uint &swizzle_log [[buffer(8)]],
|
||||
constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
step4_gemm_dequant_impl<32, 128, 512, 32, 1, 4>(
|
||||
A, B, C, M, N, K, scale_a, scale_w, swizzle_log, tiles_m, tiles_n,
|
||||
tgid, sgid, lid);
|
||||
}
|
||||
|
||||
// INT32 entry points
|
||||
kernel void int8_matmul_int32(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
constant uint &swizzle_log [[buffer(6)]],
|
||||
constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
step4_gemm_int32_impl<128, 128, 512, 32, 4, 4>(
|
||||
A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid);
|
||||
}
|
||||
|
||||
kernel void int8_matmul_int32_small(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
constant uint &swizzle_log [[buffer(6)]],
|
||||
constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
step4_gemm_int32_impl<32, 128, 512, 32, 1, 4>(
|
||||
A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid);
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
// ============================================================
|
||||
// Per-token quantization: FP16 → INT8 + float32 scale
|
||||
// Target: Apple M5, Metal 4
|
||||
//
|
||||
// Each threadgroup handles one row (one token).
|
||||
// Threads cooperate to find absmax via simdgroup reduce,
|
||||
// then quantize in parallel.
|
||||
//
|
||||
// Host dispatch:
|
||||
// threadgroup = (min(256, ceil(K/32)*32), 1, 1)
|
||||
// grid = (M, 1, 1)
|
||||
// ============================================================
|
||||
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
kernel void
|
||||
quantize_per_token(const device half *X [[buffer(0)]], // [M, K] FP16 input
|
||||
device int8_t *A [[buffer(1)]], // [M, K] INT8 output
|
||||
device float *scale [[buffer(2)]], // [M] float32 scale
|
||||
constant uint &M [[buffer(3)]],
|
||||
constant uint &K [[buffer(4)]],
|
||||
uint gid [[threadgroup_position_in_grid]], // row index
|
||||
uint lid [[thread_index_in_threadgroup]],
|
||||
uint tg_size [[threads_per_threadgroup]]) {
|
||||
if (gid >= M) {
|
||||
return;
|
||||
}
|
||||
|
||||
const device half *row_in = X + gid * K;
|
||||
device int8_t *row_out = A + gid * K;
|
||||
|
||||
// Step 1: Find local absmax
|
||||
float local_max = 0.0f;
|
||||
for (uint i = lid; i < K; i += tg_size) {
|
||||
float v = abs(float(row_in[i]));
|
||||
local_max = max(local_max, v);
|
||||
}
|
||||
|
||||
// Step 2: Simdgroup reduce max
|
||||
float sg_max = simd_max(local_max);
|
||||
|
||||
// Step 3: Threadgroup reduce across simdgroups via shared memory
|
||||
threadgroup float sg_maxes[8]; // up to 8 simdgroups (256/32)
|
||||
threadgroup float shared_scale;
|
||||
threadgroup float shared_inv_scale;
|
||||
uint sg_id = lid / 32;
|
||||
uint sg_lid = lid % 32;
|
||||
if (sg_lid == 0) {
|
||||
sg_maxes[sg_id] = sg_max;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Final reduce (first simdgroup only)
|
||||
if (sg_id == 0) {
|
||||
float row_max = 0.0f;
|
||||
uint num_sgs = (tg_size + 31) / 32;
|
||||
if (sg_lid < num_sgs) {
|
||||
row_max = sg_maxes[sg_lid];
|
||||
}
|
||||
row_max = simd_max(row_max);
|
||||
|
||||
// Compute and broadcast scale
|
||||
float s = row_max / 255.0f;
|
||||
if (s == 0.0f) {
|
||||
s = 1.0f;
|
||||
}
|
||||
|
||||
if (sg_lid == 0) {
|
||||
shared_scale = s;
|
||||
shared_inv_scale = 1.0f / s;
|
||||
// Store scale to output
|
||||
scale[gid] = s;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Step 4: All threads read broadcasted scale
|
||||
float inv_s = shared_inv_scale;
|
||||
|
||||
// Step 5: Quantize
|
||||
for (uint i = lid; i < K; i += tg_size) {
|
||||
float v = float(row_in[i]) * inv_s;
|
||||
v = clamp(round(v), -128.0f, 127.0f);
|
||||
row_out[i] = int8_t(v);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,475 @@
|
||||
// ============================================================
|
||||
// W8A8 INT8×INT8→INT32 TensorOps GEMM
|
||||
// Target: Apple M5 (G17G), Metal 4
|
||||
//
|
||||
// Variants:
|
||||
// - fused dequant: INT8×INT8→FP16, with per-token/per-channel scales
|
||||
// - raw INT32: INT8×INT8→INT32, no scale (pure integer GEMM)
|
||||
// Multi-config: large (BM=128) and small (BM=32) tiles
|
||||
// Swizzle dispatch for L2 cache locality
|
||||
//
|
||||
// matmul2d(16,32,16) via MPP cooperative_tensor
|
||||
// ============================================================
|
||||
|
||||
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
// ── NAXFrag layout constants ────────────────────────────────────
|
||||
constant constexpr short kElemsPerFrag = 8;
|
||||
constant constexpr short kElemCols = 4;
|
||||
constant constexpr short kElemRowsJump = 8;
|
||||
|
||||
// ── NAXFrag coordinate mapping ──────────────────────────────────
|
||||
inline 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};
|
||||
}
|
||||
|
||||
// ── Fragment load: device → register ────────────────────────────
|
||||
template <typename T>
|
||||
inline void nax_frag_load(thread T *dst, const device T *src, int ld, short2 sc,
|
||||
short off_m = 0, short off_n = 0) {
|
||||
src += (sc.y + off_m) * ld + (sc.x + off_n);
|
||||
for (short i = 0; i < 2; i++) {
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
dst[i * kElemCols + j] = src[(i * kElemRowsJump) * ld + j];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Fragment store: raw INT32 (no dequant) ───────────────────
|
||||
inline void nax_frag_store_int32(const thread int32_t *src, device int32_t *dst,
|
||||
int ld, short2 sc, short off_m, short off_n,
|
||||
uint M, uint N, uint m_base, uint n_base) {
|
||||
for (short i = 0; i < 2; i++) {
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
uint mi = m_base + sc.y + off_m + i * kElemRowsJump;
|
||||
uint ni = n_base + sc.x + off_n + j;
|
||||
if (mi < M && ni < N) {
|
||||
dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] =
|
||||
src[i * kElemCols + j];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Fragment store with bounds check and dequant ────────────────
|
||||
inline void nax_frag_store_dequant(const thread int32_t *src, device half *dst,
|
||||
int ld, short2 sc, short off_m, short off_n,
|
||||
uint M, uint N, uint m_base, uint n_base,
|
||||
const device float *scale_a,
|
||||
const device float *scale_w) {
|
||||
for (short i = 0; i < 2; i++) {
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
uint mi = m_base + sc.y + off_m + i * kElemRowsJump;
|
||||
uint ni = n_base + sc.x + off_n + j;
|
||||
if (mi < M && ni < N) {
|
||||
float val = float(src[i * kElemCols + j]) * scale_a[mi] * scale_w[ni];
|
||||
dst[(sc.y + off_m + i * kElemRowsJump) * ld + (sc.x + off_n + j)] =
|
||||
half(val);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Generic GEMM kernel ─────────────────────────────────────────
|
||||
// Template params: BM, BN, BK, SK, WM, WN
|
||||
// Each SG computes SM×SN = (BM/WM) × (BN/WN) output
|
||||
// SM and SN must be 32 (2×2 of 16×16 fragments)
|
||||
//
|
||||
// swizzle_log: passed via constant buffer
|
||||
// tid_y = (tgid.y << swizzle_log) + (tgid.x & ((1<<swizzle_log)-1))
|
||||
// tid_x = tgid.x >> swizzle_log
|
||||
template <int BM, int BN, int BK, int SK, int WM, int WN>
|
||||
void w8a8_gemm_impl(const device int8_t *A, const device int8_t *B,
|
||||
device half *C, uint M, uint N, uint K,
|
||||
const device float *scale_a, const device float *scale_w,
|
||||
uint swizzle_log, uint tiles_m, uint tiles_n, uint2 tgid,
|
||||
uint sgid, uint lid) {
|
||||
constexpr int SM = BM / WM; // 32
|
||||
constexpr int SN = BN / WN; // 32
|
||||
constexpr short TM = SM / 16; // 2
|
||||
constexpr short TN = SN / 16; // 2
|
||||
constexpr short TK = SK / 16; // 2
|
||||
|
||||
// Swizzle decode
|
||||
uint tid_y = (tgid.y << swizzle_log) + (tgid.x & ((1u << swizzle_log) - 1u));
|
||||
uint tid_x = tgid.x >> swizzle_log;
|
||||
|
||||
// Bounds check (swizzle can create out-of-bounds tiles)
|
||||
if (tid_x >= tiles_n || tid_y >= tiles_m) {
|
||||
return;
|
||||
}
|
||||
|
||||
short2 sc = nax_get_coord(ushort(lid));
|
||||
uint sg_row = sgid / WN;
|
||||
uint sg_col = sgid % WN;
|
||||
uint m_base = tid_y * BM + sg_row * SM;
|
||||
uint n_base = tid_x * BN + sg_col * SN;
|
||||
|
||||
const device int8_t *sg_A = A + m_base * K;
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Main K loop ─────────────────────────────────────────────
|
||||
int gemm_k_iters = int(K) / 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++) {
|
||||
nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc, short(mm * 16),
|
||||
short(kk * 16));
|
||||
}
|
||||
}
|
||||
for (short kk = 0; kk < TK; kk++) {
|
||||
for (short nn = 0; nn < TN; nn++) {
|
||||
nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), 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;
|
||||
}
|
||||
|
||||
// ── Remainder K ─────────────────────────────────────────────
|
||||
int rem_k = int(K) - gemm_k_iters * BK;
|
||||
for (int kk1 = 0; kk1 < rem_k; kk1 += 16) {
|
||||
int8_t a_frag[TM][kElemsPerFrag];
|
||||
int8_t b_frag[TN][kElemsPerFrag];
|
||||
short psk = short(max(0, rem_k - kk1));
|
||||
|
||||
for (short mm = 0; mm < TM; mm++) {
|
||||
const device int8_t *ptr = sg_A + kk1 + (sc.y + mm * 16) * K + sc.x;
|
||||
for (short i = 0; i < 2; i++) {
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
short ki = short(sc.x + j);
|
||||
a_frag[mm][i * kElemCols + j] =
|
||||
(ki < psk) ? ptr[(i * kElemRowsJump) * K + j] : int8_t(0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (short nn = 0; nn < TN; nn++) {
|
||||
const device int8_t *ptr = sg_B + kk1 * N + nn * 16 + sc.y * N + sc.x;
|
||||
for (short i = 0; i < 2; i++) {
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
short ki = short(sc.y + i * kElemRowsJump);
|
||||
b_frag[nn][i * kElemCols + j] =
|
||||
(ki < psk) ? ptr[(i * kElemRowsJump) * N + j] : int8_t(0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (short mm = 0; mm < TM; mm++) {
|
||||
for (short i = 0; i < kElemsPerFrag; i++) {
|
||||
ct_a[i] = a_frag[mm][i];
|
||||
}
|
||||
for (short i = 0; i < kElemsPerFrag; i++) {
|
||||
ct_b[i] = b_frag[0][i];
|
||||
ct_b[kElemsPerFrag + i] = b_frag[1][i];
|
||||
}
|
||||
short c0 = mm * TN, 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];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Store with fused dequant ────────────────────────────────
|
||||
device half *D = C + m_base * N + n_base;
|
||||
for (short mm = 0; mm < TM; mm++) {
|
||||
for (short nn = 0; nn < TN; nn++) {
|
||||
nax_frag_store_dequant(c_frags[mm * TN + nn], D, int(N), sc,
|
||||
short(mm * 16), short(nn * 16), M, N, m_base,
|
||||
n_base, scale_a, scale_w);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Kernel entry points
|
||||
// ============================================================
|
||||
|
||||
// ── Large tile: BM=128, BN=128, BK=512, WM=4, WN=4 ────────────
|
||||
// 16 SG, 512 threads/TG. Best for M≥128.
|
||||
kernel void w8a8_matmul_fused_dequant(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device half *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
const device float *scale_a [[buffer(6)]],
|
||||
const device float *scale_w [[buffer(7)]],
|
||||
constant uint &swizzle_log [[buffer(8)]],
|
||||
constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
w8a8_gemm_impl<128, 128, 512, 32, 4, 4>(A, B, C, M, N, K, scale_a, scale_w,
|
||||
swizzle_log, tiles_m, tiles_n, tgid,
|
||||
sgid, lid);
|
||||
}
|
||||
|
||||
// ── Small tile: BM=32, BN=128, BK=512, WM=1, WN=4 ─────────────
|
||||
// 4 SG, 128 threads/TG. Optimized for M≤64.
|
||||
kernel void w8a8_matmul_fused_dequant_small(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device half *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
const device float *scale_a [[buffer(6)]],
|
||||
const device float *scale_w [[buffer(7)]],
|
||||
constant uint &swizzle_log [[buffer(8)]],
|
||||
constant uint &tiles_m [[buffer(9)]], constant uint &tiles_n [[buffer(10)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
w8a8_gemm_impl<32, 128, 512, 32, 1, 4>(A, B, C, M, N, K, scale_a, scale_w,
|
||||
swizzle_log, tiles_m, tiles_n, tgid,
|
||||
sgid, lid);
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Raw INT32 GEMM impl (no dequant, no scales)
|
||||
// ============================================================
|
||||
|
||||
template <int BM, int BN, int BK, int SK, int WM, int WN>
|
||||
void w8a8_gemm_int32_impl(const device int8_t *A, const device int8_t *B,
|
||||
device int32_t *C, uint M, uint N, uint K,
|
||||
uint swizzle_log, uint tiles_m, uint tiles_n,
|
||||
uint2 tgid, uint sgid, uint lid) {
|
||||
constexpr int SM = BM / WM;
|
||||
constexpr int SN = BN / WN;
|
||||
constexpr short TM = SM / 16;
|
||||
constexpr short TN = SN / 16;
|
||||
constexpr short TK = SK / 16;
|
||||
|
||||
uint tid_y = (tgid.y << swizzle_log) + (tgid.x & ((1u << swizzle_log) - 1u));
|
||||
uint tid_x = tgid.x >> swizzle_log;
|
||||
if (tid_x >= tiles_n || tid_y >= tiles_m) {
|
||||
return;
|
||||
}
|
||||
|
||||
short2 sc = nax_get_coord(ushort(lid));
|
||||
uint sg_row = sgid / WN;
|
||||
uint sg_col = sgid % WN;
|
||||
uint m_base = tid_y * BM + sg_row * SM;
|
||||
uint n_base = tid_x * BN + sg_col * SN;
|
||||
|
||||
const device int8_t *sg_A = A + m_base * K;
|
||||
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) / 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++) {
|
||||
nax_frag_load(a_frags[mm][kk], sg_A + kk1, int(K), sc, short(mm * 16),
|
||||
short(kk * 16));
|
||||
}
|
||||
}
|
||||
for (short kk = 0; kk < TK; kk++) {
|
||||
for (short nn = 0; nn < TN; nn++) {
|
||||
nax_frag_load(b_frags[kk][nn], sg_B + kk1 * N, int(N), 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;
|
||||
}
|
||||
|
||||
// Remainder K
|
||||
int rem_k = int(K) - gemm_k_iters * BK;
|
||||
for (int kk1 = 0; kk1 < rem_k; kk1 += 16) {
|
||||
int8_t a_frag[TM][kElemsPerFrag];
|
||||
int8_t b_frag[TN][kElemsPerFrag];
|
||||
short psk = short(max(0, rem_k - kk1));
|
||||
for (short mm = 0; mm < TM; mm++) {
|
||||
const device int8_t *ptr = sg_A + kk1 + (sc.y + mm * 16) * K + sc.x;
|
||||
for (short i = 0; i < 2; i++) {
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
short ki = short(sc.x + j);
|
||||
a_frag[mm][i * kElemCols + j] =
|
||||
(ki < psk) ? ptr[(i * kElemRowsJump) * K + j] : int8_t(0);
|
||||
}
|
||||
}
|
||||
}
|
||||
for (short nn = 0; nn < TN; nn++) {
|
||||
const device int8_t *ptr = sg_B + kk1 * N + nn * 16 + sc.y * N + sc.x;
|
||||
for (short i = 0; i < 2; i++) {
|
||||
for (short j = 0; j < kElemCols; j++) {
|
||||
short ki = short(sc.y + i * kElemRowsJump);
|
||||
b_frag[nn][i * kElemCols + j] =
|
||||
(ki < psk) ? ptr[(i * kElemRowsJump) * N + j] : int8_t(0);
|
||||
}
|
||||
}
|
||||
}
|
||||
for (short mm = 0; mm < TM; mm++) {
|
||||
for (short i = 0; i < kElemsPerFrag; i++) {
|
||||
ct_a[i] = a_frag[mm][i];
|
||||
}
|
||||
for (short i = 0; i < kElemsPerFrag; i++) {
|
||||
ct_b[i] = b_frag[0][i];
|
||||
ct_b[kElemsPerFrag + i] = b_frag[1][i];
|
||||
}
|
||||
short c0 = mm * TN, 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];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Store raw INT32
|
||||
device int32_t *D = C + m_base * N + n_base;
|
||||
for (short mm = 0; mm < TM; mm++) {
|
||||
for (short nn = 0; nn < TN; nn++) {
|
||||
nax_frag_store_int32(c_frags[mm * TN + nn], D, int(N), sc, short(mm * 16),
|
||||
short(nn * 16), M, N, m_base, n_base);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Kernel entry points — raw INT32 output
|
||||
// ============================================================
|
||||
|
||||
kernel void int8_matmul_int32(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
constant uint &swizzle_log [[buffer(6)]],
|
||||
constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
w8a8_gemm_int32_impl<128, 128, 512, 32, 4, 4>(
|
||||
A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid);
|
||||
}
|
||||
|
||||
kernel void int8_matmul_int32_small(
|
||||
const device int8_t *A [[buffer(0)]], const device int8_t *B [[buffer(1)]],
|
||||
device int32_t *C [[buffer(2)]], constant uint &M [[buffer(3)]],
|
||||
constant uint &N [[buffer(4)]], constant uint &K [[buffer(5)]],
|
||||
constant uint &swizzle_log [[buffer(6)]],
|
||||
constant uint &tiles_m [[buffer(7)]], constant uint &tiles_n [[buffer(8)]],
|
||||
uint2 tgid [[threadgroup_position_in_grid]],
|
||||
uint sgid [[simdgroup_index_in_threadgroup]],
|
||||
uint lid [[thread_index_in_simdgroup]]) {
|
||||
w8a8_gemm_int32_impl<32, 128, 512, 32, 1, 4>(
|
||||
A, B, C, M, N, K, swizzle_log, tiles_m, tiles_n, tgid, sgid, lid);
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
// ============================================================
|
||||
// Per-token quantization: FP16 → INT8 + float32 scale
|
||||
// Target: Apple M5, Metal 4
|
||||
//
|
||||
// Each threadgroup handles one row (one token).
|
||||
// Threads cooperate to find absmax via simdgroup reduce,
|
||||
// then quantize in parallel.
|
||||
//
|
||||
// Host dispatch:
|
||||
// threadgroup = (min(256, ceil(K/32)*32), 1, 1)
|
||||
// grid = (M, 1, 1)
|
||||
// ============================================================
|
||||
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
|
||||
kernel void
|
||||
quantize_per_token(const device half *X [[buffer(0)]], // [M, K] FP16 input
|
||||
device int8_t *A [[buffer(1)]], // [M, K] INT8 output
|
||||
device float *scale [[buffer(2)]], // [M] float32 scale
|
||||
constant uint &M [[buffer(3)]],
|
||||
constant uint &K [[buffer(4)]],
|
||||
uint gid [[threadgroup_position_in_grid]], // row index
|
||||
uint lid [[thread_index_in_threadgroup]],
|
||||
uint tg_size [[threads_per_threadgroup]]) {
|
||||
if (gid >= M) {
|
||||
return;
|
||||
}
|
||||
|
||||
const device half *row_in = X + gid * K;
|
||||
device int8_t *row_out = A + gid * K;
|
||||
|
||||
// Step 1: Find local absmax
|
||||
float local_max = 0.0f;
|
||||
for (uint i = lid; i < K; i += tg_size) {
|
||||
float v = abs(float(row_in[i]));
|
||||
local_max = max(local_max, v);
|
||||
}
|
||||
|
||||
// Step 2: Simdgroup reduce max
|
||||
float sg_max = simd_max(local_max);
|
||||
|
||||
// Step 3: Threadgroup reduce across simdgroups via shared memory
|
||||
threadgroup float sg_maxes[8]; // up to 8 simdgroups (256/32)
|
||||
threadgroup float shared_scale;
|
||||
threadgroup float shared_inv_scale;
|
||||
uint sg_id = lid / 32;
|
||||
uint sg_lid = lid % 32;
|
||||
if (sg_lid == 0) {
|
||||
sg_maxes[sg_id] = sg_max;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Final reduce (first simdgroup only)
|
||||
if (sg_id == 0) {
|
||||
float row_max = 0.0f;
|
||||
uint num_sgs = (tg_size + 31) / 32;
|
||||
if (sg_lid < num_sgs) {
|
||||
row_max = sg_maxes[sg_lid];
|
||||
}
|
||||
row_max = simd_max(row_max);
|
||||
|
||||
// Compute and broadcast scale
|
||||
float s = row_max / 255.0f;
|
||||
if (s == 0.0f) {
|
||||
s = 1.0f;
|
||||
}
|
||||
|
||||
if (sg_lid == 0) {
|
||||
shared_scale = s;
|
||||
shared_inv_scale = 1.0f / s;
|
||||
// Store scale to output
|
||||
scale[gid] = s;
|
||||
}
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
// Step 4: All threads read broadcasted scale
|
||||
float inv_s = shared_inv_scale;
|
||||
|
||||
// Step 5: Quantize
|
||||
for (uint i = lid; i < K; i += tg_size) {
|
||||
float v = float(row_in[i]) * inv_s;
|
||||
v = clamp(round(v), -128.0f, 127.0f);
|
||||
row_out[i] = int8_t(v);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user