// Compare repeated draft/accept rounds against a supplied reference JSON. // Usage: test-dspark-loop #include "../src/llama-ext.h" #include "common.h" #include "llama.h" #include "speculative.h" #ifdef LLAMA_DSPARK_MARKOV_CUDA # include "dspark-markov.h" # include # include # include #endif #ifdef LLAMA_DSPARK_MARKOV_METAL # include "dspark-markov-metal.h" # include # include # include #endif #include #include #include #include #include #include #include using json = nlohmann::json; static void fail(const std::string & msg) { fprintf(stderr, "FAIL: %s\n", msg.c_str()); exit(1); } // --- synthetic "target tap feature" stand-in --------------------------- // Matches the reference fixture's // hash_u32/synth_feat/synth_bonus_token (same constants, same integer ops -- // see that file's header comment for why this is safe to duplicate rather // than share: it's a closed-form pure function of small integers, not a // stateful RNG stream, so bit-parity across languages just falls out of // using the same uint32 wraparound arithmetic). static uint32_t hash_u32(uint32_t x) { x ^= x >> 16; x *= 0x7feb352du; x ^= x >> 15; x *= 0x846ca68bu; x ^= x >> 16; return x; } static float synth_feat(int64_t pos, int64_t d) { const uint32_t h = hash_u32((uint32_t) (pos * 131071 + d * 97 + 12345)); const int32_t m = (int32_t) (h % 2000u) - 1000; // [-1000, 999] return (float) m / 500.0f; // [-2.0, 1.998] } static int32_t synth_bonus_token(int32_t round_idx, int32_t vocab_size, int32_t mask_token_id) { const uint32_t h = hash_u32((uint32_t) round_idx * 2654435761u + 999983u); int32_t v = (int32_t) (h % (uint32_t) (vocab_size - 1)); if (v == mask_token_id) { v = (v + 1) % vocab_size; } return v; } // mirrored verbatim from dspark_phase2_py_ref.py static const std::vector ACCEPT_SCHEDULE = { 7, 3, 0, 7, 5, 1, 4 }; static const std::vector PROMPT = { 1, 2, 3, 4, 5 }; static std::vector synth_feat_rows(int64_t pos_beg, int64_t n_rows, int64_t n_embd_cap) { std::vector feat((size_t) n_rows * n_embd_cap); for (int64_t i = 0; i < n_rows; ++i) { for (int64_t d = 0; d < n_embd_cap; ++d) { feat[(size_t) i * n_embd_cap + d] = synth_feat(pos_beg + i, d); } } return feat; } #ifdef LLAMA_DSPARK_MARKOV_METAL static int test_metal_markov() { int checks = 0; auto check = [&](bool good, const char * label) { if (!good) { fail(label); } ++checks; }; for (auto shape : { std::pair{ 17, 1 }, { 129, 31 }, { 257, 256 }, { 1025, 33 }, { 8193, 256 } }) { const int v = shape.first, rank = shape.second, mask = v - 1; std::vector a(size_t(v) * rank), b(a.size()), base(size_t(v) * 4); for (size_t i = 0; i < a.size(); ++i) { a[i] = float(int(hash_u32(uint32_t(i)) % 31) - 15) / 16; b[i] = float(int(hash_u32(uint32_t(i) + 987) % 31) - 15) / 16; } for (size_t i = 0; i < base.size(); ++i) { base[i] = float(int(hash_u32(uint32_t(i) + 345) % 31) - 15) / 16; } std::unique_ptr ctx( dspark_markov_metal_init(a.data(), b.data(), v, rank, mask), dspark_markov_metal_free); check(bool(ctx), "Metal context creation"); for (int slots : { 1, 4, 2 }) { std::vector output(slots); check(dspark_markov_metal_resample(ctx.get(), base.data(), 1, slots, output.data()), "Metal resample"); int previous = 1; for (int slot = 0; slot < slots; ++slot) { double best = -std::numeric_limits::infinity(); int winner = -1; for (int token = 0; token < v; ++token) { if (token == mask) { continue; } double score = base[size_t(slot) * v + token]; for (int r = 0; r < rank; ++r) { score += double(a[size_t(previous) * rank + r]) * b[size_t(token) * rank + r]; } if (score > best) { best = score; winner = token; } } check(output[slot] == winner, "Metal vs independent F64 chain oracle"); previous = winner; } } int32_t out[4]; check(!dspark_markov_metal_resample(ctx.get(), base.data(), mask, 1, out), "mask anchor rejected"); check(!dspark_markov_metal_resample(ctx.get(), base.data(), -1, 1, out), "negative anchor rejected"); check(!dspark_markov_metal_resample(ctx.get(), base.data(), v, 1, out), "out-of-range anchor rejected"); check(!dspark_markov_metal_resample(ctx.get(), base.data(), 1, 0, out), "zero slots rejected"); check(!dspark_markov_metal_resample(ctx.get(), base.data(), 1, 1, nullptr), "null output rejected"); const float saved = base[0]; for (float bad : { std::numeric_limits::quiet_NaN(), std::numeric_limits::infinity(), -std::numeric_limits::infinity() }) { base[0] = bad; check(!dspark_markov_metal_resample(ctx.get(), base.data(), 1, 4, out), "nonfinite score rejected"); } base[0] = saved; check(dspark_markov_metal_resample(ctx.get(), base.data(), 1, 4, out), "recovery after rejected round"); } std::vector zeros(129 * 31), base(129 * 4); std::unique_ptr ties( dspark_markov_metal_init(zeros.data(), zeros.data(), 129, 31, 0), dspark_markov_metal_free); check(bool(ties), "tie context"); for (int slot = 0; slot < 4; ++slot) { base[slot * 129] = 1e30f; } int32_t out[4]; check(dspark_markov_metal_resample(ties.get(), base.data(), 1, 4, out), "tie and dominant mask round"); for (int token : out) { check(token == 1, "lowest unmasked ID wins ties"); } for (int slot = 0; slot < 4; ++slot) { base[slot * 129 + 128] = 1; } check(dspark_markov_metal_resample(ties.get(), base.data(), 1, 4, out), "ragged final row round"); for (int token : out) { check(token == 128, "ragged final row can win"); } for (float excluded : { -std::numeric_limits::infinity(), std::numeric_limits::quiet_NaN() }) { base[0] = excluded; check(dspark_markov_metal_resample(ties.get(), base.data(), 1, 4, out), "excluded mask logit ignored"); } check(!dspark_markov_metal_init(zeros.data(), zeros.data(), 129, 31, 129), "invalid mask metadata"); zeros[0] = std::numeric_limits::quiet_NaN(); check(!dspark_markov_metal_init(zeros.data(), zeros.data(), 129, 31, 0), "nonfinite weights rejected"); printf("Metal Markov: %d checks passed\n", checks); return 0; } #endif #ifdef LLAMA_DSPARK_MARKOV_CUDA static int test_cuda_dcut() { int checks = 0; for (int rank : { 1, 31, 33, 256 }) { const int vocab = 67; std::vector a(vocab * rank), b(vocab * rank), base(4 * vocab); for (size_t i = 0; i < a.size(); ++i) { a[i] = ((int) (i % 11) - 5) * 0.03f; b[i] = ((int) (i % 17) - 8) * 0.02f; } for (size_t i = 0; i < base.size(); ++i) { base[i] = std::sin((float) i) * 3.0f; } for (int k = 0; k < 4; ++k) { base[k * vocab + vocab - 1] = 1000.0f; } auto * ctx = dspark_markov_cuda_init(a.data(), b.data(), vocab, rank, vocab - 1); if (!ctx) { fail("CUDA D-cut init"); } for (int n : { 4, 1, 3, 2, 4 }) { for (int mode = 0; mode < 3; ++mode) { float costs[4] = { 1, 1, 1, 1 }; if (mode == 1) { costs[1] = 20; costs[2] = 30; costs[3] = 40; } if (mode == 2) { costs[1] = 1.01f; costs[2] = 1.1f; costs[3] = 1.2f; } int32_t got[4], original[4], kept; float probs[4]; if (!dspark_markov_cuda_resample(ctx, base.data(), 5, n, original) || !dspark_markov_cuda_dcut(ctx, base.data(), 5, n, got, costs, &kept, probs)) { fail("CUDA D-cut run"); } int prev = 5, expected_keep = 1; double survival = 1, expected = 1, best_rate = -1; for (int k = 0; k < n; ++k) { std::vector logits(vocab); for (int v = 0; v < vocab; ++v) { logits[v] = base[k * vocab + v]; for (int r = 0; r < rank; ++r) { logits[v] += (double) a[prev * rank + r] * b[v * rank + r]; } } logits[vocab - 1] = -std::numeric_limits::infinity(); const int winner = std::max_element(logits.begin(), logits.end()) - logits.begin(); double sum = 0; for (double logit : logits) { sum += std::exp(logit - logits[winner]); } const double p = 1 / sum; if (got[k] != winner || got[k] != original[k] || std::abs(probs[k] - p) > 1e-5) { fail("CUDA D-cut F64 oracle"); } survival *= p; expected += survival; const double rate = expected / costs[k]; if (rate > best_rate) { best_rate = rate; expected_keep = k + 1; } prev = winner; ++checks; } if (kept != expected_keep) { fail("CUDA D-cut selector oracle"); } ++checks; } } float invalid[4] = { 1, 2, 3, 0 }; int32_t got[4], kept; if (dspark_markov_cuda_dcut(ctx, base.data(), 0, 4, got, invalid, &kept)) { fail("D-cut accepts zero cost"); } invalid[3] = std::numeric_limits::quiet_NaN(); if (dspark_markov_cuda_dcut(ctx, base.data(), 0, 4, got, invalid, &kept)) { fail("D-cut accepts NaN cost"); } checks += 2; dspark_markov_cuda_free(ctx); } printf("CUDA_DCUT_ORACLE PASS checks=%d\n", checks); return 0; } #endif int main(int argc, char ** argv) { #ifdef LLAMA_DSPARK_MARKOV_CUDA if (argc == 2 && std::string(argv[1]) == "--cuda-dcut-self-test") { return test_cuda_dcut(); } #endif #ifdef LLAMA_DSPARK_MARKOV_METAL if (argc == 2 && std::string(argv[1]) == "--metal-markov-self-test") { return test_metal_markov(); } #endif if (argc < 3) { fprintf(stderr, "usage: %s \n", argv[0]); return 1; } const std::string model_path = argv[1]; const std::string ref_path = argv[2]; std::ifstream f(ref_path); if (!f) { fail("could not open ref file: " + ref_path); } json ref; f >> ref; const int64_t n_embd_cap_ref = ref.at("n_embd_cap").get(); const int32_t block_size_ref = ref.at("block_size").get(); const int32_t vocab_size_ref = ref.at("vocab_size").get(); const int32_t mask_token_id_ref = ref.at("mask_token_id").get(); const int32_t prefill_bonus_ref = ref.at("prefill_bonus").get(); const auto & rounds_ref = ref.at("rounds"); llama_backend_init(); llama_model_params mparams = llama_model_default_params(); mparams.n_gpu_layers = 0; // CPU: deterministic, no Metal precision surprises llama_model * model = llama_model_load_from_file(model_path.c_str(), mparams); if (!model) { fail("failed to load model: " + model_path); } llama_dspark_meta meta; if (!llama_model_dspark_get_meta(model, &meta)) { fail("llama_model_dspark_get_meta failed -- not a dspark model?"); } printf( "meta: n_embd=%lld n_vocab=%lld n_capture=%lld n_embd_cap=%lld block_size=%d mask_token_id=%d " "markov_rank=%lld\n", (long long) meta.n_embd, (long long) meta.n_vocab, (long long) meta.n_capture, (long long) meta.n_embd_cap, meta.block_size, meta.mask_token_id, (long long) meta.markov_rank); if (meta.n_embd_cap != n_embd_cap_ref) { fail("n_embd_cap mismatch vs ref.json"); } if (meta.block_size != block_size_ref) { fail("block_size mismatch vs ref.json"); } if (meta.n_vocab != vocab_size_ref) { fail("vocab_size mismatch vs ref.json"); } if (meta.mask_token_id != mask_token_id_ref) { fail("mask_token_id mismatch vs ref.json"); } const int64_t n_embd_cap = meta.n_embd_cap; const int32_t block_size = meta.block_size; const int32_t vocab_size = (int32_t) meta.n_vocab; const int32_t mask_token_id = meta.mask_token_id; // size the context generously: worst-case round needs ctx_len(<= previous // block_size+1) + block_size tokens in one llama_decode call. const uint32_t n_ctx_max = (uint32_t) (block_size + 1 + block_size) + 8; llama_context_params cparams = llama_context_default_params(); cparams.n_ctx = n_ctx_max; cparams.n_batch = n_ctx_max; cparams.n_ubatch = n_ctx_max; cparams.n_seq_max = 1; cparams.no_perf = true; llama_context * ctx = llama_init_from_model(model, cparams); if (!ctx) { fail("llama_init_from_model failed"); } common_params_speculative sparams; sparams.types = { COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK }; sparams.draft.ctx_dft = ctx; // No real target model exists in this synthetic test (see file header // comment); dspark's process() path (which reads from ctx_tgt) is never // exercised here -- context rows are injected directly via // common_speculative_dspark_stage_ctx_test(). ctx_tgt only needs to be a // valid, non-null context to satisfy the impl's construction-time assert. sparams.draft.ctx_tgt = ctx; sparams.draft.n_max = block_size; sparams.draft.n_min = 0; common_speculative * spec = common_speculative_init(sparams, /* n_seq = */ 1); if (!spec) { fail("common_speculative_init returned null"); } common_speculative_begin(spec, /* seq_id = */ 0, PROMPT); // one-time prefill seeding: the whole prompt's (synthetic) tap features, // positions [0, N). { const int64_t N = (int64_t) PROMPT.size(); std::vector feat = synth_feat_rows(0, N, n_embd_cap); std::vector pos(N); for (int64_t i = 0; i < N; ++i) { pos[i] = (int32_t) i; } if (!common_speculative_dspark_stage_ctx_test(spec, 0, feat.data(), N, n_embd_cap, pos.data())) { fail("common_speculative_dspark_stage_ctx_test (prefill) failed"); } } llama_pos start = (llama_pos) PROMPT.size(); llama_token id_last = (llama_token) synth_bonus_token(-1, vocab_size, mask_token_id); if (id_last != prefill_bonus_ref) { fail("C++/python prefill bonus token disagree -- synth_bonus_token drifted"); } int32_t n_mismatch_rounds = 0; for (size_t r = 0; r < rounds_ref.size(); ++r) { const auto & rr = rounds_ref[r]; const int32_t n_accepted = rr.at("n_accepted").get(); const int64_t ctx_len_ref = rr.at("ctx_len").get(); const int64_t start_ref = rr.at("start").get(); const std::vector sampled_ref = rr.at("sampled").get>(); if ((int32_t) ACCEPT_SCHEDULE[r % ACCEPT_SCHEDULE.size()] != n_accepted) { fail("ACCEPT_SCHEDULE drifted out of sync with ref.json at round " + std::to_string(r)); } if ((int64_t) start != start_ref) { fail("start bookkeeping disagrees with ref.json at round " + std::to_string(r) + " (cpp=" + std::to_string(start) + " py=" + std::to_string(start_ref) + ")"); } common_speculative_draft_params & dp = common_speculative_get_draft_params(spec, 0); dp.drafting = true; dp.n_max = -1; dp.n_past = start; dp.id_last = id_last; dp.prompt = nullptr; // unused by dspark llama_tokens result; dp.result = &result; common_speculative_draft(spec); printf("round %zu: n_past=%d id_last=%d -> result=[", r, start, id_last); for (auto t : result) { printf("%d ", t); } printf("] expected=["); for (auto t : sampled_ref) { printf("%d ", t); } printf("]\n"); if (result.size() != (size_t) block_size) { fail("round " + std::to_string(r) + ": expected block_size=" + std::to_string(block_size) + " drafted tokens, got " + std::to_string(result.size())); } bool round_ok = true; for (int32_t k = 0; k < block_size; ++k) { if (result[k] != sampled_ref[k]) { round_ok = false; } } if (!round_ok) { n_mismatch_rounds++; fprintf(stderr, " MISMATCH at round %zu\n", r); } // stage this round's verify-capture: the target would verify the // anchor + all block_size drafted tokens in one batch (verify_length // = block_size + 1), regardless of how many end up accepted -- accept() // below trims this down to the actually-committed prefix. { std::vector feat = synth_feat_rows(start, block_size + 1, n_embd_cap); std::vector pos(block_size + 1); for (int32_t i = 0; i < block_size + 1; ++i) { pos[i] = start + i; } if (!common_speculative_dspark_stage_ctx_test(spec, 0, feat.data(), block_size + 1, n_embd_cap, pos.data())) { fail("common_speculative_dspark_stage_ctx_test (verify) failed at round " + std::to_string(r)); } } common_speculative_accept(spec, 0, (uint16_t) n_accepted); const int32_t bonus = synth_bonus_token((int32_t) r, vocab_size, mask_token_id); const int32_t bonus_ref = rr.at("bonus").get(); if (bonus != bonus_ref) { fail("C++/python bonus token disagree at round " + std::to_string(r)); } id_last = (llama_token) bonus; start = start + n_accepted + 1; GGML_UNUSED(ctx_len_ref); } common_speculative_print_stats(spec); common_speculative_free(spec); llama_free(ctx); llama_model_free(model); llama_backend_free(); if (n_mismatch_rounds > 0) { fail(std::to_string(n_mismatch_rounds) + "/" + std::to_string(rounds_ref.size()) + " rounds mismatched the Python reference"); } printf( "\nPhase 2 gate PASSED: %zu/%zu rounds token-for-token identical to the DeepSpec reference " "(cache growth/crop, block seeding, RoPE positions, sequential markov resample).\n", rounds_ref.size(), rounds_ref.size()); return 0; }