Something went wrong. Try again.
strata but optimized for m5 (and maybe more later)
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245
import subprocessfrom 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!")