diff --git a/CMakeLists.txt b/CMakeLists.txt index edf7fa73..39f25c7b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -17,4 +17,35 @@ add_compile_options(-Wall -Wextra -Wpedantic) enable_testing() -# Components are added here as they land, each with its tests. +add_library(onebit_core STATIC + engine/core/gguf.cpp + engine/core/dequant.cpp + engine/core/model_config.cpp + engine/core/parallel.cpp) +target_include_directories(onebit_core PUBLIC engine/core) +find_package(Threads REQUIRED) +target_link_libraries(onebit_core PUBLIC Threads::Threads) + +add_library(onebit_cpu STATIC engine/backends/cpu/cpu_model.cpp) +target_include_directories(onebit_cpu PUBLIC engine/backends/cpu) +target_link_libraries(onebit_cpu PUBLIC onebit_core) + +# Unit tests: synthetic data only, so they run anywhere (including CI). +foreach(t test_gguf test_dequant) + add_executable(${t} tests/${t}.cpp) + target_include_directories(${t} PRIVATE tests) + target_link_libraries(${t} PRIVATE onebit_core) + add_test(NAME ${t} COMMAND ${t}) +endforeach() + +# Golden test: the CPU reference against HF transformers fp32 logits. Needs a +# model and a golden dir (tools/golden/make_golden.py), so it is registered +# only when both are given, e.g. +# -DENGINE_GOLDEN_MODEL=/path/Qwen3-0.6B-BF16.gguf -DENGINE_GOLDEN_DIR=/path/golden +add_executable(golden_cpu tests/golden_cpu.cpp) +target_link_libraries(golden_cpu PRIVATE onebit_cpu) +set(ENGINE_GOLDEN_MODEL "" CACHE FILEPATH "GGUF model for the golden test") +set(ENGINE_GOLDEN_DIR "" CACHE PATH "Golden logits directory for the golden test") +if(ENGINE_GOLDEN_MODEL AND ENGINE_GOLDEN_DIR) + add_test(NAME golden_cpu COMMAND golden_cpu ${ENGINE_GOLDEN_MODEL} ${ENGINE_GOLDEN_DIR}) +endif() diff --git a/docs/PORTING.md b/docs/PORTING.md index d0890cc6..ace7df7d 100644 --- a/docs/PORTING.md +++ b/docs/PORTING.md @@ -5,7 +5,7 @@ and what has to be true before it lands here. Refs are branches or commits in th | Component | Source | State at source | Gate to land here | |---|---|---|---| -| CPU reference | `src/gguf_reader.cpp`, `src/tokenizer.cpp`, `tools/qwen36_full_ref.py` | works | matches HF transformers fp32 logits on Qwen3-0.6B | +| CPU reference | `src/gguf_reader.cpp`, `src/tokenizer.cpp`, `tools/qwen36_full_ref.py` | **landed** (qwen3; see docs/cpu-reference.md). Tokenizer not yet | matches HF transformers fp32 logits on Qwen3-0.6B | | Arch registry | `src/model_registry.cpp` + `Testing/census_*.json` | 569 tokens map 2,030 HF arch strings (mapping only) | data file + separate verified list | | NPU backend | `engine/npu/src/npu_engine_universal.cpp` (`I8Ctx::init_elf`) | ELF-native, matches its own baseline | golden test vs CPU reference | | NPU ELF dispatch table | branch `backup/iso-build-elf-native-2026-09-22` (`kElfDesigns`) | ELF and xclbin modes match token for token; 0 xclbin opens | same, in this engine | diff --git a/docs/cpu-reference.md b/docs/cpu-reference.md new file mode 100644 index 00000000..6c16b962 --- /dev/null +++ b/docs/cpu-reference.md @@ -0,0 +1,62 @@ +# CPU reference + +`engine/backends/cpu/` is the fp32 oracle the NPU and GPU backends are tested +against. It dequantizes every weight to fp32 at load time and runs one token at +a time in fp32. + +It accepts only architectures that have passed the golden test below. Any other +architecture is refused by name. Loading also fails if the GGUF holds a tensor +the forward pass does not use, because an unused tensor means the model has +semantics this code does not implement. + +| Architecture | Weight types | Golden test | +|---|---|---| +| `qwen3` | F32, F16, BF16, Q8_0 | Qwen3-0.6B: pass | + +## Golden test: Qwen3-0.6B + +Reference: HF transformers, fp32, eager attention, on the same checkpoint. +The test feeds the 36-token sequence (12 prompt tokens plus the reference's 24 +greedy tokens) teacher-forced and compares the next-token logits at every position. + +| Metric | Result | Gate | +|---|---|---| +| Argmax agreement | 36 / 36 | all | +| Worst per-position KL(ref ‖ cpu) | 1.3e-10 | ≤ 1e-6 | +| Worst max \|Δlogit\| | 1.0e-4 | reported only | + +To confirm the test can fail, two deliberately broken builds were run: + +- **Wrong RoPE layout** (Normal instead of Neox): 17 / 36 argmax mismatches, worst KL 7.9. +- **K-norm skipped**: 34 / 36 mismatches, worst KL 24.9. + +Both failed the gate. + +Inputs, as recorded in `tests/golden/qwen3-0.6b-capitals/meta.json`: + +- Checkpoint `Qwen/Qwen3-0.6B` at revision `c1899de289a04d12100db370d81485cdf75e47ca`. +- The GGUF is converted with llama.cpp `convert_hf_to_gguf.py --outtype bf16` (llama.cpp `e71b805`). + Its sha256 is `33d6f6c9f0dd21dd3d9c99f514498ea594119508e538aa36773993a3ffd39aff`. +- The golden logits are 21.9 MB, so they are not committed. They are regenerated + by the script below; `meta.json` records their sha256. + +Measured on Strix Halo (32 threads): 0.6 s load, 24 tok/s teacher-forced. + +## Reproduce + +```bash +python tools/golden/make_golden.py --model Qwen/Qwen3-0.6B \ + --prompt "The capital of France is Paris. The capital of Japan is" \ + --n-gen 24 --out golden/qwen3-0.6b-capitals +python llama.cpp/convert_hf_to_gguf.py --outtype bf16 --outfile Qwen3-0.6B-BF16.gguf + +cmake -B build -DENGINE_GOLDEN_MODEL=$PWD/Qwen3-0.6B-BF16.gguf \ + -DENGINE_GOLDEN_DIR=$PWD/golden/qwen3-0.6b-capitals +cmake --build build && ctest --test-dir build --output-on-failure +``` + +## Not done yet + +- **Tokenizer.** The test uses token ids from the HF tokenizer. +- **Quantized types beyond Q8_0** (Q4_K, Q6_K and others). Each needs its own dequant test first. +- **A longer-context golden** (hundreds of positions) to cover RoPE at larger angles. diff --git a/engine/backends/cpu/cpu_model.cpp b/engine/backends/cpu/cpu_model.cpp new file mode 100644 index 00000000..121459d1 --- /dev/null +++ b/engine/backends/cpu/cpu_model.cpp @@ -0,0 +1,219 @@ +#include "cpu_model.h" + +#include +#include +#include +#include +#include + +#include "dequant.h" + +namespace onebit { + +namespace { + +// KV cache length. Kept well under the training context so the reference +// stays small; raise it when a test needs longer sequences. +constexpr uint32_t kMaxCtx = 4096; + +void rms_norm(const float* x, const float* w, float* out, uint32_t n, float eps) { + double ss = 0.0; + for (uint32_t i = 0; i < n; ++i) ss += double(x[i]) * x[i]; + const float scale = 1.0f / std::sqrt(float(ss / n) + eps); + for (uint32_t i = 0; i < n; ++i) out[i] = w[i] * (x[i] * scale); +} + +void rope(float* x, uint32_t n_heads, uint32_t head_dim, uint32_t pos, float theta, RopeStyle style) { + const uint32_t half = head_dim / 2; + for (uint32_t i = 0; i < half; ++i) { + const float inv_freq = 1.0f / std::pow(theta, float(2 * i) / float(head_dim)); + const float angle = float(pos) * inv_freq; + const float c = std::cos(angle), s = std::sin(angle); + for (uint32_t h = 0; h < n_heads; ++h) { + float* v = x + h * head_dim; + const uint32_t a = style == RopeStyle::Neox ? i : 2 * i; + const uint32_t b = style == RopeStyle::Neox ? i + half : 2 * i + 1; + const float x0 = v[a], x1 = v[b]; + v[a] = x0 * c - x1 * s; + v[b] = x0 * s + x1 * c; + } + } +} + +float silu(float x) { return x / (1.0f + std::exp(-x)); } + +class Loader { +public: + explicit Loader(const GgufFile& f) : f_(f) {} + + // Dequantizes tensor `name`, checking its shape is exactly `shape`. + std::expected, std::string> take(const std::string& name, std::vector shape) { + const GgufTensor* t = f_.tensor(name); + if (!t) return std::unexpected(std::format("missing tensor {}", name)); + if (t->ne != shape) return std::unexpected(std::format("tensor {} has an unexpected shape", name)); + std::vector out(t->n_elements()); + if (auto r = dequantize(t->type, t->data, out.data(), out.size()); !r) + return std::unexpected(std::format("tensor {}: {}", name, r.error())); + used_.insert(name); + return out; + } + + bool has(const std::string& name) const { return f_.tensor(name) != nullptr; } + + std::expected check_all_used() const { + for (const GgufTensor& t : f_.tensors()) + if (!used_.contains(t.name)) return std::unexpected(std::format("tensor {} is not used by the reference", t.name)); + return {}; + } + +private: + const GgufFile& f_; + std::set used_; +}; + +} // namespace + +std::expected CpuModel::load(const std::string& gguf_path, size_t n_threads) { + auto file = GgufFile::open(gguf_path); + if (!file) return std::unexpected(file.error()); + auto cfg = read_model_config(*file); + if (!cfg) return std::unexpected(cfg.error()); + + CpuModel m; + m.cfg_ = *cfg; + const ModelConfig& c = m.cfg_; + const uint64_t E = c.n_embd, D = c.head_dim, Q = uint64_t{c.n_head} * D, KV = uint64_t{c.n_head_kv} * D; + const uint64_t F = c.n_ff, V = c.n_vocab; + + Loader ld(*file); +#define TAKE(dst, name, ...) \ + do { \ + auto r_ = ld.take(name, {__VA_ARGS__}); \ + if (!r_) return std::unexpected(std::move(r_.error())); \ + dst = std::move(*r_); \ + } while (0) + + TAKE(m.tok_embd_, "token_embd.weight", E, V); + TAKE(m.output_norm_, "output_norm.weight", E); + if (ld.has("output.weight")) TAKE(m.output_, "output.weight", E, V); + + m.layers_.resize(c.n_layer); + for (uint32_t i = 0; i < c.n_layer; ++i) { + Layer& L = m.layers_[i]; + const std::string p = std::format("blk.{}.", i); + TAKE(L.attn_norm, p + "attn_norm.weight", E); + TAKE(L.wq, p + "attn_q.weight", E, Q); + TAKE(L.wk, p + "attn_k.weight", E, KV); + TAKE(L.wv, p + "attn_v.weight", E, KV); + TAKE(L.wo, p + "attn_output.weight", Q, E); + TAKE(L.q_norm, p + "attn_q_norm.weight", D); + TAKE(L.k_norm, p + "attn_k_norm.weight", D); + TAKE(L.ffn_norm, p + "ffn_norm.weight", E); + TAKE(L.w_gate, p + "ffn_gate.weight", E, F); + TAKE(L.w_up, p + "ffn_up.weight", E, F); + TAKE(L.w_down, p + "ffn_down.weight", F, E); + } +#undef TAKE + if (auto r = ld.check_all_used(); !r) return std::unexpected(r.error()); + + m.n_ctx_ = std::min(c.n_ctx_train, kMaxCtx); + const size_t cache = size_t{c.n_layer} * m.n_ctx_ * KV; + m.k_cache_.assign(cache, 0.0f); + m.v_cache_.assign(cache, 0.0f); + m.x_.resize(E); + m.xn_.resize(E); + m.q_.resize(Q); + m.k_.resize(KV); + m.v_.resize(KV); + m.attn_.resize(Q); + m.scores_.resize(m.n_ctx_); + m.gate_.resize(F); + m.up_.resize(F); + m.logits_.resize(V); + m.pool_ = std::make_unique(n_threads); + return m; +} + +void CpuModel::reset() { n_past_ = 0; } + +void CpuModel::matvec(const std::vector& w, const float* x, float* y, uint32_t rows, uint32_t cols) { + pool_->parallel_for(rows, [&](size_t begin, size_t end) { + for (size_t r = begin; r < end; ++r) { + const float* row = w.data() + r * cols; + float acc = 0.0f; + for (uint32_t i = 0; i < cols; ++i) acc += row[i] * x[i]; + y[r] = acc; + } + }); +} + +std::expected*, std::string> CpuModel::forward(int32_t token) { + const ModelConfig& c = cfg_; + if (token < 0 || uint32_t(token) >= c.n_vocab) return std::unexpected(std::format("token {} out of range", token)); + if (n_past_ >= n_ctx_) return std::unexpected(std::format("context full ({} tokens)", n_ctx_)); + + const uint32_t E = c.n_embd, D = c.head_dim, H = c.n_head, HKV = c.n_head_kv; + const uint32_t KV = HKV * D, pos = n_past_, group = H / HKV; + const float att_scale = 1.0f / std::sqrt(float(D)); + + std::copy_n(tok_embd_.data() + size_t(token) * E, E, x_.data()); + + for (uint32_t l = 0; l < c.n_layer; ++l) { + const Layer& L = layers_[l]; + + rms_norm(x_.data(), L.attn_norm.data(), xn_.data(), E, c.rms_eps); + matvec(L.wq, xn_.data(), q_.data(), H * D, E); + matvec(L.wk, xn_.data(), k_.data(), KV, E); + matvec(L.wv, xn_.data(), v_.data(), KV, E); + for (uint32_t h = 0; h < H; ++h) rms_norm(q_.data() + h * D, L.q_norm.data(), q_.data() + h * D, D, c.rms_eps); + for (uint32_t h = 0; h < HKV; ++h) rms_norm(k_.data() + h * D, L.k_norm.data(), k_.data() + h * D, D, c.rms_eps); + rope(q_.data(), H, D, pos, c.rope_theta, c.rope_style); + rope(k_.data(), HKV, D, pos, c.rope_theta, c.rope_style); + + float* kc = k_cache_.data() + (size_t(l) * n_ctx_) * KV; + float* vc = v_cache_.data() + (size_t(l) * n_ctx_) * KV; + std::copy_n(k_.data(), KV, kc + size_t(pos) * KV); + std::copy_n(v_.data(), KV, vc + size_t(pos) * KV); + + for (uint32_t h = 0; h < H; ++h) { + const float* q = q_.data() + h * D; + const uint32_t kvh = h / group; + float mx = -INFINITY; + for (uint32_t t = 0; t <= pos; ++t) { + const float* k = kc + size_t(t) * KV + kvh * D; + float s = 0.0f; + for (uint32_t i = 0; i < D; ++i) s += q[i] * k[i]; + scores_[t] = s * att_scale; + mx = std::max(mx, scores_[t]); + } + double sum = 0.0; + for (uint32_t t = 0; t <= pos; ++t) { + scores_[t] = std::exp(scores_[t] - mx); + sum += scores_[t]; + } + float* out = attn_.data() + h * D; + std::fill_n(out, D, 0.0f); + for (uint32_t t = 0; t <= pos; ++t) { + const float p = float(scores_[t] / sum); + const float* v = vc + size_t(t) * KV + kvh * D; + for (uint32_t i = 0; i < D; ++i) out[i] += p * v[i]; + } + } + matvec(L.wo, attn_.data(), xn_.data(), E, H * D); + for (uint32_t i = 0; i < E; ++i) x_[i] += xn_[i]; + + rms_norm(x_.data(), L.ffn_norm.data(), xn_.data(), E, c.rms_eps); + matvec(L.w_gate, xn_.data(), gate_.data(), c.n_ff, E); + matvec(L.w_up, xn_.data(), up_.data(), c.n_ff, E); + for (uint32_t i = 0; i < c.n_ff; ++i) gate_[i] = silu(gate_[i]) * up_[i]; + matvec(L.w_down, gate_.data(), xn_.data(), E, c.n_ff); + for (uint32_t i = 0; i < E; ++i) x_[i] += xn_[i]; + } + + rms_norm(x_.data(), output_norm_.data(), xn_.data(), E, c.rms_eps); + matvec(output_.empty() ? tok_embd_ : output_, xn_.data(), logits_.data(), c.n_vocab, E); + ++n_past_; + return &logits_; +} + +} // namespace onebit diff --git a/engine/backends/cpu/cpu_model.h b/engine/backends/cpu/cpu_model.h new file mode 100644 index 00000000..7e5e240d --- /dev/null +++ b/engine/backends/cpu/cpu_model.h @@ -0,0 +1,64 @@ +// cpu_model.h: fp32 reference decoder on the CPU. +// +// This is the correctness oracle for the NPU and GPU backends, so it favours +// plain, checkable code over speed: every weight is dequantized to fp32 at load +// time and every op runs in fp32, one token at a time. +#pragma once + +#include +#include +#include +#include +#include + +#include "gguf.h" +#include "model_config.h" +#include "parallel.h" + +namespace onebit { + +class CpuModel { +public: + // Loads every tensor of the GGUF. Fails if a tensor is missing, has the + // wrong shape, has an unsupported type, or is present but unused (an + // unused tensor means the model has semantics this code does not model). + static std::expected load(const std::string& gguf_path, size_t n_threads = 0); + + const ModelConfig& config() const { return cfg_; } + + // Clears the KV cache. + void reset(); + + // Number of tokens in the KV cache. + uint32_t n_past() const { return n_past_; } + + // Runs one token at position n_past() and returns the next-token logits + // (config().n_vocab floats). Fails when the context is full. + std::expected*, std::string> forward(int32_t token); + + CpuModel(CpuModel&&) = default; + CpuModel& operator=(CpuModel&&) = default; + +private: + CpuModel() = default; + + struct Layer { + std::vector attn_norm, wq, wk, wv, wo, q_norm, k_norm; + std::vector ffn_norm, w_gate, w_up, w_down; + }; + + void matvec(const std::vector& w, const float* x, float* y, uint32_t rows, uint32_t cols); + + ModelConfig cfg_; + uint32_t n_ctx_ = 0; + std::vector tok_embd_, output_norm_, output_; // output_ empty when tied to tok_embd_ + std::vector layers_; + std::vector k_cache_, v_cache_; // [layer][pos][n_head_kv * head_dim] + uint32_t n_past_ = 0; + + // Scratch. + std::vector x_, xn_, q_, k_, v_, attn_, scores_, gate_, up_, logits_; + std::unique_ptr pool_; +}; + +} // namespace onebit diff --git a/engine/core/dequant.cpp b/engine/core/dequant.cpp new file mode 100644 index 00000000..75ceb1bd --- /dev/null +++ b/engine/core/dequant.cpp @@ -0,0 +1,84 @@ +#include "dequant.h" + +#include +#include +#include + +namespace onebit { + +float fp16_to_fp32(uint16_t h) { + const uint32_t sign = uint32_t{h & 0x8000u} << 16; + const uint32_t exp = (h >> 10) & 0x1fu; + uint32_t mant = h & 0x3ffu; + uint32_t bits; + if (exp == 0) { + if (mant == 0) { + bits = sign; + } else { + // Subnormal: normalise the mantissa. + int e = -1; + do { + ++e; + mant <<= 1; + } while ((mant & 0x400u) == 0); + bits = sign | uint32_t(127 - 15 - e) << 23 | (mant & 0x3ffu) << 13; + } + } else if (exp == 0x1f) { + bits = sign | 0x7f800000u | mant << 13; // inf / nan + } else { + bits = sign | (exp + (127 - 15)) << 23 | mant << 13; + } + return std::bit_cast(bits); +} + +float bf16_to_fp32(uint16_t b) { return std::bit_cast(uint32_t{b} << 16); } + +bool dequant_supported(GgmlType t) { + switch (t) { + case GgmlType::F32: + case GgmlType::F16: + case GgmlType::BF16: + case GgmlType::Q8_0: return true; + default: return false; + } +} + +namespace { + +uint16_t load_u16(const std::byte* p) { + uint16_t v; + std::memcpy(&v, p, 2); + return v; +} + +} // namespace + +std::expected dequantize(GgmlType t, const std::byte* src, float* dst, uint64_t n) { + switch (t) { + case GgmlType::F32: + std::memcpy(dst, src, n * sizeof(float)); + return {}; + case GgmlType::F16: + for (uint64_t i = 0; i < n; ++i) dst[i] = fp16_to_fp32(load_u16(src + 2 * i)); + return {}; + case GgmlType::BF16: + for (uint64_t i = 0; i < n; ++i) dst[i] = bf16_to_fp32(load_u16(src + 2 * i)); + return {}; + case GgmlType::Q8_0: { + // Block of 32: fp16 scale d, then 32 int8 quants. x = d * q. + constexpr uint64_t kBlock = 32, kBytes = 34; + if (n % kBlock != 0) return std::unexpected(std::string("Q8_0: n is not a multiple of 32")); + for (uint64_t b = 0; b < n / kBlock; ++b) { + const std::byte* blk = src + b * kBytes; + const float d = fp16_to_fp32(load_u16(blk)); + const auto* q = reinterpret_cast(blk + 2); + for (uint64_t i = 0; i < kBlock; ++i) dst[b * kBlock + i] = d * float(q[i]); + } + return {}; + } + default: + return std::unexpected(std::format("dequantize: type {} not supported yet", ggml_type_name(t))); + } +} + +} // namespace onebit diff --git a/engine/core/dequant.h b/engine/core/dequant.h new file mode 100644 index 00000000..db14f031 --- /dev/null +++ b/engine/core/dequant.h @@ -0,0 +1,23 @@ +// dequant.h: convert ggml tensor data to fp32. +#pragma once + +#include +#include +#include +#include + +#include "gguf.h" + +namespace onebit { + +float fp16_to_fp32(uint16_t h); +float bf16_to_fp32(uint16_t b); + +// Whether dequantize() handles type t. +bool dequant_supported(GgmlType t); + +// Write n fp32 values decoded from src (of type t) into dst. +// n must be a whole number of blocks for block-quantized types. +std::expected dequantize(GgmlType t, const std::byte* src, float* dst, uint64_t n); + +} // namespace onebit diff --git a/engine/core/gguf.cpp b/engine/core/gguf.cpp new file mode 100644 index 00000000..ef6ea86c --- /dev/null +++ b/engine/core/gguf.cpp @@ -0,0 +1,318 @@ +#include "gguf.h" + +#include +#include +#include +#include + +#include +#include +#include + +namespace onebit { + +std::string_view ggml_type_name(GgmlType t) { + switch (t) { + case GgmlType::F32: return "F32"; + case GgmlType::F16: return "F16"; + case GgmlType::Q4_0: return "Q4_0"; + case GgmlType::Q4_1: return "Q4_1"; + case GgmlType::Q5_0: return "Q5_0"; + case GgmlType::Q5_1: return "Q5_1"; + case GgmlType::Q8_0: return "Q8_0"; + case GgmlType::Q8_1: return "Q8_1"; + case GgmlType::Q2_K: return "Q2_K"; + case GgmlType::Q3_K: return "Q3_K"; + case GgmlType::Q4_K: return "Q4_K"; + case GgmlType::Q5_K: return "Q5_K"; + case GgmlType::Q6_K: return "Q6_K"; + case GgmlType::Q8_K: return "Q8_K"; + case GgmlType::BF16: return "BF16"; + } + return "unknown"; +} + +uint64_t GgufTensor::n_elements() const { + uint64_t n = 1; + for (uint64_t d : ne) n *= d; + return n; +} + +namespace { + +constexpr uint32_t kMagic = 0x46554747; // "GGUF" little-endian +constexpr uint64_t kDefaultAlignment = 32; +// Guards against corrupt headers asking for absurd allocations. +constexpr uint64_t kMaxCount = uint64_t{1} << 32; +constexpr uint32_t kMaxDims = 4; + +enum ValueType : uint32_t { + UINT8 = 0, INT8 = 1, UINT16 = 2, INT16 = 3, UINT32 = 4, INT32 = 5, + FLOAT32 = 6, BOOL = 7, STRING = 8, ARRAY = 9, UINT64 = 10, INT64 = 11, FLOAT64 = 12, +}; + +class Cursor { +public: + Cursor(const std::byte* base, size_t size) : base_(base), size_(size) {} + + size_t pos() const { return pos_; } + + template + std::expected read() { + if (size_ - pos_ < sizeof(T)) return std::unexpected(truncated()); + T out; + std::memcpy(&out, base_ + pos_, sizeof(T)); + pos_ += sizeof(T); + return out; + } + + std::expected read_string() { + auto len = read(); + if (!len) return std::unexpected(len.error()); + if (size_ - pos_ < *len) return std::unexpected(truncated()); + std::string s(reinterpret_cast(base_ + pos_), *len); + pos_ += *len; + return s; + } + + std::expected read_value(uint32_t type, int depth = 0) { + switch (type) { + case UINT8: return uint_value(); + case UINT16: return uint_value(); + case UINT32: return uint_value(); + case UINT64: return uint_value(); + case INT8: return int_value(); + case INT16: return int_value(); + case INT32: return int_value(); + case INT64: return int_value(); + case FLOAT32: return float_value(); + case FLOAT64: return float_value(); + case BOOL: { + auto b = read(); + if (!b) return std::unexpected(b.error()); + return GgufValue{*b != 0}; + } + case STRING: { + auto s = read_string(); + if (!s) return std::unexpected(s.error()); + return GgufValue{std::move(*s)}; + } + case ARRAY: { + if (depth > 0) return std::unexpected(std::string("nested GGUF arrays are not supported")); + auto elem_type = read(); + if (!elem_type) return std::unexpected(elem_type.error()); + auto count = read(); + if (!count) return std::unexpected(count.error()); + if (*count > kMaxCount) return std::unexpected(std::string("GGUF array too large")); + GgufArray arr; + arr.reserve(*count); + for (uint64_t i = 0; i < *count; ++i) { + auto v = read_value(*elem_type, depth + 1); + if (!v) return std::unexpected(v.error()); + arr.push_back(std::move(*v)); + } + return GgufValue{std::move(arr)}; + } + } + return std::unexpected(std::format("unknown GGUF value type {}", type)); + } + +private: + template + std::expected uint_value() { + auto x = read(); + if (!x) return std::unexpected(x.error()); + return GgufValue{uint64_t{*x}}; + } + template + std::expected int_value() { + auto x = read(); + if (!x) return std::unexpected(x.error()); + return GgufValue{int64_t{*x}}; + } + template + std::expected float_value() { + auto x = read(); + if (!x) return std::unexpected(x.error()); + return GgufValue{double{*x}}; + } + std::string truncated() const { return std::format("GGUF truncated at byte {}", pos_); } + + const std::byte* base_; + size_t size_; + size_t pos_ = 0; +}; + +// Bytes occupied by n elements of type t, or nullopt for types whose block +// size this reader does not know (the tensor is still listed, just unsized). +std::optional type_bytes(GgmlType t, uint64_t n) { + auto blocks = [&](uint64_t block, uint64_t bytes) -> std::optional { + if (n % block != 0) return std::nullopt; + return n / block * bytes; + }; + switch (t) { + case GgmlType::F32: return n * 4; + case GgmlType::F16: + case GgmlType::BF16: return n * 2; + case GgmlType::Q4_0: return blocks(32, 18); + case GgmlType::Q4_1: return blocks(32, 20); + case GgmlType::Q5_0: return blocks(32, 22); + case GgmlType::Q5_1: return blocks(32, 24); + case GgmlType::Q8_0: return blocks(32, 34); + case GgmlType::Q8_1: return blocks(32, 36); + case GgmlType::Q2_K: return blocks(256, 84); + case GgmlType::Q3_K: return blocks(256, 110); + case GgmlType::Q4_K: return blocks(256, 144); + case GgmlType::Q5_K: return blocks(256, 176); + case GgmlType::Q6_K: return blocks(256, 210); + case GgmlType::Q8_K: return blocks(256, 292); + } + return std::nullopt; +} + +} // namespace + +std::expected GgufFile::open(const std::string& path) { + int fd = ::open(path.c_str(), O_RDONLY | O_CLOEXEC); + if (fd < 0) return std::unexpected(std::format("cannot open {}: {}", path, std::strerror(errno))); + struct stat st {}; + if (fstat(fd, &st) != 0) { + ::close(fd); + return std::unexpected(std::format("cannot stat {}", path)); + } + GgufFile f; + f.size_ = static_cast(st.st_size); + void* m = mmap(nullptr, f.size_, PROT_READ, MAP_PRIVATE, fd, 0); + ::close(fd); + if (m == MAP_FAILED) return std::unexpected(std::format("cannot mmap {}", path)); + f.map_ = m; + + const auto* base = static_cast(m); + Cursor c(base, f.size_); + auto fail = [&](std::string msg) { return std::unexpected(path + ": " + msg); }; + + auto magic = c.read(); + if (!magic) return fail(magic.error()); + if (*magic != kMagic) return fail("not a GGUF file (bad magic)"); + auto version = c.read(); + if (!version) return fail(version.error()); + if (*version != 2 && *version != 3) return fail(std::format("unsupported GGUF version {}", *version)); + f.version_ = *version; + + auto n_tensors = c.read(); + if (!n_tensors) return fail(n_tensors.error()); + auto n_kv = c.read(); + if (!n_kv) return fail(n_kv.error()); + if (*n_tensors > kMaxCount || *n_kv > kMaxCount) return fail("header counts out of range"); + + for (uint64_t i = 0; i < *n_kv; ++i) { + auto key = c.read_string(); + if (!key) return fail(key.error()); + auto type = c.read(); + if (!type) return fail(type.error()); + auto value = c.read_value(*type); + if (!value) return fail(std::format("key {}: {}", *key, value.error())); + f.kv_.insert_or_assign(std::move(*key), std::move(*value)); + } + + f.tensors_.reserve(*n_tensors); + for (uint64_t i = 0; i < *n_tensors; ++i) { + GgufTensor t; + auto name = c.read_string(); + if (!name) return fail(name.error()); + t.name = std::move(*name); + auto n_dims = c.read(); + if (!n_dims) return fail(n_dims.error()); + if (*n_dims == 0 || *n_dims > kMaxDims) return fail(std::format("tensor {}: bad rank {}", t.name, *n_dims)); + for (uint32_t d = 0; d < *n_dims; ++d) { + auto dim = c.read(); + if (!dim) return fail(dim.error()); + t.ne.push_back(*dim); + } + auto type = c.read(); + if (!type) return fail(type.error()); + t.type = static_cast(*type); + auto offset = c.read(); + if (!offset) return fail(offset.error()); + t.offset = *offset; + f.tensors_.push_back(std::move(t)); + } + + uint64_t alignment = f.get_uint("general.alignment").value_or(kDefaultAlignment); + if (alignment == 0 || (alignment & (alignment - 1)) != 0) return fail("general.alignment is not a power of two"); + uint64_t data_start = (c.pos() + alignment - 1) / alignment * alignment; + + for (size_t i = 0; i < f.tensors_.size(); ++i) { + GgufTensor& t = f.tensors_[i]; + uint64_t begin = data_start + t.offset; + auto bytes = type_bytes(t.type, t.n_elements()); + if (!bytes) return fail(std::format("tensor {}: unsupported type {} or bad shape", t.name, + static_cast(t.type))); + if (begin > f.size_ || *bytes > f.size_ - begin) return fail(std::format("tensor {} extends past end of file", t.name)); + t.data = base + begin; + if (!f.by_name_.emplace(t.name, i).second) return fail(std::format("duplicate tensor {}", t.name)); + } + return f; +} + +GgufFile::GgufFile(GgufFile&& o) noexcept { *this = std::move(o); } + +GgufFile& GgufFile::operator=(GgufFile&& o) noexcept { + if (this != &o) { + release(); + map_ = std::exchange(o.map_, nullptr); + size_ = std::exchange(o.size_, 0); + version_ = o.version_; + kv_ = std::move(o.kv_); + tensors_ = std::move(o.tensors_); + by_name_ = std::move(o.by_name_); + } + return *this; +} + +GgufFile::~GgufFile() { release(); } + +void GgufFile::release() { + if (map_) munmap(map_, size_); + map_ = nullptr; + size_ = 0; +} + +const GgufValue* GgufFile::find(std::string_view key) const { + auto it = kv_.find(std::string(key)); + return it == kv_.end() ? nullptr : &it->second; +} + +const GgufTensor* GgufFile::tensor(std::string_view name) const { + auto it = by_name_.find(name); + return it == by_name_.end() ? nullptr : &tensors_[it->second]; +} + +std::optional GgufFile::get_uint(std::string_view key) const { + const GgufValue* v = find(key); + if (!v) return std::nullopt; + if (auto* u = std::get_if(&v->v)) return *u; + if (auto* i = std::get_if(&v->v); i && *i >= 0) return static_cast(*i); + return std::nullopt; +} + +std::optional GgufFile::get_float(std::string_view key) const { + const GgufValue* v = find(key); + if (!v) return std::nullopt; + if (auto* d = std::get_if(&v->v)) return *d; + return std::nullopt; +} + +std::optional GgufFile::get_string(std::string_view key) const { + const GgufValue* v = find(key); + if (!v) return std::nullopt; + if (auto* s = std::get_if(&v->v)) return *s; + return std::nullopt; +} + +const GgufArray* GgufFile::get_array(std::string_view key) const { + const GgufValue* v = find(key); + return v ? std::get_if(&v->v) : nullptr; +} + +} // namespace onebit diff --git a/engine/core/gguf.h b/engine/core/gguf.h new file mode 100644 index 00000000..08fe80e5 --- /dev/null +++ b/engine/core/gguf.h @@ -0,0 +1,94 @@ +// gguf.h: read-only GGUF v2/v3 reader over a memory-mapped file. +// +// Format reference: https://github.com/ggml-org/ggml/blob/master/docs/gguf.md +// Metadata values are decoded into GgufValue; tensor data is never copied, and +// GgufTensor::data points into the mapping, which lives as long as the GgufFile. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace onebit { + +// ggml tensor types this engine knows by id. Values match ggml.h. +enum class GgmlType : uint32_t { + F32 = 0, + F16 = 1, + Q4_0 = 2, + Q4_1 = 3, + Q5_0 = 6, + Q5_1 = 7, + Q8_0 = 8, + Q8_1 = 9, + Q2_K = 10, + Q3_K = 11, + Q4_K = 12, + Q5_K = 13, + Q6_K = 14, + Q8_K = 15, + BF16 = 30, +}; + +std::string_view ggml_type_name(GgmlType t); + +struct GgufValue; +using GgufArray = std::vector; + +struct GgufValue { + std::variant v; +}; + +struct GgufTensor { + std::string name; + GgmlType type; + std::vector ne; // ne[0] is the contiguous (fastest) dimension + uint64_t offset = 0; // relative to the start of the data section + const std::byte* data = nullptr; + + uint64_t n_elements() const; +}; + +class GgufFile { +public: + static std::expected open(const std::string& path); + + GgufFile(GgufFile&& other) noexcept; + GgufFile& operator=(GgufFile&& other) noexcept; + GgufFile(const GgufFile&) = delete; + GgufFile& operator=(const GgufFile&) = delete; + ~GgufFile(); + + uint32_t version() const { return version_; } + const std::map& metadata() const { return kv_; } + const std::vector& tensors() const { return tensors_; } + + const GgufValue* find(std::string_view key) const; + const GgufTensor* tensor(std::string_view name) const; + + // Typed accessors. Integer getters accept any integer metadata type that + // fits; they return nullopt when the key is missing or has another type. + std::optional get_uint(std::string_view key) const; + std::optional get_float(std::string_view key) const; + std::optional get_string(std::string_view key) const; + const GgufArray* get_array(std::string_view key) const; + +private: + GgufFile() = default; + void release(); + + void* map_ = nullptr; + size_t size_ = 0; + uint32_t version_ = 0; + std::map kv_; + std::vector tensors_; + std::map> by_name_; +}; + +} // namespace onebit diff --git a/engine/core/model_config.cpp b/engine/core/model_config.cpp new file mode 100644 index 00000000..a8013a76 --- /dev/null +++ b/engine/core/model_config.cpp @@ -0,0 +1,54 @@ +#include "model_config.h" + +#include + +namespace onebit { + +std::expected read_model_config(const GgufFile& f) { + ModelConfig c; + auto arch = f.get_string("general.architecture"); + if (!arch) return std::unexpected(std::string("missing general.architecture")); + c.arch = *arch; + + // Only architectures that have passed the golden test are accepted. + if (c.arch == "qwen3") { + c.rope_style = RopeStyle::Neox; + } else { + return std::unexpected(std::format("architecture '{}' is not verified on the CPU reference", c.arch)); + } + + auto u32 = [&](const std::string& suffix) -> std::expected { + auto v = f.get_uint(c.arch + "." + suffix); + if (!v) return std::unexpected(std::format("missing {}.{}", c.arch, suffix)); + if (*v == 0 || *v > UINT32_MAX) return std::unexpected(std::format("{}.{} out of range", c.arch, suffix)); + return static_cast(*v); + }; + auto set = [&](uint32_t& field, const std::string& suffix) -> std::expected { + auto v = u32(suffix); + if (!v) return std::unexpected(v.error()); + field = *v; + return {}; + }; + + for (auto [field, key] : {std::pair{&c.n_layer, "block_count"}, {&c.n_embd, "embedding_length"}, + {&c.n_head, "attention.head_count"}, {&c.n_ff, "feed_forward_length"}, + {&c.n_ctx_train, "context_length"}}) { + if (auto r = set(*field, key); !r) return std::unexpected(r.error()); + } + c.n_head_kv = u32("attention.head_count_kv").value_or(c.n_head); + c.head_dim = u32("attention.key_length").value_or(c.n_embd / c.n_head); + if (auto vl = f.get_uint(c.arch + ".attention.value_length"); vl && *vl != c.head_dim) + return std::unexpected(std::string("key_length != value_length is not supported")); + if (c.n_head % c.n_head_kv != 0) return std::unexpected(std::string("head_count is not a multiple of head_count_kv")); + + if (auto eps = f.get_float(c.arch + ".attention.layer_norm_rms_epsilon")) c.rms_eps = static_cast(*eps); + if (auto theta = f.get_float(c.arch + ".rope.freq_base")) c.rope_theta = static_cast(*theta); + + const GgufTensor* emb = f.tensor("token_embd.weight"); + if (!emb || emb->ne.size() != 2 || emb->ne[0] != c.n_embd) + return std::unexpected(std::string("token_embd.weight missing or wrong shape")); + c.n_vocab = static_cast(emb->ne[1]); + return c; +} + +} // namespace onebit diff --git a/engine/core/model_config.h b/engine/core/model_config.h new file mode 100644 index 00000000..d0713e9c --- /dev/null +++ b/engine/core/model_config.h @@ -0,0 +1,36 @@ +// model_config.h: decoder hyper-parameters read from GGUF metadata. +#pragma once + +#include +#include +#include + +#include "gguf.h" + +namespace onebit { + +enum class RopeStyle { + Normal, // rotate adjacent pairs (x[2i], x[2i+1]); llama GGUFs are permuted for this + Neox, // rotate halves (x[i], x[i + d/2]); qwen2/qwen3 +}; + +struct ModelConfig { + std::string arch; // general.architecture, e.g. "qwen3" + uint32_t n_layer = 0; + uint32_t n_embd = 0; + uint32_t n_head = 0; + uint32_t n_head_kv = 0; + uint32_t head_dim = 0; + uint32_t n_ff = 0; + uint32_t n_vocab = 0; + uint32_t n_ctx_train = 0; + float rms_eps = 1e-6f; + float rope_theta = 10000.0f; + RopeStyle rope_style = RopeStyle::Normal; +}; + +// Reads and validates the config. Fails for architectures the engine has not +// been verified on, rather than guessing their semantics. +std::expected read_model_config(const GgufFile& f); + +} // namespace onebit diff --git a/engine/core/parallel.cpp b/engine/core/parallel.cpp new file mode 100644 index 00000000..258b4d81 --- /dev/null +++ b/engine/core/parallel.cpp @@ -0,0 +1,63 @@ +#include "parallel.h" + +#include + +namespace onebit { + +ThreadPool::ThreadPool(size_t n_threads) { + if (n_threads == 0) n_threads = std::max(1u, std::thread::hardware_concurrency()); + for (size_t i = 1; i < n_threads; ++i) workers_.emplace_back([this, i] { worker_loop(i); }); +} + +ThreadPool::~ThreadPool() { + { + std::lock_guard lock(mu_); + stop_ = true; + } + start_cv_.notify_all(); + for (auto& t : workers_) t.join(); +} + +void ThreadPool::run_chunk(size_t index) { + const size_t n = job_n_, parts = size(); + const size_t begin = n * index / parts, end = n * (index + 1) / parts; + if (begin < end) (*job_)(begin, end); +} + +void ThreadPool::worker_loop(size_t index) { + size_t seen = 0; + for (;;) { + { + std::unique_lock lock(mu_); + start_cv_.wait(lock, [&] { return stop_ || generation_ != seen; }); + if (stop_) return; + seen = generation_; + } + run_chunk(index); + { + std::lock_guard lock(mu_); + if (--pending_ == 0) done_cv_.notify_one(); + } + } +} + +void ThreadPool::parallel_for(size_t n, const std::function& fn) { + if (workers_.empty() || n < 2) { + if (n > 0) fn(0, n); + return; + } + { + std::lock_guard lock(mu_); + job_ = &fn; + job_n_ = n; + pending_ = workers_.size(); + ++generation_; + } + start_cv_.notify_all(); + run_chunk(0); + std::unique_lock lock(mu_); + done_cv_.wait(lock, [&] { return pending_ == 0; }); + job_ = nullptr; +} + +} // namespace onebit diff --git a/engine/core/parallel.h b/engine/core/parallel.h new file mode 100644 index 00000000..7a660851 --- /dev/null +++ b/engine/core/parallel.h @@ -0,0 +1,43 @@ +// parallel.h: a fixed pool of worker threads running one parallel_for at a time. +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace onebit { + +class ThreadPool { +public: + // n_threads == 0 uses std::thread::hardware_concurrency(). + explicit ThreadPool(size_t n_threads = 0); + ~ThreadPool(); + ThreadPool(const ThreadPool&) = delete; + ThreadPool& operator=(const ThreadPool&) = delete; + + size_t size() const { return workers_.size() + 1; } + + // Split [0, n) into contiguous chunks and call fn(begin, end) on each, + // using the calling thread as one of the workers. Blocks until done. + // Not reentrant: fn must not call parallel_for on the same pool. + void parallel_for(size_t n, const std::function& fn); + +private: + void worker_loop(size_t index); + void run_chunk(size_t index); + + std::vector workers_; + std::mutex mu_; + std::condition_variable start_cv_; + std::condition_variable done_cv_; + const std::function* job_ = nullptr; + size_t job_n_ = 0; + size_t generation_ = 0; + size_t pending_ = 0; + bool stop_ = false; +}; + +} // namespace onebit diff --git a/tests/check.h b/tests/check.h new file mode 100644 index 00000000..dfc8daf4 --- /dev/null +++ b/tests/check.h @@ -0,0 +1,29 @@ +// check.h: minimal test assertions (no framework dependency). +#pragma once + +#include +#include + +inline int g_failures = 0; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "%s:%d: CHECK failed: %s\n", __FILE__, __LINE__, #cond); \ + ++g_failures; \ + } \ + } while (0) + +#define REQUIRE(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "%s:%d: REQUIRE failed: %s\n", __FILE__, __LINE__, #cond); \ + std::exit(1); \ + } \ + } while (0) + +inline int test_result(const char* name) { + if (g_failures) std::fprintf(stderr, "%s: %d check(s) failed\n", name, g_failures); + else std::printf("%s: ok\n", name); + return g_failures ? 1 : 0; +} diff --git a/tests/gguf_writer.h b/tests/gguf_writer.h new file mode 100644 index 00000000..e94a448e --- /dev/null +++ b/tests/gguf_writer.h @@ -0,0 +1,91 @@ +// gguf_writer.h: builds small GGUF files in memory for tests. +#pragma once + +#include +#include +#include +#include +#include +#include + +struct GgufWriter { + struct Tensor { + std::string name; + std::vector ne; + uint32_t type; + std::vector bytes; + }; + + std::vector kv; + uint64_t n_kv = 0; + std::vector tensors; + uint64_t alignment = 32; + + template + static void put(std::vector& out, T v) { + uint8_t b[sizeof(T)]; + std::memcpy(b, &v, sizeof(T)); + out.insert(out.end(), b, b + sizeof(T)); + } + static void put_str(std::vector& out, const std::string& s) { + put(out, s.size()); + out.insert(out.end(), s.begin(), s.end()); + } + + void key(const std::string& k, uint32_t type) { + put_str(kv, k); + put(kv, type); + ++n_kv; + } + void u32(const std::string& k, uint32_t v) { key(k, 4); put(kv, v); } + void i32(const std::string& k, int32_t v) { key(k, 5); put(kv, v); } + void f32(const std::string& k, float v) { key(k, 6); put(kv, v); } + void str(const std::string& k, const std::string& v) { key(k, 8); put_str(kv, v); } + void str_array(const std::string& k, const std::vector& v) { + key(k, 9); + put(kv, 8); + put(kv, v.size()); + for (const auto& s : v) put_str(kv, s); + } + void tensor_f32(const std::string& name, std::vector ne, const std::vector& data) { + Tensor t{name, std::move(ne), 0, {}}; + t.bytes.resize(data.size() * 4); + std::memcpy(t.bytes.data(), data.data(), t.bytes.size()); + tensors.push_back(std::move(t)); + } + + std::vector build() const { + std::vector out; + put(out, 0x46554747); + put(out, 3); + put(out, tensors.size()); + put(out, n_kv); + out.insert(out.end(), kv.begin(), kv.end()); + std::vector offsets; + uint64_t off = 0; + for (const auto& t : tensors) { + offsets.push_back(off); + off += (t.bytes.size() + alignment - 1) / alignment * alignment; + } + for (size_t i = 0; i < tensors.size(); ++i) { + const auto& t = tensors[i]; + put_str(out, t.name); + put(out, t.ne.size()); + for (auto d : t.ne) put(out, d); + put(out, t.type); + put(out, offsets[i]); + } + while (out.size() % alignment) out.push_back(0); + for (const auto& t : tensors) { + out.insert(out.end(), t.bytes.begin(), t.bytes.end()); + while (out.size() % alignment) out.push_back(0); + } + return out; + } + + static std::string write_temp(const std::vector& bytes, const std::string& tag) { + std::string path = std::string(std::getenv("TMPDIR") ? std::getenv("TMPDIR") : "/tmp") + "/onebit-test-" + tag + ".gguf"; + std::ofstream(path, std::ios::binary).write(reinterpret_cast(bytes.data()), bytes.size()); + return path; + } +}; diff --git a/tests/golden/qwen3-0.6b-capitals/meta.json b/tests/golden/qwen3-0.6b-capitals/meta.json new file mode 100644 index 00000000..121f603b --- /dev/null +++ b/tests/golden/qwen3-0.6b-capitals/meta.json @@ -0,0 +1,13 @@ +{ + "model": "/home/bcloud/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/c1899de289a04d12100db370d81485cdf75e47ca", + "revision": "c1899de289a04d12100db370d81485cdf75e47ca", + "prompt": "The capital of France is Paris. The capital of Japan is", + "n_prompt": 12, + "n_tokens": 36, + "vocab": 151936, + "dtype": "float32", + "attn_implementation": "eager", + "torch": "2.13.0+cpu", + "transformers": "5.13.0.dev0", + "logits_sha256": "3aaec202fa3d5e4dca1ebc053623836a57960f242ccbd34e98036aa97654820a" +} \ No newline at end of file diff --git a/tests/golden/qwen3-0.6b-capitals/tokens.txt b/tests/golden/qwen3-0.6b-capitals/tokens.txt new file mode 100644 index 00000000..8e443131 --- /dev/null +++ b/tests/golden/qwen3-0.6b-capitals/tokens.txt @@ -0,0 +1 @@ +785 6722 315 9625 374 12095 13 576 6722 315 6323 374 26194 13 576 6722 315 15344 374 21718 13 576 6722 315 17689 374 24081 13 576 6722 315 5616 374 26549 13 576 diff --git a/tests/golden_cpu.cpp b/tests/golden_cpu.cpp new file mode 100644 index 00000000..1fb78160 --- /dev/null +++ b/tests/golden_cpu.cpp @@ -0,0 +1,124 @@ +// golden_cpu: compare the CPU reference against golden logits from HF transformers. +// +// usage: golden_cpu [--max-kl X] [--threads N] +// +// golden_dir (written by tools/golden/make_golden.py) holds: +// tokens.txt whitespace-separated token ids (prompt followed by the reference's greedy continuation) +// logits.f32 float32 [n_tokens][n_vocab]: the reference's next-token logits after each token +// +// The model is fed every token teacher-forced. The test passes when the argmax +// agrees at every position and every per-position KL(ref || cpu) is <= max_kl. +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cpu_model.h" + +using namespace onebit; + +namespace { + +std::vector log_softmax(const float* x, size_t n) { + const double mx = *std::max_element(x, x + n); + double sum = 0.0; + for (size_t i = 0; i < n; ++i) sum += std::exp(double(x[i]) - mx); + const double lse = mx + std::log(sum); + std::vector out(n); + for (size_t i = 0; i < n; ++i) out[i] = double(x[i]) - lse; + return out; +} + +double kl(const float* ref, const float* got, size_t n) { + auto lp = log_softmax(ref, n), lq = log_softmax(got, n); + double d = 0.0; + for (size_t i = 0; i < n; ++i) d += std::exp(lp[i]) * (lp[i] - lq[i]); + return d; +} + +size_t argmax(const float* x, size_t n) { return size_t(std::max_element(x, x + n) - x); } + +} // namespace + +int main(int argc, char** argv) { + if (argc < 3) { + std::fprintf(stderr, "usage: %s [--max-kl X] [--threads N]\n", argv[0]); + return 2; + } + const std::string model_path = argv[1], dir = argv[2]; + double max_kl = 1e-6; + size_t threads = 0; + for (int i = 3; i + 1 < argc; i += 2) { + std::string flag = argv[i]; + if (flag == "--max-kl") max_kl = std::atof(argv[i + 1]); + else if (flag == "--threads") threads = std::strtoul(argv[i + 1], nullptr, 10); + else { + std::fprintf(stderr, "unknown flag %s\n", flag.c_str()); + return 2; + } + } + + std::vector tokens; + { + std::ifstream in(dir + "/tokens.txt"); + for (int32_t t; in >> t;) tokens.push_back(t); + } + if (tokens.empty()) { + std::fprintf(stderr, "no tokens in %s/tokens.txt\n", dir.c_str()); + return 2; + } + + auto t0 = std::chrono::steady_clock::now(); + auto model = CpuModel::load(model_path, threads); + if (!model) { + std::fprintf(stderr, "load failed: %s\n", model.error().c_str()); + return 1; + } + const size_t V = model->config().n_vocab; + auto t1 = std::chrono::steady_clock::now(); + + std::vector ref(tokens.size() * V); + { + std::ifstream in(dir + "/logits.f32", std::ios::binary); + in.read(reinterpret_cast(ref.data()), std::streamsize(ref.size() * sizeof(float))); + if (!in || in.peek() != EOF) { + std::fprintf(stderr, "logits.f32 is not %zu x %zu floats\n", tokens.size(), V); + return 2; + } + } + + double worst_kl = 0.0, worst_abs = 0.0; + size_t mismatches = 0; + for (size_t p = 0; p < tokens.size(); ++p) { + auto logits = model->forward(tokens[p]); + if (!logits) { + std::fprintf(stderr, "forward failed at %zu: %s\n", p, logits.error().c_str()); + return 1; + } + const float* r = ref.data() + p * V; + const float* g = (*logits)->data(); + const double k = kl(r, g, V); + double a = 0.0; + for (size_t i = 0; i < V; ++i) a = std::max(a, double(std::fabs(r[i] - g[i]))); + const size_t ar = argmax(r, V), ag = argmax(g, V); + if (ar != ag) ++mismatches; + worst_kl = std::max(worst_kl, k); + worst_abs = std::max(worst_abs, a); + std::printf("pos %3zu tok %6d argmax ref %6zu cpu %6zu %s KL %.3e max|dlogit| %.3e\n", p, tokens[p], ar, ag, + ar == ag ? " " : "!!", k, a); + } + auto t2 = std::chrono::steady_clock::now(); + const double load_s = std::chrono::duration(t1 - t0).count(); + const double run_s = std::chrono::duration(t2 - t1).count(); + + const bool pass = mismatches == 0 && worst_kl <= max_kl; + std::printf("\n%zu positions, argmax mismatches %zu, worst KL %.3e (gate %.1e), worst max|dlogit| %.3e\n", + tokens.size(), mismatches, worst_kl, max_kl, worst_abs); + std::printf("load %.1f s, %.2f tok/s teacher-forced\n", load_s, tokens.size() / run_s); + std::printf("%s\n", pass ? "PASS" : "FAIL"); + return pass ? 0 : 1; +} diff --git a/tests/test_dequant.cpp b/tests/test_dequant.cpp new file mode 100644 index 00000000..a2e6a703 --- /dev/null +++ b/tests/test_dequant.cpp @@ -0,0 +1,60 @@ +#include +#include +#include + +#include "check.h" +#include "dequant.h" + +using namespace onebit; + +static void test_fp16() { + CHECK(fp16_to_fp32(0x3C00) == 1.0f); + CHECK(fp16_to_fp32(0xC000) == -2.0f); + CHECK(fp16_to_fp32(0x3555) == 0.333251953125f); + CHECK(fp16_to_fp32(0x7BFF) == 65504.0f); + CHECK(fp16_to_fp32(0x0001) == std::ldexp(1.0f, -24)); // smallest subnormal + CHECK(fp16_to_fp32(0x03FF) == std::ldexp(1023.0f, -24)); // largest subnormal + CHECK(fp16_to_fp32(0x0400) == std::ldexp(1.0f, -14)); // smallest normal + CHECK(fp16_to_fp32(0x0000) == 0.0f); + CHECK(std::signbit(fp16_to_fp32(0x8000))); + CHECK(std::isinf(fp16_to_fp32(0x7C00))); + CHECK(std::isnan(fp16_to_fp32(0x7E00))); +} + +static void test_bf16() { + CHECK(bf16_to_fp32(0x3F80) == 1.0f); + CHECK(bf16_to_fp32(0xC040) == -3.0f); + CHECK(bf16_to_fp32(0x3EAB) == 0.333984375f); +} + +static void test_q8_0() { + // Two blocks: scale 0.5 then scale -1.0 (fp16), quants -128..127 pattern. + std::vector buf(68); + const uint16_t scales[2] = {0x3800, 0xBC00}; + for (int b = 0; b < 2; ++b) { + std::memcpy(buf.data() + b * 34, &scales[b], 2); + for (int i = 0; i < 32; ++i) buf[b * 34 + 2 + i] = std::byte(uint8_t(int8_t(i * 8 - 128))); + } + std::vector out(64); + REQUIRE(dequantize(GgmlType::Q8_0, buf.data(), out.data(), 64).has_value()); + CHECK(out[0] == -64.0f); + CHECK(out[31] == 0.5f * 120.0f); + CHECK(out[32] == 128.0f); + CHECK(out[63] == -120.0f); + CHECK(!dequantize(GgmlType::Q8_0, buf.data(), out.data(), 40).has_value()); +} + +static void test_unsupported() { + float out[256]; + std::byte in[256] = {}; + CHECK(!dequant_supported(GgmlType::Q4_K)); + CHECK(!dequantize(GgmlType::Q4_K, in, out, 256).has_value()); +} + +int main() { + test_fp16(); + test_bf16(); + test_q8_0(); + test_unsupported(); + return test_result("test_dequant"); +} diff --git a/tests/test_gguf.cpp b/tests/test_gguf.cpp new file mode 100644 index 00000000..e3dd74b8 --- /dev/null +++ b/tests/test_gguf.cpp @@ -0,0 +1,78 @@ +#include +#include +#include + +#include "check.h" +#include "gguf.h" +#include "gguf_writer.h" + +using namespace onebit; + +static void test_roundtrip() { + GgufWriter w; + w.str("general.architecture", "qwen3"); + w.u32("qwen3.block_count", 2); + w.i32("signed", -7); + w.f32("qwen3.rope.freq_base", 1000000.0f); + w.str_array("tokenizer.ggml.tokens", {"a", "bb", ""}); + w.tensor_f32("x.weight", {3, 2}, {1, 2, 3, 4, 5, 6}); + w.tensor_f32("y.weight", {1}, {9}); + std::string path = GgufWriter::write_temp(w.build(), "roundtrip"); + + auto f = GgufFile::open(path); + REQUIRE(f.has_value()); + CHECK(f->version() == 3); + CHECK(f->get_string("general.architecture") == "qwen3"); + CHECK(f->get_uint("qwen3.block_count") == 2u); + CHECK(!f->get_uint("signed").has_value()); // negative: not representable + CHECK(f->get_float("qwen3.rope.freq_base") == 1000000.0); + CHECK(!f->get_string("missing").has_value()); + const GgufArray* toks = f->get_array("tokenizer.ggml.tokens"); + REQUIRE(toks && toks->size() == 3); + CHECK(std::get((*toks)[1].v) == "bb"); + + const GgufTensor* x = f->tensor("x.weight"); + REQUIRE(x); + CHECK(x->type == GgmlType::F32); + CHECK(x->n_elements() == 6); + CHECK((x->ne == std::vector{3, 2})); + CHECK(reinterpret_cast(x->data)[5] == 6.0f); + const GgufTensor* y = f->tensor("y.weight"); + REQUIRE(y); + CHECK(reinterpret_cast(y->data) % 32 == 0); + CHECK(reinterpret_cast(y->data)[0] == 9.0f); + CHECK(f->tensor("z.weight") == nullptr); + std::remove(path.c_str()); +} + +static void test_rejects_bad_files() { + GgufWriter w; + w.tensor_f32("x.weight", {4}, {1, 2, 3, 4}); + auto good = w.build(); + + auto bad_magic = good; + bad_magic[0] = 'X'; + std::string p1 = GgufWriter::write_temp(bad_magic, "badmagic"); + CHECK(!GgufFile::open(p1).has_value()); + + auto truncated = good; + truncated.resize(truncated.size() - 20); // 16 data bytes + 16 padding: cut into the data + std::string p2 = GgufWriter::write_temp(truncated, "truncated"); + CHECK(!GgufFile::open(p2).has_value()); + + auto header_only = good; + header_only.resize(20); + std::string p3 = GgufWriter::write_temp(header_only, "header"); + CHECK(!GgufFile::open(p3).has_value()); + + CHECK(!GgufFile::open("/nonexistent/file.gguf").has_value()); + std::remove(p1.c_str()); + std::remove(p2.c_str()); + std::remove(p3.c_str()); +} + +int main() { + test_roundtrip(); + test_rejects_bad_files(); + return test_result("test_gguf"); +} diff --git a/tools/golden/make_golden.py b/tools/golden/make_golden.py new file mode 100644 index 00000000..21d0db17 --- /dev/null +++ b/tools/golden/make_golden.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +"""Write golden next-token logits from HF transformers (fp32, eager attention). + +usage: make_golden.py --model --prompt "..." --n-gen 16 --out + +Output (read by tests/golden_cpu.cpp): + tokens.txt prompt ids followed by the fp32 model's greedy continuation + logits.f32 float32 [n_tokens][vocab], the logits after each token (teacher-forced) + meta.json model id, revision, prompt, library versions, sha256 of logits.f32 +""" +import argparse +import hashlib +import json +import os + +import numpy as np +import torch +import transformers +from transformers import AutoModelForCausalLM, AutoTokenizer + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--prompt", required=True) + ap.add_argument("--n-gen", type=int, default=16) + ap.add_argument("--out", required=True) + args = ap.parse_args() + + torch.manual_seed(0) + tok = AutoTokenizer.from_pretrained(args.model) + model = AutoModelForCausalLM.from_pretrained( + args.model, dtype=torch.float32, attn_implementation="eager" + ).eval() + + ids = tok(args.prompt, add_special_tokens=False)["input_ids"] + with torch.no_grad(): + gen = model.generate( + torch.tensor([ids]), max_new_tokens=args.n_gen, do_sample=False, + eos_token_id=None, pad_token_id=tok.eos_token_id, + )[0].tolist() + # Teacher-forced pass over the whole sequence: logits after every token. + logits = model(torch.tensor([gen])).logits[0].float().numpy() + + os.makedirs(args.out, exist_ok=True) + with open(os.path.join(args.out, "tokens.txt"), "w") as f: + f.write(" ".join(map(str, gen)) + "\n") + blob = np.ascontiguousarray(logits, dtype=np.float32).tobytes() + with open(os.path.join(args.out, "logits.f32"), "wb") as f: + f.write(blob) + meta = { + "model": args.model, + "revision": getattr(model.config, "_commit_hash", None), + "prompt": args.prompt, + "n_prompt": len(ids), + "n_tokens": len(gen), + "vocab": int(logits.shape[1]), + "dtype": "float32", + "attn_implementation": "eager", + "torch": torch.__version__, + "transformers": transformers.__version__, + "logits_sha256": hashlib.sha256(blob).hexdigest(), + } + with open(os.path.join(args.out, "meta.json"), "w") as f: + json.dump(meta, f, indent=2) + print(json.dumps(meta, indent=2)) + print("continuation:", repr(tok.decode(gen[len(ids):]))) + + +if __name__ == "__main__": + main()