Merge f-022-specfix: first-window HIP verify hang (PLE flag before the first launch of each window graph) + 1-token prompt starts the speculative loop

CMakeLists.txt

@@ -1425,6 +1425,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(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)
 endif()
 if(STRATA_BUILD_TESTS OR STRATA_BUILD_CONVERSATION_TESTS)
   add_executable(message_boundary_test src/program/message_boundary_test.cpp)

include/strata/core/verify.hpp

@@ -530,6 +530,7 @@ private:
     cudaEvent_t df_join_[2] = {};
     cudaGraphExec_t exec_[9] = {};
     cudaGraphExec_t exec_nr_[9] = {};   // #871: the doorbell variant of a stage that is all-resident otherwise
+    bool launched_[9] = {}, launched_nr_[9] = {};   ///< this graph has been launched once (early_ple_flag: a cold first launch raises the PLE flag first)
     // fork 34a35db (STRATA_SIDE_MIXER=1): the GDN gates and the QSA indexer on a branch of the window's graph beside the
     // mixer's projections, joined before the recurrence / the attention
     cudaStream_t side_ = nullptr;

include/strata/spec/spec_start.hpp

@@ -0,0 +1,20 @@
+// include/strata/spec/spec_start.hpp - where the speculative loop of `strata generate --spec` begins, and whether it runs.
+//
+// The CLI's prompt path leaves the position the speculative loop starts at in one variable.  Until 0.2.2 it was 0 for "the
+// loop is not used", which is also the position a native pack starts the loop at when the prompt is a single token (its
+// last prompt token is window 0 at position 0): such a prompt never speculated.  -1 is the "not used" value now.
+#pragma once
+
+#include <cstdint>
+
+namespace strata::spec {
+
+inline constexpr int64_t kSpecNotStarted = -1;
+
+/// The loop runs when the prompt path named a start position (0 is a real one), tokens are still wanted and the text did
+/// not already end.
+inline constexpr bool spec_loop_runs(int64_t spec_pos, int64_t produced, int64_t max_new, bool ended) {
+    return spec_pos >= 0 && produced < max_new && !ended;
+}
+
+}  // namespace strata::spec

src/core/verify.cpp

@@ -269,6 +269,26 @@ bool sh_stream_on() {
     return on;
 }
 
+// 0.2.2: the zero-doorbell window graph has one host wait, the PLE flag (layer 1's wait kernel spins until the host has read
+// the window's PLE rows).  Raised after the launch returns, it can deadlock HIP on the very first launch of a graph: the
+// stall report of a hung first window (gfx1151, 910-token prompt, about 1 run in 6) shows the host inside hipGraphLaunch
+// waiting for a runtime lock while the runtime's worker waits for the GPU, which is spinning on this flag; the raise that
+// would free it sits behind the launch.  So the first launch of each window graph reads the rows and raises the flag BEFORE
+// the launch (the wait then passes at once); later launches keep the order that overlaps the read with layer 0 (reading
+// first costs about 1.4% decode on gfx1151).  Mode 0: never, 1: every launch, 2: the first launch of each graph.  Default:
+// 2 on HIP, 0 elsewhere (the CUDA path is unchanged).  STRATA_VERIFY_EARLY_PLE_FLAG=0 / 1 / 2 forces a mode.
+int early_ple_flag_mode() {
+    static const int mode = [] {
+        if (const char* v = std::getenv("STRATA_VERIFY_EARLY_PLE_FLAG")) return std::atoi(v) < 0 ? 0 : std::min(2, std::atoi(v));
+#if defined(STRATA_USE_HIP)
+        return 2;
+#else
+        return 0;
+#endif
+    }();
+    return mode;
+}
+
 // A one-token window always keeps its token, so its graph advances the sequence state itself (the GDN conv history
 // and recurrence state) and Verifier::commit launches no commit graph after it (eddoursul's fork, F7).  CUDA only:
 // HIP keeps the commit graph after every window (STRATA_ONE_TOKEN_COMMIT=0 does too).
@@ -2732,6 +2752,20 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool,
 #if !defined(STRATA_USE_HIP) && !defined(STRATA_HIP_GFX906)
     std::optional<HostGraphLaunch> launch_o;   // (not set when the previous stage launched this window already)
 #endif
+    const bool test_stall = g_test_stall > 0 && windows + 1 == g_test_stall;   // #267 test hook (off: false)
+    bool flag_raised = false;   // the PLE flag went up before the launch (early_ple_flag_mode)
+    bool& graph_launched = ar_off_ ? launched_nr_[T] : launched_[T];
+    if (!pre && ar_on() && !test_stall && (early_ple_flag_mode() == 1 || (early_ple_flag_mode() == 2 && !graph_launched))) {
+        if (do_ple) {
+            const Clock::time_point tp = Clock::now();
+            if (!ss.ple.table->gather_batch(ple_rows, (size_t) T, h_ple_, err)) return false;
+            std::atomic_thread_fence(std::memory_order_seq_cst);
+            _mm_sfence();
+            ms_host += ms_since(tp);
+        }
+        *(volatile uint32_t*) h_flag_ = 1;
+        flag_raised = true;
+    }
     if (!pre) {
         if (trace_h_ != nullptr) std::memset(trace_h_, 0, trace_n_ * 8);   // #649: this window's breadcrumbs only
         std::atomic_thread_fence(std::memory_order_seq_cst);
@@ -2756,6 +2790,7 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool,
         const cudaError_t le = cudaGraphLaunch(ar_off_ ? exec_nr_[T] : exec_[T], cs_);
 #endif
         trace_ev("LAUNCHED", -1, -1, (int64_t) le);
+        graph_launched = true;
         if (le != cudaSuccess) { err = std::string("verify: launch: ") + cudaGetErrorString(le); return false; }
 #if !defined(STRATA_USE_HIP) && !defined(STRATA_HIP_GFX906)
         (void) launch_o->query();
@@ -2793,9 +2828,10 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool,
     const int G = groups_[T] > 0 ? groups_[T] : 1;
     const int gtb[2] = {0, (T + 1) / 2}, gte[2] = {G == 2 ? (T + 1) / 2 : T, T};
     const int64_t steps = (le_ - lb_) * G;
-    const bool test_stall = g_test_stall > 0 && windows + 1 == g_test_stall;   // #267 test hook (off: false)
     if (ar_on() && !test_stall) {
-        if (do_ple) {
+        if (flag_raised) {
+            // up before the launch (early_ple_flag_mode)
+        } else if (do_ple) {
             const Clock::time_point tp = Clock::now();
             if (!ss.ple.table->gather_batch(ple_rows, (size_t) T, h_ple_, err)) return false;
             std::atomic_thread_fence(std::memory_order_seq_cst);

src/program/generate.cpp

@@ -84,6 +84,7 @@
 #include "strata/spec/draft_policy.hpp"
 #include "strata/spec/draft_source.hpp"
 #include "strata/spec/suffix_drafter.hpp"
+#include "strata/spec/spec_start.hpp"
 #include "strata/program/batch_groups.hpp"
 #include "strata/kernels/cvec.hpp"
 #include "strata/core/progress.hpp"
@@ -15534,7 +15535,7 @@ int main(int argc, char** argv) {
     // ---- plan v0.3 P5: the prompt's conditioning positions [0, n_prompt - 1) in batched chunks.  The token loop
     // then starts at the last prompt position, whose prediction is the first generated token.
     int64_t pos_start = 0;
-    int64_t spec_pos = 0;   // plan v0.3 P6: where the speculative loop starts (0 = not used)
+    int64_t spec_pos = strata::spec::kSpecNotStarted;   // plan v0.3 P6: where the speculative loop starts (0 is a real position: a 1-token prompt on a native pack)
     strata::prefill::Prefill prefill;
     // #1595: not by default when the cache holds every expert (nothing to take)
     strata::prefill::Prefill::arm_cpu_share(!multi_gpu && !o.no_pool,
@@ -15924,7 +15925,7 @@ int main(int argc, char** argv) {
     // round emits (accepted drafts + 1) tokens.  `commit` keeps the state of the tokens that were emitted.
     const bool ended = o.stop_eos && !produced.empty() &&
                        std::find(o.eos_ids.begin(), o.eos_ids.end(), (int64_t) produced.back()) != o.eos_ids.end();
-    if (spec_pos > 0 && (int64_t) produced.size() < o.max_new && !ended) {
+    if (strata::spec::spec_loop_runs(spec_pos, (int64_t) produced.size(), o.max_new, ended)) {
         auto read_ids = [](const std::string& path, std::vector<int64_t>& ids) {
             std::ifstream in(path);
             std::string text((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>());
@@ -16225,8 +16226,32 @@ int main(int argc, char** argv) {
         }
         const double pool_ms0 = drive.cpu_ms;
         const int64_t misses0 = drive.d.multi_misses, entries0 = drive.d.multi_entries;
+        // STRATA_SPEC_WATCH=<seconds> (diagnostics, off by default): the serve watchdog for this loop.  A run whose verify
+        // windows and prompt chunks stop beating for that long prints the stage it is in and the stall report, then ends.
+        if (const char* sw = std::getenv("STRATA_SPEC_WATCH"); sw != nullptr && std::atoi(sw) > 0) {
+            const int limit = std::atoi(sw);
+            strata::core::progress().busy.store(true);
+            std::thread([limit] {
+                strata::core::Progress& pr = strata::core::progress();
+                uint64_t last = pr.beats.load(), ticks_at = pr.ticks.load();
+                auto since = std::chrono::steady_clock::now();
+                for (;;) {
+                    std::this_thread::sleep_for(std::chrono::seconds(1));
+                    const auto now = std::chrono::steady_clock::now();
+                    const uint64_t b = pr.beats.load();
+                    if (b != last) { last = b; ticks_at = pr.ticks.load(); since = now; continue; }
+                    if (now - since < std::chrono::seconds(limit)) continue;
+                    std::fprintf(stderr, "strata generate: no progress for %d s (%s)\n", limit, stage_text().c_str());
+                    stall_report(stderr, pr.ticks.load() - ticks_at);
+                    strata::core::release_gpu_waits(stderr);
+                    std::fflush(stderr);
+                    std::abort();
+                }
+            }).detach();
+        }
         while ((int64_t) produced.size() < max_new) {
             const Clock::time_point t0 = Clock::now();
+            strata::core::progress_at("spec loop: choosing the window", (int64_t) produced.size());
             int T = S_mtp;
             if (use_mtp && o.spec_min_p > 0.0) {
                 T = 1;
@@ -16338,6 +16363,7 @@ int main(int argc, char** argv) {
             bool adapt_ok = true;
             if (!drive.d.usage.empty() && ((rounds + 1) % o.adapt_every) == 0)
                 adapt_thr = std::thread([&] { strata::platform::name_this_thread("strata-adapt"); strata::aux_cpus::pin_current_thread(); adapt_ok = adapt(); });
+            strata::core::progress_at("spec loop: commit", (int64_t) T);
             if (!ver.commit(a + 1, err)) {
                 if (adapt_thr.joinable()) adapt_thr.join();
                 std::fprintf(stderr, "strata generate: %s\n", err.c_str());

src/spec/spec_start_test.cpp

@@ -0,0 +1,22 @@
+// src/spec/spec_start_test.cpp - the speculative loop's start position: 0 is a real position (a 1-token prompt on a native pack).
+#include "strata/spec/spec_start.hpp"
+
+#include <cstdio>
+
+using namespace strata::spec;
+
+int main() {
+    int fail = 0;
+    auto check = [&](bool ok, const char* what) {
+        if (!ok) { std::printf("FAIL: %s\n", what); ++fail; }
+    };
+    check(!spec_loop_runs(kSpecNotStarted, 0, 100, false), "no start position: the loop is not used");
+    check(spec_loop_runs(0, 0, 100, false), "position 0 (a 1-token prompt on a native pack) starts the loop");
+    check(spec_loop_runs(910, 0, 100, false), "a long prompt's start position");
+    check(spec_loop_runs(1, 0, 100, false), "position 1 (a 1-token prompt on the per-token path)");
+    check(!spec_loop_runs(0, 100, 100, false), "nothing left to generate");
+    check(!spec_loop_runs(0, 3, 100, true), "the text already ended");
+    check(kSpecNotStarted < 0, "the sentinel is not a position");
+    std::printf("spec_start tests: %s\n", fail ? "FAILED" : "OK");
+    return fail ? 1 : 0;
+}

src/spec/suffix_drafter_test.cpp

@@ -128,6 +128,54 @@ void unit_tests() {
         check(d.propose(3, out) == 0, "propose after propose_after");
         check(d.propose_after(nullptr, 0, 3, out) == d.propose(3, out), "propose_after with no drafts == propose");
     }
+    {   // randomized soundness (run under ASan/UBSan in CI-like builds): small vocabularies repeat a lot, tiny tables fill up,
+        // propose_after is interleaved; every proposal must be a real earlier continuation of the matched suffix
+        uint64_t rng = 0x243F6A8885A308D3ull;
+        auto next = [&rng] { rng ^= rng << 13; rng ^= rng >> 7; rng ^= rng << 17; return rng; };
+        long proposals = 0, bad = 0;
+        for (int round = 0; round < 300; ++round) {
+            const int vocab = 2 + (int) (next() % 6), min_m = 3 + (int) (next() % 3);
+            const size_t cap = 8 + (size_t) (next() % 200);   // often far below the history: the table fills
+            SuffixDrafter d(min_m, 4 + (int) (next() % 40), cap);
+            std::vector<int32_t> hist;
+            for (int step = 0; step < 400; ++step) {
+                const int burst = 1 + (int) (next() % 3);
+                for (int b = 0; b < burst; ++b) {
+                    const int32_t t = hist.size() > 8 && (next() % 3) == 0 ? hist[hist.size() - 1 - (size_t) (next() % 8)]
+                                                                           : (int32_t) (next() % (uint64_t) vocab);
+                    hist.push_back(t);
+                    d.append(t);
+                }
+                const int max_k = 1 + (int) (next() % 8);
+                int32_t out[16] = {};
+                int32_t extra[3] = {};
+                const int n_extra = (int) (next() % 4);
+                for (int e = 0; e < n_extra; ++e) extra[e] = (int32_t) (next() % (uint64_t) vocab);
+                const int k = n_extra > 0 && (next() % 2) ? d.propose_after(extra, n_extra, max_k, out) : d.propose(max_k, out);
+                if (k < 0 || k > max_k) { ++bad; continue; }
+                if (d.size() != hist.size()) { ++bad; continue; }
+                if (k == 0) continue;
+                ++proposals;
+                const int m = d.last_match();
+                if (m < min_m) { ++bad; continue; }
+                // the same query through propose() on the history (no extra tokens) must agree with the invariant too
+                bool found = false;
+                for (size_t p = 0; p + 1 < hist.size() + (size_t) n_extra && !found; ++p) {
+                    // candidate continuation hist[p+1 ..] after a suffix of length m that ends at p
+                    bool match = true;
+                    for (int t = 0; t < k && match; ++t) {
+                        const size_t at = p + 1 + (size_t) t;
+                        const int32_t v = at < hist.size() ? hist[at] : (at - hist.size() < (size_t) n_extra ? extra[at - hist.size()] : -1);
+                        match = v == out[t];
+                    }
+                    found = match;
+                }
+                if (!found) ++bad;
+            }
+        }
+        check(bad == 0, "randomized proposals are real continuations within bounds");
+        check(proposals > 1000, "randomized test exercised proposals");
+    }
     std::printf("suffix_drafter unit tests: %s\n", g_fail ? "FAILED" : "OK");
 }
 }  // namespace