peer tier: STRATA_ADAPT_OVERLAP=1 with --peer-device (the peer's copies land while the next window runs)

In peer mode the adaptive overlap was refused, so every decode window first waited for the previous round's copies. With the overlap the peer
tier issues its copies in layer order, signals each layer's landing the way the primary does (a stream-ordered D2H copy of the layer number into
a pinned int, or host functions), its residency table takes the swaps at the window's start (flip_pending), and a layer waits for a tier's copies
only when that tier swapped in it (AdaptGate: peer_landed, peer_has, per-tier last). The wait before a window clears the peer's marks with the
primary's. Without the variable nothing changes.

Ported from #2720 (bulhenvaster), rebased by hand onto the current gate (the watchdog label stays). Not ported: STRATA_PEER_LAUNCH_IN_POOL
(the author measured no gain on his build), the new decode timing line and the documentation paragraph.

(cherry picked from commit e78b64a7348d685dfd11799fdb395661f2120436)

include/strata/core/peer_experts.hpp

@@ -20,6 +20,7 @@
 
 #include <cuda_runtime.h>
 
+#include <atomic>
 #include <cstdint>
 #include <string>
 #include <utility>
@@ -72,15 +73,32 @@ public:
     /// swapped-in expert becomes resident at `apply_pending`.  Call only between verify windows.
     bool adapt(const float* usage, const int32_t* res0, int max_swaps, std::string& err);
     void apply_pending(bool wait);
+    /// STRATA_ADAPT_OVERLAP=1 with a peer: adapt() then issues its copies in layer order and raises `*landed` to each
+    /// layer as that layer's copies land (host functions on the refill stream), marking `(*has)[layer]` and `*last`
+    /// first; the verify window holds such a layer's plan back until then (AdaptGate).  Call once after open().
+    /// `copy_signal`: the landing is told by a stream-ordered D2H copy of the layer number into a pinned int instead
+    /// (no callback thread to schedule beside the pool's spinning workers); that int is returned, null when host
+    /// functions are used.
+    volatile int32_t* set_overlap_gate(std::atomic<int32_t>* landed, std::vector<uint8_t>* has, int32_t* last,
+                                       bool copy_signal);
+    /// The pending swaps enter the residency table now, their copies still landing (the gate protects every read);
+    /// apply_pending(true) after the window still waits for them and clears the list.
+    void flip_pending();
+    bool has_pending() const { return !pending_.empty(); }
 
     int64_t resident() const { return resident_; }
     /// the last launch writes its rows into the host rows itself (the caller must not touch them)
     bool launched_direct() const { return launched_direct_; }
+    /// whether a launch given host rows writes them itself (STRATA_PEER_DIRECT, default on): then the caller need not
+    /// zero a peer row at all, whether the launch has happened yet or not
+    static bool direct_mode();
     double gib() const { return cache_.gib(); }
     int64_t entries() const { return entries_; }      ///< routed entries the peer computed
     int64_t experts() const { return experts_; }      ///< distinct (layer, expert) the peer computed
     int64_t swaps() const { return swaps_; }
     double ms_wait = 0;                               ///< host time spent in finish() waiting for the peer
+    double ms_launch = 0;                             ///< host time spent in launch() (grouping, copies, kernel launches)
+    int64_t launches = 0;                             ///< layer launches that gave the peer at least one row
 
 private:
     int device_ = -1;
@@ -90,6 +108,13 @@ private:
     ExpertSource* src_ = nullptr;
     std::vector<int32_t> res_;                        ///< [n_layers * n_expert] -> peer slot or -1
     std::vector<std::pair<int32_t, int32_t>> pending_;   ///< (residency index, slot) once the copies land
+    struct Tick { std::atomic<int32_t>* counter; int32_t layer; };
+    std::atomic<int32_t>* gate_landed_ = nullptr;        ///< set_overlap_gate: the counter, per-layer marks and the last layer
+    std::vector<uint8_t>* gate_has_ = nullptr;
+    int32_t* gate_last_ = nullptr;
+    std::vector<Tick> ticks_;                            ///< one per layer (stable storage for the host functions)
+    int32_t* d_gate_vals_ = nullptr;                     ///< device [n_layers]: the layer numbers, the source of the signal copies
+    volatile int32_t* h_gate_ = nullptr;                 ///< pinned int those copies write (null: host functions)
     int64_t resident_ = 0, entries_ = 0, experts_ = 0, swaps_ = 0;
     cudaStream_t stream_ = nullptr, refill_ = nullptr;
     cudaEvent_t refill_ev_ = nullptr;

include/strata/core/verify.hpp

@@ -48,10 +48,11 @@ namespace strata::core {
 /// bytes arrive layer by layer: `run` holds a layer's plan (the first thing that lets the GPU read the layer's expert
 /// slots) back until `landed >= layer` for a layer with a swap (`has[layer]`).  Set `active` only around one window.
 struct AdaptGate {
-    struct Tick { AdaptGate* gate; int32_t layer; };
+    struct Tick { std::atomic<int32_t>* counter; int32_t layer; };
     std::atomic<int32_t> landed{-1};      ///< the highest layer whose swap copies have landed (n_layers: all of them)
     std::atomic<bool> active{false};
     std::vector<uint8_t> has;             ///< per layer: a swap copies into this layer's slots
+    int32_t last = -1;                    ///< the highest layer with a swap this round (-1: none)
     std::vector<Tick> ticks;              ///< host-function contexts, one per layer and the final one (stable storage)
     const int32_t* dvals = nullptr;       ///< device [0..n_layers]: the source of those copies
     volatile int32_t* hland = nullptr;   ///< pinned host int the copy stream writes (a D2H copy of the layer number), instead of a host function; null: host functions
@@ -61,6 +62,19 @@ struct AdaptGate {
         const int32_t b = *hland;
         return a > b ? a : b;
     }
+    /// The peer tier's copies (--peer-device with STRATA_ADAPT_OVERLAP=1): its own counter, raised on its refill stream
+    /// in layer order the same two ways (`peer_hland`: the pinned int its D2H copies write; null: host functions); a
+    /// layer with a peer swap waits for it as well.
+    std::atomic<int32_t> peer_landed{-1};
+    std::vector<uint8_t> peer_has;
+    int32_t peer_last = -1;
+    volatile int32_t* peer_hland = nullptr;
+    int32_t peer_reached() const {
+        const int32_t a = peer_landed.load(std::memory_order_acquire);
+        if (peer_hland == nullptr) return a;
+        const int32_t b = *peer_hland;
+        return a > b ? a : b;
+    }
 };
 
 class NativeHead;

src/core/peer_experts.cpp

@@ -109,6 +109,8 @@ void PeerExperts::close() {
         if (h_x_) cudaFreeHost(h_x_);
         if (h_out_) cudaFreeHost(h_out_);
         if (h_meta_) cudaFreeHost(h_meta_);
+        if (d_gate_vals_) cudaFree(d_gate_vals_);
+        if (h_gate_) cudaFreeHost((void*) h_gate_);
         if (refill_ev_) cudaEventDestroy(refill_ev_);
         if (stream_) cudaStreamDestroy(stream_);
         if (refill_) cudaStreamDestroy(refill_);
@@ -117,6 +119,8 @@ void PeerExperts::close() {
     d_meta_ = h_meta_ = d_scratch_ = nullptr;
     d_q8_ = nullptr;
     d_sx_ = d_sh_ = nullptr;
+    d_gate_vals_ = nullptr;
+    h_gate_ = nullptr;
     stream_ = refill_ = nullptr;
     refill_ev_ = nullptr;
     device_ = -1;
@@ -242,10 +246,19 @@ bool PeerExperts::open(int device, const std::vector<std::pair<int32_t, int32_t>
     return true;
 }
 
+bool PeerExperts::direct_mode() {
+    static const bool v = [] { const char* e = std::getenv("STRATA_PEER_DIRECT"); return e == nullptr || std::atoi(e) != 0; }();
+    return v;
+}
+
 bool PeerExperts::launch(int64_t layer, const float* x, const int32_t* ids, int64_t n_tok, int64_t k,
                          const int32_t* kind, std::string& err, float* out) {
-    static const bool direct_env = [] { const char* v = std::getenv("STRATA_PEER_DIRECT"); return v == nullptr || std::atoi(v) != 0; }();
-    const bool direct = direct_env && out != nullptr;
+    const bool direct = direct_mode() && out != nullptr;
+    struct Timer {
+        double& acc;
+        std::chrono::steady_clock::time_point t0 = std::chrono::steady_clock::now();
+        ~Timer() { acc += std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t0).count(); }
+    } timer{ms_launch};
     launched_direct_ = false;
     launched_rows_ = 0;
     row_of_.clear();
@@ -316,6 +329,7 @@ bool PeerExperts::launch(int64_t layer, const float* x, const int32_t* ids, int6
 
 bool PeerExperts::finish(float* out, std::string& err) {
     if (launched_rows_ == 0) return true;
+    ++launches;
     const auto t0 = std::chrono::steady_clock::now();
     {
         On on(device_);
@@ -362,20 +376,85 @@ bool PeerExperts::adapt(const float* usage, const int32_t* res0, int max_swaps,
     if (swaps.empty()) return true;
     std::sort(swaps.begin(), swaps.end(), [](const Swap& a, const Swap& b) { return a.gain > b.gain; });
     if ((int) swaps.size() > max_swaps) swaps.resize((size_t) max_swaps);
+    const bool gated = gate_landed_ != nullptr;
+    if (gated) {   // the copies go in layer order and each layer's landing raises the gate's counter
+        std::stable_sort(swaps.begin(), swaps.end(), [](const Swap& a, const Swap& b) { return a.layer < b.layer; });
+        gate_landed_->store(-1, std::memory_order_release);
+        if (h_gate_ != nullptr) *h_gate_ = -1;
+        gate_has_->assign((size_t) n_layers_, 0);
+        *gate_last_ = -1;
+    }
     On on(device_);
+    int32_t issued_layer = -1;   // the layer whose copies are on the stream and not yet ticked
+    auto tick = [&](int32_t layer) -> bool {
+        if ((h_gate_ != nullptr
+                 ? cudaMemcpyAsync((void*) h_gate_, d_gate_vals_ + layer, 4, cudaMemcpyDeviceToHost, refill_)
+                 : cudaLaunchHostFunc(refill_, [](void* p) {
+                auto* t = (Tick*) p;
+                int32_t cur = t->counter->load(std::memory_order_relaxed);
+                while (cur < t->layer && !t->counter->compare_exchange_weak(cur, t->layer, std::memory_order_release)) {}
+            }, &ticks_[(size_t) layer])) == cudaSuccess)
+            return true;
+        (void) cudaGetLastError();
+        cudaStreamSynchronize(refill_);   // no callback: wait here, the gate opens at once
+        gate_landed_->store((int32_t) n_layers_, std::memory_order_release);
+        return false;
+    };
     for (const Swap& s : swaps) {
         const size_t in = (size_t) (s.layer * n_expert_ + s.in), out = (size_t) (s.layer * n_expert_ + s.out);
         const int32_t slot = res_[out];
         const uint8_t* b = src_->blob(s.layer, s.in);
         if (slot < 0 || b == nullptr) continue;
+        if (gated && issued_layer >= 0 && issued_layer != s.layer) tick(issued_layer);
         if (!ck(cudaMemcpyAsync(cache_.device_slot(slot), b, (size_t) lay.blob_bytes(s.layer), cudaMemcpyHostToDevice,
                                 refill_), "refill", err))
             return false;
+        if (gated) {
+            (*gate_has_)[(size_t) s.layer] = 1;
+            *gate_last_ = std::max(*gate_last_, s.layer);
+            issued_layer = s.layer;
+        }
         res_[out] = kNotResident;              // evicted now: the CPU computes it meanwhile
         pending_.emplace_back((int32_t) in, slot);
         ++swaps_;
     }
-    return ck(cudaEventRecord(refill_ev_, refill_), "refill event", err);
+    if (gated && issued_layer >= 0) tick(issued_layer);
+    if (!ck(cudaEventRecord(refill_ev_, refill_), "refill event", err)) return false;
+    // #2188: hand the copies and their landing signals to the GPU now (WDDM batches a stream's work until a driver
+    // call flushes it, and the verify window only spins on the pinned int; Linux submits at once: a no-op)
+    if (gated && h_gate_ != nullptr && issued_layer >= 0) (void) cudaStreamQuery(refill_);
+    return true;
+}
+
+volatile int32_t* PeerExperts::set_overlap_gate(std::atomic<int32_t>* landed, std::vector<uint8_t>* has, int32_t* last,
+                                                bool copy_signal) {
+    gate_landed_ = landed;
+    gate_has_ = has;
+    gate_last_ = last;
+    ticks_.resize((size_t) n_layers_);
+    for (int64_t l = 0; l < n_layers_; ++l) ticks_[(size_t) l] = {landed, (int32_t) l};
+    if (copy_signal && h_gate_ == nullptr) {
+        On on(device_);
+        std::vector<int32_t> iota((size_t) n_layers_);
+        for (size_t i = 0; i < iota.size(); ++i) iota[i] = (int32_t) i;
+        int32_t* vals = nullptr;
+        void* hp = nullptr;
+        if (cudaMalloc((void**) &vals, iota.size() * 4) == cudaSuccess &&
+            cudaMemcpy(vals, iota.data(), iota.size() * 4, cudaMemcpyHostToDevice) == cudaSuccess &&
+            cudaHostAlloc(&hp, 64, cudaHostAllocPortable) == cudaSuccess) {
+            *(volatile int32_t*) hp = -1;
+            d_gate_vals_ = vals;
+            h_gate_ = (volatile int32_t*) hp;
+        } else {   // host functions then
+            (void) cudaGetLastError();
+            if (vals) cudaFree(vals);
+        }
+    }
+    return h_gate_;
+}
+
+void PeerExperts::flip_pending() {
+    for (const auto& [i, slot] : pending_) res_[(size_t) i] = slot;
 }
 
 void PeerExperts::apply_pending(bool wait) {

src/core/verify.cpp

@@ -3148,8 +3148,13 @@ void Verifier::accumulate_profile_inner(const unsigned long long* stamps) {
 bool Verifier::adapt_gate_wait(int32_t layer, std::string& err) {
     const Clock::time_point t0 = Clock::now();
     uint32_t spins = 0;
+    // each tier's counter climbs to the last layer it swaps this round, so an all-layers wait asks for that layer
+    const AdaptGate& ag = *adapt_gate_;
+    const int32_t want = std::min(layer, ag.last), want_peer = std::min(layer, ag.peer_last);
+    const bool wait_main = ag.last >= 0 && ((size_t) layer >= ag.has.size() || ag.has[(size_t) layer] != 0);
+    const bool wait_peer = ag.peer_last >= 0 && ((size_t) layer >= ag.peer_has.size() || ag.peer_has[(size_t) layer] != 0);
     bool announced = false;   // the watchdog's label is set only once the wait really starts: no cost when the copies landed
-    while (adapt_gate_->reached() < layer) {
+    while ((wait_main && ag.reached() < want) || (wait_peer && ag.peer_reached() < want_peer)) {
         if (!announced) {
             progress_at("verify window: waiting for the adaptive tier's copies of layer", layer);
             announced = true;
@@ -3478,7 +3483,8 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool,
         if (g_trace) trace_ev("RANG", k, l, (int64_t) std::chrono::duration_cast<std::chrono::microseconds>(b - a).count());
         VDBG("layer %lld rang\n", (long long) l);
         if (adapt_gate_ != nullptr && adapt_gate_->active.load(std::memory_order_acquire) && l >= 0 &&
-            (size_t) l < adapt_gate_->has.size() && adapt_gate_->has[(size_t) l] != 0 &&
+            (((size_t) l < adapt_gate_->has.size() && adapt_gate_->has[(size_t) l] != 0) ||
+             ((size_t) l < adapt_gate_->peer_has.size() && adapt_gate_->peer_has[(size_t) l] != 0)) &&
             !adapt_gate_wait((int32_t) l, err)) return false;
         cur_layer_ = want - 1;
         set_plan_slot(grp);

src/program/generate.cpp

@@ -12809,9 +12809,10 @@ int main(int argc, char** argv) {
         int64_t aov_flips = 0;
         // #2010: the overlap is opt-in on every backend now, so the off state is the default and has to be visible too.
         // The predicates are side-effect free, so they are read in both states: they name why it could not run.
+        // A peer tier (--peer-device) joins the overlap: measured on one rig so far (RTX 5060 Ti x8 + RTX 3060, IQ3_XXS),
+        // where the blocking wait took 3-8 ms of a 27-43 ms window.
         const bool aov_want = adapt_overlap();
         const char* aov_why = multi_gpu || !stages.empty() ? "not on a layer split"
-                              : peer.valid() ? "not with --peer-device"
                               : remote_opt ? "not with --remote-expert-opt"
                               : o.adapt_async ? "not with --adapt-async"
                               : o.pipeline_windows > 0 ? "not with --pipeline-windows"
@@ -12828,8 +12829,9 @@ int main(int argc, char** argv) {
         if (aov_want) {
             if (aov_why == nullptr) {
                 agate.ticks.resize((size_t) g.n_layers + 1);
-                for (int64_t l = 0; l <= g.n_layers; ++l) agate.ticks[(size_t) l] = {&agate, (int32_t) l};
+                for (int64_t l = 0; l <= g.n_layers; ++l) agate.ticks[(size_t) l] = {&agate.landed, (int32_t) l};
                 agate.has.assign((size_t) g.n_layers, 0);
+                agate.peer_has.assign((size_t) g.n_layers, 0);
                 if (!ver.set_adapt_gate(&agate)) aov_why = "the device-planned layers (STRATA_VERIFY_DEVICE_PLAN) are on";
                 // #2188: on Windows the default is the host function.  Under WDDM a stream's work sits in the driver's
                 // command buffer until something flushes it, and nothing does while the host only spins on the pinned
@@ -12859,11 +12861,16 @@ int main(int argc, char** argv) {
                         std::fprintf(stderr, "strata serve: adaptive overlap signal: stream-ordered copies (STRATA_ADAPT_SIGNAL=host: a host function)\n");
                     } else { (void) cudaGetLastError(); }
                 }
+                // the peer tier's copies land over its own link and are gated the same way, with the same kind of signal
+                if (aov_why == nullptr && peer.valid())
+                    agate.peer_hland = peer.set_overlap_gate(&agate.peer_landed, &agate.peer_has, &agate.peer_last, sig_copy);
             }
             aov = aov_why == nullptr;
             if (aov)
                 std::fprintf(stderr, "strata serve: adaptive overlap (STRATA_ADAPT_OVERLAP=0 turns it off): the adaptive tier's copies land while the next window runs "
-                                     "(same tokens, a layer waits for its own copies)\n");
+                                     "(same tokens, a layer waits for its own copies)%s\n",
+                             !peer.valid() ? "" : agate.peer_hland != nullptr ? "; the peer tier's copies too, over its own link (stream-ordered copies)"
+                                                                              : "; the peer tier's copies too, over its own link (host functions)");
             else if (std::getenv("STRATA_ADAPT_OVERLAP") != nullptr)   // the env var was set: say what stopped it
                 std::fprintf(stderr, "strata serve: STRATA_ADAPT_OVERLAP=1 is off: %s\n", aov_why);
         } else
@@ -12881,7 +12888,15 @@ int main(int argc, char** argv) {
         };
         auto apply_pending = [&](bool wait) {
             if (peer.valid()) peer.apply_pending(wait);
-            if (pending.empty()) return;
+            if (pending.empty()) {
+                // a round where only the peer swapped: its copies have landed (or, without the wait, are still in
+                // flight), and the gate closes with them
+                if (!(peer.valid() && peer.has_pending())) {
+                    agate.active.store(false, std::memory_order_release);
+                    aflipped = false;
+                }
+                return;
+            }
             if (wait) cudaEventSynchronize(adapt_ev);
             else if (cudaEventQuery(adapt_ev) != cudaSuccess) return;
             for (auto& st : stages)
@@ -13051,7 +13066,12 @@ int main(int argc, char** argv) {
                 agate.landed.store(-1, std::memory_order_release);
                 if (agate.hland != nullptr) *agate.hland = -1;
                 agate.has.assign((size_t) g.n_layers, 0);
-                for (const Swap& s : swaps) agate.has[(size_t) s.layer] = 1;
+                agate.last = -1;
+                for (const Swap& s : swaps) { agate.has[(size_t) s.layer] = 1; agate.last = std::max(agate.last, s.layer); }
+                // the peer's marks of the previous round go too (its adapt below sets this round's, if it swaps anything)
+                agate.peer_landed.store(-1, std::memory_order_release);
+                agate.peer_has.assign((size_t) g.n_layers, 0);
+                agate.peer_last = -1;
             }
             // the copies back read the cache of the card that owns each layer (see resident_stage_swaps)
             auto home_of = [&](int64_t layer) -> SwapHome {
@@ -13137,8 +13157,8 @@ int main(int argc, char** argv) {
                          ? cudaMemcpyAsync((void*) agate.hland, agate.dvals + s.layer, 4, cudaMemcpyDeviceToHost, adapt_stream)
                          : cudaLaunchHostFunc(adapt_stream, [](void* p) {
                         auto* t = (strata::core::AdaptGate::Tick*) p;
-                        int32_t cur = t->gate->landed.load(std::memory_order_relaxed);
-                        while (cur < t->layer && !t->gate->landed.compare_exchange_weak(cur, t->layer, std::memory_order_release)) {}
+                        int32_t cur = t->counter->load(std::memory_order_relaxed);
+                        while (cur < t->layer && !t->counter->compare_exchange_weak(cur, t->layer, std::memory_order_release)) {}
                     }, &agate.ticks[(size_t) s.layer])) != cudaSuccess) {
                     (void) cudaGetLastError();
                     cudaStreamSynchronize(adapt_stream);   // no callback: wait here, the gate opens at once
@@ -18932,7 +18952,7 @@ int main(int argc, char** argv) {
                 const bool asw_skip = first_window || merged;   // the session's tier choice counts steady decode windows only
                 double asw_wait_ms = 0;
                 if (adapt_nowait()) apply_pending(false);
-                else if (aov && !pending.empty() && !aflipped) {
+                else if (aov && (!pending.empty() || (peer.valid() && peer.has_pending())) && !aflipped) {
                     // 0.1.43 STRATA_ADAPT_OVERLAP: the tables take the swaps now, exactly where the wait would have ended;
                     // the copies still landing are waited for per layer inside the window (AdaptGate)
                     // (resident RAM mode: the staged evicted experts look committed to the planner at once, the RAM
@@ -18941,6 +18961,7 @@ int main(int argc, char** argv) {
                         apply_pending(true);   // not possible here: the plain wait
                     } else {
                         for (const auto& [i, slot] : pending) host_res[(size_t) i] = slot;
+                        if (peer.valid()) peer.flip_pending();   // its copies land over its own link, gated the same way
                         res_upload();
                         aflipped = true;
                         ++aov_flips;