fix: speculative verify with --mtp-draft > 2 commits wrong tokens

ds4.c

@@ -20157,11 +20157,14 @@ static bool metal_graph_verify_suffix_tops(
                                        g->spec_logits,
                                        DS4_N_VOCAB) != 0;
         } else if (top_rows) {
+            /* top-1 of each of the top_rows rows: n_tokens=top_rows, top_k=1.
+             * The order is transposed vs the indexer-score callers; a swap
+             * silently scores row 0's runner-ups instead of each row. */
             ok = ds4_gpu_indexer_topk_tensor(g->comp_selected,
                                                g->spec_logits,
                                                DS4_N_VOCAB,
-                                               1,
-                                               top_rows) != 0;
+                                               top_rows,
+                                               1) != 0;
         }
     }
     if (ok) ok = ds4_gpu_end_commands() != 0;

tests/ds4_test.c

@@ -87,6 +87,9 @@ static void test_restore_canonical_streaming_prefill(
 
 static ds4_engine *test_open_engine(bool quality) {
     ds4_engine *engine = NULL;
+    /* DS4_TEST_MTP loads the MTP head on the fast engine so the speculative
+     * verify regression can reuse it; draft=4 hits the multi-row verify path. */
+    const char *mtp = getenv("DS4_TEST_MTP");
     ds4_engine_options opt = {
         .model_path = test_model_path(),
 #ifdef __APPLE__
@@ -103,6 +106,8 @@ static ds4_engine *test_open_engine(bool quality) {
             test_env_gib("DS4_TEST_SSD_STREAMING_CACHE_GB"),
         .ssd_streaming_preload_experts =
             test_env_u32("DS4_TEST_SSD_STREAMING_PRELOAD_EXPERTS"),
+        .mtp_path = (mtp && mtp[0] && !quality) ? mtp : NULL,
+        .mtp_draft_tokens = (mtp && mtp[0] && !quality) ? 4 : 0,
     };
     TEST_ASSERT(ds4_engine_open(&engine, &opt) == 0);
     return engine;
@@ -1795,6 +1800,153 @@ static void test_tool_call_quality(void) {
     test_close_engine(true);
 }
 
+/* Greedy speculative decode: capture committed tokens and the largest accepted
+ * chunk, so the caller can confirm the multi-row verify path actually ran. */
+static bool test_mtp_capture_speculative(ds4_engine *engine, const ds4_tokens *prompt,
+                                         int max_tokens, int *out, int *out_len,
+                                         int *max_chunk) {
+    *out_len = 0;
+    *max_chunk = 0;
+    ds4_session *session = NULL;
+    TEST_ASSERT(ds4_session_create(&session, engine, 32768) == 0);
+    if (!session) return false;
+
+    char err[160];
+    bool ok = ds4_session_sync(session, prompt, err, sizeof(err)) == 0;
+    TEST_ASSERT(ok);
+
+    const int eos = ds4_token_eos(engine);
+    int n = 0;
+    bool stop = false;
+    while (ok && !stop && n < max_tokens) {
+        const int token = ds4_session_argmax(session);
+        if (token == eos) break;
+
+        int toks[17]; /* base token + draft depth, which the engine clamps to 16 */
+        const int ntok = ds4_session_eval_speculative_argmax(
+            session, token, max_tokens - n, eos, toks,
+            (int)(sizeof(toks) / sizeof(toks[0])), err, sizeof(err));
+        if (ntok < 0) { ok = false; TEST_ASSERT(false); break; }
+        if (ntok > *max_chunk) *max_chunk = ntok;
+
+        for (int j = 0; j < ntok; j++) {
+            if (toks[j] == eos) { stop = true; break; }
+            out[n++] = toks[j];
+            if (n >= max_tokens) { stop = true; break; }
+        }
+    }
+
+    *out_len = n;
+    ds4_session_free(session);
+    return ok;
+}
+
+/* Replay toks[] through plain decode and return the largest gap between a
+ * position's argmax logit and the committed token's logit.  Correct speculation
+ * commits (near-)argmax tokens (gap ~0); a mis-committed token gives a big gap. */
+static bool test_mtp_worst_argmax_gap(ds4_engine *engine, const ds4_tokens *prompt,
+                                      const int *toks, int n,
+                                      float *worst_gap, int *worst_at) {
+    *worst_gap = 0.0f;
+    *worst_at = -1;
+    ds4_session *session = NULL;
+    TEST_ASSERT(ds4_session_create(&session, engine, 32768) == 0);
+    if (!session) return false;
+
+    char err[160];
+    bool ok = ds4_session_sync(session, prompt, err, sizeof(err)) == 0;
+    TEST_ASSERT(ok);
+
+    for (int i = 0; ok && i < n; i++) {
+        ds4_token_score best, cur;
+        ok = ds4_session_top_logprobs(session, &best, 1) >= 1 &&
+             ds4_session_token_logprob(session, toks[i], &cur) == 1;
+        TEST_ASSERT(ok);
+        if (!ok) break;
+
+        const float gap = best.logit - cur.logit;
+        if (gap > *worst_gap) { *worst_gap = gap; *worst_at = i; }
+        if (ds4_session_eval(session, toks[i], err, sizeof(err)) != 0) { ok = false; TEST_ASSERT(false); break; }
+    }
+
+    ds4_session_free(session);
+    return ok;
+}
+
+/* Verbatim-copy task: keeps the model confident (a mis-committed token shows as
+ * a large argmax gap) and draft acceptance high (so the multi-row verify path is
+ * exercised across the generation). */
+static const char *test_mtp_copy_prompt(void) {
+    return
+        "Reproduce the following C code EXACTLY, character for character, "
+        "inside a single code block and output nothing else:\n\n"
+        "```c\n"
+        "static uint32_t clamp_u32(uint32_t v, uint32_t lo, uint32_t hi) {\n"
+        "    if (v < lo) return lo;\n"
+        "    if (v > hi) return hi;\n"
+        "    return v;\n"
+        "}\n"
+        "\n"
+        "static uint32_t ring_advance(uint32_t pos, uint32_t cap) {\n"
+        "    uint32_t next = pos + 1u;\n"
+        "    return next >= cap ? 0u : next;\n"
+        "}\n"
+        "\n"
+        "static int scratch_init(scratch *s, uint32_t ctx_size) {\n"
+        "    if (ctx_size == 0u) ctx_size = 1u;\n"
+        "    s->ctx_size = ctx_size;\n"
+        "    s->comp_cap = ctx_size / 4u + 2u;\n"
+        "    s->rows = clamp_u32(s->comp_cap, 1u, 4096u);\n"
+        "    s->head = 0u;\n"
+        "    return s->rows > 0u ? 0 : -1;\n"
+        "}\n"
+        "```\n";
+}
+
+#define TEST_MTP_MAXGEN 256
+
+/* Regression for the swapped top-k arguments in metal_graph_verify_suffix_tops
+ * at draft depth > 2.  Replays the committed speculative tokens through plain
+ * decode and requires each to be a (near-)argmax: that is the verify invariant,
+ * and unlike comparing token streams it tolerates the near-greedy tie
+ * divergences.  Needs an MTP head, so it self-skips without DS4_TEST_MTP. */
+static void test_mtp_verify_depth(void) {
+    ds4_engine *engine = test_get_engine(false);
+    if (!engine || !ds4_engine_has_mtp(engine)) {
+        fprintf(stderr, "ds4-test: mtp-verify-depth skipped (set DS4_TEST_MTP to an MTP GGUF)\n");
+        return;
+    }
+    TEST_ASSERT(ds4_engine_mtp_draft_tokens(engine) > 2);
+
+    ds4_tokens prompt = {0};
+    ds4_chat_begin(engine, &prompt);
+    ds4_chat_append_message(engine, &prompt, "user", test_mtp_copy_prompt());
+    ds4_chat_append_assistant_prefix(engine, &prompt, DS4_THINK_NONE);
+    TEST_ASSERT(prompt.len > 0);
+
+    int *spec = malloc((size_t)TEST_MTP_MAXGEN * sizeof(*spec));
+    TEST_ASSERT(spec != NULL);
+    if (spec && prompt.len > 0) {
+        int nspec = 0, max_chunk = 0;
+        const bool ok_spec = test_mtp_capture_speculative(engine, &prompt, TEST_MTP_MAXGEN,
+                                                          spec, &nspec, &max_chunk);
+        TEST_ASSERT(ok_spec);
+        TEST_ASSERT(max_chunk > 1);  /* multi-token chunks committed: the multi-row path ran */
+        TEST_ASSERT(nspec > 128);    /* enough output to surface the bug, incl. a spurious-EOS truncation */
+
+        float worst_gap = 0.0f;
+        int worst_at = -1;
+        const bool ok_check = test_mtp_worst_argmax_gap(engine, &prompt, spec, nspec,
+                                                        &worst_gap, &worst_at);
+        TEST_ASSERT(ok_check);
+        fprintf(stderr, "ds4-test: mtp-verify-depth nspec=%d max_chunk=%d worst_argmax_gap=%.3f at=%d\n",
+                nspec, max_chunk, worst_gap, worst_at);
+        TEST_ASSERT(worst_gap <= 2.0f);  /* correct: ~0; bug: ~21 on the reference model */
+    }
+
+    free(spec);
+    ds4_tokens_free(&prompt);
+}
 #endif
 
 static void test_server_unit_group(void) {
@@ -1820,6 +1972,7 @@ static const ds4_test_entry test_entries[] = {
     {"--metal-kernels", "metal-kernels", "isolated Metal kernel numeric regressions", test_metal_kernel_group},
     {"--metal-tensor-equivalence", "metal-tensor-equivalence", "fast/quality Metal prompt-logit and greedy equivalence", test_metal_mpp_equivalence},
     {"--streaming-decode-prefill-correctness", "streaming-decode-prefill-correctness", "streaming decode-style cold prefill drift and repeatability", test_streaming_decode_prefill_correctness},
+    {"--mtp-verify-depth", "mtp-verify-depth", "MTP speculative verify commits autoregressive-identical tokens at draft depth > 2", test_mtp_verify_depth},
 #endif
     {"--server", "server", "server parser/rendering/cache unit tests", test_server_unit_group},
 };