spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Quantization-aware-training kernel: fake-quantize a weight matrix row by4// row (per-output-channel scales). One thread per row, two strided passes5// over the row (reduce scale, then quantize). Rows are weight-matrix sized6// (C = in_features, a few thousand max) and this runs once per layer per7// forward, so a simple thread-per-row layout is plenty.8//9// Backward is the straight-through estimator (identity), so there is no10// backward kernel — autograd routes dout straight into the master weight's11// grad.12#include <metal_stdlib>13using namespace metal;1415// p = (C, mode); mode 0 = int8 (absmax/127), mode 1 = ternary BitNet-style16// (scale = mean|w|, values in {-s, 0, +s}).17kernel void fake_quant_f32(device const float* W   [[buffer(0)]],18                           device float*       OUT [[buffer(1)]],19                           constant uint2&     p   [[buffer(2)]],20                           uint row [[thread_position_in_grid]]) {21    const uint C = p.x;22    device const float* w = W + ulong(row) * C;23    device float* out = OUT + ulong(row) * C;2425    if (p.y == 0) { // int8: symmetric absmax26        float amax = 0.0f;27        for (uint j = 0; j < C; ++j) amax = max(amax, fabs(w[j]));28        const float s = max(amax / 127.0f, 1e-12f);29        for (uint j = 0; j < C; ++j) out[j] = rint(w[j] / s) * s;30    } else {        // ternary: BitNet b1.58 absmean31        float asum = 0.0f;32        for (uint j = 0; j < C; ++j) asum += fabs(w[j]);33        const float s = max(asum / float(C), 1e-12f);34        for (uint j = 0; j < C; ++j)35            out[j] = clamp(rint(w[j] / s), -1.0f, 1.0f) * s;36    }37}38