import subprocess
from pathlib import Path
subprocess.run(["git", "-C", "/content/llama.cpp", "checkout", "tools/imatrix/imatrix.cpp"], check=True)
p = Path("/content/llama.cpp/tools/imatrix/imatrix.cpp")
orig = p.read_text()
hook_code = """
// === REAP_TRACE_COLLECTOR_V2 ===
struct ReapTraceCollector {
static constexpr int N_LAYERS = 48;
static constexpr int N_EXPERTS = 512;
static constexpr int TOP_M = 64;
double prob_sum[N_LAYERS][N_EXPERTS] = {{0.0}};
double reap_l2_sum[N_LAYERS][N_EXPERTS] = {{0.0}};
double reap_l2_sq_sum[N_LAYERS][N_EXPERTS] = {{0.0}};
int64_t topk_count[N_LAYERS][N_EXPERTS] = {{0}};
int64_t nonfinite_drops = 0;
std::vector<int32_t> last_topk_ids[N_LAYERS];
int64_t last_topk_n_used[N_LAYERS] = {0};
int64_t last_topk_n_tokens[N_LAYERS] = {0};
int64_t trace_tokens_written = 0;
int64_t max_trace_tokens = 16384;
int32_t cur_ubatch_tokens = 0;
std::vector<uint16_t> ubatch_ids;
std::vector<uint16_t> ubatch_logits_f16;
FILE * trace_fp = nullptr;
void init_if_needed() {
if (trace_fp == nullptr && trace_tokens_written < max_trace_tokens) {
const char * tpath = getenv("Q4RT_OUT");
if (!tpath) tpath = "/content/routing_trace.q4rt";
trace_fp = fopen(tpath, "wb");
if (trace_fp) {
const char magic[4] = {'Q', '4', 'R', 'T'};
uint16_t ver = 1, nl = N_LAYERS, ne = N_EXPERTS, tm = TOP_M;
uint32_t ntok_placeholder = 0;
fwrite(magic, 1, 4, trace_fp);
fwrite(&ver, 2, 1, trace_fp);
fwrite(&nl, 2, 1, trace_fp);
fwrite(&ne, 2, 1, trace_fp);
fwrite(&tm, 2, 1, trace_fp);
fwrite(&ntok_placeholder, 4, 1, trace_fp);
fflush(trace_fp);
}
}
}
void on_logits(int il, const float * data, int64_t n_expert, int64_t n_tokens) {
if (il < 0 || il >= N_LAYERS || n_expert != N_EXPERTS) return;
init_if_needed();
if (il == 0 && trace_tokens_written < max_trace_tokens) {
cur_ubatch_tokens = (int32_t)std::min<int64_t>(n_tokens, max_trace_tokens - trace_tokens_written);
ubatch_ids.assign((size_t)cur_ubatch_tokens * N_LAYERS * TOP_M, 0);
ubatch_logits_f16.assign((size_t)cur_ubatch_tokens * N_LAYERS * TOP_M, 0);
}
std::pair<float, int> pairs[N_EXPERTS];
double exps[N_EXPERTS];
for (int64_t row = 0; row < n_tokens; ++row) {
const float * z = data + row * n_expert;
float max_z = -1e30f;
for (int e = 0; e < N_EXPERTS; ++e) {
float v = std::isfinite(z[e]) ? z[e] : -1e30f;
if (v > max_z) max_z = v;
pairs[e] = { v, e };
}
double sum_exp = 0.0;
for (int e = 0; e < N_EXPERTS; ++e) {
exps[e] = std::exp((double)pairs[e].first - (double)max_z);
sum_exp += exps[e];
}
double inv_sum = sum_exp > 0.0 ? (1.0 / sum_exp) : 0.0;
for (int e = 0; e < N_EXPERTS; ++e) {
prob_sum[il][e] += exps[e] * inv_sum;
}
if (row < cur_ubatch_tokens && trace_tokens_written < max_trace_tokens) {
std::partial_sort(pairs, pairs + TOP_M, pairs + N_EXPERTS,
[](const std::pair<float, int> & a, const std::pair<float, int> & b) {
return (a.first != b.first) ? (a.first > b.first) : (a.second < b.second);
});
size_t base = ((size_t)row * N_LAYERS + (size_t)il) * TOP_M;
for (int m = 0; m < TOP_M; ++m) {
ubatch_ids[base + m] = (uint16_t)pairs[m].second;
ubatch_logits_f16[base + m] = ggml_fp32_to_fp16(pairs[m].first);
}
}
}
if (il == N_LAYERS - 1 && cur_ubatch_tokens > 0 && trace_fp) {
fseek(trace_fp, 0, SEEK_END);
for (int32_t row = 0; row < cur_ubatch_tokens; ++row) {
for (int l = 0; l < N_LAYERS; ++l) {
size_t base = ((size_t)row * N_LAYERS + (size_t)l) * TOP_M;
fwrite(&ubatch_ids[base], sizeof(uint16_t), TOP_M, trace_fp);
fwrite(&ubatch_logits_f16[base], sizeof(uint16_t), TOP_M, trace_fp);
}
}
trace_tokens_written += cur_ubatch_tokens;
cur_ubatch_tokens = 0;
uint32_t ntok = (uint32_t)trace_tokens_written;
fseek(trace_fp, 12, SEEK_SET);
fwrite(&ntok, 4, 1, trace_fp);
fseek(trace_fp, 0, SEEK_END);
fflush(trace_fp);
}
}
void on_topk(int il, const char * raw_ids, const ggml_tensor * t) {
if (il < 0 || il >= N_LAYERS) return;
const int64_t n_used = t->ne[0];
const int64_t n_tokens = t->ne[1];
last_topk_n_used[il] = n_used;
last_topk_n_tokens[il] = n_tokens;
last_topk_ids[il].resize((size_t)(n_used * n_tokens));
for (int64_t row = 0; row < n_tokens; ++row) {
for (int64_t idx = 0; idx < n_used; ++idx) {
int32_t ex = *(const int32_t *)(raw_ids + row * t->nb[1] + idx * t->nb[0]);
last_topk_ids[il][(size_t)(row * n_used + idx)] = ex;
}
}
}
void on_weighted_out(int il, const char * weighted_data, const ggml_tensor * t) {
if (il < 0 || il >= N_LAYERS) return;
const int64_t ne0 = t->ne[0]; // 2560
const int64_t n_used = t->ne[1]; // 10
const int64_t n_tokens = t->ne[2];
if (last_topk_n_used[il] != n_used || last_topk_n_tokens[il] != n_tokens) {
GGML_ABORT("ReapTraceCollector: topk shape mismatch at layer %d", il);
}
const auto & ids_vec = last_topk_ids[il];
for (int64_t row = 0; row < n_tokens; ++row) {
for (int64_t idx = 0; idx < n_used; ++idx) {
const int32_t ex = ids_vec[(size_t)(row * n_used + idx)];
GGML_ASSERT(ex >= 0 && ex < N_EXPERTS);
const float * y = (const float *)(weighted_data + idx * t->nb[1] + row * t->nb[2]);
double sq = 0.0;
for (int64_t d = 0; d < ne0; ++d) {
double v = (double)y[d];
sq += v * v;
}
if (std::isfinite(sq)) {
topk_count[il][ex]++;
reap_l2_sq_sum[il][ex] += sq;
reap_l2_sum[il][ex] += std::sqrt(sq);
} else {
nonfinite_drops++;
}
}
}
}
void save_json_only() {
int64_t tot_c = 0;
for (int e = 0; e < N_EXPERTS; ++e) tot_c += topk_count[0][e];
if (tot_c == 0) return;
const char * spath = getenv("REAP_OUT");
if (!spath) spath = "/content/reap_saliency.json";
std::string tmp_path = std::string(spath) + ".tmp";
FILE * fp = fopen(tmp_path.c_str(), "w");
if (!fp) return;
fprintf(fp, "{\\n \\"trace_tokens_written\\": %lld,\\n \\"nonfinite_drops\\": %lld,\\n \\"layers\\": {\\n",
(long long)trace_tokens_written, (long long)nonfinite_drops);
for (int l = 0; l < N_LAYERS; ++l) {
fprintf(fp, " \\"%d\\": {\\n", l);
fprintf(fp, " \\"count\\": [");
for (int e = 0; e < N_EXPERTS; ++e) fprintf(fp, "%s%lld", e ? "," : "", (long long)topk_count[l][e]);
fprintf(fp, "],\\n \\"prob_sum\\": [");
for (int e = 0; e < N_EXPERTS; ++e) fprintf(fp, "%s%.6f", e ? "," : "", prob_sum[l][e]);
fprintf(fp, "],\\n \\"reap_l2_sum\\": [");
for (int e = 0; e < N_EXPERTS; ++e) fprintf(fp, "%s%.6f", e ? "," : "", reap_l2_sum[l][e]);
fprintf(fp, "],\\n \\"reap_l2_sq_sum\\": [");
for (int e = 0; e < N_EXPERTS; ++e) fprintf(fp, "%s%.6f", e ? "," : "", reap_l2_sq_sum[l][e]);
fprintf(fp, "]\\n }%s\\n", (l + 1 < N_LAYERS) ? "," : "");
}
fprintf(fp, " }\\n}\\n");
fclose(fp);
rename(tmp_path.c_str(), spath);
}
~ReapTraceCollector() {
save_json_only();
if (trace_fp) {
fclose(trace_fp);
trace_fp = nullptr;
}
}
};
static ReapTraceCollector g_reap;
static std::vector<char> g_reap_buf;
"""
target_ask = """ if (ask) {
if (t->op == GGML_OP_MUL_MAT_ID) { return true; } // collect all indirect matrix multiplications"""
replacement_ask = """ if (ask) {
if ((t->op == GGML_OP_MUL_MAT && strncmp(t->name, "ffn_moe_logits-", 15) == 0) ||
(t->op == GGML_OP_MUL && strncmp(t->name, "ffn_moe_weighted-", 17) == 0)) { return true; }
if (t->op == GGML_OP_MUL_MAT_ID) { return true; } // collect all indirect matrix multiplications"""
target_lock = """ std::lock_guard<std::mutex> lock(m_mutex);"""
replacement_lock = """ std::lock_guard<std::mutex> lock(m_mutex);
if (t->op == GGML_OP_MUL_MAT && strncmp(t->name, "ffn_moe_logits-", 15) == 0 && t->type == GGML_TYPE_F32) {
int il = atoi(t->name + 15);
size_t nb = ggml_nbytes(t);
g_reap_buf.resize(nb);
ggml_backend_tensor_get(t, g_reap_buf.data(), 0, nb);
g_reap.on_logits(il, (const float *)g_reap_buf.data(), t->ne[0], t->ne[1]);
return true;
}
if (t->op == GGML_OP_MUL && strncmp(t->name, "ffn_moe_weighted-", 17) == 0 && t->type == GGML_TYPE_F32) {
int il = atoi(t->name + 17);
const ggml_tensor * down_t = t->src[0];
GGML_ASSERT(down_t && down_t->op == GGML_OP_MUL_MAT_ID && down_t->src[2]);
const ggml_tensor * ids = down_t->src[2];
size_t nb_ids = ggml_nbytes(ids);
g_reap_buf.resize(nb_ids);
ggml_backend_tensor_get(ids, g_reap_buf.data(), 0, nb_ids);
g_reap.on_topk(il, g_reap_buf.data(), ids);
size_t nb = ggml_nbytes(t);
g_reap_buf.resize(nb);
ggml_backend_tensor_get(t, g_reap_buf.data(), 0, nb);
g_reap.on_weighted_out(il, g_reap_buf.data(), t);
return true;
}"""
patched = orig.replace("bool IMatrixCollector::collect_imatrix(", hook_code + "\nbool IMatrixCollector::collect_imatrix(", 1)
patched = patched.replace(target_ask, replacement_ask, 1)
patched = patched.replace(target_lock, replacement_lock, 1)
patched = patched.replace("save_imatrix();", "save_imatrix(); g_reap.save_json_only();")
p.write_text(patched)
print("Patched imatrix.cpp v2 successfully!")