bonsai.cpp/tests/test-dspark-real-eval.cpp

// Phase 3 (real-target) eval harness for the dspark block-draft loop
// (common_speculative_impl_draft_dspark, common/speculative.cpp).
//
// tests/test-dspark-loop.cpp gates loop MECHANICS (cache crop, RoPE
// positions, sequential markov resample) against a tiny synthetic
// checkpoint and synthetic target-tap features -- it never touches a real
// target model. This program is the first harness that drives the SAME
// public common_speculative_* API against a target and the REAL
// r16b drafter GGUF end to end, to get a first real accept-rate/tau signal.
//
// Why this file exists (there is no pre-existing CLI/server path):
//   - `--spec-type draft-dspark` parses (common/arg.cpp), but neither
//     examples/speculative-simple/speculative-simple.cpp nor
//     tools/server/server-context.cpp ever calls llama_set_capture_layers()
//     on the target context. dspark's process() (see common/speculative.cpp)
//     needs the target's multi-layer tap captured via
//     llama_get_embeddings_capture_ith(), which returns null unless capture
//     is engaged AND logits/output were requested for every row. Phase 2's
//     scope was the loop implementation + its own synthetic-target gate, not
//     this CLI/server wiring -- see DSPARK_FWD_PLAN.md.
//   - the drafter's target_layers array (the layer-id list to pass to
//     llama_set_capture_layers) has no public getter: llama_model's own
//     string-KV cache explicitly SKIPS array-typed GGUF keys (see
//     llama_model_base::load_hparams in src/llama-model.cpp, "if (type ==
//     GGUF_TYPE_ARRAY) continue;"), so llama_model_meta_val_str() can never
//     see it. This harness instead reads the drafter GGUF's own
//     "<arch>.dspark.target_layers" key directly via the low-level gguf.h
//     C API (no core-file changes needed).
//
// The per-round draft/verify/accept loop below mirrors
// examples/speculative-simple/speculative-simple.cpp's target-verify pattern
// (common_sampler_sample_and_accept_n against a greedy/temp=0 target
// sampler) but is NOT a drop-in generalization of that file: dspark differs
// in two structural ways that file doesn't handle --
//   1. it needs common_speculative_process() called explicitly on every
//      target batch (prefill AND each round's verify batch) so the tap
//      capture actually gets consumed -- see the header comment above
//      common_speculative_need_embd_capture() in common/speculative.h.
//   2. it must NEVER llama_decode() the verify batch against ctx_dft (that
//      file does this for ordinary draft-model types); dspark's drafter
//      cache is advanced entirely inside common_speculative_draft() itself
//      via the out-of-band llama_set_dspark_ctx() staging, and re-decoding
//      the verify batch's token ids through the drafter's own embedding
//      table would be meaningless (see src/models/dspark.cpp).
//
// usage: test-dspark-real-eval <target.gguf> <drafter.gguf> [n_predict=96] [n_gpu_layers=999] [dataset.jsonl] [n_prompts=24] [n_max_override=block_size]
//
// n_max_override (1..block_size) caps the per-round draft length dp.n_max.
// The dspark impl itself always produces a full block_size block (its
// markov-resample loop is fixed-length and its drafter cache is cropped back
// to `start` inside draft() regardless of acceptance), but the generic
// common_speculative_draft() dispatcher truncates *dp.result to dp.n_max
// after the impl returns -- so capping here shrinks only the VERIFY batch,
// leaving the drafter's per-round cost unchanged (the honest semantics for a
// block-diffusion drafter).
//
// when dataset.jsonl is given, prompts are loaded DeepSpec-style (one
// {"turns": [...]} object per line, single turn only) and run through the
// target's own tokenizer.chat_template (via common_chat_templates_apply)
// instead of this file's built-in plain-text-continuation PROMPTS below --
// for a run directly comparable to the DeepSpec/Python-side numbers, which
// apply the same chat template.

#include "../src/llama-ext.h"
#include "chat.h"
#include "common.h"
#include "gguf.h"
#include "llama.h"
#include "nlohmann/json.hpp"
#include "sampling.h"
#include "speculative.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <string>
#include <vector>

[[noreturn]] static void fail(const std::string & msg) {
    fprintf(stderr, "FAIL: %s\n", msg.c_str());
    exit(1);
}

struct prompt_spec {
    std::string                  category;
    std::string                  text;
    std::vector<common_chat_msg> messages = {};
};

// JSONL prompt format:
// one JSON object per line, {"turns": ["<single user turn>"]}. Loading real
// DeepSpec prompts (instead of this file's own hand-written set below) and
// running them through the target's own chat template is what makes a run
// directly comparable to the DeepSpec/Python-side accept-rate numbers,
// rather than this harness's own plain-text-continuation methodology.
static std::vector<prompt_spec> load_deepspec_jsonl(const std::string & path, int limit) {
    std::ifstream f(path);
    if (!f) {
        fail("failed to open dataset file: " + path);
    }

    std::vector<prompt_spec> out;
    std::string              line;
    while (std::getline(f, line) && (limit <= 0 || (int) out.size() < limit)) {
        if (line.empty()) {
            continue;
        }
        nlohmann::json j = nlohmann::json::parse(line);
        if (j.contains("messages")) {
            prompt_spec prompt{ "clean-heldout", "", {} };
            for (const auto & message : j.at("messages")) {
                common_chat_msg msg;
                msg.role    = message.at("role").get<std::string>();
                msg.content = message.at("content").get<std::string>();
                prompt.messages.push_back(std::move(msg));
            }
            // Exclude the reference answer, but preserve preceding conversation turns.
            if (!prompt.messages.empty() && prompt.messages.back().role == "assistant") {
                prompt.messages.pop_back();
            }
            if (prompt.messages.empty() || prompt.messages.back().role != "user") {
                fail("held-out prompt must end in a user turn after removing its reference answer");
            }
            out.push_back(std::move(prompt));
        } else {
            out.push_back({ "deepspec", j.at("turns").at(0).get<std::string>() });
        }
    }

    return out;
}

// 24 held-out prompts across 3 categories. Plain-text continuations (no
// chat template applied) so this measures raw free-running rollout accept
// behavior, same spirit as the DeepSpec free-running alpaca/arena-hard eval
// this is being compared against.
static const std::vector<prompt_spec> PROMPTS = {
    { "code",      "def is_prime(n):\n    \"\"\"Return True if n is a prime number, else False.\"\"\"\n"               },
    { "code",
     "import heapq\n\ndef k_smallest(nums, k):\n    \"\"\"Return the k smallest elements of nums, sorted.\"\"\"\n"     },
    { "code",
     "class LRUCache:\n    \"\"\"A least-recently-used cache with fixed capacity.\"\"\"\n    def __init__(self, "
     "capacity: int):\n"                                                                                               },
    { "code",
     "// Reverse a singly linked list in place.\nstruct Node { int val; Node* next; };\nNode* reverse(Node* head) "
     "{\n"                                                                                                             },
    { "code",
     "def merge_intervals(intervals):\n    \"\"\"Given a list of [start, end] intervals, merge all overlapping "
     "ones.\"\"\"\n"                                                                                                   },
    { "code",
     "#include <vector>\n#include <algorithm>\n// Binary search for the first index where arr[i] >= target.\nint "
     "lower_bound_idx(std::vector<int>& arr, int target) {\n"                                                          },
    { "code",
     "def quicksort(arr):\n    \"\"\"Sort arr in place using the quicksort algorithm.\"\"\"\n    if len(arr) <= 1:\n  "
     "      return arr\n"                                                                                              },
    { "code",      "-- SQL: return the top 3 highest-paid employees per department.\nSELECT\n"                         },
    { "chat",      "Q: What's the difference between a list and a tuple in Python?\nA:"                                },
    { "chat",      "Q: Can you explain how photosynthesis works, in simple terms?\nA:"                                 },
    { "chat",      "Q: I have $500 and want to invest it for 5 years. What are some options to consider?\nA:"          },
    { "chat",      "Q: Write a short, friendly email declining a meeting invite because of a scheduling conflict.\nA:" },
    { "chat",      "Q: What are three practical tips for improving sleep quality?\nA:"                                 },
    { "chat",      "Q: Explain the plot of Romeo and Juliet in two sentences.\nA:"                                     },
    { "chat",      "Q: My laptop fan is very loud under light load. What should I check first?\nA:"                    },
    { "chat",      "Q: Summarize the main causes of the French Revolution in a short paragraph.\nA:"                   },
    { "reasoning",
     "Q: A train leaves city A at 60 mph and another leaves city B (300 miles away) at 90 mph, heading toward each "
     "other. How long until they meet?\nA: Let's think step by step."                                                  },
    { "reasoning",
     "Q: If all bloops are razzies and all razzies are lazzies, are all bloops definitely lazzies? Explain.\nA:"       },
    { "reasoning",
     "Q: A store marks up an item by 40% then offers a 25% discount off the marked-up price. Is the final price "
     "higher or lower than the original? By how much?\nA: Let's think step by step."                                   },
    { "reasoning", "Q: Why is the sky blue during the day but red/orange at sunset?\nA:"                               },
    { "reasoning",
     "Q: You have 8 identical-looking balls, one of which is heavier. Using a balance scale only twice, how do you "
     "find the heavier ball?\nA: Let's think step by step."                                                            },
    { "reasoning",
     "Q: Which is a better estimate of the number of piano tuners in a large city: 50, 500, or 5000? Explain your "
     "reasoning.\nA:"                                                                                                  },
    { "reasoning",
     "Q: Two coworkers, Alice and Bob, always tell the truth or always lie. Alice says \"Bob always lies.\" What can "
     "you conclude?\nA: Let's think step by step."                                                                     },
    { "reasoning",
     "Q: A recipe calls for 3/4 cup of sugar for 12 cookies. How much sugar is needed for 30 cookies?\nA: Let's think "
     "step by step."                                                                                                   },
};

// dspark ships no tokenizer / vocab of its own (see conversion/dspark.py and
// src/models/dspark.cpp): the drafter GGUF's target_layers array is only
// discoverable by reading the file's own metadata directly, not via any
// llama_model_* accessor -- see the file header comment.
static std::vector<int32_t> read_dspark_target_layers(const std::string & drafter_path) {
    struct gguf_init_params gp   = { /* .no_alloc = */ true, /* .ctx = */ nullptr };
    gguf_context *          gctx = gguf_init_from_file(drafter_path.c_str(), gp);
    if (gctx == nullptr) {
        fail("gguf_init_from_file failed for " + drafter_path);
    }

    const int64_t arch_kid = gguf_find_key(gctx, "general.architecture");
    if (arch_kid < 0) {
        fail(drafter_path + ": missing general.architecture key");
    }
    const std::string arch = gguf_get_val_str(gctx, arch_kid);

    const std::string key = arch + ".dspark.target_layers";
    const int64_t     kid = gguf_find_key(gctx, key.c_str());
    if (kid < 0) {
        fail(drafter_path + ": missing GGUF key " + key);
    }
    if (gguf_get_kv_type(gctx, kid) != GGUF_TYPE_ARRAY) {
        fail(key + " is not an array-typed KV");
    }

    const enum gguf_type arr_type = gguf_get_arr_type(gctx, kid);
    const size_t         n        = gguf_get_arr_n(gctx, kid);
    const void *         data     = gguf_get_arr_data(gctx, kid);

    std::vector<int32_t> out(n);
    for (size_t i = 0; i < n; i++) {
        switch (arr_type) {
            case GGUF_TYPE_INT32:
                out[i] = ((const int32_t *) data)[i];
                break;
            case GGUF_TYPE_UINT32:
                out[i] = (int32_t) ((const uint32_t *) data)[i];
                break;
            case GGUF_TYPE_INT64:
                out[i] = (int32_t) ((const int64_t *) data)[i];
                break;
            case GGUF_TYPE_UINT64:
                out[i] = (int32_t) ((const uint64_t *) data)[i];
                break;
            default:
                fail(key + ": unexpected array element gguf_type " + std::to_string((int) arr_type));
        }
    }

    gguf_free(gctx);
    return out;
}

int main(int argc, char ** argv) {
    if (argc < 3) {
        fprintf(stderr,
                "usage: %s <target.gguf> <drafter.gguf> [n_predict=96] [n_gpu_layers=999] [dataset.jsonl] "
                "[n_prompts=24] [n_max_override=block_size]\n",
                argv[0]);
        return 1;
    }
    const std::string target_path   = argv[1];
    const std::string drafter_path  = argv[2];
    const int         n_predict_max = argc > 3 ? std::atoi(argv[3]) : 96;
    const int         n_gpu_layers  = argc > 4 ? std::atoi(argv[4]) : 999;
    const std::string dataset_path  = argc > 5 ? argv[5] : "";
    const int         n_prompts     = argc > 6 ? std::atoi(argv[6]) : 24;
    const int         n_max_arg     = argc > 7 ? std::atoi(argv[7]) : 0;  // 0 == use the drafter's baked block_size
    const char *      oracle_env    = std::getenv("DSPARK_PARITY_ORACLE_ROWS");
    char *            oracle_end    = nullptr;
    const long        oracle_value  = oracle_env ? std::strtol(oracle_env, &oracle_end, 10) : 0;
    if (oracle_env && (oracle_end == oracle_env || *oracle_end || oracle_value < 1 || oracle_value > 9)) {
        fail("DSPARK_PARITY_ORACLE_ROWS must be 1..9");
    }
    const int  oracle_rows       = (int) oracle_value;
    size_t     oracle_mismatches = 0;
    const auto env_size          = [](const char * name, long fallback, long minimum, long maximum) {
        const char * value = std::getenv(name);
        if (!value) {
            return fallback;
        }
        char *     end    = nullptr;
        const long parsed = std::strtol(value, &end, 10);
        if (end == value || *end || parsed < minimum || parsed > maximum) {
            fail(std::string("invalid ") + name);
        }
        return parsed;
    };
    const long eval_context = env_size("DSPARK_EVAL_CONTEXT", 4096, 128, 131072);
    const long eval_batch   = env_size("DSPARK_EVAL_BATCH", 2048, 16, 2048);

    std::vector<int32_t> target_layers = read_dspark_target_layers(drafter_path);
    printf("dspark target_layers (%zu):", target_layers.size());
    for (auto l : target_layers) {
        printf(" %d", l);
    }
    printf("\n");

    llama_backend_init();

    // --- target ---
    llama_model_params mparams_tgt = llama_model_default_params();
    mparams_tgt.n_gpu_layers       = n_gpu_layers;
    llama_model * model_tgt        = llama_model_load_from_file(target_path.c_str(), mparams_tgt);
    if (!model_tgt) {
        fail("failed to load target model: " + target_path);
    }
    const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt);

    // Both models are loaded before allocating contexts.
    llama_model_params mparams_dft = llama_model_default_params();
    mparams_dft.n_gpu_layers       = n_gpu_layers;
    const char * shared_head       = std::getenv("LLAMA_DSPARK_SHARED_HEAD");
    if (shared_head && std::strcmp(shared_head, "1") == 0) {
        mparams_dft.dspark_head_source = model_tgt;
    }
    llama_model * model_dft = llama_model_load_from_file(drafter_path.c_str(), mparams_dft);
    if (!model_dft) {
        fail("failed to load drafter model: " + drafter_path);
    }

    llama_dspark_meta meta;
    if (!llama_model_dspark_get_meta(model_dft, &meta)) {
        fail("drafter GGUF does not look like a dspark model (missing dspark.*.block_size KV)");
    }
    if ((int64_t) target_layers.size() != meta.n_capture) {
        fail("target_layers count (" + std::to_string(target_layers.size()) + ") != meta.n_capture (" +
             std::to_string(meta.n_capture) + ")");
    }

    printf(
        "dspark 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);

    llama_context_params cparams_tgt = llama_context_default_params();
    cparams_tgt.n_ctx                = eval_context;
    cparams_tgt.n_batch              = eval_batch;
    cparams_tgt.n_ubatch             = eval_batch;
    cparams_tgt.n_seq_max            = 1;
    cparams_tgt.no_perf              = true;
    // dspark verifies a whole draft block against the target in one
    // llama_decode(), then crops the target's cache back to the accepted
    // length with a PARTIAL llama_memory_seq_rm(). On a hybrid GDN/attention
    // target that partial removal only succeeds if the recurrent-state
    // rollback ring (n_rs_seq) was sized up front -- see
    // common_params_speculative::need_n_rs_seq() and
    // llama_memory_recurrent::seq_rm()'s per-token snapshot index path.
    cparams_tgt.n_rs_seq             = (uint32_t) meta.block_size;

    llama_context * ctx_tgt = llama_init_from_model(model_tgt, cparams_tgt);
    if (!ctx_tgt) {
        fail("failed to create target context");
    }

    if (llama_model_is_hybrid(model_tgt) && llama_n_rs_seq(ctx_tgt) == 0) {
        fail(
            "target is a hybrid GDN/attention model but its context has "
            "n_rs_seq=0 -- the post-verify partial crop would silently "
            "no-op instead of rolling back the recurrent state (see "
            "llama_memory_hybrid::seq_rm)");
    }

    auto ar_params         = cparams_tgt;
    ar_params.n_rs_seq     = 0;
    llama_context * ctx_ar = llama_init_from_model(model_tgt, ar_params);
    if (!ctx_ar) {
        fail("failed to create plain AR context");
    }
    printf("BASELINE: separate AR context n_rs_seq=%u, capture off; shared target weights\n", llama_n_rs_seq(ctx_ar));

    const bool                use_chat_template = !dataset_path.empty();
    common_chat_templates_ptr tmpls;
    std::vector<prompt_spec>  prompts;
    if (use_chat_template) {
        tmpls   = common_chat_templates_init(model_tgt, "");
        prompts = load_deepspec_jsonl(dataset_path, n_prompts);
        printf("loaded %zu prompts from %s, chat-templated\n", prompts.size(), dataset_path.c_str());
    } else {
        prompts = PROMPTS;
    }

    llama_context_params cparams_dft  = llama_context_default_params();
    cparams_dft.n_ctx                 = eval_context;
    cparams_dft.n_batch               = eval_batch;
    cparams_dft.n_ubatch              = eval_batch;
    cparams_dft.n_seq_max             = 1;
    cparams_dft.no_perf               = true;
    cparams_dft.n_outputs_max         = meta.block_size;
    cparams_dft.n_outputs_max_per_seq = meta.block_size;

    llama_context * ctx_dft = llama_init_from_model(model_dft, cparams_dft);
    if (!ctx_dft) {
        fail("failed to create drafter context");
    }

    // engage the target-layer tap capture ONCE, permanently, on the target
    // context (see file header comment -- this is the missing piece no
    // existing CLI/server path wires up for dspark).
    llama_set_capture_layers(ctx_tgt, target_layers.data(), target_layers.size());

    common_params_speculative sparams;
    sparams.types         = { COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK };
    sparams.draft.ctx_dft = ctx_dft;
    sparams.draft.ctx_tgt = ctx_tgt;
    sparams.draft.n_max   = n_max_arg > 0 ? n_max_arg : meta.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");
    }

    if (!common_speculative_need_embd_capture(spec)) {
        fail("expected the dspark impl to report need_embd_capture() == true");
    }

    common_params_sampling sparams_smpl;
    sparams_smpl.temp = 0.0f;  // greedy / deterministic target verification
    sparams_smpl.seed = 42;

    const llama_seq_id seq_id     = 0;
    const int32_t      block_size = meta.block_size;

    // runtime draft-length cap (see file header). the drafter still produces
    // (and pays for) a full block_size block every round; only the first
    // n_draft tokens are kept and verified.
    int32_t n_draft = n_max_arg > 0 ? n_max_arg : block_size;
    if (n_draft < 1 || n_draft > block_size) {
        fail("n_max_override must be in [1, block_size=" + std::to_string(block_size) + "], got " +
             std::to_string(n_draft));
    }
    printf("draft length per round: n_max=%d (block_size=%d)\n", n_draft, block_size);

    llama_batch batch_tgt = llama_batch_init((int32_t) llama_n_batch(ctx_tgt), 0, 1);

    int64_t total_drafted = 0, total_accepted = 0, total_rounds = 0, total_predicted = 0;
    int64_t total_ar_predicted = 0;
    double  total_ar_seconds = 0.0, total_sp_seconds = 0.0;

    // accept-by-depth: for 1-based draft position i, depth_reached[i] counts
    // rounds where position i was reached (i.e. positions 1..i-1 were all
    // accepted and the draft was at least i long), depth_accepted[i] counts
    // rounds where it was also accepted. depth_accepted[i]/depth_reached[i]
    // is the conditional per-depth accept rate d(i).
    std::vector<int64_t> depth_reached(block_size + 1, 0);
    std::vector<int64_t> depth_accepted(block_size + 1, 0);

    struct cat_stats {
        int64_t drafted = 0, accepted = 0, rounds = 0;
    };

    std::vector<std::string> cat_names;
    std::vector<cat_stats>   cat_stats_v;

    // DSPARK_PARITY_CONTINUE=1: record a parity failure, drop that prompt from every total, keep going.
    // Default (unset) keeps the hard refusal.
    const bool parity_continue =
        std::getenv("DSPARK_PARITY_CONTINUE") && std::strcmp(std::getenv("DSPARK_PARITY_CONTINUE"), "1") == 0;
    std::vector<size_t> parity_failed_prompts;
    for (size_t pi = 0; pi < prompts.size(); ++pi) {
        const auto & ps = prompts[pi];

        llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, 0, -1);

        std::string text = ps.text;
        if (use_chat_template) {
            common_chat_templates_inputs cinputs;
            if (ps.messages.empty()) {
                cinputs.messages.push_back({ "user", ps.text, {}, {}, "", "", "" });
            } else {
                cinputs.messages = ps.messages;
            }
            cinputs.enable_thinking       = false;
            cinputs.add_generation_prompt = true;
            text                          = common_chat_templates_apply(tmpls.get(), cinputs).prompt;
        }

        std::vector<llama_token> inp =
            common_tokenize(ctx_tgt, text, /* add_special = */ true, /* parse_special = */ true);
        if (inp.size() < 2) {
            fprintf(stderr, "skipping prompt %zu: too short after tokenization\n", pi);
            continue;
        }
        if (inp.size() + (size_t) n_predict_max + (size_t) block_size + 8 > llama_n_ctx(ctx_tgt)) {
            fail("prompt plus output exceeds target context capacity");
        }
        printf("PROMPT: index=%zu input_tokens=%zu context_capacity=%u batch=%ld\n", pi, inp.size(),
               llama_n_ctx(ctx_tgt), eval_batch);

        llama_token              id_last = inp.back();
        std::vector<llama_token> prompt_tgt(inp.begin(), inp.end() - 1);

        std::swap(ctx_tgt, ctx_ar);

        // === vanilla AR baseline (measured first, same prompt, same target
        // context) -- capture is disabled so this pays no dspark tap-capture
        // overhead, i.e. it is genuinely comparable to plain decoding, not
        // "dspark plumbing with drafting turned off". ===
        llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, 0, -1);
        llama_set_capture_layers(ctx_tgt, nullptr, 0);

        const auto t_ar0 = std::chrono::steady_clock::now();

        const auto prefill = [&](const llama_tokens & tokens, bool capture) {
            for (size_t begin = 0; begin < tokens.size();) {
                const size_t end = std::min(tokens.size(), begin + (size_t) eval_batch);
                common_batch_clear(batch_tgt);
                for (size_t i = begin; i < end; ++i) {
                    common_batch_add(batch_tgt, tokens[i], (llama_pos) i, { seq_id },
                                     capture || i + 1 == tokens.size());
                }
                if (llama_decode(ctx_tgt, batch_tgt) != 0) {
                    fail("prefill chunk failed");
                }
                if (capture && !common_speculative_process(spec, batch_tgt)) {
                    fail("prefill capture failed");
                }
                begin = end;
            }
        };
        prefill(inp, false);

        common_sampler_ptr              smpl_ar(common_sampler_init(model_tgt, sparams_smpl));
        const int                       n_vocab = llama_vocab_n_tokens(vocab_tgt);
        std::vector<std::vector<float>> ar_logits;
        const auto                      record_ar_logits = [&]() {
            if (oracle_rows) {
                const float * logits = llama_get_logits_ith(ctx_tgt, -1);
                ar_logits.emplace_back(logits, logits + n_vocab);
            }
        };
        record_ar_logits();
        llama_token  ar_cur    = common_sampler_sample(smpl_ar.get(), ctx_tgt, -1);
        llama_tokens ar_tokens = { ar_cur };
        common_sampler_accept(smpl_ar.get(), ar_cur, /* accept_grammar = */ true);

        int        ar_n_past      = (int) prompt_tgt.size() + 1;  // position of ar_cur
        int        ar_n_predicted = 1;
        bool       ar_has_eos     = llama_vocab_is_eog(vocab_tgt, ar_cur);
        const auto t_ar_decode    = std::chrono::steady_clock::now();

        while (ar_n_predicted < n_predict_max && !ar_has_eos) {
            common_batch_clear(batch_tgt);
            common_batch_add(batch_tgt, ar_cur, (llama_pos) ar_n_past, { seq_id }, /* logits = */ true);
            if (llama_decode(ctx_tgt, batch_tgt) != 0) {
                fail("AR decode failed at prompt " + std::to_string(pi));
            }

            record_ar_logits();
            ar_cur = common_sampler_sample(smpl_ar.get(), ctx_tgt, -1);
            ar_tokens.push_back(ar_cur);
            common_sampler_accept(smpl_ar.get(), ar_cur, /* accept_grammar = */ true);
            ar_n_past++;
            ar_n_predicted++;
            if (llama_vocab_is_eog(vocab_tgt, ar_cur)) {
                ar_has_eos = true;
            }
        }

        const double ar_seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - t_ar0).count();
        const double ar_decode_seconds =
            std::chrono::duration<double>(std::chrono::steady_clock::now() - t_ar_decode).count();
        const double ar_tok_per_sec = ar_n_predicted / ar_seconds;

        if (oracle_rows) {
            const char * same_path_env = std::getenv("DSPARK_ORACLE_SAME_PATH");
            const char * kld_env       = std::getenv("DSPARK_ORACLE_KLD");
            for (const char * value : { same_path_env, kld_env }) {
                if (value && std::strcmp(value, "0") && std::strcmp(value, "1")) {
                    fail("oracle diagnostic flags must be 0 or 1");
                }
            }
            const bool same_path   = same_path_env && std::strcmp(same_path_env, "1") == 0;
            const bool measure_kld = kld_env && std::strcmp(kld_env, "1") == 0;
            if (!same_path) {
                std::swap(ctx_tgt, ctx_ar);
            }
            llama_set_capture_layers(ctx_tgt, nullptr, 0);
            // Teacher force the AR stream. No drafter, rejection, or partial rollback runs here.
            if (!llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, 0, -1)) {
                fail("oracle reset failed");
            }
            const char * split_env = std::getenv("DSPARK_PARITY_ORACLE_SPLIT_PREFILL");
            if (split_env && std::string(split_env) != "0" && std::string(split_env) != "1") {
                fail("split-prefill oracle flag must be 0 or 1");
            }
            const bool split_prefill = split_env && std::string(split_env) == "1";
            fprintf(stderr, "ORACLE split_prefill=%d\n", (int) split_prefill);
            prefill(split_prefill ? prompt_tgt : inp, false);
            size_t     mismatches = 0;
            float      max_abs    = 0.0f;
            double     kld_sum = 0.0, kld_max = 0.0;
            size_t     kld_count  = 0;
            const auto compare    = [&](size_t output_index, int row) {
                const float * logits    = llama_get_logits_ith(ctx_tgt, row);
                const auto &  reference = ar_logits[output_index];
                const int     top       = (int) (std::max_element(logits, logits + n_vocab) - logits);
                for (int v = 0; v < n_vocab; ++v) {
                    if (!std::isfinite(logits[v]) || !std::isfinite(reference[v])) {
                        fail("oracle non-finite logit");
                    }
                    max_abs = std::max(max_abs, std::abs(logits[v] - reference[v]));
                }
                if (measure_kld) {
                    const double ref_max  = *std::max_element(reference.begin(), reference.end());
                    const double test_max = logits[top];
                    double       ref_sum = 0.0, test_sum = 0.0, weighted_delta = 0.0;
                    for (int v = 0; v < n_vocab; ++v) {
                        const double weight = std::exp((double) reference[v] - ref_max);
                        ref_sum += weight;
                        test_sum += std::exp((double) logits[v] - test_max);
                        weighted_delta += weight * ((double) reference[v] - logits[v]);
                    }
                    double kld = weighted_delta / ref_sum + std::log(test_sum / ref_sum) + test_max - ref_max;
                    if (!std::isfinite(kld) || kld < -1e-10) {
                        fail("invalid oracle KLD");
                    }
                    kld = std::max(0.0, kld);
                    kld_sum += kld;
                    kld_max = std::max(kld_max, kld);
                    ++kld_count;
                    if (top != ar_tokens[output_index]) {
                        fprintf(stderr,
                                "ORACLE probabilities: prompt=%zu token=%zu kld=%.12g p_ref_ar=%.9g p_ref_other=%.9g "
                                "p_test_ar=%.9g p_test_other=%.9g\n",
                                pi, output_index, kld, std::exp(reference[ar_tokens[output_index]] - ref_max) / ref_sum,
                                std::exp(reference[top] - ref_max) / ref_sum,
                                std::exp(logits[ar_tokens[output_index]] - test_max) / test_sum, 1.0 / test_sum);
                    }
                }
                if (top != ar_tokens[output_index]) {
                    ++mismatches;
                    fprintf(
                        stderr,
                        "ORACLE mismatch: prompt=%zu token=%zu rows=%d AR=%d block=%d ar_margin=%g block_margin=%g\n",
                        pi, output_index, oracle_rows, ar_tokens[output_index], top,
                        reference[ar_tokens[output_index]] - reference[top],
                        logits[top] - logits[ar_tokens[output_index]]);
                }
            };
            if (split_prefill) {
                for (size_t start = 0; start < ar_tokens.size();) {
                    const size_t count = std::min((size_t) oracle_rows, ar_tokens.size() - start);
                    common_batch_clear(batch_tgt);
                    for (size_t j = 0; j < count; ++j) {
                        const size_t      output_index = start + j;
                        const llama_token token        = output_index == 0 ? inp.back() : ar_tokens[output_index - 1];
                        common_batch_add(batch_tgt, token, (llama_pos) (inp.size() - 1 + output_index), { seq_id },
                                         true);
                    }
                    if (llama_decode(ctx_tgt, batch_tgt) != 0) {
                        fail("split-prefill oracle decode failed");
                    }
                    for (size_t j = 0; j < count; ++j) {
                        compare(start + j, (int) j);
                    }
                    start += count;
                }
            } else {
                compare(0, -1);
                for (size_t start = 0; start + 1 < ar_tokens.size();) {
                    const size_t count = std::min((size_t) oracle_rows, ar_tokens.size() - 1 - start);
                    common_batch_clear(batch_tgt);
                    for (size_t j = 0; j < count; ++j) {
                        common_batch_add(batch_tgt, ar_tokens[start + j], (llama_pos) (inp.size() + start + j),
                                         { seq_id }, true);
                    }
                    if (llama_decode(ctx_tgt, batch_tgt) != 0) {
                        fail("oracle decode failed");
                    }
                    for (size_t j = 0; j < count; ++j) {
                        compare(start + j + 1, (int) j);
                    }
                    start += count;
                }
            }
            fprintf(stderr, "ORACLE summary: prompt=%zu rows=%d tokens=%zu mismatches=%zu max_abs_logit=%g\n", pi,
                    oracle_rows, ar_tokens.size(), mismatches, max_abs);
            if (measure_kld) {
                fprintf(stderr,
                        "ORACLE KLD: prompt=%zu rows=%d split_prefill=%d same_path=%d count=%zu mean=%.12g max=%.12g\n",
                        pi, oracle_rows, (int) split_prefill, (int) same_path, kld_count,
                        kld_count ? kld_sum / kld_count : 0.0, kld_max);
            }
            if (same_path) {
                std::swap(ctx_tgt, ctx_ar);
            }
            oracle_mismatches += mismatches;
            continue;  // Diagnostic timings include logit copies and are not benchmarks.
        }

        std::swap(ctx_tgt, ctx_ar);

        // === DSpark pass (existing logic below, now timed) ===
        llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, 0, -1);
        llama_set_capture_layers(ctx_tgt, target_layers.data(), target_layers.size());

        const auto t_sp0 = std::chrono::steady_clock::now();

        common_speculative_begin(spec, seq_id, prompt_tgt);

        common_sampler_ptr smpl(common_sampler_init(model_tgt, sparams_smpl));

        // manual prefill with per-row logits requested (llama_batch_get_one()
        // only requests the last row -- dspark's process() needs a capture
        // row for EVERY prompt position, see common/speculative.cpp).
        prefill(prompt_tgt, true);

        int          n_past      = (int) prompt_tgt.size();  // == position of id_last
        int          n_predicted = 0;
        llama_tokens spec_tokens;
        bool         has_eos = false;

        int64_t              prompt_drafted = 0, prompt_accepted = 0, prompt_rounds = 0;
        const auto           t_decode = std::chrono::steady_clock::now();
        int64_t              draft_us = 0, verify_us = 0;
        int64_t              steady_us = 0, steady_rounds = 0;
        std::vector<int64_t> retained_counts((size_t) n_draft + 1, 0);

        while (n_predicted < n_predict_max && !has_eos) {
            const auto   t_round = std::chrono::steady_clock::now();
            llama_tokens draft;

            common_speculative_draft_params & dp = common_speculative_get_draft_params(spec, seq_id);
            dp.drafting                          = true;
            dp.n_max                             = n_draft;
            dp.n_past                            = n_past;
            dp.id_last                           = id_last;
            dp.prompt                            = nullptr;  // unused by dspark
            dp.result                            = &draft;

            const auto t_draft = std::chrono::steady_clock::now();
            common_speculative_draft(spec);
            draft_us +=
                std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::steady_clock::now() - t_draft)
                    .count();

            if (draft.empty()) {
                fprintf(stderr, "warn: empty draft at prompt %zu, n_past=%d -- stopping this prompt early\n", pi,
                        n_past);
                break;
            }
            if (draft.size() >= retained_counts.size()) {
                fail("retained depth exceeds configured cap");
            }
            retained_counts[draft.size()]++;

            // target verify batch: [id_last, draft0, draft1, ..., draftN-1],
            // matching examples/speculative-simple/speculative-simple.cpp.
            common_batch_clear(batch_tgt);
            common_batch_add(batch_tgt, id_last, (llama_pos) n_past, { seq_id }, /* logits = */ true);
            for (size_t i = 0; i < draft.size(); ++i) {
                common_batch_add(batch_tgt, draft[i], (llama_pos) (n_past + 1 + (int) i), { seq_id },
                                 /* logits = */ true);
            }

            const auto t_verify = std::chrono::steady_clock::now();
            if (llama_decode(ctx_tgt, batch_tgt) != 0) {
                fail("verify decode failed at prompt " + std::to_string(pi));
            }
            llama_synchronize(ctx_tgt);
            verify_us +=
                std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::steady_clock::now() - t_verify)
                    .count();
            // NOTE: unlike examples/speculative-simple.cpp's generic draft-model
            // path, dspark must NOT llama_decode() this batch against ctx_dft --
            // its drafter cache was already advanced inside common_speculative_draft()
            // above via the out-of-band dspark-ctx staging (see file header comment).
            if (!common_speculative_process(spec, batch_tgt)) {
                fail("common_speculative_process (verify) failed at prompt " + std::to_string(pi));
            }

            auto ids = common_sampler_sample_and_accept_n(smpl.get(), ctx_tgt, draft);
            if (ids.empty()) {
                fail("common_sampler_sample_and_accept_n returned empty");
            }

            const uint16_t n_accepted = (uint16_t) (ids.size() - 1);
            common_speculative_accept(spec, seq_id, n_accepted);

            // accept-by-depth: position i (1-based) is reached iff positions
            // 1..i-1 were all accepted, i.e. i <= n_accepted + 1.
            for (int32_t i = 1; i <= (int32_t) draft.size(); ++i) {
                if (i <= (int32_t) n_accepted + 1) {
                    depth_reached[i]++;
                }
                if (i <= (int32_t) n_accepted) {
                    depth_accepted[i]++;
                }
            }

            prompt_drafted += (int64_t) draft.size();
            prompt_accepted += n_accepted;
            prompt_rounds += 1;

            // total newly committed tokens this round == ids.size() (n_accepted
            // draft tokens + exactly one new sample: either the mismatch
            // replacement or, on full acceptance, the bonus token) --
            // mirrors speculative-simple.cpp's n_past bookkeeping exactly.
            n_past += (int) ids.size();

            for (size_t i = 0; i < ids.size(); ++i) {
                if (n_predicted == n_predict_max) {
                    break;
                }
                id_last = ids[i];
                spec_tokens.push_back(id_last);
                n_predicted++;
                if (llama_vocab_is_eog(vocab_tgt, id_last)) {
                    has_eos = true;
                    break;
                }
            }

            // drop the rejected tail of this round's verify batch from the
            // target's KV cache (dspark's own drafter-side cache was already
            // cropped inside draft() itself). must not ignore failure here --
            // on a hybrid GDN/attention target this is a bounded partial
            // rollback of the recurrent state, and a silently ignored no-op
            // would leave every round's rejected draft tail permanently
            // baked into the recurrent state instead of failing loudly.
            if (!llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, n_past, -1)) {
                fail("target recurrent-state rollback failed");
            }
            if (prompt_rounds > 2) {
                steady_us +=
                    std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::steady_clock::now() - t_round)
                        .count();
                ++steady_rounds;
            }
        }

        const double sp_seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - t_sp0).count();
        const double decode_seconds =
            std::chrono::duration<double>(std::chrono::steady_clock::now() - t_decode).count();
        printf(
            "DCUT_TIMING prompt=%zu rounds=%lld emitted=%d decode_s=%.9f ar_decode_s=%.9f ar_decode_tokens=%d "
            "draft_us=%lld verify_us=%lld steady_us=%lld steady_rounds=%lld depths=",
            pi, (long long) prompt_rounds, n_predicted, decode_seconds, ar_decode_seconds, ar_n_predicted - 1,
            (long long) draft_us, (long long) verify_us, (long long) steady_us, (long long) steady_rounds);
        for (size_t d = 1; d < retained_counts.size(); ++d) {
            printf("%s%lld", d == 1 ? "" : ",", (long long) retained_counts[d]);
        }
        printf("\n");
        const double sp_tok_per_sec = n_predicted / sp_seconds;
        const double speedup        = sp_tok_per_sec / ar_tok_per_sec;
        if (spec_tokens != ar_tokens) {
            size_t first = 0;
            while (first < spec_tokens.size() && first < ar_tokens.size() && spec_tokens[first] == ar_tokens[first]) {
                ++first;
            }
            fprintf(stderr, "AR parity FAILED: prompt=%zu first_difference=%zu AR=%d spec=%d\n", pi, first,
                    first < ar_tokens.size() ? ar_tokens[first] : -1,
                    first < spec_tokens.size() ? spec_tokens[first] : -1);
            if (parity_continue) {
                parity_failed_prompts.push_back(pi);
                printf("[%2zu] PARITY_FAILED excluded from totals (first_difference=%zu)\n", pi, first);
                continue;
            }
            fail("refusing performance claim for non-equivalent output");
        }
        printf("AR parity PASSED: prompt=%zu tokens=%zu\n", pi, ar_tokens.size());

        total_drafted += prompt_drafted;
        total_accepted += prompt_accepted;
        total_rounds += prompt_rounds;
        total_predicted += n_predicted;

        total_ar_predicted += ar_n_predicted;
        total_ar_seconds += ar_seconds;
        total_sp_seconds += sp_seconds;

        {
            bool found = false;
            for (size_t ci = 0; ci < cat_names.size(); ++ci) {
                if (cat_names[ci] == ps.category) {
                    cat_stats_v[ci].drafted += prompt_drafted;
                    cat_stats_v[ci].accepted += prompt_accepted;
                    cat_stats_v[ci].rounds += prompt_rounds;
                    found = true;
                    break;
                }
            }
            if (!found) {
                cat_names.push_back(ps.category);
                cat_stats_v.push_back({ prompt_drafted, prompt_accepted, prompt_rounds });
            }
        }

        const double accept_rate = prompt_drafted > 0 ? (double) prompt_accepted / (double) prompt_drafted : 0.0;
        const double tau         = prompt_rounds > 0 ? (double) prompt_accepted / (double) prompt_rounds + 1.0 : 0.0;

        printf(
            "[%2zu][%-9s] n_predicted=%-4d rounds=%-4lld drafted=%-5lld accepted=%-5lld accept=%.3f tau=%.3f "
            "ar_tok_s=%.2f sp_tok_s=%.2f speedup=%.3f\n",
            pi, ps.category.c_str(), n_predicted, (long long) prompt_rounds, (long long) prompt_drafted,
            (long long) prompt_accepted, accept_rate, tau, ar_tok_per_sec, sp_tok_per_sec, speedup);
        fflush(stdout);
    }

    if (!oracle_rows) {
        printf("\n=== per-category ===\n");
        for (size_t ci = 0; ci < cat_names.size(); ++ci) {
            const auto & cs          = cat_stats_v[ci];
            const double accept_rate = cs.drafted > 0 ? (double) cs.accepted / (double) cs.drafted : 0.0;
            const double tau         = cs.rounds > 0 ? (double) cs.accepted / (double) cs.rounds + 1.0 : 0.0;
            printf("%-9s: rounds=%-5lld drafted=%-6lld accepted=%-6lld accept=%.4f tau=%.4f\n", cat_names[ci].c_str(),
                   (long long) cs.rounds, (long long) cs.drafted, (long long) cs.accepted, accept_rate, tau);
        }

        printf("\n=== accept-by-depth (conditional: given position i was reached) ===\n");
        printf("%-6s %-9s %-9s %s\n", "depth", "reached", "accepted", "accept_i");
        for (int32_t i = 1; i <= n_draft; ++i) {
            const double r = depth_reached[i] > 0 ? (double) depth_accepted[i] / (double) depth_reached[i] : 0.0;
            printf("%-6d %-9lld %-9lld %.4f\n", i, (long long) depth_reached[i], (long long) depth_accepted[i], r);
        }

        common_speculative_print_stats(spec);

        const double accept_rate_all = total_drafted > 0 ? (double) total_accepted / (double) total_drafted : 0.0;
        const double tau_all         = total_rounds > 0 ? (double) total_accepted / (double) total_rounds + 1.0 : 0.0;

        // aggregate tok/s: total tokens over total wall time across all prompts,
        // not a naive mean of per-prompt ratios (which would over-weight short
        // prompts) -- same convention as the per-category accept/tau rollup above.
        const double ar_tok_per_sec_all = total_ar_seconds > 0.0 ? total_ar_predicted / total_ar_seconds : 0.0;
        const double sp_tok_per_sec_all = total_sp_seconds > 0.0 ? total_predicted / total_sp_seconds : 0.0;
        const double speedup_all        = ar_tok_per_sec_all > 0.0 ? sp_tok_per_sec_all / ar_tok_per_sec_all : 0.0;

        if (parity_continue) {
            printf("PARITY_CLEAN prompts=%zu failed=%zu failed_ids=", prompts.size() - parity_failed_prompts.size(),
                   parity_failed_prompts.size());
            for (size_t k = 0; k < parity_failed_prompts.size(); ++k) {
                printf("%s%zu", k ? "," : "", parity_failed_prompts[k]);
            }
            printf("\n");
        }
        printf(
            "\n=== OVERALL: prompts=%zu n_predicted=%lld rounds=%lld drafted=%lld accepted=%lld accept=%.4f tau=%.4f "
            "ar_tok_s=%.2f sp_tok_s=%.2f speedup=%.3f ===\n",
            prompts.size() - parity_failed_prompts.size(), (long long) total_predicted, (long long) total_rounds,
            (long long) total_drafted, (long long) total_accepted, accept_rate_all, tau_all, ar_tok_per_sec_all,
            sp_tok_per_sec_all, speedup_all);
    }

    llama_batch_free(batch_tgt);
    common_speculative_free(spec);
    llama_free(ctx_dft);
    llama_free(ctx_tgt);
    llama_free(ctx_ar);
    llama_model_free(model_dft);
    llama_model_free(model_tgt);
    llama_backend_free();

    return oracle_mismatches ? 2 : 0;
}