DFlash2: the MTP chain's cost window and q lists apply only to candidates the MTP drew

With DFlash2 as the drafter the MTP stays loaded as the fallback (use_mtp stays true), and three decisions tested use_mtp where they meant
"the MTP chain drew this round's candidates":
- STRATA_SPEC_COST=1 (a GB10 default): T = clamp(chain_rows, 1, S_mtp), but chain_rows is only written by the MTP chain callback, so with DFlash2
  and no --draft-conf the window stayed at one row (CLI and serve; code_py 175 vs 258 tok/s on an RTX PRO 6000, serve 166-230 vs 283-397).
- STRATA_SPEC_PROB=1, sampled requests: ver.set_spec_q(mtp.spec_q()) judged DFlash2's deterministic candidates against the MTP's stale q lists
  (CLI: 107 of 137 windows), and serve's gated path was not declined for DFlash2.
- the lookup choice read mtp.prob() for the same reason.
spec/window_rule.hpp holds the window rule and the q rule, both loops call it. The default path is unchanged: window_rule_test compares the new
rule with a copy of the old code over 9216 inputs (only the DFlash2 + cost rule + no --draft-conf cases differ). The CLI stats also print the
window histogram for any rule that sets the width and, when STRATA_SPEC_PROB judged windows, their count and source.

CMakeLists.txt

@@ -1896,6 +1896,9 @@ if(STRATA_BUILD_TESTS)
   add_executable(draft_policy_test src/spec/draft_policy_test.cpp)
   target_link_libraries(draft_policy_test PRIVATE strata_spec)
   add_test(NAME draft_policy_test COMMAND draft_policy_test)
+  add_executable(window_rule_test src/spec/window_rule_test.cpp)   # which rule picks the window / whose q judges it, by the active draft source
+  target_include_directories(window_rule_test PRIVATE include)
+  add_test(NAME window_rule_test COMMAND window_rule_test)
   add_executable(spec_start_test src/spec/spec_start_test.cpp)
   target_include_directories(spec_start_test PRIVATE include)
   add_test(NAME spec_start_test COMMAND spec_start_test)

include/strata/spec/window_rule.hpp

@@ -0,0 +1,53 @@
+// include/strata/spec/window_rule.hpp - which rule picks a round's verify window, by the round's ACTIVE draft source.
+//
+// `use_mtp` means "the MTP layer is loaded"; with DFlash2 it stays true because the MTP is the fallback.  The decisions below
+// read what the MTP chain just produced (the cost-aware rule's `chain_rows`, the draft probabilities), so they apply only when
+// the MTP chain is the source of this round's candidates.  With DFlash2 as the source, `chain_rows` is never written (the chain
+// callback does not run) and would pin the window at one row.  Header only, no device types: the CLI loop and the serve loop of
+// every backend call it.
+#pragma once
+
+#include <algorithm>
+
+namespace strata::spec {
+
+struct WindowRuleIn {
+    bool use_mtp = false;        ///< the MTP layer is loaded (a fallback beside DFlash2 counts)
+    bool df2_active = false;     ///< DFlash2 drew this round's candidates (CLI: use_df2; serve: this request's drafter)
+    bool spec_cost = false;      ///< STRATA_SPEC_COST=1: the cost-aware MTP window
+    bool first_window = false;   ///< the prompt's last token alone: the caller sets its own width afterwards
+    int chain_rows = 1;          ///< the window the cost rule kept in the MTP chain just drafted (MTP only)
+    int cap = 1;                 ///< S_mtp: the widest window of this source
+    int start = 1;               ///< the width before any rule (CLI: cap; serve: cap with the MTP, 1 without)
+    double spec_min_p = 0.0;     ///< --spec-min-p: stop at the first draft whose probability is below it (0: off)
+    double draft_conf = 0.0;     ///< --draft-conf: DFlash2's window ends where the product of the probabilities falls below it (0: off)
+    const float* dprob = nullptr;   ///< the drafts' probabilities, dprob[j] for draft j (the row j + 1 of the window)
+};
+
+/// The cost-aware MTP window applies only to the MTP chain's own candidates.
+inline bool cost_window_applies(bool use_mtp, bool spec_cost, bool df2_active, bool first_window) {
+    return use_mtp && spec_cost && !df2_active && !first_window;
+}
+
+/// The MTP's q lists (STRATA_SPEC_PROB) are the proposal law only for candidates the MTP drew.
+inline bool mtp_q_applies(bool use_mtp, bool df2_active, bool mtp_prob, bool from_lookup, bool first_window, int mtp_rows) {
+    return use_mtp && !df2_active && mtp_prob && !from_lookup && !first_window && mtp_rows > 1;
+}
+
+inline int choose_window(const WindowRuleIn& in) {
+    int T = in.start;
+    if (cost_window_applies(in.use_mtp, in.spec_cost, in.df2_active, in.first_window)) {
+        T = std::clamp(in.chain_rows, 1, in.cap);
+    } else if (in.use_mtp && in.spec_min_p > 0.0 && in.dprob != nullptr) {
+        T = 1;
+        while (T < in.cap && in.dprob[T - 1] >= (float) in.spec_min_p) ++T;
+    }
+    if (in.df2_active && in.draft_conf > 0.0 && in.dprob != nullptr) {   // DFlash2's own confidence: every draft so far holds
+        T = 1;
+        double chain = 1.0;
+        while (T < in.cap && (chain *= (double) in.dprob[T - 1]) >= in.draft_conf) ++T;
+    }
+    return T;
+}
+
+}  // namespace strata::spec

src/program/generate.cpp

@@ -98,6 +98,7 @@
 #include "strata/program/draft_kv_plan.hpp"
 #include "strata/program/message_boundary.hpp"
 #include "strata/spec/draft_policy.hpp"
+#include "strata/spec/window_rule.hpp"
 #include "strata/spec/draft_source.hpp"
 #include "strata/spec/suffix_drafter.hpp"
 #include "strata/spec/spec_start.hpp"
@@ -16401,7 +16402,7 @@ int main(int argc, char** argv) {
             if (use_mtp) mtp.set_draft_sampling(req_sp);   // STRATA_SPEC_COUPLED=1: sampled drafts (a no-op otherwise)
             const bool gated_proposal = use_mtp && mtp.prob() && strata::core::spec_gate_pick() &&
                 gated_proposal_usable(o.pipeline_windows < 2 && o.batch == 0 && n_stages == 1 && !multi_gpu && drive.d.remote_count == 0,
-                                      o.lookup_chain, o.adapt_swaps, req_pcie_frac != 0.0 || o.pcie_frac_auto, spec_cost, S,
+                                      o.lookup_chain, o.adapt_swaps, req_pcie_frac != 0.0 || o.pcie_frac_auto, spec_cost || req_df2, S,
                                       req_spec_min_p);
             {   // aux_cpus=0|1: where the non-pool, non-host threads run for this request (-1 = the start-up --aux-cpus); moves
                 // threads only.  An enabled placement is swept at every request: it catches threads created since.
@@ -17399,19 +17400,14 @@ int main(int argc, char** argv) {
                     if (first_window && merge_T > 0) T = merge_T;   // STRATA_ADMIT_MERGE_FIRST: the prompt's last tokens ride in this window
                 } else {
                     T = use_mtp ? S_mtp : 1;   // no --mtp: one token a round unless a lookup draft fires
-                    if (use_mtp && spec_cost && !first_window) {
-                        T = std::clamp(chain_rows, 1, S_mtp);
-                    } else if (use_mtp && req_spec_min_p > 0.0) {
-                        T = 1;
-                        while (T < S_mtp && dprob[(size_t) T - 1] >= (float) req_spec_min_p) ++T;
+                    {   // the window rule of this request's ACTIVE source (spec/window_rule.hpp), as in the CLI loop: the MTP chain's cost rule
+                        // and probabilities do not govern DFlash2's candidates, though the MTP stays loaded beside it
+                        strata::spec::WindowRuleIn wr;
+                        wr.use_mtp = use_mtp; wr.df2_active = req_df2; wr.spec_cost = spec_cost; wr.first_window = first_window;
+                        wr.chain_rows = chain_rows; wr.cap = S_mtp; wr.start = T; wr.spec_min_p = req_spec_min_p;
+                        wr.draft_conf = o.draft_conf; wr.dprob = dprob.data();
+                        T = strata::spec::choose_window(wr);
                     }
-#if defined(STRATA_HAVE_DFLASH2)
-                    if (req_df2 && o.draft_conf > 0.0 && !first_window) {   // --draft-conf in serve as in the CLI loop: the chance that every draft so far holds
-                        T = 1;
-                        double chain = 1.0;
-                        while (T < S_mtp && (chain *= (double) dprob[(size_t) T - 1]) >= o.draft_conf) ++T;
-                    }
-#endif
                     if (first_window) T = merge_T > 0 ? merge_T : 1;
                     // a repeat of earlier context (prompt lookup) where the MTP's own first guess agrees: the policy takes it
                     // when its expected tokens per ms, from the measured acceptance and window costs, beat the MTP window's
@@ -17425,7 +17421,7 @@ int main(int argc, char** argv) {
                         }
                         // a probabilistic window must not condition the source on the draw: the policy copy fixed at the request's start
                         // chooses, and the lookup is taken without asking whether the MTP's first guess agrees
-                        const bool probabilistic = use_mtp && mtp.prob();
+                        const bool probabilistic = use_mtp && !req_df2 && mtp.prob();   // DFlash2's candidates are not draws from the MTP's q
                         if (k > 0 && (probabilistic || !use_mtp || sbuf[0] == drafts[0])) {
                             const strata::spec::DraftPolicy::Pick pk = probabilistic
                                 ? source_policy.choose(S_mtp, k, sfx_match) : policy.choose(T, k, sfx_match);
@@ -17649,7 +17645,7 @@ int main(int argc, char** argv) {
                 // a suffix window and the lookup chain's tail are point masses and keep the exact-match rule
                 // (never on a first window: a merged one holds prompt tokens, not drafts)
                 if (gated_proposal) ver.set_spec_proposal(proposal.n_q ? proposal.mu.data() : nullptr, proposal.n_q);
-                else if (use_mtp && mtp.prob() && !from_sfx && !first_window && T_mtp > 1) ver.set_spec_q(mtp.spec_q(), T_mtp - 1);
+                else if (strata::spec::mtp_q_applies(use_mtp, req_df2, use_mtp && mtp.prob(), from_sfx, first_window, T_mtp)) ver.set_spec_q(mtp.spec_q(), T_mtp - 1);   // only candidates the MTP drew
                 const auto decode_pool_fn = use_mtp && mtp.prob() ? invariant_win_pool_fn : win_pool_fn;
                 exit_if_gpu_fatal("");   // (the test hook counts windows here)
                 if (!ver.run(T, window.data(), p, decode_pool_fn, win_pool_user, outv.data(), err) || drive.d.failed) {
@@ -19265,19 +19261,14 @@ int main(int argc, char** argv) {
                 T = proposal.T; from_sfx = proposal.from_sfx;
                 std::copy(dprob.begin(), dprob.begin() + (std::ptrdiff_t) o.spec, dwin.begin());
             } else {
-            if (use_mtp && spec_cost && !first_window) {
-                T = std::clamp(chain_rows, 1, S_mtp);
-            } else if (use_mtp && o.spec_min_p > 0.0) {
-                T = 1;
-                while (T < S_mtp && dprob[(size_t) T - 1] >= (float) o.spec_min_p) ++T;
+            {   // the window rule of the round's ACTIVE source (spec/window_rule.hpp): the MTP chain's cost rule and probabilities do not
+                // govern DFlash2's candidates, though the MTP stays loaded beside it as the fallback (use_mtp)
+                strata::spec::WindowRuleIn wr;
+                wr.use_mtp = use_mtp; wr.df2_active = spec_uses_df2; wr.spec_cost = spec_cost; wr.first_window = first_window;
+                wr.chain_rows = chain_rows; wr.cap = S_mtp; wr.start = T; wr.spec_min_p = o.spec_min_p;
+                wr.draft_conf = o.draft_conf; wr.dprob = dprob.data();
+                T = strata::spec::choose_window(wr);
             }
-#if defined(STRATA_HAVE_DFLASH2)
-            if (use_df2 && o.draft_conf > 0.0) {   // DFlash2's own confidence: the chance that every draft so far holds, from the selector's probabilities
-                T = 1;
-                double chain = 1.0;
-                while (T < S_mtp && (chain *= (double) dprob[(size_t) T - 1]) >= o.draft_conf) ++T;
-            }
-#endif
             if (first_window) T = 1;
             std::copy(dprob.begin(), dprob.begin() + (std::ptrdiff_t) o.spec, dwin.begin());
             if (!t_cycle.empty() && !first_window) T = t_cycle[(size_t) (rounds % (int64_t) t_cycle.size())];
@@ -19286,7 +19277,7 @@ int main(int argc, char** argv) {
                 sfx_match = sfx.last_match();
                 // a probabilistic window must not condition the source on the draw: the policy copy fixed at the start chooses, and the
                 // lookup is taken without asking whether the MTP's first guess agrees
-                const bool probabilistic = use_mtp && mtp.prob();
+                const bool probabilistic = use_mtp && !spec_uses_df2 && mtp.prob();   // DFlash2's candidates are not draws from the MTP's q
                 if (k > 0 && (probabilistic || !use_mtp || sbuf[0] == drafts[0])) {
                     const strata::spec::DraftPolicy::Pick pk = probabilistic
                         ? source_policy.choose(S_mtp, k, sfx_match) : policy.choose(T, k, sfx_match);
@@ -19339,7 +19330,9 @@ int main(int argc, char** argv) {
             if (adapt_nowait()) apply_pending(false);
                 else if (pending.empty() || ++pending_age >= adapt_lag()) apply_pending(true);
             if (gated_proposal) ver.set_spec_proposal(proposal.n_q ? proposal.mu.data() : nullptr, proposal.n_q);
-            else if (use_mtp && mtp.prob() && !from_sfx && T_mtp > 1 && o.spec_corrupt <= 0) ver.set_spec_q(mtp.spec_q(), T_mtp - 1);
+            // the MTP's q lists judge only candidates the MTP drew: DFlash2's are a deterministic pick (exact-match path)
+            else if (strata::spec::mtp_q_applies(use_mtp, spec_uses_df2, use_mtp && mtp.prob(), from_sfx, first_window, T_mtp) && o.spec_corrupt <= 0)
+                ver.set_spec_q(mtp.spec_q(), T_mtp - 1);
             const auto decode_pool_fn = use_mtp && mtp.prob() ? &drive_pool_multi_invariant : &drive_pool_multi;
             if (!ver.run(T, window.data(), p, decode_pool_fn, &drive, outv.data(), err)) {
                 // The graph of a window size is captured the first time it is used, and the draft source took VRAM: when the capture runs
@@ -19472,11 +19465,14 @@ int main(int argc, char** argv) {
         if (!follow.empty())
             std::printf("%-24s %lld of %lld emitted tokens differ from the argmax\n", "follow",
                         (long long) follow_differ, (long long) follow_emitted);
-        if (o.spec_min_p > 0.0) {
+        if (o.spec_min_p > 0.0 || o.draft_conf > 0.0 || spec_cost) {   // (any rule that sets the width)
             std::printf("%-24s", "window sizes");
             for (size_t i = 1; i < window_hist.size(); ++i) std::printf(" T%zu:%lld", i, (long long) window_hist[i]);
-            std::printf("  (min draft probability %.2f)\n", o.spec_min_p);
+            std::printf("  (min draft probability %.2f, draft-conf %.2f, cost rule %d)\n", o.spec_min_p, o.draft_conf, spec_cost ? 1 : 0);
         }
+        if (ver.spec_windows > 0)   // STRATA_SPEC_PROB: windows judged against the MTP's q lists (only an MTP-drawn window should be)
+            std::printf("%-24s %lld windows judged against q lists (source: %s)\n", "probabilistic accept", (long long) ver.spec_windows,
+                        spec_uses_df2 ? "DFlash2" : "MTP");
         if (o.suffix_draft > 0)
             std::printf("%-24s %lld windows, drafts accepted %lld of %lld\n", "suffix drafts", (long long) sfx_windows,
                         (long long) sfx_ok, (long long) sfx_drafts);

src/spec/window_rule_test.cpp

@@ -0,0 +1,108 @@
+// src/spec/window_rule_test.cpp - which rule picks the verify window, and whose q lists judge the candidates (spec/window_rule.hpp).
+//
+// DFlash2 runs with the MTP layer loaded beside it as the fallback (use_mtp stays true).  Before this rule existed, the cost-aware
+// window (STRATA_SPEC_COST=1, a GB10 default) read `chain_rows`, which only the MTP chain writes, so DFlash2 candidates were
+// verified one row at a time; and the MTP's q lists (STRATA_SPEC_PROB=1) were attached to DFlash2's windows.  The tests:
+//   1. the table of the fixed cases (DFlash2 + cost rule with and without --draft-conf / --spec-min-p);
+//   2. the DEFAULT PATH IS UNCHANGED: for every input the old inline code handled correctly (MTP as the source, any flags;
+//      DFlash2 with --draft-conf > 0 or with the cost rule off) a copy of the old code returns the same window;
+//   3. the q rule.
+#include "strata/spec/window_rule.hpp"
+
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+
+using namespace strata::spec;
+
+namespace {
+int g_fail = 0;
+void check(bool ok, const char* what) {
+    std::printf("  %-78s %s\n", what, ok ? "ok" : "FAIL");
+    if (!ok) ++g_fail;
+}
+
+// the inline code of 93970e0a2 (CLI loop; the serve loop differs only in the start value and in first_window, which both callers
+// overwrite afterwards)
+int old_rule(const WindowRuleIn& in) {
+    int T = in.start;
+    if (in.use_mtp && in.spec_cost && !in.first_window) {
+        T = std::clamp(in.chain_rows, 1, in.cap);
+    } else if (in.use_mtp && in.spec_min_p > 0.0) {
+        T = 1;
+        while (T < in.cap && in.dprob[T - 1] >= (float) in.spec_min_p) ++T;
+    }
+    if (in.df2_active && in.draft_conf > 0.0) {
+        T = 1;
+        double chain = 1.0;
+        while (T < in.cap && (chain *= (double) in.dprob[T - 1]) >= in.draft_conf) ++T;
+    }
+    return T;
+}
+}  // namespace
+
+int main() {
+    std::printf("window_rule_test\n");
+    const float p[8] = {0.9f, 0.8f, 0.3f, 0.9f, 0.9f, 0.9f, 0.9f, 0.9f};
+    {   // 1. the table
+        WindowRuleIn in;
+        in.use_mtp = true; in.cap = 6; in.start = 6; in.dprob = p; in.chain_rows = 1;
+        in.df2_active = true; in.spec_cost = true;
+        check(choose_window(in) == 6, "DFlash2 + cost rule, no min-p, no conf: the full cap (was 1)");
+        in.spec_min_p = 0.5;
+        check(choose_window(in) == 3, "DFlash2 + cost rule + --spec-min-p 0.5: stops at the weak draft (was 1)");
+        in.spec_min_p = 0.0; in.draft_conf = 0.35;
+        check(choose_window(in) == 3, "DFlash2 + cost rule + --draft-conf 0.35: the chain product rule (0.9*0.8*0.3=0.216)");
+        in.draft_conf = 0.0; in.first_window = true;
+        check(choose_window(in) == 6, "first window: left to the caller (it sets 1 / merge_T after)");
+        in = WindowRuleIn{};
+        in.use_mtp = true; in.cap = 4; in.start = 4; in.dprob = p; in.spec_cost = true; in.chain_rows = 2;
+        check(choose_window(in) == 2, "MTP + cost rule: the chain's rows");
+        in.chain_rows = 9;
+        check(choose_window(in) == 4, "MTP + cost rule: clamped to the cap");
+        in.spec_cost = false; in.spec_min_p = 0.5;
+        check(choose_window(in) == 3, "MTP + min-p 0.5: the first weak draft ends the window");
+        in.spec_min_p = 0.0;
+        check(choose_window(in) == 4, "MTP, no rule: the full cap");
+        in.use_mtp = false; in.start = 1;
+        check(choose_window(in) == 1, "no MTP (serve): one row unless a lookup fires");
+    }
+    {   // 2. exhaustive over the small input space; random probabilities
+        std::srand(7);
+        long n = 0, same = 0, differ_df2_cost = 0, differ_other = 0;
+        for (int use_mtp = 0; use_mtp < 2; ++use_mtp)
+            for (int df2 = 0; df2 < 2; ++df2)
+                for (int cost = 0; cost < 2; ++cost)
+                    for (int first = 0; first < 2; ++first)
+                        for (int cap = 1; cap <= 8; ++cap)
+                            for (int rows = 1; rows <= 8; ++rows)
+                                for (double mp : {0.0, 0.25, 0.5})
+                                    for (double dc : {0.0, 0.15, 0.35}) {
+                                        float pr[8];
+                                        for (float& x : pr) x = (float) (std::rand() % 101) / 100.0f;
+                                        WindowRuleIn in;
+                                        in.use_mtp = use_mtp; in.df2_active = df2; in.spec_cost = cost; in.first_window = first;
+                                        in.chain_rows = rows; in.cap = cap; in.start = use_mtp ? cap : 1; in.spec_min_p = mp;
+                                        in.draft_conf = dc; in.dprob = pr;
+                                        ++n;
+                                        const int a = old_rule(in), b = choose_window(in);
+                                        if (a == b) { ++same; continue; }
+                                        // a difference is allowed only where the old rule used the MTP chain's rows for DFlash2's candidates
+                                        if (use_mtp && df2 && cost && !first && dc == 0.0) ++differ_df2_cost; else ++differ_other;
+                                    }
+        std::printf("  %ld inputs: %ld identical, %ld differ (DFlash2 + cost rule, no --draft-conf), %ld differ elsewhere\n", n, same, differ_df2_cost, differ_other);
+        check(differ_other == 0, "default path unchanged: no input outside the buggy case changes its window");
+        check(differ_df2_cost > 0, "the buggy case does change (the test sees it)");
+    }
+    {   // 3. the q rule
+        check(mtp_q_applies(true, false, true, false, false, 4), "MTP drew the window, sampled, probabilistic: q attached");
+        check(!mtp_q_applies(true, true, true, false, false, 4), "DFlash2 drew the window: no MTP q (exact-match path)");
+        check(!mtp_q_applies(true, false, true, true, false, 4), "a lookup window: point masses, no q");
+        check(!mtp_q_applies(true, false, true, false, true, 4), "first window: no q");
+        check(!mtp_q_applies(true, false, true, false, false, 1), "one row: no q");
+        check(!mtp_q_applies(true, false, false, false, false, 4), "greedy / mode off: no q");
+        check(!mtp_q_applies(false, false, true, false, false, 4), "no MTP: no q");
+    }
+    std::printf("%s\n", g_fail == 0 ? "window_rule_test: ok" : "window_rule_test: FAILED");
+    return g_fail == 0 ? 0 : 1;
+}