Layer split: pipelined verify windows for one conversation (--pipeline-windows, opt-in)

docs/MULTI_GPU.md

@@ -168,6 +168,45 @@ The Coder on an RTX 5080 + RTX 3090 (Ryzen 9 9950X3D), 32K context; details in
   3090 pair slower (68 / 90 tok/s decode): every extra card costs its own round per window.
 - More cards pay off when the model's routed experts do not fit the faster ones.
 
+## One conversation with both cards busy (`--pipeline-windows`, opt-in)
+
+With a split the cards take turns on a verify window: the first card runs its layers and hands off, then waits while
+the last card runs the rest, the head and the draft. `--pipeline-windows 2` lets the first card start the **next**
+window while the last card still verifies this one. The next window is a guess: that this window is accepted whole
+and that its bonus token is the one the draft layer predicts (the draft layer is run on through the drafts of the
+window in flight). When the guess holds, half of the next window is already done; when it does not, the first card
+puts its state back (a copy of its recurrent state taken while the window ran) and the next window is built from the
+real tokens. A guessed window is only started when the draft layer's estimate says it is likely to be kept. Windows
+copied from earlier context (`--suffix-draft`) are guessed past as well, which is where edits that copy text gain
+most. `--pipeline-windows 1` overlaps only the short prompt reads that go through the verify windows
+(`--short-read`).
+
+Two cards, exactly two stages, `--serve`. In the config:
+
+```
+"args": [ ..., "--pipeline-windows", "2" ],
+"layer_split": "20"
+```
+
+- **Cost**: a second verify window on each card, 160 MiB more kept out of each card's expert cache, plus two copies
+  of the first card's recurrent state (about 3 MiB per GDN layer it runs) on the first card with `2`.
+- **Same text**: the last card only ever runs windows that are verified, and every window row computes what it
+  would in any other window, so the tokens are the serial loop's. With `STRATA_IQ_MT_MIN=1 --pcie-frac 0
+  --adapt-every 0` the greedy output is identical bit for bit to the serial loop's with the same expert caches. The pipeline keeps
+  its VRAM out of the caches, so against a run without the flag a few experts move from a card to the CPU, which
+  rounds them differently, and a near-tie can flip (a serial run given the same caches through `--vram-reserve-mib`
+  matches it exactly).
+- **Off, with one line in the log saying why**, with `--batch` slots, `--peer-device`, the helper caches
+  (`--expert-cache-device1..3`, `--remote-expert-opt`), a split into three or more stages or onto one GPU
+  (`--split-device 0`), or no draft layer. A request with repetition penalties (`penalty_last_n`) or coupled
+  draft sampling decodes serially.
+- **Measured** (Swift 1.5 IQ3_XXS, 160K context, q4_0 KV, the stock draft layer, RTX 4060 Ti (layers 0-19) +
+  RTX 5080 (20-47), i9-14900KF, 32 GB of RAM with the resident RAM mode on the split (#848); greedy, 500 tokens, two
+  interleaved pairs of three rounds, decode tok/s): Python code 90.4 -> 103.1, C code 76.6 -> 81.6, English prose
+  63.0 -> 71.1, Italian prose 41.6 -> 46.3, a copy-heavy edit (a 5 KB file back with a rename) 90.5 -> 117.9; mean
+  72.4 -> 84.0 (+16%). It pays when the windows are GPU-bound: with the experts read through the OS file cache
+  (`--mmap-experts` on that 32 GB PC) the file reads dominate and it measured no faster.
+
 ## Several conversations at once
 
 With a layer split, `--batch N --batch-groups G --trim-stage-weights` decodes several conversations together and

include/strata/core/expert_source.hpp

@@ -483,6 +483,9 @@ class FileExpertSource : public ExpertSource {
     uint8_t* exchange_buffer(int64_t q) const;
     /// Requires `has_resident(layer, in)`, `!has_resident(layer, out)` and `exchange_buffer(q)` holding out's blob.
     bool stage_exchange(int64_t layer, int64_t in, int64_t out, int64_t q);
+    /// --pipeline-windows: size the exchange table now, so a `stage_exchange` on the adaptive tier's thread never
+    /// reallocates it under a concurrent `blob` (the pool reads it while windows are in flight).
+    void prepare_overrides() { if (override_.empty()) override_.assign((size_t) blobs_, nullptr); }
     /// After the GPU copies of every staged swap have landed.  Returns how many exchanges were applied.
     int64_t commit_exchanges();
     int64_t exchanges() const { return exchanges_; }

include/strata/core/mtp.hpp

@@ -48,6 +48,7 @@ class MtpDrafter {
     void set_prompt_len(int64_t n) { prompt_len_ = n; }
     /// At most this many drafts per round (below max_t - 1): a window longer than the MTP's comes from elsewhere.
     void set_max_drafts(int k) { max_drafts_ = k; }
+    int max_t() const { return max_t_; }
     uint64_t vram_bytes() const { return vram_; }
     /// The draft layer's K/V state (read-only: --serve's STRATA_STATE_HASH check hashes it)
     const QsaState& kv_state() const { return st_; }
@@ -61,8 +62,14 @@ class MtpDrafter {
     bool bind(const WeightTable& wt, const NativeHead* head, const float* window_R, std::string& err);
 
     /// Prompt cells [cell0, cell0 + n): residual rows `R_rows` (device, hc*n_embd each) and `next_tokens` (host,
-    /// the token at position cell+1).  Runs in batches of up to max_t rows.
-    bool prefill(const float* R_rows, const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err);
+    /// the token at position cell+1).  Runs in batches of up to max_t rows.  `sync` false: returns without waiting
+    /// for the drafter's stream (the caller orders on it: `stream()`).
+    bool prefill(const float* R_rows, const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err,
+                 bool sync = true);
+    /// --pipeline-windows: capture the prefill graphs and size its input records for windows of up to max_t rows now,
+    /// so a later prefill(sync = false) of a window neither captures nor allocates (both wait for the device).
+    bool prepare_prefill(std::string& err);
+    cudaStream_t stream() const { return cs_; }
 
     /// One round: catch-up over T cells from `p` (rows = the window's final residuals, `tokens` = the window's
     /// argmaxes: row t pairs R_{p+t} with the token at p+t+1), then the draft chain from row `a` (the last
@@ -85,6 +92,31 @@ class MtpDrafter {
     void set_ple_session(SessionState* s) { ple_ss_ = s; }
     bool coupled() const { return coupled_active_; }
 
+    // ---- --pipeline-windows 2: the chain as one asynchronous launch (the round and its steps back to back, no host
+    // wait), its first `n_force` steps fed the given tokens instead of their own picks (teacher forcing: the drafts of
+    // a window already in flight), each output read once the event after its step has completed.
+    /// Before `load`: the round/step graphs get the forcing kernel (a no-op while its token is -1), and the drafter's
+    /// stream the highest priority (STRATA_MTP_PRIORITY=0: the default priority).
+    void set_force_capture(bool on) { force_on_ = on; }
+    /// The rows the next round reads (null: the bound `window_R`), copied into the bound buffer before the round.
+    void set_source_R(const float* rows) { src_R_ = rows; }
+    /// Capture the round graphs for T = 1..max_t and every step graph now (a capture syncs the drafter's stream).
+    bool prepare_chain(std::string& err);
+    /// The catch-up over T cells from `p` (as draft()), then `n_out` outputs from row `a`: output j is the chain's
+    /// own pick, and step j (1 <= j <= n_force) takes force[j - 1] as its input token instead of output j - 1.
+    /// `n_early`: an event after each of the first `n_early` outputs (chain_outputs_ready).
+    bool chain_launch(int T, const int32_t* tokens, int64_t p, int a, const int32_t* force, int n_force, int n_out,
+                      int n_early, std::string& err);
+    /// 1: every output is in chain_tok()/chain_prob(); 0: still running; -1: an error.
+    int chain_poll(std::string& err);
+    /// How many of the chain's first outputs (up to `n_early`; all of them once the chain is done) have landed in
+    /// chain_tok()/chain_prob(); -1: an error.  A window's size is often decided by its first outputs (a draft below
+    /// spec_min_p ends it), so the next launch need not wait for all of them.
+    int chain_outputs_ready(std::string& err);
+    bool chain_live() const { return chain_live_; }
+    const int32_t* chain_tok() const { return chain_tok_; }
+    const float* chain_prob() const { return chain_prob_; }
+
     double ms_draft = 0, ms_prefill = 0;
     int64_t rounds = 0;
 
@@ -132,6 +164,16 @@ class MtpDrafter {
     int max_t_ = 0;
     int device_ = -1;   ///< the device `load` ran on: the public calls switch to it (layer split)
     int max_drafts_ = 1 << 30;
+    // --pipeline-windows 2 (chain_launch)
+    bool stage_source_R(int T, std::string& err);   ///< set_source_R's rows into the bound buffer
+    bool force_on_ = false, chain_live_ = false;
+    const float* src_R_ = nullptr;
+    int32_t *h_force_ = nullptr, *m_force_ = nullptr;   ///< the forced tokens (mapped), -1 = the step's own pick
+    cudaEvent_t ev_chain_ = nullptr;
+    cudaEvent_t ev_step_[8] = {};                       ///< after each of the first n_early outputs
+    int steps_seen_ = 0, chain_n_ = 0, chain_early_ = 0;
+    int32_t chain_tok_[8] = {};
+    float chain_prob_[8] = {};
     int64_t n_vocab_ = 0;
     uint64_t vram_ = 0;
     cudaStream_t cs_ = nullptr;

include/strata/core/progress.hpp

@@ -39,6 +39,8 @@ inline std::atomic<DiagFn>& release_gpu_fn() { static std::atomic<DiagFn> f{null
 inline void release_gpu_waits(std::FILE* f) {
     if (auto fn = release_gpu_fn().load()) fn(f);
 }
+/// --pipeline-windows: the pipelined loop's windows and the drafter's chain (set while that loop runs)
+inline std::atomic<DiagFn>& diag_pipeline_fn() { static std::atomic<DiagFn> f{nullptr}; return f; }
 
 inline Progress& progress() {
     static Progress p;

include/strata/core/verify.hpp

@@ -180,6 +180,45 @@ class Verifier {
     /// false with `err` when it failed.  Free when nothing is pending.
     bool wait_commit(std::string& err);
 
+    // ---- PIPELINED WINDOWS (--pipeline-windows, a layer split on two GPUs).  One conversation's windows with the
+    // stages overlapped: stage 0 runs window K+1 while stage 1 still runs window K.  The same window as `run`, driven
+    // without blocking the host, so one host thread keeps a window in flight on each stage (the batch pipeline's
+    // pattern, batch_launch / batch_poll, for one sequence with drafts).  Two verifiers per stage (one per window
+    // parity) share the stage's stream and each has its own hand-off.  `pl_launch` stages and launches (it never
+    // captures: `capture_all` first, with nothing in flight), `service` serves the layers whose doorbells have rung
+    // and returns at once, `done` polls the window's completion, `pl_finish` reads the picks, `pl_commit_async`
+    // queues the commit.  Nothing chains to `next_`: the caller drives every stage.
+    /// The compute stream to use instead of a private one (the two verifiers of one stage share it).  Before `init`.
+    void set_stream(cudaStream_t s) { ext_stream_ = s; }
+    /// Every layer copies the token rows to the host (doorbell_publish), not only the layers with a routed expert
+    /// outside the VRAM tier by the device's residency table: with windows in flight the adaptive tier marks an
+    /// expert evicted on the host (the pool then computes it on the CPU) before the device table follows.  Also turns
+    /// the device-planned layers (E-6) off.  Before `init`.
+    void set_always_publish(bool on) { always_publish_ = on; }
+    cudaStream_t stream() const { return cs_; }
+    int device() const { return device_; }
+    /// Capture every window size and the commit graph now (a capture syncs the stream: never with a window in flight).
+    bool capture_all(std::string& err);
+    bool pl_launch(int T, const int32_t* tokens, int64_t pos0, std::string& err);
+    /// Stage a window ahead of its launch (positions, and the PLE rows from `ple_prev` = the two tokens before the
+    /// window as they WILL be, their pages prefetched).  A later `pl_launch` of the same window (T, pos0, tokens, and
+    /// `ss.ple_prev` equal to `ple_prev` by then) skips the staging; anything else stages again.
+    bool prestage(int T, const int32_t* tokens, int64_t pos0, const int32_t ple_prev[2], std::string& err);
+    /// 1: every layer served; 0: the GPU has not reached the next layer yet; -1: an error (`err`).
+    int service(PoolMultiFn pool, void* user, std::string& err);
+    bool in_flight() const { return fl_active_; }
+    /// The window's graph (and its profile copy) completed; false while it runs.  An error sets `err`.
+    bool done(std::string& err);
+    /// After `done`: the profile, the last stage's host sampling and picks (`out` may be null on an earlier stage).
+    bool pl_finish(int32_t* out, std::string& err);
+    /// The commit without a host sync; `ss.ple_prev` advances now (host side).  A second call for the same window
+    /// (after its state was restored) replays it with another count.
+    bool pl_commit_async(int n_keep, std::string& err);
+    /// Fold another verifier's counters and GPU profile into this one's (the two verifiers of one stage report once).
+    void absorb_stats(Verifier& o);
+    /// The watchdog's line for a pipelined verifier: in flight, layers served, the GPU's ring and flags, its events.
+    void diag_pipelined(std::FILE* f, const char* name) const;
+
     /// Measurement hook (STRATA_LOGPOS): after run(), write one line per row t of the last window's head -
     /// "pos target logprob top top_logprob hit extra_logprob target_logprob_without_extra" - where row t is the
     /// distribution at pos0 + t, targets[t] is the token at pos0 + t + 1, extra_logprob is the log-probability of
@@ -242,6 +281,19 @@ class Verifier {
     int64_t last_pos_b_[8] = {};
     bool capture_batch(const int* rows, int S, int hbase, std::string& err);
     void collect_profile();   ///< STRATA_VERIFY_PROFILE: add the last window's stamps to prof_sum_
+    void accumulate_profile(const unsigned long long* stamps);   ///< one window's stamps (host copy) into prof_sum_
+    // pipelined windows (pl_launch ...)
+    cudaStream_t ext_stream_ = nullptr;   ///< set_stream: the stage's shared stream (not destroyed here)
+    bool always_publish_ = false;
+    void pl_stage(int T, const int32_t* tokens, int64_t pos0, const int32_t ple_prev[2]);
+    cudaEvent_t ev_done_ = nullptr, ev_commit_ = nullptr;
+    unsigned long long* prof_pin_ = nullptr;   ///< pinned host copy of the stamps (pipelined windows)
+    bool fl_active_ = false, fl_prof_ = false, fl_ple_ = false, commit_live_ = false, pl_prestaged_ = false;
+    int fl_T_ = 0;
+    int64_t fl_k_ = 0, fl_total_ = 0;
+    double fl_since_ms_ = 0, fl_flush_ms_ = 0, fl_launch_ms_ = 0;
+    int32_t pl_prev_[2] = {-1, -1};
+    std::vector<uint32_t> pl_ple_rows_;   ///< the window's PLE rows (T x PLE_N_HEADS), gathered when layer 0 is served
     bool capture_commit_batch(const int* rows, int S, int hbase, std::string& err);
     bool stage_batch(const int* rows, int S, int hbase, const int32_t* tokens, const int64_t* pos, std::string& err);
     strata::kernels::SamplerParams sampling_ = [] {

include/strata/kernels/verify_kernels.hpp

@@ -83,6 +83,9 @@ void ident_hits(const int32_t* ids, int n, int32_t* slot, int32_t* dst, int32_t*
 void mtp_select(const float* R_src, int64_t R_stride, const int32_t* ids, const int32_t* row_dev, float* R_dst,
                 int32_t* tok_dst, int32_t* out, int j, void* stream, const float* probs = nullptr,
                 float* out_p = nullptr);
+/// --pipeline-windows, the drafter's chain teacher forced: `*tok = force[j]` when force[j] >= 0 (`force` is mapped
+/// host memory, read when the kernel runs), else `*tok` is left as it is.
+void force_token(int32_t* tok, const int32_t* force, int j, void* stream);
 /// dst row i = src row ids[i] (row_bytes each, multiple of 16), for n rows.
 void gather_rows(const uint8_t* src, int64_t row_bytes, const int32_t* ids, int64_t n, uint8_t* dst, void* stream);
 /// ids[t] = table[ids[t]] for n entries (a subset index back to a token id).