wehub-resource-sync
88 lines
2.5 KiB
Metal
88 lines
2.5 KiB
Metal
// ============================================================
|
|
// 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);
|
|
}
|
|
}
|