// vendor/strata/metal/src/heal_k256_q2_0.cpp - Phase 3 Signed-d AWQ/Lloyd-Max Q2_0 Expert Healer +
// Real-Trace Layerwise Scale Calibration + Monotonic Shared-Expert Ridge Solver for K256_reap_l2_sum
//
// Modes:
// 1. heal_k256_q2_0 --quantize <in-metal.gguf> <out-healed-metal.gguf> [gain_gu=0.50] [gain_d=0.30] [cal_iq3s_pfx]
// Quantizes all 48 x 256 routed experts to Signed-d AWQ/Lloyd-Max Q2_0 (unlocking the sign bit of fp16 d
// for a 0-bit polarity switch between {-|d|, 0, +|d|, +2|d|} and {-2|d|, -|d|, 0, +|d|}), weighted by
// tempered real activation RMS (E[mixed_d^2])^0.5 (beta=0.50) on W_gate/W_up and analytical
// fourth-root importance (sigma_g * sigma_u)^0.5 ~ (E[h_c^2])^0.25 (beta=0.25) on W_down
// when cal_iq3s_pfx is provided (falling back to hc_ffn_norm RMS if omitted), and saves <out>.experts.d_bak.bin.
// 2. heal_k256_q2_0 --calibrate-and-heal <in-metal.gguf> <out-healed-metal.gguf> <cal_iq3s_pfx> <yr_q2_bin> <out_writers_bin> [gamma=0.75] [eta_shd=0.15] [ridge_rel=6.0]
// Calibrates per-layer W_down fp16 d scale alpha_l = cos_l^(1-gamma) / norm_ratio_l directly from the real
// multi-domain teacher trace (all t_cal tokens x all 256 routed experts, eliminating Gaussian depth overshoot)
// and solves the ridge regression (full-rank 640/640 on 46/48 layers; L0=430, L2=617) for blk.0..47.ffn_down_shexp.weight with
// monotonic Q8_0 / IQ4_NL block quantization, writing the 48 modified ffn_down_shexp tensors to <out_writers_bin>.
#ifndef ACCELERATE_NEW_LAPACK
#define ACCELERATE_NEW_LAPACK
#endif
#include "gguf_lite.h"
#include "ref_quant.h"
#include <Accelerate/Accelerate.h>
#include <algorithm>
#include <atomic>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fcntl.h>
#include <string>
#include <sys/stat.h>
#include <thread>
#include <unistd.h>
#include <vector>
static constexpr int NL = 48;
static constexpr int NE = 256;
static constexpr int N = 2560;
static constexpr int D = 640;
static constexpr uint64_t ALIGN = 16384;
static constexpr uint64_t Q2_SLICE_BYTES = 460800; // 2560 * 640 / 64 * 18
static constexpr size_t N_BLOCKS = Q2_SLICE_BYTES / 18;
static constexpr int8_t KVALUES_IQ4NL[16] = {
-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113
};
static void require_io(bool ok, const char* msg) {
if (!ok) {
std::perror(msg);
std::exit(2);
}
}
static void pread_exact(int fd, void* buf, size_t n, off_t off, const char* msg) {
uint8_t* p = static_cast<uint8_t*>(buf);
size_t got = 0;
while (got < n) {
ssize_t r = ::pread(fd, p + got, n - got, off + (off_t)got);
require_io(r > 0, msg);
got += (size_t)r;
}
}
static void pwrite_exact(int fd, const void* buf, size_t n, off_t off, const char* msg) {
const uint8_t* p = static_cast<const uint8_t*>(buf);
size_t put = 0;
while (put < n) {
ssize_t w = ::pwrite(fd, p + put, n - put, off + (off_t)put);
require_io(w > 0, msg);
put += (size_t)w;
}
}
static uint64_t fnv1a_64(const void* data, size_t nbytes) {
const uint8_t* p = static_cast<const uint8_t*>(data);
uint64_t h = 14695981039346656037ULL;
for (size_t i = 0; i < nbytes; ++i) {
h ^= (uint64_t)p[i];
h *= 1099511628211ULL;
}
return h;
}
static void quantize_q2_0_signed_awq_fast(
const float* x, size_t n_rows, size_t n_cols, const float* col_imp, uint8_t* out, float gain_exp
) {
const size_t blocks_per_row = n_cols / 64;
for (size_t r = 0; r < n_rows; ++r) {
for (size_t b = 0; b < blocks_per_row; ++b) {
const float* xb = x + r * n_cols + b * 64;
const float* wb = col_imp ? (col_imp + b * 64) : nullptr;
float wxb[64], wi_arr[64];
float amax = 0.0f, norm_x2 = 0.0f;
for (int i = 0; i < 64; ++i) {
const float xi = xb[i];
const float wi = wb ? wb[i] : 1.0f;
wi_arr[i] = wi;
wxb[i] = wi * xi;
amax = std::max(amax, std::fabs(xi));
norm_x2 += wxb[i] * xi;
}
int best_q[64] = {0};
uint16_t d16 = 0;
if (amax > 0.0f && norm_x2 > 0.0f) {
float best_score = -1.0f, best_qw = 0.0f, best_qq = 1.0f;
static constexpr float seeds[10] = {
0.42f, 0.52f, 0.64f, 0.78f, 0.94f,
-0.42f, -0.52f, -0.64f, -0.78f, -0.94f
};
for (int s = 0; s < 10; ++s) {
float d = amax * seeds[s];
for (int iter = 0; iter < 2; ++iter) {
if (std::fabs(d) <= 1e-12f) break;
const float inv = 1.0f / d;
float qw = 0.0f, qq = 0.0f;
int q_cur[64];
for (int i = 0; i < 64; ++i) {
int q = (int)std::lrintf(xb[i] * inv);
q = std::max(-1, std::min(2, q));
q_cur[i] = q;
const float fq = (float)q;
qw += fq * wxb[i];
qq += wi_arr[i] * (fq * fq);
}
if (qq <= 0.0f) break;
const float d_new = qw / qq;
if ( (d > 0.0f ? d_new <= 0.0f : d_new >= 0.0f) ) break;
const float score = (qw * qw) / qq;
if (score > best_score) {
best_score = score;
best_qw = qw;
best_qq = qq;
std::memcpy(best_q, q_cur, sizeof(best_q));
}
d = d_new;
}
}
float d_ls = best_qw / best_qq;
float r2 = best_score / std::max(norm_x2, 1e-20f);
float gain = (gain_exp == 0.0f) ? 1.0f : std::pow(std::max(r2, 0.25f), -0.5f * gain_exp);
d16 = ref::f32_to_f16(d_ls * gain);
}
uint8_t* dst = out + (r * blocks_per_row + b) * 18;
std::memcpy(dst, &d16, 2);
for (int j = 0; j < 16; ++j) {
uint8_t byte = 0;
for (int k = 0; k < 4; ++k) {
const uint8_t code = (uint8_t)(best_q[j * 4 + k] + 1);
byte |= (uint8_t)(code << (2 * k));
}
dst[2 + j] = byte;
}
}
}
}
static void quantize_iq4nl_monotonic(const uint8_t* raw_old, const float* w_target, size_t n_elems, uint8_t* raw_new) {
const size_t n_blks = n_elems / 32;
std::vector<float> w_old_deq(32);
for (size_t b = 0; b < n_blks; ++b) {
const uint8_t* bo = raw_old + b * 18;
uint8_t* bn = raw_new + b * 18;
std::memcpy(bn, bo, 18);
ref::dequant(ref::IQ4_NL, bo, w_old_deq.data(), 32);
const float* wt = w_target + b * 32;
float err_old = 0.0f;
for (int i = 0; i < 32; ++i) {
float diff = w_old_deq[i] - wt[i];
err_old += diff * diff;
}
uint16_t d_old_u;
std::memcpy(&d_old_u, bo, 2);
float d_old = ref::f16_to_f32(d_old_u);
if (std::fabs(d_old) < 1e-12f) continue;
float best_err = err_old;
uint8_t best_blk[18];
std::memcpy(best_blk, bo, 18);
// Search 25 fine scale multipliers around d_old so IQ4_NL preserves small ridge updates
for (int s = -12; s <= 12; ++s) {
float mult = 1.0f + 0.015f * (float)s;
uint16_t d_cand_u = ref::f32_to_f16(d_old * mult);
float d_cand = ref::f16_to_f32(d_cand_u);
uint8_t idx[32];
float err_c = 0.0f;
for (int i = 0; i < 32; ++i) {
float best_w_err = 1e30f;
uint8_t best_q = 0;
for (uint8_t q = 0; q < 16; ++q) {
float v = d_cand * (float)KVALUES_IQ4NL[q];
float de = (v - wt[i]) * (v - wt[i]);
if (de < best_w_err) {
best_w_err = de;
best_q = q;
}
}
idx[i] = best_q;
err_c += best_w_err;
}
if (err_c + 1e-11f < best_err) {
best_err = err_c;
std::memcpy(best_blk, &d_cand_u, 2);
for (int j = 0; j < 16; ++j) {
best_blk[2 + j] = (uint8_t)((idx[j] & 0x0F) | ((idx[j + 16] & 0x0F) << 4));
}
}
}
std::memcpy(bn, best_blk, 18);
}
}
static void quantize_q8_0_monotonic(const uint8_t* raw_old, const float* w_target, size_t n_elems, uint8_t* raw_new) {
const size_t n_blks = n_elems / 32;
std::vector<float> w_old_deq(32);
for (size_t b = 0; b < n_blks; ++b) {
const uint8_t* bo = raw_old + b * 34;
uint8_t* bn = raw_new + b * 34;
std::memcpy(bn, bo, 34);
ref::dequant(ref::Q8_0, bo, w_old_deq.data(), 32);
const float* wt = w_target + b * 32;
float err_old = 0.0f, amax = 0.0f;
for (int i = 0; i < 32; ++i) {
float diff = w_old_deq[i] - wt[i];
err_old += diff * diff;
amax = std::max(amax, std::fabs(wt[i]));
}
uint16_t d_cand_u = ref::f32_to_f16(amax / 127.0f);
float d_cand = ref::f16_to_f32(d_cand_u);
if (d_cand <= 0.0f) continue;
int8_t qs[32];
float err_c = 0.0f;
for (int i = 0; i < 32; ++i) {
int q = (int)std::lrintf(wt[i] / d_cand);
q = std::max(-127, std::min(127, q));
qs[i] = (int8_t)q;
float diff = d_cand * (float)q - wt[i];
err_c += diff * diff;
}
if (err_c + 1e-11f < err_old) {
std::memcpy(bn, &d_cand_u, 2);
std::memcpy(bn + 2, qs, 32);
}
}
}
static int run_quantize_mode(int argc, char** argv) {
if (argc < 4) {
std::fprintf(stderr, "usage: heal_k256_q2_0 --quantize <in-metal.gguf> <out-healed-metal.gguf> [gain_gu=0.50] [gain_d=0.30] [cal_iq3s_pfx]\n");
return 2;
}
const std::string in_gguf = argv[2];
const std::string out_gguf = argv[3];
const float gain_gu = (argc > 4) ? std::atof(argv[4]) : 0.50f;
const float gain_d = (argc > 5) ? std::atof(argv[5]) : 0.30f;
const std::string cal_pfx = (argc > 6) ? argv[6] : "";
const std::string in_exp = in_gguf + ".experts";
const std::string out_exp = out_gguf + ".experts";
const std::string bak_d_path = out_exp + ".d_bak.bin";
gguf::File f;
std::string err;
if (!f.open(in_gguf, err)) {
std::fprintf(stderr, "gguf open failed: %s\n", err.c_str());
return 1;
}
// Check if real teacher activation dump is available for tempered RMS importance
bool use_real_act = false;
size_t t_cal = 0;
int fd_cal_x = -1;
if (!cal_pfx.empty()) {
struct stat st{};
std::string mix_path = cal_pfx + ".mixed";
if (::stat(mix_path.c_str(), &st) == 0 && st.st_size > 0 && st.st_size % ((size_t)NL * N * sizeof(float)) == 0) {
t_cal = (size_t)st.st_size / ((size_t)NL * N * sizeof(float));
fd_cal_x = ::open(mix_path.c_str(), O_RDONLY);
require_io(fd_cal_x >= 0, "open cal_pfx.mixed");
use_real_act = (t_cal >= 32);
}
}
std::vector<std::vector<float>> layer_s_in(NL, std::vector<float>(N, 1.0f));
std::vector<std::vector<float>> layer_s2_raw(NL, std::vector<float>(N, 1.0f));
if (use_real_act) {
std::vector<float> x_l(t_cal * N);
for (int l = 0; l < NL; ++l) {
pread_exact(fd_cal_x, x_l.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread cal_pfx.mixed");
for (int d = 0; d < N; ++d) {
double ms2 = 0.0;
for (size_t t = 0; t < t_cal; ++t) {
double v = x_l[t * N + d];
ms2 += v * v;
}
double e_x2 = std::max(ms2 / (double)t_cal, 1e-8);
layer_s2_raw[l][d] = (float)e_x2;
// Tempered real activation RMS (beta = 0.5: sqrt(E[mixed_d^2]))
layer_s_in[l][d] = (float)std::sqrt(e_x2);
}
}
::close(fd_cal_x);
std::printf("Using tempered real activation RMS (T=%zu, beta=0.50) + per-expert (sigma_g*sigma_u)^0.5 (beta=0.25) for Signed-d Q2_0\n", t_cal);
} else {
for (int l = 0; l < NL; ++l) {
const std::string tname = "blk." + std::to_string(l) + ".hc_ffn_norm.weight";
const gguf::Tensor* tn = f.find(tname);
if (tn && tn->type == 0) {
const float* wn = reinterpret_cast<const float*>(static_cast<const uint8_t*>(f.map) + tn->offset);
for (int d = 0; d < N; ++d) {
float ms2 = 0.0f;
for (int c = 0; c < 4; ++c) {
const float v = wn[c * N + d];
ms2 += v * v;
}
layer_s_in[l][d] = std::sqrt(std::max(ms2 * 0.25f, 1e-12f));
}
}
}
}
const size_t hlen = std::min<size_t>(f.file_bytes, 16 * 1024 * 1024);
std::vector<uint8_t> hdr(hlen);
std::memcpy(hdr.data(), f.map, hlen);
for (int l = 0; l < NL; ++l) {
for (const char* role : {"gate", "up", "down"}) {
std::string name = "blk." + std::to_string(l) + ".ffn_" + role + "_exps.weight";
const uint64_t nlen = name.size();
std::vector<uint8_t> pat(8 + nlen);
std::memcpy(pat.data(), &nlen, 8);
std::memcpy(pat.data() + 8, name.data(), nlen);
auto it = std::search(hdr.begin(), hdr.end(), pat.begin(), pat.end());
require_io(it != hdr.end(), ("find tensor in GGUF header: " + name).c_str());
size_t pos = (size_t)(it - hdr.begin()) + pat.size();
uint32_t nd = 0;
std::memcpy(&nd, hdr.data() + pos, 4);
pos += 4 + (size_t)nd * 8;
uint32_t new_type = 42;
std::memcpy(hdr.data() + pos, &new_type, 4);
}
}
int ofd = ::open(out_gguf.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644);
require_io(ofd >= 0, "open out_gguf");
fcntl(ofd, F_NOCACHE, 1);
pwrite_exact(ofd, hdr.data(), hlen, 0, "pwrite out_gguf hdr");
size_t pos = hlen;
const size_t chunk = 64ull << 20;
while (pos < f.file_bytes) {
const size_t n = std::min(chunk, f.file_bytes - pos);
pwrite_exact(ofd, (const uint8_t*)f.map + pos, n, (off_t)pos, "pwrite out_gguf body");
pos += n;
}
::fsync(ofd);
::close(ofd);
int ifd = ::open(in_exp.c_str(), O_RDONLY);
require_io(ifd >= 0, "open in_exp");
fcntl(ifd, F_NOCACHE, 1);
uint64_t in_tab[NL * 4];
pread_exact(ifd, in_tab, sizeof(in_tab), 64, "pread in_tab");
const uint64_t exp_data0 = ALIGN;
uint64_t out_offs[NL];
uint64_t at = exp_data0;
const uint64_t out_exp_bytes = 3 * Q2_SLICE_BYTES;
const uint64_t out_layer_bytes = (uint64_t)NE * out_exp_bytes;
for (int l = 0; l < NL; ++l) {
at = (at + ALIGN - 1) & ~(ALIGN - 1);
out_offs[l] = at;
at += out_layer_bytes;
}
int efd = ::open(out_exp.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644);
require_io(efd >= 0, "open out_exp");
fcntl(efd, F_NOCACHE, 1);
uint8_t strx_hdr[ALIGN];
std::memset(strx_hdr, 0, sizeof(strx_hdr));
std::memcpy(strx_hdr, "STRX", 4);
uint32_t v1 = 1, nl_u = NL, ne_u = NE;
uint64_t src_sz = f.file_bytes;
std::memcpy(strx_hdr + 4, &v1, 4);
std::memcpy(strx_hdr + 8, &nl_u, 4);
std::memcpy(strx_hdr + 12, &ne_u, 4);
std::memcpy(strx_hdr + 16, &src_sz, 8);
for (int l = 0; l < NL; ++l) {
uint64_t row[4] = {out_offs[l], Q2_SLICE_BYTES, Q2_SLICE_BYTES, Q2_SLICE_BYTES};
std::memcpy(strx_hdr + 64 + l * 32, row, 32);
}
pwrite_exact(efd, strx_hdr, ALIGN, 0, "pwrite strx_hdr");
std::vector<uint16_t> ds_all((size_t)NL * NE * N_BLOCKS);
const size_t n_elems = (size_t)N * D;
auto t0 = std::chrono::steady_clock::now();
for (int l = 0; l < NL; ++l) {
const std::string b = "blk." + std::to_string(l) + ".";
int tg = f.find(b + "ffn_gate_exps.weight")->type;
int tu = f.find(b + "ffn_up_exps.weight")->type;
int td = f.find(b + "ffn_down_exps.weight")->type;
const uint64_t in_off = in_tab[l * 4 + 0];
const uint64_t gs = in_tab[l * 4 + 1];
const uint64_t us = in_tab[l * 4 + 2];
const uint64_t ds = in_tab[l * 4 + 3];
const uint64_t in_exp_step = gs + us + ds;
const size_t in_layer_bytes = (size_t)NE * in_exp_step;
std::vector<uint8_t> in_buf(in_layer_bytes);
pread_exact(ifd, in_buf.data(), in_layer_bytes, (off_t)in_off, "pread in_layer");
std::vector<uint8_t> out_buf(out_layer_bytes);
std::atomic<int> next_e{0};
const int n_threads = 10;
std::vector<std::thread> threads;
for (int tid = 0; tid < n_threads; ++tid) {
threads.emplace_back([&]() {
std::vector<float> f_gate(n_elems), f_up(n_elems), f_down(n_elems), s_down(D, 1.0f);
const float* s2_l = layer_s2_raw[l].data();
for (;;) {
int e = next_e.fetch_add(1);
if (e >= NE) break;
const uint8_t* src = in_buf.data() + (size_t)e * in_exp_step;
uint8_t* dst = out_buf.data() + (size_t)e * out_exp_bytes;
if (tg == ref::Q2_0) {
std::memcpy(dst, src, Q2_SLICE_BYTES);
if (use_real_act && td != ref::Q2_0) ref::dequant(tg, src, f_gate.data(), n_elems);
} else {
ref::dequant(tg, src, f_gate.data(), n_elems);
quantize_q2_0_signed_awq_fast(f_gate.data(), D, N, layer_s_in[l].data(), dst, gain_gu);
}
if (tu == ref::Q2_0) {
std::memcpy(dst + Q2_SLICE_BYTES, src + gs, Q2_SLICE_BYTES);
if (use_real_act && td != ref::Q2_0) ref::dequant(tu, src + gs, f_up.data(), n_elems);
} else {
ref::dequant(tu, src + gs, f_up.data(), n_elems);
quantize_q2_0_signed_awq_fast(f_up.data(), D, N, layer_s_in[l].data(), dst + Q2_SLICE_BYTES, gain_gu);
}
uint8_t* qd_dst = dst + 2 * Q2_SLICE_BYTES;
if (td == ref::Q2_0) {
std::memcpy(qd_dst, src + gs + us, Q2_SLICE_BYTES);
} else {
ref::dequant(td, src + gs + us, f_down.data(), n_elems);
const float* d_imp = nullptr;
if (use_real_act) {
for (int c = 0; c < D; ++c) {
const float* rg = f_gate.data() + (size_t)c * N;
const float* ru = f_up.data() + (size_t)c * N;
double vg = 0.0, vu = 0.0;
for (int d = 0; d < N; ++d) {
const double s2d = s2_l[d];
vg += (double)rg[d] * rg[d] * s2d;
vu += (double)ru[d] * ru[d] * s2d;
}
s_down[c] = (float)std::pow(std::max(vg * vu, 1e-16), 0.25);
}
d_imp = s_down.data();
}
quantize_q2_0_signed_awq_fast(f_down.data(), N, D, d_imp, qd_dst, gain_d);
}
uint16_t* d_bak = ds_all.data() + ((size_t)l * NE + e) * N_BLOCKS;
for (size_t blk = 0; blk < N_BLOCKS; ++blk) {
std::memcpy(&d_bak[blk], qd_dst + blk * 18, 2);
}
}
});
}
for (auto& th : threads) th.join();
pwrite_exact(efd, out_buf.data(), out_buf.size(), (off_t)out_offs[l], "pwrite out_layer");
if ((l + 1) % 8 == 0 || l == NL - 1) {
double el = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();
std::printf(" Quantized layer %d/%d to Signed-d AWQ Q2_0 (%.1fs elapsed)\n", l + 1, NL, el);
std::fflush(stdout);
}
}
FILE* fb = std::fopen(bak_d_path.c_str(), "wb");
require_io(fb != nullptr, "fopen d_bak.bin");
require_io(std::fwrite(ds_all.data(), sizeof(uint16_t), ds_all.size(), fb) == ds_all.size(), "fwrite d_bak.bin");
std::fclose(fb);
// Save .quant.meta fingerprint over (gain_gu, gain_d, cal_pfx, t_cal)
const std::string qmeta_path = out_exp + ".quant.meta";
FILE* fqm = std::fopen(qmeta_path.c_str(), "wb");
require_io(fqm != nullptr, ("fopen " + qmeta_path).c_str());
uint32_t gu_bits = 0, gd_bits = 0;
std::memcpy(&gu_bits, &gain_gu, 4);
std::memcpy(&gd_bits, &gain_d, 4);
const uint64_t pfx_hash = fnv1a_64(cal_pfx.data(), cal_pfx.size());
const uint64_t qmeta[3] = {((uint64_t)gu_bits << 32) | gd_bits, pfx_hash, (uint64_t)t_cal};
require_io(std::fwrite(qmeta, sizeof(uint64_t), 3, fqm) == 3, "fwrite quant.meta");
std::fclose(fqm);
::fsync(efd);
::close(efd);
::close(ifd);
return 0;
}
static int run_calibrate_and_heal_mode(int argc, char** argv) {
if (argc < 7) {
std::fprintf(stderr, "usage: heal_k256_q2_0 --calibrate-and-heal <in-metal.gguf> <out-healed-metal.gguf> <cal_iq3s_pfx> <yr_q2_bin> <out_writers_bin> [gamma=0.75] [eta_shd=0.15] [ridge_rel=6.0]\n");
return 2;
}
const std::string orig_gguf = argv[2];
const std::string healed_gguf = argv[3];
const std::string cal_pfx = argv[4];
const std::string yr_q2_bin = argv[5];
const std::string out_writers_bin = argv[6];
const float gamma = (argc > 7) ? std::atof(argv[7]) : 0.75f;
const float eta_shd = (argc > 8) ? std::atof(argv[8]) : 0.15f;
const float ridge_rel = (argc > 9) ? std::atof(argv[9]) : 6.0f;
const std::string healed_exp = healed_gguf + ".experts";
const std::string bak_d_path = healed_exp + ".d_bak.bin";
gguf::File gf;
std::string err;
if (!gf.open(orig_gguf, err)) {
std::fprintf(stderr, "gguf open failed: %s\n", err.c_str());
return 1;
}
int fd_heal = ::open(healed_exp.c_str(), O_RDWR);
require_io(fd_heal >= 0, "open healed_exp");
uint64_t tab_heal[NL * 4];
pread_exact(fd_heal, tab_heal, sizeof(tab_heal), 64, "pread tab_heal");
FILE* fb = std::fopen(bak_d_path.c_str(), "rb");
require_io(fb != nullptr, ("fopen " + bak_d_path).c_str());
std::vector<uint16_t> ds_all((size_t)NL * NE * N_BLOCKS);
require_io(std::fread(ds_all.data(), sizeof(uint16_t), ds_all.size(), fb) == ds_all.size(), "fread d_bak.bin");
std::fclose(fb);
// Infer t_cal dynamically from (cal_pfx + ".mixed") file size and cross-check .yr, .ids, .wts, .hh sizes
struct stat st_mix{}, st_yr{}, st_ids{}, st_wts{}, st_hh{};
require_io(::stat((cal_pfx + ".mixed").c_str(), &st_mix) == 0 && st_mix.st_size > 0, "stat cal_pfx.mixed");
require_io(st_mix.st_size % ((size_t)NL * N * sizeof(float)) == 0, "cal_pfx.mixed size divisible by NL*N*4");
const size_t t_cal = (size_t)st_mix.st_size / ((size_t)NL * N * sizeof(float));
require_io(::stat((cal_pfx + ".yr").c_str(), &st_yr) == 0 && (size_t)st_yr.st_size == (size_t)NL * t_cal * N * sizeof(float), "cal_pfx.yr size match");
require_io(::stat((cal_pfx + ".ids").c_str(), &st_ids) == 0 && (size_t)st_ids.st_size == (size_t)NL * t_cal * 10 * sizeof(int32_t), "cal_pfx.ids size match");
require_io(::stat((cal_pfx + ".wts").c_str(), &st_wts) == 0 && (size_t)st_wts.st_size == (size_t)NL * t_cal * 10 * sizeof(float), "cal_pfx.wts size match");
require_io(::stat((cal_pfx + ".hh").c_str(), &st_hh) == 0 && (size_t)st_hh.st_size == (size_t)NL * t_cal * D * sizeof(uint16_t), "cal_pfx.hh size match");
int fd_mix = ::open((cal_pfx + ".mixed").c_str(), O_RDONLY);
int fd_yr_ref = ::open((cal_pfx + ".yr").c_str(), O_RDONLY);
require_io(fd_mix >= 0 && fd_yr_ref >= 0, "open cal_pfx .mixed/.yr");
// Compute cache fingerprint over d_bak.bin + L0 & L47 of .mixed + L0/L24/L47 gate/up slices of healed_exp + t_cal
std::vector<float> mix_l0_l47(2 * t_cal * N);
pread_exact(fd_mix, mix_l0_l47.data(), t_cal * N * sizeof(float), 0, "pread mix_l0");
pread_exact(fd_mix, mix_l0_l47.data() + t_cal * N, t_cal * N * sizeof(float), (off_t)((size_t)(NL - 1) * t_cal * N * sizeof(float)), "pread mix_l47");
std::vector<uint8_t> gu_sample(3 * 2 * Q2_SLICE_BYTES);
pread_exact(fd_heal, gu_sample.data() + 0 * 2 * Q2_SLICE_BYTES, 2 * Q2_SLICE_BYTES, (off_t)tab_heal[0 * 4 + 0], "pread gu_l0");
pread_exact(fd_heal, gu_sample.data() + 1 * 2 * Q2_SLICE_BYTES, 2 * Q2_SLICE_BYTES, (off_t)tab_heal[24 * 4 + 0], "pread gu_l24");
pread_exact(fd_heal, gu_sample.data() + 2 * 2 * Q2_SLICE_BYTES, 2 * Q2_SLICE_BYTES, (off_t)tab_heal[47 * 4 + 0], "pread gu_l47");
const uint64_t dbak_hash = fnv1a_64(ds_all.data(), ds_all.size() * sizeof(uint16_t));
const uint64_t mix_hash = fnv1a_64(mix_l0_l47.data(), mix_l0_l47.size() * sizeof(float));
const uint64_t gu_hash = fnv1a_64(gu_sample.data(), gu_sample.size());
const uint64_t want_meta[4] = {dbak_hash, mix_hash, gu_hash, (uint64_t)t_cal};
const std::string meta_path = yr_q2_bin + ".meta";
bool cache_valid = false;
struct stat st_q2{};
if (::stat(yr_q2_bin.c_str(), &st_q2) == 0 && (size_t)st_q2.st_size == (size_t)NL * t_cal * N * sizeof(float)) {
FILE* fm = std::fopen(meta_path.c_str(), "rb");
if (fm) {
uint64_t got_meta[4] = {0, 0, 0, 0};
if (std::fread(got_meta, sizeof(uint64_t), 4, fm) == 4 &&
std::memcmp(got_meta, want_meta, sizeof(want_meta)) == 0) {
cache_valid = true;
}
std::fclose(fm);
}
}
if (!cache_valid) {
int fd_ids = ::open((cal_pfx + ".ids").c_str(), O_RDONLY);
int fd_wts = ::open((cal_pfx + ".wts").c_str(), O_RDONLY);
int fd_q2w = ::open(yr_q2_bin.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644);
require_io(fd_ids >= 0 && fd_wts >= 0 && fd_q2w >= 0, "open cal_pfx .ids/.wts or yr_q2_bin");
std::atomic<int> nl_q2{0};
std::vector<std::thread> q2_ths;
for (int tid = 0; tid < 8; ++tid) {
q2_ths.emplace_back([&]() {
const size_t n_elems = (size_t)N * D;
std::vector<float> x(t_cal * N), yr_q2(t_cal * N, 0.0f);
std::vector<int32_t> ids(t_cal * 10);
std::vector<float> wts(t_cal * 10);
std::vector<uint8_t> in_h((size_t)NE * 3 * Q2_SLICE_BYTES);
std::vector<float> wg(n_elems), wu(n_elems), wd(n_elems);
for (;;) {
int l = nl_q2.fetch_add(1);
if (l >= NL) break;
pread_exact(fd_mix, x.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread mix");
pread_exact(fd_ids, ids.data(), t_cal * 10 * sizeof(int32_t), (off_t)((size_t)l * t_cal * 10 * sizeof(int32_t)), "pread ids");
pread_exact(fd_wts, wts.data(), t_cal * 10 * sizeof(float), (off_t)((size_t)l * t_cal * 10 * sizeof(float)), "pread wts");
pread_exact(fd_heal, in_h.data(), in_h.size(), (off_t)tab_heal[l * 4 + 0], "pread in_h");
for (int e = 0; e < NE; ++e) {
uint8_t* qd_s = in_h.data() + ((size_t)e * 3 + 2) * Q2_SLICE_BYTES;
const uint16_t* d_b = ds_all.data() + ((size_t)l * NE + e) * N_BLOCKS;
for (size_t blk = 0; blk < N_BLOCKS; ++blk) std::memcpy(qd_s + blk * 18, &d_b[blk], 2);
}
std::fill(yr_q2.begin(), yr_q2.end(), 0.0f);
std::vector<std::vector<std::pair<int, float>>> exp_toks(NE);
for (size_t t = 0; t < t_cal; ++t) {
for (int k = 0; k < 10; ++k) {
int e = ids[t * 10 + k];
float w = wts[t * 10 + k];
if (e >= 0 && e < NE && w > 0.0f) exp_toks[e].push_back({(int)t, w});
}
}
std::vector<float> x_e, g_e, u_e, y_e;
for (int e = 0; e < NE; ++e) {
const auto& hits = exp_toks[e];
if (hits.empty()) continue;
int m = (int)hits.size();
x_e.resize((size_t)m * N); g_e.resize((size_t)m * D); u_e.resize((size_t)m * D); y_e.resize((size_t)m * N);
for (int i = 0; i < m; ++i) std::memcpy(x_e.data() + (size_t)i * N, x.data() + (size_t)hits[i].first * N, N * sizeof(float));
const uint8_t* rh = in_h.data() + (size_t)e * 3 * Q2_SLICE_BYTES;
ref::dequant(ref::Q2_0, rh, wg.data(), n_elems);
ref::dequant(ref::Q2_0, rh + Q2_SLICE_BYTES, wu.data(), n_elems);
ref::dequant(ref::Q2_0, rh + 2 * Q2_SLICE_BYTES, wd.data(), n_elems);
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m, D, N, 1.0f, x_e.data(), N, wg.data(), N, 0.0f, g_e.data(), D);
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m, D, N, 1.0f, x_e.data(), N, wu.data(), N, 0.0f, u_e.data(), D);
for (int i = 0; i < m * D; ++i) { float gv = g_e[i]; g_e[i] = (gv / (1.0f + std::exp(-gv))) * u_e[i]; }
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m, N, D, 1.0f, g_e.data(), D, wd.data(), D, 0.0f, y_e.data(), N);
for (int i = 0; i < m; ++i) {
float* dst = yr_q2.data() + (size_t)hits[i].first * N;
const float* src = y_e.data() + (size_t)i * N;
float w = hits[i].second;
for (int c = 0; c < N; ++c) dst[c] += w * src[c];
}
}
pwrite_exact(fd_q2w, yr_q2.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pwrite yr_q2");
}
});
}
for (auto& th : q2_ths) th.join();
::close(fd_q2w); ::close(fd_wts); ::close(fd_ids);
FILE* fm = std::fopen(meta_path.c_str(), "wb");
require_io(fm != nullptr, "fopen yr_q2_bin.meta");
require_io(std::fwrite(want_meta, sizeof(uint64_t), 4, fm) == 4, "fwrite yr_q2_bin.meta");
std::fclose(fm);
}
int fd_yr_q2 = ::open(yr_q2_bin.c_str(), O_RDONLY);
require_io(fd_yr_q2 >= 0, "open yr_q2_bin");
// Part A (Real-Trace Layerwise Scale Calibration) + Part B (Full-Rank Shared-Expert Ridge Regression)
FILE* fout = std::fopen(out_writers_bin.c_str(), "wb");
require_io(fout != nullptr, "fopen out_writers_bin");
uint32_t n_tensors = (eta_shd > 0.0f) ? NL : 0;
require_io(std::fwrite(&n_tensors, sizeof(uint32_t), 1, fout) == 1, "fwrite n_tensors");
auto write_record = [&](const std::string& name, uint32_t gtype, const void* data, uint32_t nbytes) {
uint32_t nlen = (uint32_t)name.size();
require_io(std::fwrite(&nlen, 4, 1, fout) == 1, "fwrite nlen");
require_io(std::fwrite(name.data(), 1, nlen, fout) == nlen, "fwrite name");
require_io(std::fwrite(>ype, 4, 1, fout) == 1, "fwrite gtype");
require_io(std::fwrite(&nbytes, 4, 1, fout) == 1, "fwrite nbytes");
require_io(std::fwrite(data, 1, nbytes, fout) == nbytes, "fwrite data");
};
std::vector<float> layer_alpha(NL, 1.0f);
std::vector<uint8_t> slice(Q2_SLICE_BYTES);
double tot_ref2 = 0.0, tot_dy2_unscaled = 0.0, tot_dy2_before = 0.0, tot_dy2_after = 0.0;
for (int l = 0; l < NL; ++l) {
std::vector<float> x(t_cal * N), yr_r(t_cal * N), yr_q(t_cal * N), dy(t_cal * N);
pread_exact(fd_mix, x.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread mix");
pread_exact(fd_yr_ref, yr_r.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread yr_ref");
pread_exact(fd_yr_q2, yr_q.data(), t_cal * N * sizeof(float), (off_t)((size_t)l * t_cal * N * sizeof(float)), "pread yr_q2");
double sum_ref2 = 0.0, sum_q2 = 0.0, sum_dot = 0.0;
for (size_t i = 0; i < t_cal * N; ++i) {
double r = yr_r[i], q = yr_q[i];
sum_ref2 += r * r;
sum_q2 += q * q;
sum_dot += r * q;
double du = r - q;
tot_dy2_unscaled += du * du;
}
tot_ref2 += sum_ref2;
float alpha_l = 1.0f;
if (gamma >= 0.0f && sum_ref2 > 0.0 && sum_q2 > 0.0) {
double norm_ratio = std::sqrt(sum_q2 / sum_ref2);
double cos_sim = std::max(sum_dot / (std::sqrt(sum_ref2) * std::sqrt(sum_q2)), 0.1);
alpha_l = (float)(std::pow(cos_sim, 1.0 - (double)gamma) / norm_ratio);
}
layer_alpha[l] = alpha_l;
// Apply alpha_l to W_down fp16 d headers in healed_exp
uint64_t off_h = tab_heal[l * 4 + 0];
for (int e = 0; e < NE; ++e) {
const uint16_t* d_bak = ds_all.data() + ((size_t)l * NE + e) * N_BLOCKS;
uint64_t d_off = off_h + (uint64_t)e * (3 * Q2_SLICE_BYTES) + 2 * Q2_SLICE_BYTES;
pread_exact(fd_heal, slice.data(), Q2_SLICE_BYTES, (off_t)d_off, "pread slice");
for (size_t b = 0; b < N_BLOCKS; ++b) {
float d = ref::f16_to_f32(d_bak[b]);
uint16_t d_new = ref::f32_to_f16(d * alpha_l);
std::memcpy(slice.data() + b * 18, &d_new, 2);
}
pwrite_exact(fd_heal, slice.data(), Q2_SLICE_BYTES, (off_t)d_off, "pwrite slice");
}
for (size_t i = 0; i < t_cal * N; ++i) {
dy[i] = yr_r[i] - alpha_l * yr_q[i];
tot_dy2_before += (double)dy[i] * dy[i];
}
if (eta_shd <= 0.0f) {
tot_dy2_after = tot_dy2_before;
continue;
}
const std::string b = "blk." + std::to_string(l) + ".";
const auto* tn_shg = gf.find(b + "ffn_gate_shexp.weight");
const auto* tn_shu = gf.find(b + "ffn_up_shexp.weight");
const auto* tn_shd = gf.find(b + "ffn_down_shexp.weight");
const auto* tn_sg = gf.find(b + "ffn_gate_inp_shexp.weight");
std::vector<float> w_shg(D * N), w_shu(D * N), w_shd(N * D), w_sg(N);
ref::dequant(tn_shg->type, (const uint8_t*)gf.map + tn_shg->offset, w_shg.data(), D * N);
ref::dequant(tn_shu->type, (const uint8_t*)gf.map + tn_shu->offset, w_shu.data(), D * N);
ref::dequant(tn_shd->type, (const uint8_t*)gf.map + tn_shd->offset, w_shd.data(), N * D);
std::memcpy(w_sg.data(), (const uint8_t*)gf.map + tn_sg->offset, N * sizeof(float));
std::vector<float> g_sh(t_cal * D), u_sh(t_cal * D), z_sh(t_cal * D), sg_l(t_cal);
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (int)t_cal, D, N, 1.0f, x.data(), N, w_shg.data(), N, 0.0f, g_sh.data(), D);
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (int)t_cal, D, N, 1.0f, x.data(), N, w_shu.data(), N, 0.0f, u_sh.data(), D);
cblas_sgemv(CblasRowMajor, CblasNoTrans, (int)t_cal, N, 1.0f, x.data(), N, w_sg.data(), 1, 0.0f, sg_l.data(), 1);
for (size_t t = 0; t < t_cal; ++t) {
float sig = 1.0f / (1.0f + std::exp(-sg_l[t]));
for (int c = 0; c < D; ++c) {
float gv = g_sh[t * D + c];
z_sh[t * D + c] = sig * (gv / (1.0f + std::exp(-gv))) * u_sh[t * D + c];
}
}
std::vector<float> G(D * D, 0.0f), B(D * N, 0.0f);
cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, D, D, (int)t_cal, 1.0f, z_sh.data(), D, z_sh.data(), D, 0.0f, G.data(), D);
cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, D, N, (int)t_cal, 1.0f, z_sh.data(), D, dy.data(), N, 0.0f, B.data(), N);
double trace_G = 0.0;
for (int c = 0; c < D; ++c) trace_G += G[c * D + c];
float lam = (float)(ridge_rel * (trace_G / D));
for (int c = 0; c < D; ++c) G[c * D + c] += lam;
std::vector<double> Gd(D * D), L_chol(D * D, 0.0);
for (int i = 0; i < D * D; ++i) Gd[i] = G[i];
for (int i = 0; i < D; ++i) {
for (int j = 0; j <= i; ++j) {
double s = Gd[i * D + j];
for (int k = 0; k < j; ++k) s -= L_chol[i * D + k] * L_chol[j * D + k];
if (i == j) L_chol[i * D + j] = std::sqrt(std::max(s, 1e-12));
else L_chol[i * D + j] = s / L_chol[j * D + j];
}
}
std::vector<float> w_shd_target = w_shd;
for (int r = 0; r < N; ++r) {
double y_fwd[D], x_bwd[D];
for (int i = 0; i < D; ++i) {
double s = B[i * N + r];
for (int k = 0; k < i; ++k) s -= L_chol[i * D + k] * y_fwd[k];
y_fwd[i] = s / L_chol[i * D + i];
}
for (int i = D - 1; i >= 0; --i) {
double s = y_fwd[i];
for (int k = i + 1; k < D; ++k) s -= L_chol[k * D + i] * x_bwd[k];
x_bwd[i] = s / L_chol[i * D + i];
}
for (int c = 0; c < D; ++c) {
w_shd_target[r * D + c] += eta_shd * (float)x_bwd[c];
}
}
const uint8_t* raw_shd_old = (const uint8_t*)gf.map + tn_shd->offset;
std::vector<uint8_t> raw_shd_new(tn_shd->bytes);
if (tn_shd->type == ref::Q8_0) {
quantize_q8_0_monotonic(raw_shd_old, w_shd_target.data(), N * D, raw_shd_new.data());
} else {
quantize_iq4nl_monotonic(raw_shd_old, w_shd_target.data(), N * D, raw_shd_new.data());
}
std::vector<float> w_shd_new_deq(N * D), delta_w(N * D), pred_dy(t_cal * N);
ref::dequant(tn_shd->type, raw_shd_new.data(), w_shd_new_deq.data(), N * D);
for (int i = 0; i < N * D; ++i) delta_w[i] = w_shd_new_deq[i] - w_shd[i];
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, (int)t_cal, N, D, 1.0f, z_sh.data(), D, delta_w.data(), D, 0.0f, pred_dy.data(), N);
for (size_t i = 0; i < t_cal * N; ++i) {
double rem = (double)dy[i] - (double)pred_dy[i];
tot_dy2_after += rem * rem;
}
write_record(b + "ffn_down_shexp.weight", (uint32_t)tn_shd->type, raw_shd_new.data(), (uint32_t)raw_shd_new.size());
}
std::fclose(fout);
::fsync(fd_heal);
::close(fd_heal);
::close(fd_mix);
::close(fd_yr_ref);
::close(fd_yr_q2);
std::printf("Calibrated 48 layers on real trace (T=%zu, gamma=%.2f, L0=%.3f L24=%.3f L47=%.3f) & solved %u ffn_down_shexp tensors (MoE MSE %.2f -> %.2f -> %.2f, %.2f%%)\n",
t_cal, gamma, layer_alpha[0], layer_alpha[24], layer_alpha[47], n_tensors,
tot_dy2_unscaled, tot_dy2_before, tot_dy2_after, (tot_dy2_after / tot_dy2_unscaled - 1.0) * 100.0);
return 0;
}
int main(int argc, char** argv) {
if (argc >= 2 && std::string(argv[1]) == "--quantize") {
return run_quantize_mode(argc, argv);
}
if (argc >= 2 && std::string(argv[1]) == "--calibrate-and-heal") {
return run_calibrate_and_heal_mode(argc, argv);
}
if (argc >= 3) {
std::vector<char*> args = {argv[0], (char*)"--quantize"};
for (int i = 1; i < argc; ++i) args.push_back(argv[i]);
return run_quantize_mode((int)args.size(), args.data());
}
std::fprintf(stderr, "usage: heal_k256_q2_0 [--quantize | --calibrate-and-heal] ...\n");
return 2;
}