clef-flash-mixed-GGUF / hsdump.cpp
Jakevin's picture
v1.0: measured-KL mixed GGUF (3.30 GB)
fba639c verified
Raw History Blame Contribute Delete
7.15 kB
// Dump per-token final-norm hidden states from a text-only qwen35 GGUF.
//
// llama.cpp b11407 and later already export these two C++ functions from
// src/llama-ext.h (final RMSNorm / t_h_nextn). They are not declared in the
// installed C header llama.h, and llama-ext.h is not installed with the dylib.
// The declarations below match the exported symbols. No source patch is required
// on b11407+. An older tree that lacks the symbols needs a newer llama.cpp,
// not a patch against this file.
//
// llama_set_embeddings_nextn(ctx, true, masked=false)
// masked=false: one row per token in the batch, ignoring batch.logits.
// embeddings stays false. batch.logits is 1 only on the last token of each
// ubatch so the unused LM head does not reserve n_vocab rows per token.
//
// tokens.bin: int32 nrec, then (int32 ntok, int32 ids[ntok]) * nrec
// stdout: (int32 ntok, int32 n_embd, float h[ntok * n_embd]) * nrec
//
// c++ -std=c++17 -O2 \
// -I llama.cpp/include -I llama.cpp/ggml/include \
// hsdump.cpp -o hsdump \
// -L /path/to/llama.cpp/build/bin -lllama \
// -Wl,-rpath,/path/to/llama.cpp/build/bin
// hsdump model.gguf tokens.bin [--ngl N] [--ctx N] [--ubatch N]
#include "llama.h"
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>
// llama-ext.h is C++ and not installed with the dylib. These two match b11407's exported symbols.
void llama_set_embeddings_nextn(struct llama_context * ctx, bool value, bool masked);
float * llama_get_embeddings_nextn(struct llama_context * ctx);
static void read_exact(FILE * f, void * dst, size_t n, const char * what) {
auto * p = (uint8_t *) dst;
size_t got = 0;
while (got < n) {
size_t r = std::fread(p + got, 1, n - got, f);
if (r == 0) {
std::fprintf(stderr, "short read (%s)\n", what);
std::exit(2);
}
got += r;
}
}
static void write_exact(const void * src, size_t n) {
auto * p = (const uint8_t *) src;
size_t put = 0;
while (put < n) {
size_t w = std::fwrite(p + put, 1, n - put, stdout);
if (w == 0) {
std::fprintf(stderr, "short stdout\n");
std::exit(2);
}
put += w;
}
}
struct Rec {
std::vector<llama_token> ids;
};
static std::vector<Rec> load_tokens(const char * path, int32_t & max_len) {
FILE * f = std::fopen(path, "rb");
if (!f) {
std::perror(path);
std::exit(2);
}
int32_t nrec = 0;
read_exact(f, &nrec, 4, "nrec");
if (nrec < 0 || nrec > 1000000) {
std::fprintf(stderr, "bad nrec %d\n", nrec);
std::exit(2);
}
std::vector<Rec> recs(nrec);
max_len = 0;
for (int32_t i = 0; i < nrec; ++i) {
int32_t n = 0;
read_exact(f, &n, 4, "ntok");
if (n <= 0 || n > 1000000) {
std::fprintf(stderr, "bad ntok %d at record %d\n", n, i);
std::exit(2);
}
recs[i].ids.resize(n);
read_exact(f, recs[i].ids.data(), (size_t) n * sizeof(llama_token), "ids");
if (n > max_len) max_len = n;
}
std::fclose(f);
return recs;
}
int main(int argc, char ** argv) {
if (argc < 3) {
std::fprintf(stderr, "usage: hsdump model.gguf tokens.bin [--ngl N] [--ctx N] [--ubatch N]\n");
return 2;
}
const char * model = argv[1];
const char * tokens_path = argv[2];
int ngl = 99;
int ctx_arg = 0;
int ubatch = 512;
for (int i = 3; i < argc; ++i) {
if (!std::strcmp(argv[i], "--ngl") && i + 1 < argc) ngl = std::atoi(argv[++i]);
else if (!std::strcmp(argv[i], "--ctx") && i + 1 < argc) ctx_arg = std::atoi(argv[++i]);
else if (!std::strcmp(argv[i], "--ubatch") && i + 1 < argc) ubatch = std::atoi(argv[++i]);
else {
std::fprintf(stderr, "unknown arg %s\n", argv[i]);
return 2;
}
}
int32_t max_len = 0;
std::vector<Rec> recs = load_tokens(tokens_path, max_len);
const int n_ctx = ctx_arg > 0 ? ctx_arg : max_len + 8;
if (max_len > n_ctx) {
std::fprintf(stderr, "record length %d exceeds n_ctx %d\n", max_len, n_ctx);
return 2;
}
std::fprintf(stderr, "records %zu max_len %d n_ctx %d ubatch %d ngl %d\n",
recs.size(), max_len, n_ctx, ubatch, ngl);
llama_backend_init();
llama_model_params mp = llama_model_default_params();
mp.n_gpu_layers = ngl;
llama_model * mdl = llama_model_load_from_file(model, mp);
if (!mdl) {
std::fprintf(stderr, "failed to load %s\n", model);
return 1;
}
llama_context_params cp = llama_context_default_params();
cp.n_ctx = (uint32_t) n_ctx;
cp.n_batch = (uint32_t) ubatch;
cp.n_ubatch = (uint32_t) ubatch;
cp.n_seq_max = 1;
cp.embeddings = false;
cp.pooling_type = LLAMA_POOLING_TYPE_NONE;
cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
llama_context * ctx = llama_init_from_model(mdl, cp);
if (!ctx) {
std::fprintf(stderr, "failed to init context\n");
return 1;
}
llama_set_embeddings_nextn(ctx, true, /*masked=*/false);
const int n_embd = llama_model_n_embd(mdl);
std::fprintf(stderr, "n_embd %d\n", n_embd);
llama_batch batch = llama_batch_init(ubatch, 0, 1);
std::vector<float> hidden;
for (size_t ri = 0; ri < recs.size(); ++ri) {
const auto & ids = recs[ri].ids;
const int32_t n = (int32_t) ids.size();
hidden.assign((size_t) n * (size_t) n_embd, 0.f);
llama_memory_clear(llama_get_memory(ctx), true);
for (int32_t pos = 0; pos < n; ) {
const int32_t ntok = (n - pos) < ubatch ? (n - pos) : ubatch;
batch.n_tokens = ntok;
for (int32_t i = 0; i < ntok; ++i) {
batch.token[i] = ids[pos + i];
batch.pos[i] = pos + i;
batch.n_seq_id[i] = 1;
batch.seq_id[i][0] = 0;
batch.logits[i] = (i == ntok - 1) ? 1 : 0;
}
const int rc = llama_decode(ctx, batch);
if (rc != 0) {
std::fprintf(stderr, "decode failed record %zu pos %d rc %d\n", ri, pos, rc);
return 1;
}
float * h = llama_get_embeddings_nextn(ctx);
if (!h) {
std::fprintf(stderr, "null nextn embeddings record %zu pos %d\n", ri, pos);
return 1;
}
std::memcpy(hidden.data() + (size_t) pos * (size_t) n_embd, h,
(size_t) ntok * (size_t) n_embd * sizeof(float));
pos += ntok;
}
int32_t hdr[2] = { n, n_embd };
write_exact(hdr, sizeof(hdr));
write_exact(hidden.data(), hidden.size() * sizeof(float));
std::fflush(stdout);
if ((ri + 1) % 10 == 0 || ri + 1 == recs.size()) {
std::fprintf(stderr, "dumped %zu/%zu\n", ri + 1, recs.size());
}
}
llama_batch_free(batch);
llama_free(ctx);
llama_model_free(mdl);
llama_backend_free();
std::fprintf(stderr, "HSDUMP_OK\n");
return 0;
}