#define ACCELERATE_NEW_LAPACK
#include "vendor/strata/metal/src/gguf_lite.h"
#include "vendor/strata/metal/src/ref_quant.h"
#include <Accelerate/Accelerate.h>
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <map>
#include <set>
#include <string>
#include <vector>
static constexpr int N = 2560, HC = 4, LR = 320, F = 640, V = 248320;
static std::vector<char> rdall(const std::string& p) {
FILE* f = std::fopen(p.c_str(), "rb");
if (!f) { std::perror(p.c_str()); std::exit(2); }
std::fseek(f, 0, SEEK_END);
long n = std::ftell(f);
std::fseek(f, 0, SEEK_SET);
std::vector<char> b(n);
if (std::fread(b.data(), 1, n, f) != (size_t)n) { std::perror("fread"); std::exit(2); }
std::fclose(f);
return b;
}
static std::vector<int> rdtoks(const std::string& p) {
std::vector<int> v;
FILE* f = std::fopen(p.c_str(), "r");
if (!f) { std::perror(p.c_str()); std::exit(2); }
int x;
while (std::fscanf(f, "%d", &x) == 1) v.push_back(x);
std::fclose(f);
return v;
}
struct TraceData {
int T = 0;
std::vector<int> toks;
std::vector<float> R; // [T, HC, N]
std::vector<float> inj; // [T, HC]
std::vector<float> z; // [T, F] = sigmoid(sgv) * hh
std::vector<float> mix0; // [T, N] at dW = 0
std::vector<float> gam; // [T, N] diagonal Jacobian d(mix)/d(d) at dW = 0
};
static void compute_mix_batch(
int T,
const float* R,
const float* inj,
const float* d_in, // [T, N] or nullptr
const std::vector<float>& wn,
const std::vector<float>& wdn,
const std::vector<float>& wup,
float* mix_out,
float* gam_out = nullptr
) {
std::vector<float> xn((size_t)T * HC * N);
std::vector<float> rs_all((size_t)T * HC);
std::vector<float> gg_all((size_t)T * HC);
for (int t = 0; t < T; ++t) {
const float* dt = d_in ? &d_in[(size_t)t * N] : nullptr;
for (int c = 0; c < HC; ++c) {
float gg = 2.0f / (1.0f + std::exp(-inj[t * HC + c] / 4.0f));
gg_all[t * HC + c] = gg;
const float* rc = &R[((size_t)t * HC + c) * N];
float* xnc = &xn[((size_t)t * HC + c) * N];
double ss = 0.0;
for (int j = 0; j < N; ++j) {
float v = rc[j] + (dt ? gg * dt[j] : 0.0f);
xnc[j] = v;
ss += (double)v * v;
}
float rs = 1.0f / std::sqrt((float)(ss / N) + 1e-6f);
rs_all[t * HC + c] = rs;
const float* wnc = &wn[c * N];
for (int j = 0; j < N; ++j) xnc[j] *= wnc[j] * rs;
}
}
std::vector<float> lo((size_t)T * LR);
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, T, LR, HC * N, 0.25f, xn.data(), HC * N, wdn.data(), HC * N, 0.0f, lo.data(), LR);
for (size_t i = 0; i < lo.size(); ++i) {
float z = lo[i];
lo[i] = z / (1.0f + std::exp(-z));
}
std::vector<float> gup((size_t)T * HC * N);
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, T, HC * N, LR, 1.0f, lo.data(), LR, wup.data(), LR, 0.0f, gup.data(), HC * N);
for (int t = 0; t < T; ++t) {
float* mix = &mix_out[(size_t)t * N];
float* gm = gam_out ? &gam_out[(size_t)t * N] : nullptr;
for (int j = 0; j < N; ++j) {
float m = 0.0f, gs = 0.0f;
for (int c = 0; c < HC; ++c) {
float s = gup[((size_t)t * HC + c) * N + j];
float gate = 1.0f / (1.0f + std::exp(-s));
m += xn[((size_t)t * HC + c) * N + j] * gate;
if (gm) gs += gg_all[t * HC + c] * rs_all[t * HC + c] * wn[c * N + j] * gate;
}
mix[j] = m * 0.25f;
if (gm) gm[j] = gs * 0.25f;
}
}
}
static void load_trace(
const std::string& pfx,
const std::string& tok_file,
const std::vector<float>& wn,
const std::vector<float>& wdn,
const std::vector<float>& wup,
TraceData& tr
) {
tr.toks = rdtoks(tok_file);
auto Rb = rdall(pfx + ".R_post47");
auto Ib = rdall(pfx + ".inj47");
auto Hb = rdall(pfx + ".hh47");
auto Sb = rdall(pfx + ".sgv47");
tr.T = (int)(Rb.size() / (HC * N * sizeof(float)));
tr.R.resize((size_t)tr.T * HC * N);
tr.inj.resize((size_t)tr.T * HC);
std::memcpy(tr.R.data(), Rb.data(), Rb.size());
std::memcpy(tr.inj.data(), Ib.data(), Ib.size());
const __fp16* hh = (const __fp16*)Hb.data();
const float* sg = (const float*)Sb.data();
tr.z.resize((size_t)tr.T * F);
for (int t = 0; t < tr.T; ++t) {
float sgs = 1.0f / (1.0f + std::exp(-sg[t]));
for (int k = 0; k < F; ++k) tr.z[(size_t)t * F + k] = sgs * (float)hh[(size_t)t * F + k];
}
tr.mix0.resize((size_t)tr.T * N);
tr.gam.resize((size_t)tr.T * N);
compute_mix_batch(tr.T, tr.R.data(), tr.inj.data(), nullptr, wn, wdn, wup, tr.mix0.data(), tr.gam.data());
}
struct FtpoRow {
int split, phrase_id, t_pos, rej_id;
std::vector<int> alt_ids;
};
struct ColRow {
int split, ban_idx, t_pos, target_id;
};
struct SepHoRow {
int phrase_id, t_pos, rej_id;
};
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;
for (int i = 0; i < 32; ++i) {
float diff = w_old_deq[i] - wt[i];
err_old += diff * diff;
}
uint8_t cand_blk[34];
ref::quantize_q8_0(wt, 32, cand_blk);
std::vector<float> w_cand_deq(32);
ref::dequant(ref::Q8_0, cand_blk, w_cand_deq.data(), 32);
float err_c = 0.0f;
for (int i = 0; i < 32; ++i) {
float diff = w_cand_deq[i] - wt[i];
err_c += diff * diff;
}
if (err_c + 1e-11f < err_old) std::memcpy(bn, cand_blk, 34);
}
}
int main() {
gguf::File g;
std::string err;
if (!g.open("models/K256_reap_l2_sum/K256-q2_0-healed-metal.gguf", err)) return 1;
auto deq = [&](const char* n, std::vector<float>& out) {
auto* t = g.find(n);
size_t ne = 1;
for (auto d : t->shape) ne *= d;
out.resize(ne);
ref::dequant(t->type, (const uint8_t*)g.map + t->offset, out.data(), ne);
};
auto q8rt = [&](std::vector<float>& w) {
std::vector<uint8_t> q(w.size() / 32 * 34);
ref::quantize_q8_0(w.data(), w.size(), q.data());
ref::dequant(ref::Q8_0, q.data(), w.data(), w.size());
};
std::vector<float> wn, wdn, wup, wb;
deq("output_hc_norm.weight", wn);
deq("output_hc_down.weight", wdn); q8rt(wdn);
deq("output_hc_up.weight", wup); q8rt(wup);
deq("blk.47.ffn_down_shexp.weight", wb);
auto* t_sd47 = g.find("blk.47.ffn_down_shexp.weight");
const uint8_t* raw_sd47_orig = (const uint8_t*)g.map + t_sd47->offset;
TraceData tr_cal, tr_p4, tr_col, tr_ho;
load_trace("models/K256_reap_l2_sum/cal_cache/cal_p3_1024", "data/m5_val_prompts/cal_disjoint_1024tok.txt", wn, wdn, wup, tr_cal);
load_trace("models/K256_reap_l2_sum/cal_cache/p4_ftpo_sft_1024", "data/m5_val_prompts/phase4_ftpo_sft_1024tok.txt", wn, wdn, wup, tr_p4);
load_trace("models/K256_reap_l2_sum/cal_cache/p4_collateral_1640", "data/m5_val_prompts/phase4_collateral_1024tok.txt", wn, wdn, wup, tr_col);
load_trace("models/K256_reap_l2_sum/cal_cache/p4_ftpo_heldout_1024", "data/m5_val_prompts/phase4_ftpo_heldout_1024tok.txt", wn, wdn, wup, tr_ho);
FILE* fm = std::fopen("/tmp/p4_meta_table.txt", "r");
int n_ftpo = 0; std::fscanf(fm, "%d", &n_ftpo);
std::vector<FtpoRow> ftpo(n_ftpo);
std::set<int> all_p4_slop_set, train_p4_slop_set;
for (int i = 0; i < n_ftpo; ++i) {
int na = 0;
std::fscanf(fm, "%d %d %d %d %d", &ftpo[i].split, &ftpo[i].phrase_id, &ftpo[i].t_pos, &ftpo[i].rej_id, &na);
ftpo[i].alt_ids.resize(na);
for (int k = 0; k < na; ++k) std::fscanf(fm, "%d", &ftpo[i].alt_ids[k]);
all_p4_slop_set.insert(ftpo[i].t_pos);
if (ftpo[i].split == 0) train_p4_slop_set.insert(ftpo[i].t_pos);
}
int n_col = 0; std::fscanf(fm, "%d", &n_col);
std::vector<ColRow> col(n_col);
for (int i = 0; i < n_col; ++i) std::fscanf(fm, "%d %d %d %d", &col[i].split, &col[i].ban_idx, &col[i].t_pos, &col[i].target_id);
int n_sep = 0; std::fscanf(fm, "%d", &n_sep);
std::vector<SepHoRow> sep_ho(n_sep);
std::set<int> sep_ho_set;
for (int i = 0; i < n_sep; ++i) {
std::fscanf(fm, "%d %d %d", &sep_ho[i].phrase_id, &sep_ho[i].t_pos, &sep_ho[i].rej_id);
sep_ho_set.insert(sep_ho[i].t_pos);
}
std::fclose(fm);
auto* tout = g.find("output.weight");
const size_t rb = tout->bytes / V;
auto get_wout_row = [&](int vid, std::vector<float>& row) {
row.resize(N);
ref::dequant(tout->type, (const uint8_t*)g.map + tout->offset + (size_t)vid * rb, row.data(), N);
};
// Notice: in tr_p4 (positions 0..460, the train multi-turn chat region), every token NOT inall_p4_slop_set
// is a real multi-turn chat background token (user turns, <im_start>, <im_end>, <think>, </think>, and ordinary assistant words)!
// Weighting those multi-turn chat background tokens strongly in G_base (e.g. w=60.0) AND putting their means
// in the zero-inner-product constraints forces v_pid . z[t] ~ 0 on ALL ordinary multi-turn chat positions!
const int K_PHR = 7, K_COL = 6;
std::vector<double> G_base((size_t)F * F, 0.0);
auto accum_z_to = [&](std::vector<double>& Gmat, const float* zt, double w) {
for (int r = 0; r < F; ++r) {
double zr = w * (double)zt[r];
for (int c = r; c < F; ++c) Gmat[(size_t)r * F + c] += zr * (double)zt[c];
}
};
for (int t = 0; t < tr_cal.T; ++t) accum_z_to(G_base, &tr_cal.z[(size_t)t * F], 1.0);
for (int t = 0; t < 480; ++t) accum_z_to(G_base, &tr_col.z[(size_t)t * F], 4.0);
std::vector<double> M_chat_bg(F, 0.0), M_chat_imend(F, 0.0), M_cal_bg(F, 0.0);
std::vector<float> gam_chat_bg(N, 0.0f);
int n_chat_bg = 0, n_chat_imend = 0;
for (int t = 0; t < tr_p4.T; ++t) {
if (all_p4_slop_set.count(t)) continue;
// Only use the train half of tr_p4 (t < 310 or t >= 460) so heldout_14 context (310..460) is never touched!
if (t >= 310 && t < 460) continue;
const float* zt = &tr_p4.z[(size_t)t * F];
if (t % 2 == 0) accum_z_to(G_base, zt, 2.0);
int next_tok = tr_p4.toks[t + 1];
if (next_tok == 248046 || next_tok == 248045 || next_tok == 248068 || next_tok == 248069) {
for (int k = 0; k < F; ++k) M_chat_imend[k] += (double)zt[k];
n_chat_imend++;
} else {
for (int k = 0; k < F; ++k) M_chat_bg[k] += (double)zt[k];
const float* gt = &tr_p4.gam[(size_t)t * N];
for (int j = 0; j < N; ++j) gam_chat_bg[j] += gt[j];
n_chat_bg++;
}
}
for (int k = 0; k < F; ++k) {
if (n_chat_bg > 0) M_chat_bg[k] /= n_chat_bg;
if (n_chat_imend > 0) M_chat_imend[k] /= n_chat_imend;
}
if (n_chat_bg > 0) for (int j = 0; j < N; ++j) gam_chat_bg[j] /= (float)n_chat_bg;
for (int t = 0; t < tr_cal.T; ++t) {
for (int k = 0; k < F; ++k) M_cal_bg[k] += (double)tr_cal.z[(size_t)t * F + k] / tr_cal.T;
}
std::vector<double> M_phr((size_t)K_PHR * F, 0.0), M_col_m((size_t)K_COL * F, 0.0);
std::vector<int> cnt_phr(K_PHR, 0), cnt_col_m(K_COL, 0);
int phrase_rej_id[K_PHR] = {};
std::vector<std::vector<int>> phrase_alts(K_PHR);
std::vector<float> phrase_gam((size_t)K_PHR * N, 0.0f);
for (const auto& fr : ftpo) {
if (fr.split == 0) {
int pid = fr.phrase_id;
phrase_rej_id[pid] = fr.rej_id;
const float* zt = &tr_p4.z[(size_t)fr.t_pos * F];
const float* gt = &tr_p4.gam[(size_t)fr.t_pos * N];
for (int k = 0; k < F; ++k) M_phr[(size_t)pid * F + k] += (double)zt[k];
for (int j = 0; j < N; ++j) phrase_gam[(size_t)pid * N + j] += gt[j];
cnt_phr[pid]++;
accum_z_to(G_base, zt, 8.0);
}
}
for (const auto& cr : col) {
if (cr.split == 0) {
int b = cr.ban_idx;
const float* zt = &tr_col.z[(size_t)cr.t_pos * F];
for (int k = 0; k < F; ++k) M_col_m[(size_t)b * F + k] += (double)zt[k];
cnt_col_m[b]++;
accum_z_to(G_base, zt, 40.0);
}
}
for (int p = 0; p < K_PHR; ++p) if (cnt_phr[p] > 0) {
for (int k = 0; k < F; ++k) M_phr[(size_t)p * F + k] /= cnt_phr[p];
for (int j = 0; j < N; ++j) phrase_gam[(size_t)p * N + j] /= cnt_phr[p];
}
for (int b = 0; b < K_COL; ++b) if (cnt_col_m[b] > 0) {
for (int k = 0; k < F; ++k) M_col_m[(size_t)b * F + k] /= cnt_col_m[b];
}
// Full pass over output.weight for ALL 4 traces: tr_cal (1022), tr_p4 (1022), tr_ho (1022), col (207)
const int M_cal = tr_cal.T, M_p4 = tr_p4.T, M_ho = tr_ho.T, M_col = (int)col.size();
const int M_base = M_cal + M_p4 + M_ho + M_col;
std::vector<float> X_base((size_t)M_base * N);
for (int t = 0; t < M_cal; ++t) std::memcpy(&X_base[(size_t)t * N], &tr_cal.mix0[(size_t)t * N], N * sizeof(float));
for (int t = 0; t < M_p4; ++t) std::memcpy(&X_base[(size_t)(M_cal + t) * N], &tr_p4.mix0[(size_t)t * N], N * sizeof(float));
for (int t = 0; t < M_ho; ++t) std::memcpy(&X_base[(size_t)(M_cal + M_p4 + t) * N], &tr_ho.mix0[(size_t)t * N], N * sizeof(float));
for (int i = 0; i < M_col; ++i) std::memcpy(&X_base[(size_t)(M_cal + M_p4 + M_ho + i) * N], &tr_col.mix0[(size_t)col[i].t_pos * N], N * sizeof(float));
const int CH = 8192;
std::vector<float> Wc((size_t)CH * N), L_base((size_t)M_base * CH);
std::vector<float> mx_b(M_base, -1e30f), se_b(M_base, 0.0f);
std::vector<int> argmax_b(M_base, -1);
std::vector<float> y_tgt_b(M_base, 0.0f);
std::vector<double> w_bar_d(N, 0.0);
std::vector<float> phrase_train_logits((size_t)K_PHR * V, 0.0f);
for (int v0 = 0; v0 < V; v0 += CH) {
int nv = std::min(CH, V - v0);
ref::dequant(tout->type, (const uint8_t*)g.map + tout->offset + (size_t)v0 * rb, Wc.data(), (size_t)nv * N);
for (int r = 0; r < nv; ++r) {
const float* wr = &Wc[(size_t)r * N];
for (int j = 0; j < N; ++j) w_bar_d[j] += (double)wr[j];
}
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, M_base, nv, N, 1.0f, X_base.data(), N, Wc.data(), N, 0.0f, L_base.data(), CH);
for (int m = 0; m < M_base; ++m) {
const float* l = &L_base[(size_t)m * CH];
float cm = mx_b[m];
int best_k = -1;
for (int k = 0; k < nv; ++k) {
if (l[k] > cm) { cm = l[k]; best_k = v0 + k; }
}
float s = se_b[m] * std::exp(mx_b[m] - cm);
for (int k = 0; k < nv; ++k) s += std::exp(l[k] - cm);
if (best_k >= 0) argmax_b[m] = best_k;
mx_b[m] = cm; se_b[m] = s;
}
for (int t = 0; t < M_ho; ++t) {
int tgt = tr_ho.toks[t + 1];
if (tgt >= v0 && tgt < v0 + nv) y_tgt_b[M_cal + M_p4 + t] = L_base[(size_t)(M_cal + M_p4 + t) * CH + (tgt - v0)];
}
for (const auto& fr : ftpo) {
if (fr.split == 0) {
const float* l = &L_base[(size_t)(M_cal + fr.t_pos) * CH];
float* dst = &phrase_train_logits[(size_t)fr.phrase_id * V + v0];
for (int k = 0; k < nv; ++k) dst[k] += l[k] / (float)cnt_phr[fr.phrase_id];
}
}
}
std::vector<float> lse_b(M_base);
for (int m = 0; m < M_base; ++m) lse_b[m] = mx_b[m] + std::log(se_b[m]);
std::vector<float> w_bar(N);
for (int j = 0; j < N; ++j) w_bar[j] = (float)(w_bar_d[j] / V);
// Exclude banned words, special tokens, punctuation, and ultra-common function words (like 279 " the", 760 "Let") from phrase_alts
// so boosting phrase_alts[pid] NEVER perturbs ordinary chat tokens when v_z < 0!
// Explicit clean natural non-common alternatives per phrase (from train_logits / grammar):
phrase_alts[0] = {1151, 3839}; // "end to": " start" (1151), " End" (3839)
phrase_alts[1] = {1204, 1332}; // "exactly": " how" (1204), " where" (1332)
phrase_alts[2] = {1204, 1332}; // "here's": " how" (1204), " where" (1332)
phrase_alts[3] = {383, 4237}; // "rather": " on" (383), " instead" (4237)
phrase_alts[4] = {1238, 1518}; // "right": " after" (1238), " before" (1518)
phrase_alts[5] = {4087, 3255}; // "short": " answer" (4087), " story" (3255)
phrase_alts[6] = {12125, 2019}; // "what's": " happening" (12125), " going" (2019)
int ban_ids_arr[6] = {809, 1056, 1092, 2243, 3023, 3404};
const char* ban_names[6] = {" end", " than", " what", " version", " away", " actually"};
auto find_ban_idx = [&](int rid) {
for (int i = 0; i < 6; ++i) if (ban_ids_arr[i] == rid) return i;
return 0;
};
// Solve per-phrase v_pid in R^640 with exact orthogonality to:
// - M_col_m[bidx] (collateral mean for that word)
// - M_cal_bg (calibration mean)
// - tr_p4.z[705] (SFT anchor for pid=0,3)
std::vector<double> V_cls((size_t)K_PHR * F, 0.0);
for (int pid = 0; pid < K_PHR; ++pid) {
int bidx = find_ban_idx(phrase_rej_id[pid]);
std::vector<double> G_b = G_base;
double trG_base = 0.0;
for (int k = 0; k < F; ++k) trG_base += G_base[(size_t)k * F + k];
double ridge = 1.20 * (trG_base / F);
double col_w = (pid == 0 || pid == 3 || pid == 4 || pid == 6) ? 12000.0 : 6500.0;
for (const auto& cr : col) {
if (cr.split == 0) {
double w_cr = (cr.ban_idx == bidx) ? col_w : 400.0;
accum_z_to(G_b, &tr_col.z[(size_t)cr.t_pos * F], w_cr);
}
}
for (int r = 0; r < F; ++r) {
G_b[(size_t)r * F + r] += ridge;
for (int c = 0; c < r; ++c) G_b[(size_t)r * F + c] = G_b[(size_t)c * F + r];
}
std::vector<std::vector<double>> c_rows;
c_rows.push_back(std::vector<double>(&M_phr[(size_t)pid * F], &M_phr[(size_t)(pid + 1) * F]));
if (pid == 1) c_rows.push_back(std::vector<double>(&M_phr[2 * F], &M_phr[3 * F]));
if (pid == 2) c_rows.push_back(std::vector<double>(&M_phr[1 * F], &M_phr[2 * F]));
c_rows.push_back(std::vector<double>(&M_col_m[(size_t)bidx * F], &M_col_m[(size_t)(bidx + 1) * F]));
c_rows.push_back(M_cal_bg);
if (pid == 0 || pid == 3) {
c_rows.push_back(std::vector<double>(&tr_p4.z[705 * F], &tr_p4.z[706 * F]));
}
int ks = (int)c_rows.size();
std::vector<double> Q_sub((size_t)ks * F, 0.0);
for (int i = 0; i < ks; ++i) for (int k = 0; k < F; ++k) Q_sub[(size_t)i * F + k] = c_rows[i][k];
char uplo = 'L';
int n_lap = F, nrhs = ks, lda = F, ldb = F, info = 0;
dposv_(&uplo, &n_lap, &nrhs, G_b.data(), &lda, Q_sub.data(), &ldb, &info);
std::vector<double> S_sub((size_t)ks * ks, 0.0);
for (int i = 0; i < ks; ++i) {
for (int j = 0; j < ks; ++j) {
double s = 0.0;
for (int k = 0; k < F; ++k) s += c_rows[i][k] * Q_sub[(size_t)j * F + k];
S_sub[(size_t)i * ks + j] = s;
}
S_sub[(size_t)i * ks + i] += 1e-6;
}
std::vector<double> rhs(ks, 0.0);
rhs[0] = 1.0;
int nr = 1, info_s = 0;
dposv_(&uplo, &ks, &nr, S_sub.data(), &ks, rhs.data(), &ks, &info_s);
for (int i = 0; i < ks; ++i) {
for (int k = 0; k < F; ++k) V_cls[(size_t)pid * F + k] += rhs[i] * Q_sub[(size_t)i * F + k];
}
}
// Output-space u_pid in R^2560:
// Constrain u_pid against:
// 0: a_rej -> -delta_rej
// 1: a_cho -> +delta_cho (natural alternatives)
// 2: a_im_end (248046) -> -0.15 * delta_rej (slight suppression of <|im_end|> so it NEVER rises!)
// 3: a_im_start (248045) -> 0.0
// 4: a_think (248068) -> 0.0
// 5: a_ethink (248069) -> 0.0
// 6: a_eos (248044) -> 0.0
int spec_ids[11] = {248046, 248045, 248068, 248069, 248044, 279, 760, 11, 13, 198, 3437};
std::vector<float> w_spec[11];
for (int s = 0; s < 11; ++s) get_wout_row(spec_ids[s], w_spec[s]);
double phrase_scales[7] = {3.7, 5.4, 5.4, 4.2, 4.1, 5.5, 3.5};
for (double mult : {0.88}) {
std::vector<float> dW((size_t)N * F, 0.0f);
for (int pid = 0; pid < K_PHR; ++pid) {
double delta_rej = mult * phrase_scales[pid];
std::vector<float> w_r, wc1, wc2;
get_wout_row(phrase_rej_id[pid], w_r);
get_wout_row(phrase_alts[pid][0], wc1);
get_wout_row(phrase_alts[pid][1], wc2);
const float* gm = &phrase_gam[(size_t)pid * N];
double d_c1 = (pid == 0) ? (4.50 * delta_rej) : ((pid == 3) ? (0.65 * delta_rej) : (0.16 * delta_rej));
double d_c2 = (pid == 0) ? (3.00 * delta_rej) : ((pid == 3) ? (0.40 * delta_rej) : (0.10 * delta_rej));
double d_imend = (pid == 0) ? (-0.28 * delta_rej) : ((pid == 3) ? (-0.26 * delta_rej) : (-0.02 * delta_rej));
const int NU = 15;
std::vector<std::vector<double>> A(NU, std::vector<double>(N));
for (int j = 0; j < N; ++j) {
A[0][j] = (double)gm[j] * (double)w_r[j];
A[1][j] = (double)gm[j] * (double)wc1[j];
A[2][j] = (double)gm[j] * (double)wc2[j];
A[3][j] = (double)gm[j] * (double)w_bar[j];
for (int s = 0; s < 11; ++s) {
A[4 + s][j] = (double)gm[j] * (double)w_spec[s][j];
}
}
std::vector<double> G_u(NU * NU, 0.0);
for (int i = 0; i < NU; ++i) {
for (int k = 0; k < NU; ++k) {
double dot = 0.0;
for (int j = 0; j < N; ++j) dot += A[i][j] * A[k][j];
G_u[i * NU + k] = dot;
}
G_u[i * NU + i] += 1e-6;
}
std::vector<double> rhs_u = {-delta_rej, +d_c1, +d_c2, 0.0, d_imend, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0};
char uplo = 'L';
int nu_lap = NU, nr = 1, info_u = 0;
dposv_(&uplo, &nu_lap, &nr, G_u.data(), &nu_lap, rhs_u.data(), &nu_lap, &info_u);
for (int j = 0; j < N; ++j) {
double u_j = 0.0;
for (int i = 0; i < NU; ++i) u_j += rhs_u[i] * A[i][j];
for (int k = 0; k < F; ++k) dW[(size_t)j * F + k] += (float)u_j * (float)V_cls[(size_t)pid * F + k];
}
}
std::vector<float> w_target((size_t)N * F);
for (size_t i = 0; i < w_target.size(); ++i) w_target[i] = wb[i] + dW[i];
std::vector<uint8_t> raw_q8((size_t)N * F / 32 * 34);
quantize_q8_0_monotonic(raw_sd47_orig, w_target.data(), w_target.size(), raw_q8.data());
std::vector<float> w_q8_deq((size_t)N * F), dq((size_t)N * F);
FILE* fq = std::fopen("/tmp/blk47_v3_q8.bin", "wb");
std::fwrite(raw_q8.data(), 1, raw_q8.size(), fq);
std::fclose(fq);
ref::dequant(ref::Q8_0, raw_q8.data(), w_q8_deq.data(), w_q8_deq.size());
double dn = 0, bn = 0;
for (size_t i = 0; i < dq.size(); ++i) {
dq[i] = w_q8_deq[i] - wb[i];
dn += (double)dq[i] * dq[i];
bn += (double)wb[i] * wb[i];
}
TraceData* trs[4] = {&tr_cal, &tr_p4, &tr_ho, &tr_col};
std::vector<float> mix_new[4];
for (int ti = 0; ti < 4; ++ti) {
int T = trs[ti]->T;
std::vector<float> d_t((size_t)T * N);
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, T, N, F, 1.0f, trs[ti]->z.data(), F, dq.data(), F, 0.0f, d_t.data(), N);
mix_new[ti].resize((size_t)T * N);
compute_mix_batch(T, trs[ti]->R.data(), trs[ti]->inj.data(), d_t.data(), wn, wdn, wup, mix_new[ti].data(), nullptr);
}
std::vector<float> X_new((size_t)M_base * N);
for (int t = 0; t < M_cal; ++t) std::memcpy(&X_new[(size_t)t * N], &mix_new[0][(size_t)t * N], N * sizeof(float));
for (int t = 0; t < M_p4; ++t) std::memcpy(&X_new[(size_t)(M_cal + t) * N], &mix_new[1][(size_t)t * N], N * sizeof(float));
for (int t = 0; t < M_ho; ++t) std::memcpy(&X_new[(size_t)(M_cal + M_p4 + t) * N], &mix_new[2][(size_t)t * N], N * sizeof(float));
for (int i = 0; i < M_col; ++i) std::memcpy(&X_new[(size_t)(M_cal + M_p4 + M_ho + i) * N], &mix_new[3][(size_t)col[i].t_pos * N], N * sizeof(float));
std::vector<float> L_new((size_t)M_base * CH), mx_n(M_base, -1e30f), se_n(M_base, 0.0f);
std::vector<int> argmax_n(M_base, -1);
std::vector<float> y_tgt_n(M_base, 0.0f);
std::vector<double> dot_sum(M_base, 0.0);
for (int v0 = 0; v0 < V; v0 += CH) {
int nv = std::min(CH, V - v0);
ref::dequant(tout->type, (const uint8_t*)g.map + tout->offset + (size_t)v0 * rb, Wc.data(), (size_t)nv * N);
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, M_base, nv, N, 1.0f, X_base.data(), N, Wc.data(), N, 0.0f, L_base.data(), CH);
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, M_base, nv, N, 1.0f, X_new.data(), N, Wc.data(), N, 0.0f, L_new.data(), CH);
for (int m = 0; m < M_base; ++m) {
const float* lb = &L_base[(size_t)m * CH];
const float* ln = &L_new[(size_t)m * CH];
float cm = mx_n[m];
int best_k = -1;
for (int k = 0; k < nv; ++k) {
if (ln[k] > cm) { cm = ln[k]; best_k = v0 + k; }
}
float s = se_n[m] * std::exp(mx_n[m] - cm);
double ds = 0.0;
float lseb = lse_b[m];
for (int k = 0; k < nv; ++k) {
s += std::exp(ln[k] - cm);
float pb = std::exp(lb[k] - lseb);
ds += (double)pb * ((double)lb[k] - (double)ln[k]);
}
if (best_k >= 0) argmax_n[m] = best_k;
mx_n[m] = cm; se_n[m] = s;
dot_sum[m] += ds;
}
for (int t = 0; t < M_ho; ++t) {
int tgt = tr_ho.toks[t + 1];
if (tgt >= v0 && tgt < v0 + nv) y_tgt_n[M_cal + M_p4 + t] = L_new[(size_t)(M_cal + M_p4 + t) * CH + (tgt - v0)];
}
}
std::vector<float> lse_n(M_base);
for (int m = 0; m < M_base; ++m) lse_n[m] = mx_n[m] + std::log(se_n[m]);
// Check bars (f) and (g) on held-out chat doc (tr_ho, M_ho = 1022):
int flips_to_spec_ho = 0, flips_to_spec_p4 = 0, flips_to_spec_slop21 = 0;
auto is_spec = [](int tok) { return tok == 248046 || tok == 248045 || tok == 248068 || tok == 248069 || tok == 248044; };
double ord_dlp_ho = 0.0, ho_kld_sum = 0.0, nonslop_kld_sum = 0.0;
int n_ord_ho = 0, flips_nonslop_ho = 0;
for (int t = 0; t < M_ho; ++t) {
int m = M_cal + M_p4 + t;
int tgt = tr_ho.toks[t + 1];
double kld_t = dot_sum[m] - (double)lse_b[m] + (double)lse_n[m];
ho_kld_sum += kld_t;
if (!sep_ho_set.count(t)) nonslop_kld_sum += kld_t;
if (!is_spec(tgt) && !is_spec(argmax_b[m]) && is_spec(argmax_n[m])) {
flips_to_spec_ho++;
if (!sep_ho_set.count(t)) {
flips_nonslop_ho++;
std::printf(" [NONSLOP FLIP t=%d] tgt=%d argmax_b=%d argmax_n=%d\n", t, tgt, argmax_b[m], argmax_n[m]);
}
}
if (!sep_ho_set.count(t) && tgt != 248046) {
double dlp = ((double)y_tgt_n[m] - lse_n[m]) - ((double)y_tgt_b[m] - lse_b[m]);
ord_dlp_ho += dlp;
n_ord_ho++;
}
}
for (int t = 0; t < M_p4; ++t) {
int m = M_cal + t;
int tgt = tr_p4.toks[t + 1];
if (!is_spec(tgt) && !is_spec(argmax_b[m]) && is_spec(argmax_n[m])) {
flips_to_spec_p4++;
std::printf(" [SFT FLIP t=%d is_slop=%d] tgt=%d argmax_b=%d argmax_n=%d\n", t, (int)all_p4_slop_set.count(t), tgt, argmax_b[m], argmax_n[m]);
}
}
// Check FTPO Train (28), Heldout_14, and Separate-Doc Heldout_21 (plus p(<|im_end|>) at slop21!)
double p_tr_b = 0, p_tr_n = 0; int n_tr = 0;
double p_ho_b = 0, p_ho_n = 0; int n_ho = 0;
for (const auto& fr : ftpo) {
int m = M_cal + fr.t_pos;
std::vector<float> wr; get_wout_row(fr.rej_id, wr);
float yr_b = 0, yr_n = 0;
for (int j = 0; j < N; ++j) { yr_b += wr[j] * X_base[(size_t)m * N + j]; yr_n += wr[j] * X_new[(size_t)m * N + j]; }
double pr_b = std::exp((double)yr_b - lse_b[m]);
double pr_n = std::exp((double)yr_n - lse_n[m]);
if (fr.split == 0) { p_tr_b += pr_b; p_tr_n += pr_n; n_tr++; }
else { p_ho_b += pr_b; p_ho_n += pr_n; n_ho++; }
}
double p_sep_b = 0, p_sep_n = 0, pim_sep_b = 0, pim_sep_n = 0;
for (int i = 0; i < n_sep; ++i) {
int m = M_cal + M_p4 + sep_ho[i].t_pos;
if (false) {
std::vector<std::pair<float, int>> top5;
for (int v0 = 0; v0 < V; v0 += CH) {
int nv = std::min(CH, V - v0);
ref::dequant(tout->type, (const uint8_t*)g.map + tout->offset + (size_t)v0 * rb, Wc.data(), (size_t)nv * N);
for (int k = 0; k < nv; ++k) {
float yv = 0;
const float* wr = &Wc[(size_t)k * N];
for (int j = 0; j < N; ++j) yv += wr[j] * X_base[(size_t)m * N + j];
if (yv > 12.0f) top5.push_back({yv, v0 + k});
}
}
std::sort(top5.rbegin(), top5.rend());
std::printf(" slop21[%2d] pid=%d t=%3d rej=%4d: ", i, sep_ho[i].phrase_id, sep_ho[i].t_pos, sep_ho[i].rej_id);
for (int k = 0; k < std::min<int>(5, top5.size()); ++k) std::printf("%d(%.2f) ", top5[k].second, top5[k].first);
std::printf("\n");
}
if (!is_spec(argmax_b[m]) && is_spec(argmax_n[m])) flips_to_spec_slop21++;
std::vector<float> wr; get_wout_row(sep_ho[i].rej_id, wr);
float yr_b = 0, yr_n = 0, yim_b = 0, yim_n = 0;
for (int j = 0; j < N; ++j) {
yr_b += wr[j] * X_base[(size_t)m * N + j];
yr_n += wr[j] * X_new[(size_t)m * N + j];
yim_b += w_spec[0][j] * X_base[(size_t)m * N + j];
yim_n += w_spec[0][j] * X_new[(size_t)m * N + j];
}
p_sep_b += std::exp((double)yr_b - lse_b[m]);
p_sep_n += std::exp((double)yr_n - lse_n[m]);
pim_sep_b += std::exp((double)yim_b - lse_b[m]);
pim_sep_n += std::exp((double)yim_n - lse_n[m]);
}
// Per-word Held-Out Collateral breakdown (n=142)
double w_sum[6] = {}, w_min[6] = {1e9, 1e9, 1e9, 1e9, 1e9, 1e9};
int w_cnt[6] = {}, w_loss10[6] = {}, tot_loss10 = 0, nc_ho = 0;
double dlp_ho_all = 0;
for (int i = 0; i < M_col; ++i) {
if (col[i].split == 0) continue;
int m = M_cal + M_p4 + M_ho + i;
std::vector<float> wt; get_wout_row(col[i].target_id, wt);
float yt_b = 0, yt_n = 0;
for (int j = 0; j < N; ++j) { yt_b += wt[j] * X_base[(size_t)m * N + j]; yt_n += wt[j] * X_new[(size_t)m * N + j]; }
double dlp = ((double)yt_n - lse_n[m]) - ((double)yt_b - lse_b[m]);
int b = col[i].ban_idx;
w_sum[b] += dlp; w_cnt[b]++; w_min[b] = std::min(w_min[b], dlp);
if (dlp < -0.10) { w_loss10[b]++; tot_loss10++; }
dlp_ho_all += dlp; nc_ho++;
}
std::printf("\n=== v3.2 mult=%.2f (||dW||/||W||=%.4f) ===\n", mult, std::sqrt(dn / bn));
std::printf(" FTPO TRAIN (n=%d): %.5f -> %.5f (-%.2f%%)\n", n_tr, p_tr_b / n_tr, p_tr_n / n_tr, (1.0 - p_tr_n / p_tr_b) * 100.0);
std::printf(" FTPO HELDOUT_DOC (n=%d): %.5f -> %.5f (-%.2f%%) | p(<|im_end|>): %.5f -> %.5f\n", n_sep, p_sep_b / n_sep, p_sep_n / n_sep, (1.0 - p_sep_n / p_sep_b) * 100.0, pim_sep_b / n_sep, pim_sep_n / n_sep);
std::printf(" FTPO HELDOUT_14 (n=%d): %.5f -> %.5f (-%.2f%%)\n", n_ho, p_ho_b / n_ho, p_ho_n / n_ho, (1.0 - p_ho_n / p_ho_b) * 100.0);
std::printf(" BAR (f) Special-Token Flips: ho_nonslop=%d (bar 0) | slop21=%d (bar 0) | ho_tot=%d | sft_doc=%d\n", flips_nonslop_ho, flips_to_spec_slop21, flips_to_spec_ho, flips_to_spec_p4);
std::printf(" BAR (g) Held-Out Chat Doc (n=%d): ord_dlp=%+.5f (bar >= -0.01) | nonslop_KLD=%.5f | whole_KLD=%.5f\n", n_ord_ho, ord_dlp_ho / n_ord_ho, nonslop_kld_sum / (M_ho - n_sep), ho_kld_sum / M_ho);
std::printf(" HELDOUT COLLATERAL overall (n=%d): mean=%+.5f | loss>0.1: %d/%d (%.1f%%, bar <= 10.0%%)\n",
nc_ho, dlp_ho_all / nc_ho, tot_loss10, nc_ho, 100.0 * tot_loss10 / nc_ho);
for (int b = 0; b < 6; ++b) {
std::printf(" word %-10s (n=%2d): mean=%+.5f (bar >= -0.05) | worst=%+.4f | <-0.1: %d\n",
ban_names[b], w_cnt[b], w_sum[b] / std::max(1, w_cnt[b]), w_min[b], w_loss10[b]);
}
}
return 0;
}