// Phase 1 gate for the dspark forward graph (src/models/dspark.cpp).
//
// Tier 1 (--tier1): load a real dspark GGUF, feed SYNTHETIC (deterministic,
// in-process) tap features + a draft block, assert the forward pass produces
// finite, sane-shaped logits. Catches wiring bugs cheaply before worrying
// about numerical correctness.
//
// Tier 2 (--tier2 <ref.bin>): feed the SAME tap features / draft tokens /
// positions used by a reference implementation, and diff
// this program's logits against the reference logits dumped in ref.bin.
//
// ref.bin layout (little-endian):
// int32 n_ctx_rows, int32 n_embd_cap, int32 block_size, int32 vocab_size
// int32[n_ctx_rows] ctx_pos
// float32[n_ctx_rows * n_embd_cap] ctx_feat (row-major)
// int32[block_size] draft_token_ids
// int32[block_size] draft_pos
// float32[block_size * vocab_size] ref_logits (row-major)
#include "../src/llama-ext.h"
#include "../src/llama-graph.h"
#include "../src/llama-model.h"
#include "common.h"
#include "ggml-backend.h"
#include "llama.h"
#include "speculative.h"
#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <functional>
#include <map>
#include <random>
#include <string>
#include <vector>
#ifdef _WIN32
// the Windows CRT has no POSIX setenv/unsetenv
static int setenv(const char * name, const char * value, int overwrite) {
if (!overwrite && std::getenv(name)) {
return 0;
}
return _putenv_s(name, value);
}
static int unsetenv(const char * name) {
return _putenv_s(name, ""); // an empty value removes the variable
}
#endif
static std::map<std::string, std::vector<float>> trace_values;
static bool trace_node(ggml_tensor * tensor, bool ask, void *) {
const std::string name = tensor->name;
const bool selected = name == "dspark_fc" || name == "dspark_hidden_norm" || name == "dspark_log_snr_fc1" ||
name == "dspark_log_snr_fc2" || name == "dspark_draft_embd_snr" ||
name == "dspark_log_snr_expanded" || name == "dspark_balanced_output" ||
name == "result_norm" || name.find("l_out-") == 0 || name.find("dspark_corr_") == 0;
if (ask) {
return selected && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor);
}
auto & values = trace_values[name];
values.resize(ggml_nelements(tensor));
ggml_backend_tensor_get(tensor, values.data(), 0, values.size() * sizeof(float));
return true;
}
[[noreturn]] static void fail(const std::string & msg) {
fprintf(stderr, "FAIL: %s\n", msg.c_str());
exit(1);
}
// dspark ships no tokenizer of its own (tied to the TARGET's vocab), so
// llama_vocab_n_tokens() is 0 for a dspark GGUF -- the real vocab width only
// exists as token_embd.weight's own shape (see src/models/dspark.cpp's
// load_arch_tensors for the mirror-image fix on the loader side).
static int64_t n_vocab_from_model(const llama_model * model) {
const struct ggml_tensor * t = model->get_tensor("token_embd.weight");
if (!t) {
fail("token_embd.weight not found in loaded model");
}
return t->ne[1];
}
struct dspark_meta {
int64_t n_embd = 0;
int64_t n_vocab = 0;
int32_t block_size = 0;
int32_t mask_token_id = 0;
int32_t n_capture = 0;
};
// pull the handful of dspark hparams we need out of GGUF metadata via the
// generic llama_model_meta_* accessors (arch-prefixed keys).
static dspark_meta read_dspark_meta(const llama_model * model) {
dspark_meta m;
m.n_embd = llama_model_n_embd(model);
m.n_vocab = n_vocab_from_model(model);
char buf[256];
if (llama_model_meta_val_str(model, "dspark.dspark.block_size", buf, sizeof(buf)) > 0) {
m.block_size = atoi(buf);
}
if (llama_model_meta_val_str(model, "dspark.dspark.mask_token_id", buf, sizeof(buf)) > 0) {
m.mask_token_id = atoi(buf);
}
const ggml_tensor * fc = model->get_tensor("dspark.fc.weight");
if (!fc || m.n_embd <= 0 || fc->ne[0] % m.n_embd != 0) {
fail("invalid target-feature projection width");
}
m.n_capture = (int32_t) (fc->ne[0] / m.n_embd);
if (m.n_capture <= 0) {
fail("no target layers captured");
}
return m;
}
static llama_context * make_ctx(llama_model * model, uint32_t n_ctx) {
llama_context_params cparams = llama_context_default_params();
cparams.n_ctx = n_ctx;
cparams.n_batch = n_ctx;
cparams.n_ubatch = n_ctx;
cparams.n_seq_max = 1;
cparams.n_outputs_max_per_seq = 0;
cparams.no_perf = true;
if (std::getenv("DSPARK_TEST_NO_FA")) {
cparams.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
}
if (std::getenv("DSPARK_TEST_TRACE") || std::getenv("DSPARK_TEST_ORACLE")) {
cparams.cb_eval = trace_node;
}
llama_context * ctx = llama_init_from_model(model, cparams);
if (!ctx) {
fail("llama_init_from_model failed");
}
// dspark attention is fully non-causal (attention_mask=None, is_causal=False
// in the reference) -- see src/models/dspark.cpp header comment.
llama_set_causal_attn(ctx, false);
return ctx;
}
// runs one dspark forward call: n_ctx_rows context rows (dummy token ids,
// real positions + tap features staged via llama_set_dspark_ctx) followed by
// block_size draft-block rows (real token ids). Returns the block_size*n_vocab
// logits copied out of the context.
static std::vector<float> run_forward(llama_context * ctx,
int64_t n_ctx_rows,
int64_t n_embd_cap,
const std::vector<float> & ctx_feat,
const std::vector<int32_t> & ctx_pos,
int32_t block_size,
const std::vector<int32_t> & draft_tokens,
const std::vector<int32_t> & draft_pos,
int64_t n_vocab,
std::vector<llama_token> * sampled = nullptr,
int32_t output_rows = -1) {
if ((int64_t) ctx_feat.size() != n_ctx_rows * n_embd_cap) {
fail("ctx_feat size mismatch");
}
if ((int64_t) ctx_pos.size() != n_ctx_rows) {
fail("ctx_pos size mismatch");
}
if ((int64_t) draft_tokens.size() != block_size) {
fail("draft_tokens size mismatch");
}
if ((int64_t) draft_pos.size() != block_size) {
fail("draft_pos size mismatch");
}
if (output_rows == -1) {
output_rows = block_size;
}
if (output_rows < 1 || output_rows > block_size) {
fail("invalid output row count");
}
llama_set_dspark_ctx(ctx, ctx_feat.data(), n_ctx_rows, n_embd_cap, ctx_pos.data());
const int32_t n_tokens = (int32_t) n_ctx_rows + block_size;
llama_batch batch = llama_batch_init(n_tokens, 0, 1);
batch.n_tokens = n_tokens;
for (int32_t i = 0; i < (int32_t) n_ctx_rows; ++i) {
batch.token[i] = 0; // dummy: unused (context rows never join the trunk/embedding path)
batch.pos[i] = ctx_pos[i];
batch.n_seq_id[i] = 1;
batch.seq_id[i][0] = 0;
batch.logits[i] = 0;
}
for (int32_t j = 0; j < block_size; ++j) {
const int32_t i = (int32_t) n_ctx_rows + j;
batch.token[i] = draft_tokens[j];
batch.pos[i] = draft_pos[j];
batch.n_seq_id[i] = 1;
batch.seq_id[i][0] = 0;
batch.logits[i] = j < output_rows;
}
int32_t rc;
try {
rc = llama_decode(ctx, batch);
} catch (...) {
llama_batch_free(batch);
llama_set_dspark_ctx(ctx, nullptr, 0, 0, nullptr);
throw;
}
llama_batch_free(batch);
if (rc != 0) {
fail("llama_decode returned " + std::to_string(rc));
}
llama_set_dspark_ctx(ctx, nullptr, 0, 0, nullptr); // clear staged context
if (sampled) {
sampled->clear();
for (int32_t i = 0; i < output_rows; ++i) {
sampled->push_back(llama_get_sampled_token_ith(ctx, i - output_rows));
}
return {};
}
float * logits = llama_get_logits(ctx);
if (!logits) {
fail("llama_get_logits returned null");
}
return std::vector<float>(logits, logits + (size_t) output_rows * n_vocab);
}
static std::vector<float> correction_reference(const llama_model * model,
const float * hidden,
const std::vector<int32_t> & ids,
const dspark_meta & meta) {
if (!hidden) {
fail("missing trunk embeddings for correction oracle");
}
auto read = [&](const char * name) {
const ggml_tensor * tensor = model->get_tensor(name);
if (!tensor || (tensor->type != GGML_TYPE_F32 && tensor->type != GGML_TYPE_BF16)) {
fail("correction oracle requires F32 or BF16 tensors");
}
std::vector<float> data(ggml_nelements(tensor));
if (tensor->type == GGML_TYPE_BF16) {
std::vector<ggml_bf16_t> packed(data.size());
ggml_backend_tensor_get(tensor, packed.data(), 0, packed.size() * sizeof(ggml_bf16_t));
ggml_bf16_to_fp32_row(packed.data(), data.data(), data.size());
} else {
ggml_backend_tensor_get(tensor, data.data(), 0, data.size() * sizeof(float));
}
return data;
};
auto norm = [](std::vector<float> value, const std::vector<float> & weight) {
double square = 0;
for (float x : value) {
square += double(x) * x;
}
const float scale = 1.0f / std::sqrt(float(square / value.size()) + 1e-6f);
for (size_t i = 0; i < value.size(); ++i) {
value[i] *= scale * weight[i];
}
return value;
};
auto mv = [](const std::vector<float> & weight, const std::vector<float> & x) {
std::vector<float> result(weight.size() / x.size(), 0);
for (size_t row = 0; row < result.size(); ++row) {
double sum = 0;
for (size_t col = 0; col < x.size(); ++col) {
sum += double(weight[row * x.size() + col]) * x[col];
}
result[row] = float(sum);
}
return result;
};
const auto embeddings = read("token_embd.weight");
const auto hn = read("dspark.correction_hidden_norm.weight");
const auto en = read("dspark.correction_embed_norm.weight");
const auto gate = read("dspark.correction_gate.weight");
const auto up = read("dspark.correction_up.weight");
const auto down = read("dspark.correction_down.weight");
const auto head = model->get_tensor("output.weight") ? read("output.weight") : embeddings;
const auto a = read("dspark.markov_head_a.weight");
const auto b = read("dspark.markov_head_b.weight");
const size_t rank = a.size() / meta.n_vocab;
std::vector<float> result;
int32_t previous = ids[0];
for (size_t slot = 0; slot < ids.size(); ++slot) {
std::vector<float> h(hidden + slot * meta.n_embd, hidden + (slot + 1) * meta.n_embd);
std::vector<float> e(embeddings.begin() + previous * meta.n_embd,
embeddings.begin() + (previous + 1) * meta.n_embd);
auto input = norm(h, hn);
e = norm(e, en);
input.insert(input.end(), e.begin(), e.end());
auto g = mv(gate, input);
auto u = mv(up, input);
for (size_t i = 0; i < g.size(); ++i) {
g[i] = g[i] / (1 + std::exp(-g[i])) * u[i];
}
auto delta = mv(down, g);
for (size_t i = 0; i < h.size(); ++i) {
h[i] += delta[i];
}
auto logits = mv(head, h);
std::vector<float> latent(a.begin() + ids[slot] * rank, a.begin() + (ids[slot] + 1) * rank);
auto bias = mv(b, latent);
for (size_t i = 0; i < logits.size(); ++i) {
logits[i] += bias[i];
}
logits[meta.mask_token_id] = -1e6f;
previous = int32_t(std::max_element(logits.begin(), logits.end()) - logits.begin());
result.insert(result.end(), logits.begin(), logits.end());
}
return result;
}
static int run_tier1(const std::string & model_path, bool metal = false, std::vector<float> * output = nullptr) {
printf("=== Tier 1: synthetic tap features, wiring/shape/finiteness check ===\n");
llama_model_params mparams = llama_model_default_params();
mparams.n_gpu_layers = metal ? 999 : 0;
ggml_backend_dev_t devices[] = { nullptr, nullptr };
if (metal) {
ggml_backend_reg_t registry = ggml_backend_reg_by_name("MTL");
if (registry && ggml_backend_reg_dev_count(registry) > 0) {
devices[0] = ggml_backend_reg_dev_get(registry, 0);
}
if (!devices[0]) {
fail("Metal device unavailable; refusing CPU fallback");
}
mparams.devices = devices;
}
llama_model_ptr head_source;
if (const char * source_path = std::getenv("DSPARK_TEST_SHARED_HEAD_SOURCE")) {
head_source.reset(llama_model_load_from_file(source_path, mparams));
if (!head_source) {
fail("could not load shared head source");
}
mparams.dspark_head_source = head_source.get();
}
llama_model * model = llama_model_load_from_file(model_path.c_str(), mparams);
if (!model) {
fail("failed to load model: " + model_path);
}
if (head_source) {
if (model->output != head_source->output || model->output->buffer != head_source->output->buffer) {
fail("shared head did not borrow the source tensor and buffer");
}
printf("Shared head identity PASSED: tensor and backend buffer are borrowed\n");
}
dspark_meta meta = read_dspark_meta(model);
printf("n_embd=%lld n_vocab=%lld block_size=%d mask_token_id=%d n_capture=%d\n", (long long) meta.n_embd,
(long long) meta.n_vocab, meta.block_size, meta.mask_token_id, meta.n_capture);
if (meta.block_size <= 0) {
fail("could not read dspark.dspark.block_size from GGUF");
}
const int64_t n_ctx_rows = 6;
const int64_t n_embd_cap = (int64_t) meta.n_capture * meta.n_embd;
const int32_t block_size = meta.block_size;
llama_context * ctx = make_ctx(model, (uint32_t) (n_ctx_rows + block_size));
// deterministic synthetic tap features (fixed seed, no external data needed).
std::mt19937 rng(1234);
std::normal_distribution<float> dist(0.0f, 2.0f);
std::vector<float> ctx_feat((size_t) n_ctx_rows * n_embd_cap);
for (auto & v : ctx_feat) {
v = dist(rng);
}
std::vector<int32_t> ctx_pos(n_ctx_rows);
for (int64_t i = 0; i < n_ctx_rows; ++i) {
ctx_pos[i] = (int32_t) i;
}
std::vector<int32_t> draft_tokens(block_size, meta.mask_token_id);
draft_tokens[0] = (int32_t) std::min<int64_t>(1000, meta.n_vocab - 1);
if (draft_tokens[0] == meta.mask_token_id) {
draft_tokens[0] = 0;
}
std::vector<int32_t> draft_pos(block_size);
for (int32_t i = 0; i < block_size; ++i) {
draft_pos[i] = (int32_t) (n_ctx_rows + i);
}
std::vector<float> logits =
run_forward(ctx, n_ctx_rows, n_embd_cap, ctx_feat, ctx_pos, block_size, draft_tokens, draft_pos, meta.n_vocab);
if (output) {
*output = logits;
}
llama_dspark_meta contract;
if (llama_model_dspark_get_meta(model, &contract) && contract.graph_corrected && std::getenv("DSPARK_TEST_TRACE")) {
for (int32_t slot = 0; slot < block_size; ++slot) {
const auto suffix = "-" + std::to_string(slot);
const auto & before = trace_values.at("dspark_corr_unmasked" + suffix);
const auto & after = trace_values.at("dspark_corr_masked" + suffix);
if (before.size() != size_t(meta.n_vocab) || after.size() != before.size()) {
fail("mask suppression trace shape mismatch");
}
for (size_t token = 0; token < after.size(); ++token) {
const float expected = token == size_t(meta.mask_token_id) ? -1e6f : before[token];
if (after[token] != expected) {
fprintf(stderr, "Mask suppression changed slot=%d token=%zu: before=%g after=%g\n", slot, token,
before[token], after[token]);
llama_free(ctx);
llama_model_free(model);
fail("mask suppression must preserve every non-mask logit exactly");
}
}
}
printf("Mask suppression invariance PASSED: backend=%s slots=%d\n", metal ? "Metal" : "CPU", block_size);
}
if (llama_model_dspark_get_meta(model, &contract) && contract.graph_corrected &&
std::getenv("DSPARK_TEST_ORACLE")) {
auto reference = correction_reference(model, trace_values.at("result_norm").data(), draft_tokens, meta);
float error = 0;
for (size_t i = 0; i < logits.size(); ++i) {
error = std::max(error, std::fabs(logits[i] - reference[i]));
if (std::fabs(logits[i] - reference[i]) > 1e-3f + 1e-2f * std::fabs(reference[i])) {
fprintf(stderr,
"Correction mismatch: backend=%s slot=%zu token=%zu actual=%g reference=%g abs_error=%g\n",
metal ? "Metal" : "CPU", size_t(i / meta.n_vocab), size_t(i % meta.n_vocab), logits[i],
reference[i], std::fabs(logits[i] - reference[i]));
llama_free(ctx);
llama_model_free(model);
fail("hidden correction differs from scalar CPU oracle");
}
}
printf("Independent correction oracle PASSED: max_abs_error=%g\n", error);
}
size_t n_nonfinite = 0;
float min_v = logits[0], max_v = logits[0];
for (float v : logits) {
if (!std::isfinite(v)) {
n_nonfinite++;
}
min_v = std::min(min_v, v);
max_v = std::max(max_v, v);
}
printf("logits: count=%zu min=%g max=%g non_finite=%zu\n", logits.size(), min_v, max_v, n_nonfinite);
for (int32_t p = 0; p < block_size; ++p) {
const float * row = logits.data() + (size_t) p * meta.n_vocab;
int64_t argmax = 0;
for (int64_t v = 1; v < meta.n_vocab; ++v) {
if (row[v] > row[argmax]) {
argmax = v;
}
}
printf(" pos %d: argmax token_id=%lld logit=%g\n", p, (long long) argmax, row[argmax]);
}
llama_free(ctx);
llama_model_free(model);
if (n_nonfinite > 0) {
fail("Tier 1 FAILED: non-finite logits present");
}
if (logits.size() != (size_t) block_size * meta.n_vocab) {
fail("Tier 1 FAILED: unexpected logits size");
}
if (const char * path = std::getenv("DSPARK_TEST_LOGITS_FILE")) {
std::ofstream dump(path, std::ios::binary);
dump.write(reinterpret_cast<const char *>(logits.data()), logits.size() * sizeof(float));
if (!dump) {
fail("could not write synthetic forward logits");
}
}
printf("Tier 1 PASSED: finite, correctly-shaped logits (%d x %lld)\n", block_size, (long long) meta.n_vocab);
return 0;
}
static int run_tier2(const std::string & model_path, const std::string & ref_path) {
printf("=== Tier 2: real drafter weights, deterministic tap features, diff vs DeepSpec reference ===\n");
std::ifstream f(ref_path, std::ios::binary);
if (!f) {
fail("could not open ref file: " + ref_path);
}
int32_t n_ctx_rows_i, n_embd_cap_i, block_size, vocab_size;
f.read((char *) &n_ctx_rows_i, 4);
f.read((char *) &n_embd_cap_i, 4);
f.read((char *) &block_size, 4);
f.read((char *) &vocab_size, 4);
if (!f) {
fail("ref file truncated (header)");
}
const int64_t n_ctx_rows = n_ctx_rows_i;
const int64_t n_embd_cap = n_embd_cap_i;
std::vector<int32_t> ctx_pos(n_ctx_rows);
f.read((char *) ctx_pos.data(), n_ctx_rows * sizeof(int32_t));
std::vector<float> ctx_feat((size_t) n_ctx_rows * n_embd_cap);
f.read((char *) ctx_feat.data(), ctx_feat.size() * sizeof(float));
std::vector<int32_t> draft_tokens(block_size);
f.read((char *) draft_tokens.data(), block_size * sizeof(int32_t));
std::vector<int32_t> draft_pos(block_size);
f.read((char *) draft_pos.data(), block_size * sizeof(int32_t));
std::vector<float> ref_logits((size_t) block_size * vocab_size);
f.read((char *) ref_logits.data(), ref_logits.size() * sizeof(float));
if (!f) {
fail("ref file truncated (payload)");
}
printf("ref: n_ctx_rows=%lld n_embd_cap=%lld block_size=%d vocab_size=%d\n", (long long) n_ctx_rows,
(long long) n_embd_cap, block_size, vocab_size);
llama_model_params mparams = llama_model_default_params();
mparams.n_gpu_layers = 0;
llama_model * model = llama_model_load_from_file(model_path.c_str(), mparams);
if (!model) {
fail("failed to load model: " + model_path);
}
const int64_t n_vocab = n_vocab_from_model(model);
if (n_vocab != vocab_size) {
fail("vocab size mismatch between GGUF (" + std::to_string(n_vocab) + ") and reference (" +
std::to_string(vocab_size) + ")");
}
llama_context * ctx = make_ctx(model, (uint32_t) (n_ctx_rows + block_size));
std::vector<float> logits =
run_forward(ctx, n_ctx_rows, n_embd_cap, ctx_feat, ctx_pos, block_size, draft_tokens, draft_pos, n_vocab);
llama_free(ctx);
llama_model_free(model);
// --- diff ---
double sum_abs_diff = 0.0, max_abs_diff = 0.0;
size_t n_nonfinite = 0;
int32_t argmax_matches = 0;
double top5_overlap_sum = 0.0;
for (int32_t p = 0; p < block_size; ++p) {
const float * a = logits.data() + (size_t) p * n_vocab; // C++
const float * b = ref_logits.data() + (size_t) p * n_vocab; // python
int64_t argmax_a = 0, argmax_b = 0;
for (int64_t v = 1; v < n_vocab; ++v) {
if (a[v] > a[argmax_a]) {
argmax_a = v;
}
if (b[v] > b[argmax_b]) {
argmax_b = v;
}
const double d = std::fabs((double) a[v] - (double) b[v]);
sum_abs_diff += d;
max_abs_diff = std::max(max_abs_diff, d);
if (!std::isfinite(a[v])) {
n_nonfinite++;
}
}
if (argmax_a == argmax_b) {
argmax_matches++;
}
// top-5 overlap (set intersection size / 5)
std::vector<int64_t> idx_a(n_vocab), idx_b(n_vocab);
for (int64_t v = 0; v < n_vocab; ++v) {
idx_a[v] = v;
idx_b[v] = v;
}
std::partial_sort(idx_a.begin(), idx_a.begin() + 5, idx_a.end(),
[&](int64_t x, int64_t y) { return a[x] > a[y]; });
std::partial_sort(idx_b.begin(), idx_b.begin() + 5, idx_b.end(),
[&](int64_t x, int64_t y) { return b[x] > b[y]; });
std::vector<int64_t> top5_a(idx_a.begin(), idx_a.begin() + 5), top5_b(idx_b.begin(), idx_b.begin() + 5);
std::sort(top5_a.begin(), top5_a.end());
std::sort(top5_b.begin(), top5_b.end());
std::vector<int64_t> inter;
std::set_intersection(top5_a.begin(), top5_a.end(), top5_b.begin(), top5_b.end(), std::back_inserter(inter));
top5_overlap_sum += inter.size() / 5.0;
printf(" pos %d: cpp_argmax=%lld py_argmax=%lld cpp_top1_logit=%g py_top1_logit=%g top5_overlap=%d/5\n", p,
(long long) argmax_a, (long long) argmax_b, a[argmax_a], b[argmax_b], (int) inter.size());
}
const double mean_abs_diff = sum_abs_diff / ((double) block_size * n_vocab);
const double argmax_match_rate = (double) argmax_matches / block_size;
const double mean_top5_overlap = top5_overlap_sum / block_size;
printf("\n--- Tier 2 summary ---\n");
printf("mean_abs_diff=%.6g max_abs_diff=%.6g non_finite=%zu\n", mean_abs_diff, max_abs_diff, n_nonfinite);
printf("argmax_match_rate=%.3f (%d/%d) mean_top5_overlap=%.3f\n", argmax_match_rate, argmax_matches, block_size,
mean_top5_overlap);
if (n_nonfinite > 0) {
fail("Tier 2 FAILED: non-finite logits");
}
return 0;
}
static bool report_exact_vector(const char * label,
const std::vector<float> & reference,
const std::vector<float> & actual) {
if (reference.empty() || reference.size() != actual.size()) {
fail(std::string(label) + ": empty or mismatched shape");
}
size_t differences = 0;
float max_error = 0;
for (size_t i = 0; i < reference.size(); ++i) {
if (!std::isfinite(reference[i]) || !std::isfinite(actual[i])) {
fail(std::string(label) + ": nonfinite value");
}
differences += std::memcmp(&reference[i], &actual[i], sizeof(float)) != 0;
max_error = std::max(max_error, std::fabs(reference[i] - actual[i]));
}
printf("A/B %s: count=%zu differing_bits_elements=%zu max_abs_error=%g\n", label, actual.size(), differences,
max_error);
return differences == 0;
}
static int run_snr_ab(const std::string & model_path, bool metal) {
std::vector<float> baseline, candidate;
// Fresh contexts prevent graph reuse from masking the environment switch.
if (setenv("DSPARK_TEST_TRACE", "1", 1) || setenv("LLAMA_DSPARK_SNR_TWO_ROWS", "0", 1)) {
fail("setenv failed");
}
trace_values.clear();
if (run_tier1(model_path, metal, &baseline)) {
fail("baseline forward failed");
}
const auto reference = trace_values;
if (!reference.count("dspark_draft_embd_snr") || reference.count("dspark_log_snr_expanded")) {
fail("baseline must run unfused conditioning");
}
if (setenv("LLAMA_DSPARK_SNR_TWO_ROWS", "1", 1)) {
fail("setenv failed");
}
trace_values.clear();
if (run_tier1(model_path, metal, &candidate)) {
fail("candidate forward failed");
}
if (!trace_values.count("dspark_log_snr_expanded") || !trace_values.count("dspark_draft_embd_snr") ||
!trace_values.count("dspark_log_snr_fc2") ||
trace_values.at("dspark_log_snr_fc2").size() >= reference.at("dspark_log_snr_fc2").size()) {
fail("candidate did not exercise two-row conditioning; no vacuous pass");
}
bool exact = report_exact_vector("conditioned embeddings", reference.at("dspark_draft_embd_snr"),
trace_values.at("dspark_draft_embd_snr"));
exact &= report_exact_vector("expanded conditioning", reference.at("dspark_log_snr_fc2"),
trace_values.at("dspark_log_snr_expanded"));
exact &= report_exact_vector("full logits", baseline, candidate);
if (!exact) {
fail("strict bitwise parity failed; keep candidate disabled");
}
printf("SNR A/B strict parity PASSED (synthetic features; not a serving acceptance benchmark)\n");
return 0;
}
static int run_output_ab(const std::string & model_path, bool metal) {
std::vector<float> baseline, candidate;
if (setenv("DSPARK_TEST_TRACE", "1", 1) || setenv("LLAMA_DSPARK_BALANCED_OUTPUT", "0", 1)) {
fail("setenv failed");
}
trace_values.clear();
if (run_tier1(model_path, metal, &baseline)) {
fail("baseline forward failed");
}
if (trace_values.count("dspark_balanced_output")) {
fail("baseline used balanced output");
}
if (setenv("LLAMA_DSPARK_BALANCED_OUTPUT", "1", 1)) {
fail("setenv failed");
}
trace_values.clear();
if (run_tier1(model_path, metal, &candidate)) {
fail("candidate forward failed");
}
if (!trace_values.count("dspark_balanced_output")) {
fail("candidate branch not exercised");
}
if (!report_exact_vector("balanced output logits", baseline, candidate)) {
fail("output parity failed");
}
printf("Balanced output strict parity PASSED (not a serving benchmark)\n");
return 0;
}
static int run_runtime_ab(const std::string & model_path, const char * backend = nullptr, bool bench = false) {
auto mp = llama_model_default_params();
mp.n_gpu_layers = backend ? 999 : 0;
ggml_backend_dev_t devices[] = { nullptr, nullptr };
if (backend) {
auto registry = ggml_backend_reg_by_name(backend);
if (!registry || ggml_backend_reg_dev_count(registry) == 0) {
fail("requested GPU backend unavailable");
}
devices[0] = ggml_backend_reg_dev_get(registry, 0);
mp.devices = devices;
}
auto * model = llama_model_load_from_file(model_path.c_str(), mp);
if (!model) {
fail("model load failed");
}
const auto meta = read_dspark_meta(model);
llama_dspark_meta contract;
if (!llama_model_dspark_get_meta(model, &contract) || !contract.graph_corrected) {
fail("runtime A/B requires a graph-corrected drafter");
}
const int64_t width = meta.n_capture * meta.n_embd;
int checks = 0;
const char * head_env = std::getenv("DSPARK_TEST_HEAD_TOP_K");
const int bench_top_k = head_env ? std::stoi(head_env) : 4096;
for (int top_k : (bench ? std::vector<int>{ bench_top_k } : std::vector<int>{ 0, 32 })) {
if (top_k > meta.n_vocab) {
continue;
}
setenv("DSPARK_HEAD_TOP_K", std::to_string(top_k).c_str(), 1);
for (int depth : { 1, 2, 4, 8 }) {
if (bench && depth != meta.block_size) {
continue;
}
if (depth > meta.block_size) {
continue;
}
std::vector<float> expected;
for (int variant : { 0, 4, 1, 2, 3 }) {
if (variant == 3 && depth != meta.block_size) {
continue;
}
setenv("LLAMA_DSPARK_REUSE_MASK", variant == 0 || variant == 4 ? "0" : "1", 1);
const bool ids_only = variant == 2 || variant == 4;
auto * ctx = make_ctx(model, 6 + depth);
auto * chain = llama_sampler_chain_init(llama_sampler_chain_default_params());
llama_sampler_chain_add(chain, llama_sampler_init_greedy());
if (ids_only && !llama_set_sampler(ctx, 0, chain)) {
fail("sampler offload failed");
}
common_speculative * spec = nullptr;
if (variant == 3) {
setenv("LLAMA_DSPARK_GREEDY_IDS", "1", 1);
unsetenv("DSPARK_FORWARD_ROWS");
unsetenv("DSPARK_CORRECTION_ROWS");
common_params_speculative params;
params.types = { COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK };
params.draft.ctx_dft = ctx;
params.draft.ctx_tgt = ctx;
params.draft.n_max = depth;
params.draft.n_min = 0;
spec = common_speculative_init(params, 1);
if (!spec) {
fail("common speculative init failed");
}
}
std::vector<double> elapsed;
for (int rep = 0; rep < (bench ? 5 : 3); ++rep) {
llama_memory_clear(llama_get_memory(ctx), true);
std::vector<float> features(6 * width);
for (size_t i = 0; i < features.size(); ++i) {
features[i] = std::sin(float(i));
}
std::vector<int32_t> positions = { 0, 1, 2, 3, 4, 5 };
std::vector<int32_t> tokens(depth, meta.mask_token_id), draft_pos(depth);
tokens[0] = 0;
for (int i = 0; i < depth; ++i) {
draft_pos[i] = 6 + i;
}
std::vector<llama_token> ids;
std::vector<float> logits;
const auto begin = std::chrono::steady_clock::now();
if (spec) {
common_speculative_begin(spec, 0, {});
if (!common_speculative_dspark_stage_ctx_test(spec, 0, features.data(), 6, width,
positions.data())) {
fail("staging failed");
}
auto & dp = common_speculative_get_draft_params(spec, 0);
dp.drafting = true;
dp.n_max = -1;
dp.n_past = 6;
dp.id_last = tokens[0];
dp.result = &ids;
common_speculative_draft(spec);
if (ids.size() != size_t(depth)) {
fail("common draft length mismatch");
}
} else {
logits = run_forward(ctx, 6, width, features, positions, depth, tokens, draft_pos, meta.n_vocab,
ids_only ? &ids : nullptr);
}
if (!ids_only && !spec) {
for (int i = 0; i < depth; ++i) {
const auto row = logits.begin() + i * meta.n_vocab;
ids.push_back(std::max_element(row, row + meta.n_vocab) - row);
}
}
const auto end = std::chrono::steady_clock::now();
if (bench && rep >= 2) {
elapsed.push_back(std::chrono::duration<double, std::milli>(end - begin).count());
}
if (variant == 0) {
expected = logits;
}
if (variant == 1 && !report_exact_vector("mask reuse logits", expected, logits)) {
fail("mask reuse changed logits");
}
if (variant >= 2) {
for (int i = 0; i < depth; ++i) {
const auto first = expected.begin() + i * meta.n_vocab;
const auto best = std::max_element(first, first + meta.n_vocab) - first;
if (ids[i] != best) {
fail("graph greedy ID mismatch");
}
}
}
++checks;
}
if (bench) {
printf("RUNTIME_BENCH backend=%s variant=%d depth=%d top_k=%d ms=%.6f,%.6f,%.6f\n",
backend ? backend : "CPU", variant, depth, top_k, elapsed[0], elapsed[1], elapsed[2]);
fflush(stdout);
}
if (variant == 1 && depth > 1) {
llama_memory_clear(llama_get_memory(ctx), true);
std::vector<float> features(6 * width, 0.0f);
std::vector<int32_t> positions = { 0, 1, 2, 3, 4, 5 };
std::vector<int32_t> tokens(depth, meta.mask_token_id), draft_pos(depth);
tokens[0] = tokens[1] = 0;
for (int i = 0; i < depth; ++i) {
draft_pos[i] = 6 + i;
}
bool rejected = false;
try {
run_forward(ctx, 6, width, features, positions, depth, tokens, draft_pos, meta.n_vocab);
} catch (const std::runtime_error & error) {
rejected = std::string(error.what()).find("anchor followed by MASK") != std::string::npos;
}
if (!rejected) {
fail("mask reuse accepted non-MASK input");
}
}
if (spec) {
common_speculative_free(spec);
}
llama_set_sampler(ctx, 0, nullptr);
llama_sampler_free(chain);
llama_free(ctx);
}
}
}
llama_model_free(model);
printf("Runtime A/B PASSED: %d forwards, mask reuse and greedy IDs (synthetic taps, backend=%s)\n", checks,
backend ? backend : "CPU");
return 0;
}
// Restore caller settings even when a graph validation throws.
struct scoped_test_env {
std::map<std::string, std::pair<bool, std::string>> saved;
scoped_test_env(std::initializer_list<const char *> names) {
for (const char * name : names) {
const char * value = std::getenv(name);
saved[name] = { value != nullptr, value ? value : "" };
}
}
~scoped_test_env() {
for (const auto & entry : saved) {
if (entry.second.first) {
setenv(entry.first.c_str(), entry.second.second.c_str(), 1);
} else {
unsetenv(entry.first.c_str());
}
}
}
};
static int run_prefix_ab(const std::string & model_path, const char * backend = nullptr) {
llama_dspark_ctx staged;
staged.n_ctx_rows = 6;
staged.n_embd_cap = 32;
llm_graph_params previous{};
previous.dspark_ctx = &staged;
previous.dspark_has_context = true;
previous.dspark_ctx_rows = staged.n_ctx_rows;
previous.dspark_ctx_width = staged.n_embd_cap;
auto next = previous;
if (!previous.allow_reuse(next)) {
fail("identical graph params did not reuse");
}
staged.n_ctx_rows = 7;
next.dspark_ctx_rows = staged.n_ctx_rows;
if (previous.allow_reuse(next)) {
fail("mutable context row count reused stale topology");
}
next = previous;
next.dspark_has_context = false;
if (previous.allow_reuse(next)) {
fail("context presence change reused stale topology");
}
next = previous;
next.dspark_ctx_width = 64;
if (previous.allow_reuse(next)) {
fail("context width change reused stale topology");
}
printf("Graph reuse shape contracts PASSED\n");
scoped_test_env environment({ "LLAMA_DSPARK_CORRECTION_PREFIX", "LLAMA_DSPARK_GREEDY_IDS",
"LLAMA_DSPARK_REUSE_MASK", "DSPARK_HEAD_TOP_K", "DSPARK_FORWARD_ROWS",
"DSPARK_CORRECTION_ROWS" });
unsetenv("LLAMA_DSPARK_CORRECTION_PREFIX");
unsetenv("DSPARK_FORWARD_ROWS");
unsetenv("DSPARK_CORRECTION_ROWS");
setenv("LLAMA_DSPARK_GREEDY_IDS", "0", 1);
setenv("LLAMA_DSPARK_REUSE_MASK", "0", 1);
auto mp = llama_model_default_params();
mp.n_gpu_layers = backend ? 999 : 0;
ggml_backend_dev_t devices[] = { nullptr, nullptr };
if (backend) {
auto registry = ggml_backend_reg_by_name(backend);
if (!registry || !ggml_backend_reg_dev_count(registry)) {
fail("requested backend unavailable");
}
devices[0] = ggml_backend_reg_dev_get(registry, 0);
mp.devices = devices;
}
auto * model = llama_model_load_from_file(model_path.c_str(), mp);
if (!model) {
fail("prefix model load failed");
}
{
auto * ctx = make_ctx(model, 16);
const int32_t valid[] = { 0 };
llama_set_capture_layers(ctx, valid, 1);
for (const std::vector<int32_t> & invalid : {
std::vector<int32_t>{ 0, -1 },
std::vector<int32_t>{ 0, (int32_t) model->hparams.n_layer() },
std::vector<int32_t>(LLAMA_MAX_LAYERS + 1, 0)
}) {
bool rejected = false;
try {
llama_set_capture_layers(ctx, invalid.data(), invalid.size());
} catch (const std::invalid_argument &) {
rejected = true;
}
if (!rejected || llama_get_n_capture(ctx) != 1) {
fail("invalid capture list was accepted or changed the previous registration");
}
}
llama_free(ctx);
printf("Capture registration contracts PASSED: 3 invalid lists rejected atomically\n");
}
const auto meta = read_dspark_meta(model);
llama_dspark_meta contract;
if (!llama_model_dspark_get_meta(model, &contract) || !contract.graph_corrected || meta.block_size < 1) {
fail("prefix A/B requires graph-corrected block drafter");
}
const int64_t width = meta.n_capture * meta.n_embd;
std::vector<float> features(6 * width);
for (size_t i = 0; i < features.size(); ++i) {
features[i] = std::sin(float(i));
}
std::vector<int32_t> positions = { 0, 1, 2, 3, 4, 5 };
std::vector<int32_t> tokens(meta.block_size, meta.mask_token_id), draft_pos(meta.block_size);
tokens[0] = 0;
for (int i = 0; i < meta.block_size; ++i) {
draft_pos[i] = 6 + i;
}
const char * top_k_env = std::getenv("DSPARK_TEST_HEAD_TOP_K");
const int restricted = top_k_env ? std::stoi(top_k_env) : std::min<int64_t>(32, meta.n_vocab);
int checks = 0;
for (int top_k : { 0, restricted }) {
if (top_k < 0 || top_k > meta.n_vocab) {
fail("invalid test head top-k");
}
setenv("DSPARK_HEAD_TOP_K", std::to_string(top_k).c_str(), 1);
unsetenv("LLAMA_DSPARK_CORRECTION_PREFIX");
auto * baseline = make_ctx(model, 6 + meta.block_size);
std::vector<float> reference;
for (int rep = 0; rep < 3; ++rep) {
llama_memory_clear(llama_get_memory(baseline), true);
auto values =
run_forward(baseline, 6, width, features, positions, meta.block_size, tokens, draft_pos, meta.n_vocab);
if (rep == 0) {
reference = values;
} else if (!report_exact_vector("full-block replay", reference, values)) {
fail("baseline replay changed logits");
}
}
llama_free(baseline);
for (int keep : { 1, 2, 3, 4, 6, 8 }) {
if (keep > meta.block_size) {
continue;
}
setenv("LLAMA_DSPARK_CORRECTION_PREFIX", std::to_string(keep).c_str(), 1);
const std::vector<float> expected(reference.begin(), reference.begin() + size_t(keep) * meta.n_vocab);
std::vector<llama_token> expected_ids;
for (int i = 0; i < keep; ++i) {
const auto row = expected.begin() + i * meta.n_vocab;
expected_ids.push_back(std::max_element(row, row + meta.n_vocab) - row);
}
for (int variant : { 0, 1, 2 }) {
setenv("LLAMA_DSPARK_GREEDY_IDS", variant == 2 ? "1" : "0", 1);
auto * ctx = make_ctx(model, 6 + meta.block_size);
llama_sampler * chain = nullptr;
common_speculative * spec = nullptr;
if (variant == 1) {
chain = llama_sampler_chain_init(llama_sampler_chain_default_params());
llama_sampler_chain_add(chain, llama_sampler_init_greedy());
if (!llama_set_sampler(ctx, 0, chain)) {
fail("prefix sampler offload failed");
}
}
if (variant == 2) {
common_params_speculative params;
params.types = { COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK };
params.draft.ctx_dft = ctx;
params.draft.ctx_tgt = ctx;
params.draft.n_max = keep;
params.draft.n_min = 0;
spec = common_speculative_init(params, 1);
if (!spec) {
fail("prefix common runner init failed");
}
}
for (int rep = 0; rep < 3; ++rep) {
llama_memory_clear(llama_get_memory(ctx), true);
std::vector<llama_token> ids;
if (spec) {
common_speculative_begin(spec, 0, {});
if (!common_speculative_dspark_stage_ctx_test(spec, 0, features.data(), 6, width,
positions.data())) {
fail("prefix common staging failed");
}
auto & dp = common_speculative_get_draft_params(spec, 0);
dp.drafting = true;
dp.n_max = keep;
dp.n_past = 6;
dp.id_last = tokens[0];
dp.result = &ids;
common_speculative_draft(spec);
} else {
auto values = run_forward(ctx, 6, width, features, positions, meta.block_size, tokens,
draft_pos, meta.n_vocab, variant == 1 ? &ids : nullptr, keep);
if (variant == 0 && !report_exact_vector("full-trunk prefix logits", expected, values)) {
fail("prefix changed retained logits");
}
}
if (variant && ids != expected_ids) {
fail("prefix changed retained IDs or length");
}
++checks;
}
if (spec) {
common_speculative_free(spec);
}
if (chain) {
llama_set_sampler(ctx, 0, nullptr);
llama_sampler_free(chain);
}
llama_free(ctx);
}
printf("PREFIX_AB backend=%s full_block=%d keep=%d top_k=%d exact=1\n", backend ? backend : "CPU",
meta.block_size, keep, top_k);
}
}
// Invalid contracts must fail before exposing a shortened output as a full block.
int negative_checks = 0;
setenv("DSPARK_HEAD_TOP_K", "0", 1);
setenv("LLAMA_DSPARK_GREEDY_IDS", "0", 1);
const auto expect_rejection = [&](const std::string & label, const std::function<void()> & action) {
bool rejected = false;
try {
action();
} catch (const std::runtime_error & error) {
rejected = true;
printf("PREFIX_REJECT case=%s reason=%s\n", label.c_str(), error.what());
}
if (!rejected) {
fail("prefix accepted invalid contract: " + label);
}
++negative_checks;
};
for (const std::string & invalid : { std::string("0"), std::to_string(meta.block_size + 1), std::string("nope") }) {
unsetenv("LLAMA_DSPARK_CORRECTION_PREFIX");
auto * ctx = make_ctx(model, 6 + meta.block_size);
setenv("LLAMA_DSPARK_CORRECTION_PREFIX", invalid.c_str(), 1);
expect_rejection("invalid-" + invalid, [&] {
run_forward(ctx, 6, width, features, positions, meta.block_size, tokens, draft_pos, meta.n_vocab);
});
llama_free(ctx);
}
const int keep = std::min<int>(4, meta.block_size);
const auto make_spec = [&](llama_context * ctx, int n_max, int n_seq) {
common_params_speculative params;
params.types = { COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK };
params.draft.ctx_dft = ctx;
params.draft.ctx_tgt = ctx;
params.draft.n_max = n_max;
params.draft.n_min = 0;
return common_speculative_init(params, n_seq);
};
for (const std::string & mode : { std::string("n-max-mismatch"), std::string("multiple-sequences") }) {
unsetenv("LLAMA_DSPARK_CORRECTION_PREFIX");
auto * ctx = make_ctx(model, 6 + meta.block_size);
setenv("LLAMA_DSPARK_CORRECTION_PREFIX", std::to_string(keep).c_str(), 1);
expect_rejection(mode, [&] {
auto * spec =
make_spec(ctx, mode == "n-max-mismatch" ? keep + 1 : keep, mode == "multiple-sequences" ? 2 : 1);
if (spec) {
common_speculative_free(spec);
}
});
llama_free(ctx);
}
for (const std::string & mode : { std::string("changed-after-init"), std::string("unset-after-init") }) {
unsetenv("LLAMA_DSPARK_CORRECTION_PREFIX");
auto * ctx = make_ctx(model, 6 + meta.block_size);
setenv("LLAMA_DSPARK_CORRECTION_PREFIX", std::to_string(keep).c_str(), 1);
auto * spec = make_spec(ctx, keep, 1);
if (!spec) {
fail("valid prefix initialization failed in negative test");
}
common_speculative_begin(spec, 0, {});
if (!common_speculative_dspark_stage_ctx_test(spec, 0, features.data(), 6, width, positions.data())) {
fail("negative prefix staging failed");
}
std::vector<llama_token> ids;
auto & dp = common_speculative_get_draft_params(spec, 0);
dp.drafting = true;
dp.n_max = keep;
dp.n_past = 6;
dp.id_last = tokens[0];
dp.result = &ids;
if (mode == "changed-after-init") {
setenv("LLAMA_DSPARK_CORRECTION_PREFIX", std::to_string(keep == 1 ? 2 : 1).c_str(), 1);
} else {
unsetenv("LLAMA_DSPARK_CORRECTION_PREFIX");
}
expect_rejection(mode, [&] { common_speculative_draft(spec); });
common_speculative_free(spec);
llama_free(ctx);
}
if (keep < meta.block_size) {
unsetenv("LLAMA_DSPARK_CORRECTION_PREFIX");
auto * ctx = make_ctx(model, 6 + meta.block_size);
setenv("LLAMA_DSPARK_CORRECTION_PREFIX", std::to_string(keep).c_str(), 1);
expect_rejection("full-output-mask-with-prefix", [&] {
run_forward(ctx, 6, width, features, positions, meta.block_size, tokens, draft_pos, meta.n_vocab);
});
llama_free(ctx);
}
printf("Prefix negative contracts PASSED: %d rejected cases\n", negative_checks);
llama_model_free(model);
printf("Prefix A/B PASSED: %d forwards, full-trunk logits and IDs exact\n", checks);
return 0;
}
int main(int argc, char ** argv) {
if (argc < 3) {
fprintf(stderr,
"usage: %s <model.gguf> --tier1 | --tier1-metal | --snr-ab-cpu | --snr-ab-metal | --output-ab-cpu | "
"--output-ab-metal | --runtime-ab-cpu | --runtime-ab-metal | --runtime-ab-cuda | "
"--runtime-bench-metal | --runtime-bench-cuda | --prefix-ab-cpu | --prefix-ab-metal | --prefix-ab-cuda "
"| --tier2 <ref.bin>\n",
argv[0]);
return 1;
}
ggml_backend_load_all();
llama_backend_init();
const std::string model_path = argv[1];
const std::string mode = argv[2];
int rc;
if (mode == "--prefix-ab-cpu" || mode == "--prefix-ab-metal" || mode == "--prefix-ab-cuda") {
rc = run_prefix_ab(model_path,
mode == "--prefix-ab-cpu" ? nullptr : (mode == "--prefix-ab-metal" ? "MTL" : "CUDA"));
} else if (mode == "--runtime-ab-cpu") {
rc = run_runtime_ab(model_path);
} else if (mode == "--runtime-ab-metal" || mode == "--runtime-ab-cuda" || mode == "--runtime-bench-metal" ||
mode == "--runtime-bench-cuda") {
rc = run_runtime_ab(model_path, mode.find("metal") != std::string::npos ? "MTL" : "CUDA",
mode.find("bench") != std::string::npos);
} else if (mode == "--output-ab-cpu" || mode == "--output-ab-metal") {
rc = run_output_ab(model_path, mode == "--output-ab-metal");
} else if (mode == "--snr-ab-cpu" || mode == "--snr-ab-metal") {
rc = run_snr_ab(model_path, mode == "--snr-ab-metal");
} else if (mode == "--tier1") {
rc = run_tier1(model_path);
} else if (mode == "--tier1-metal") {
std::vector<float> cpu, metal;
run_tier1(model_path, false, &cpu);
const auto cpu_trace = trace_values;
trace_values.clear();
rc = run_tier1(model_path, true, &metal);
for (const auto & item : cpu_trace) {
const auto & actual = trace_values.at(item.first);
if (actual.size() != item.second.size()) {
fail("trace shape mismatch");
}
float error = 0;
for (size_t i = 0; i < actual.size(); ++i) {
error = std::max(error, std::fabs(actual[i] - item.second[i]));
}
printf("trace %s count=%zu max_abs_error=%g\n", item.first.c_str(), actual.size(), error);
}
if (cpu.size() != metal.size()) {
fail("CPU/Metal output shapes differ");
}
const size_t slots =
cpu_trace.count("dspark_corr_masked-0") ? cpu.size() / cpu_trace.at("dspark_corr_masked-0").size() : 0;
if (slots) {
const size_t vocab = cpu.size() / slots;
size_t matches = 0;
for (size_t slot = 0; slot < slots; ++slot) {
auto a = cpu.begin() + slot * vocab;
auto b = metal.begin() + slot * vocab;
matches += std::max_element(a, a + vocab) - a == std::max_element(b, b + vocab) - b;
}
printf("CPU/Metal draft argmax matches: %zu/%zu (separate from full-logit tolerance)\n", matches, slots);
}
float max_error = 0.0f;
for (size_t i = 0; i < cpu.size(); ++i) {
const float error = std::fabs(cpu[i] - metal[i]);
max_error = std::max(max_error, error);
if (error > 1e-3f + 1e-2f * std::fabs(cpu[i])) {
fail("CPU/Metal logit tolerance exceeded (atol=1e-3, rtol=1e-2)");
}
}
printf("CPU/Metal full-logit comparison PASSED: count=%zu max_abs_error=%g\n", cpu.size(), max_error);
} else if (mode == "--tier2") {
if (argc < 4) {
fail("--tier2 requires a ref.bin path");
}
rc = run_tier2(model_path, argv[3]);
} else {
fail("unknown mode: " + mode);
rc = 1;
}
llama_backend_free();
return rc;
}