multi-GPU: second GPU as an opt-in expert-cache tier (--peer-device), resubmit of #229 on v0.1.36

CMakeLists.txt

@@ -340,7 +340,7 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
   # about arithmetic, and this one knows the ORDER the arithmetic happens in.
   add_library(strata_engine STATIC src/core/layer.cpp src/core/session.cpp src/core/expert_source.cpp src/core/remote_experts.cpp
     src/core/expert_cache.cpp src/core/native_head.cpp src/core/native_dense.cpp src/core/verify.cpp src/core/mtp.cpp
-    src/core/conversation_snapshot.cpp src/core/conversation_state.cpp src/core/conversation_memory.cpp)
+    src/core/conversation_snapshot.cpp src/core/conversation_state.cpp src/core/conversation_memory.cpp src/core/peer_experts.cpp)
   target_include_directories(strata_engine PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
   target_link_libraries(strata_engine PUBLIC strata_core strata_kernels strata_kernels_cpu ${_strata_gpu_runtime_target})
   if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/core/file_expert_source_test.cpp")

docs/SECOND_GPU.md

@@ -83,3 +83,42 @@ short decode requests.
 The existing warning about the CUDA0 expert-cache GPU hit path still applies:
 its outputs diverge from cache-off runs. Treat performance as experimental
 until the generated tokens have been validated.
+
+## Peer tier (`--peer-device`)
+
+`--peer-device N` puts a second adaptive expert cache on CUDA device N. It
+takes the ranked pairs CUDA0's cache does not hold, as many as fit. The peer
+computes the rows of its own experts, for decode windows and for prompt
+chunks; the activations and the results cross NVLink or another P2P path. The
+tier adapts while the server runs, like the primary cache. It is an
+alternative to the CUDA1-3 caches above, not a third tier beside them.
+
+- `--peer-device N` (default off): enable the tier on CUDA device N (N >= 1).
+- `--peer-reserve-mib M` (default 600): leave M MiB free on the peer card; the
+  cache takes what remains. The prompt-path buffers need this headroom.
+- `--peer-slots N` (default 0): cap the tier at N experts; 0 = as many as fit.
+- `--peer-adapt-swaps N` (default -1): swaps per adaptive round on the peer;
+  -1 uses the primary's `--adapt-swaps`.
+- `--peer-prefill-rows N` (default -1): the share of each prompt chunk's rows
+  the peer computes; -1 is half of chunk x top-k, 0 keeps prompt rows on the
+  primary.
+
+`--peer-device` requires `--expert-profile` and an enabled expert cache, and
+the device must be visible; it refuses otherwise. It also refuses
+`--layer-split` (a different second-GPU mode: use one or the other) and
+`--expert-cache-device1..3` (the peer tier already caches experts on that card):
+
+    strata generate: --peer-device cannot be combined with --layer-split (use one or the other)
+    strata generate: --peer-device cannot be combined with --expert-cache-device1..3 (the peer tier already caches experts there)
+
+Without `--peer-device` the binary is unchanged; its output is byte-identical
+to the release. With `--peer-device` and the same expert set split across the
+two cards, the generated tokens are byte-identical to the single-GPU run under
+the exactness gate.
+
+With a peer the prompt path keeps the MMQ path; the fused int8 prompt path is
+not yet combined with the peer's rows. Mapped host buffers gain
+`cudaHostAllocPortable` only with a peer, since only then does a second
+context write them. The tier size is manual for now (`--peer-reserve-mib`,
+`--peer-slots`); automatic sizing on small cards wants the buffer lending of
+#216 and is a follow-up.

include/strata/core/expert_source.hpp

@@ -43,6 +43,7 @@ struct ExpertLayout;
 
 namespace strata::core {
 
+class PeerExperts;   // multi-GPU: the second GPU's expert tier (peer_experts.hpp)
 class RemoteExperts;
 struct LoadStats;
 
@@ -336,6 +337,9 @@ struct ExpertDispatch {
     std::vector<float> usage;
     int64_t multi_misses = 0;      ///< distinct (layer, expert) pairs the CPU computed in verify windows
     int64_t multi_entries = 0;     ///< routed (token, expert) entries the CPU served in verify windows
+    /// Multi-GPU: the second GPU's tier.  Its experts are computed there instead of on the CPU (kind 2).
+    PeerExperts* peer = nullptr;
+    int64_t peer_entries = 0;      ///< routed entries the peer served in verify windows
     /// Set when `dispatch` could not produce an answer.  The loop itself has no error channel, so this is
     /// where a source failure surfaces: the driver checks it after `session_loop` returns rather than the
     /// engine computing from a half-filled `parts`.

include/strata/core/peer_experts.hpp

@@ -0,0 +1,106 @@
+// include/strata/core/peer_experts.hpp - multi-GPU: a second expert tier on another GPU.
+//
+// The primary GPU (device 0) keeps everything it had: the dense weights, the KV cache, the MTP layer, its own
+// adaptive expert tier and the captured verify graph.  This class adds an expert tier on a SECOND card.  From
+// the verify graph's point of view nothing changes: for each layer the pool thread decides which routed experts
+// the primary computes (its VRAM tier, the PCIe share), and every remaining expert is a "CPU row" the graph reads
+// from mapped host memory after the pool raises its flag.  The peer takes the rows whose experts it holds: the
+// pool thread launches them on the peer right after publishing the primary's plan, the CPU computes the rest at
+// the same time, and the peer's rows are written into the same mapped rows before the pool returns.
+//
+// The peer computes with the SAME kernels as the primary's hit path (q8_1 activations + native_expert_grouped),
+// from the same activation values, so an expert gives bit-identical rows on either card.
+//
+// The tier follows the conversation like the primary's: `adapt` swaps the most-routed experts that neither card
+// holds into the peer's least-used slots (per layer, so a slot keeps its layer's blob size).
+#pragma once
+
+#include "strata/core/expert_cache.hpp"
+#include "strata/core/expert_source.hpp"
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace strata::core {
+
+/// multi-GPU: set once at startup when a peer device is configured (--peer-device N).  The mapped host buffers
+/// the peer card writes into are then allocated with cudaHostAllocPortable; without a peer the flag stays off.
+void set_peer_portable(bool on);
+bool peer_portable();
+
+class PeerExperts {
+public:
+    PeerExperts() = default;
+    ~PeerExperts();
+    PeerExperts(const PeerExperts&) = delete;
+    PeerExperts& operator=(const PeerExperts&) = delete;
+
+    /// Opens the tier on `device`: fills it with the first pairs of `ranked` that `primary` does not hold, as many
+    /// as fit in the card's free memory minus `reserve_mib` (`max_slots` > 0 caps the count).
+    bool open(int device, const std::vector<std::pair<int32_t, int32_t>>& ranked, const ExpertCache& primary,
+              ExpertSource& src, int64_t n_layers, int64_t n_expert, int reserve_mib, int64_t max_slots,
+              std::string& err);
+    void close();
+    bool valid() const { return device_ >= 0; }
+    int device() const { return device_; }
+
+    /// Whether the peer holds (layer, expert) right now.
+    bool has(int64_t layer, int64_t expert) const { return res_[(size_t) (layer * n_expert_ + expert)] >= 0; }
+
+    /// The device address of (layer, expert)'s blob on the peer, or null when it is not resident.
+    const uint8_t* slot_ptr(int64_t layer, int64_t expert) {
+        const int32_t sl = res_[(size_t) (layer * n_expert_ + expert)];
+        return sl >= 0 ? cache_.device_slot(sl) : nullptr;
+    }
+    /// Whether the primary (device 0) and the peer can read each other's memory directly (NVLink / PCIe P2P).
+    bool p2p() const { return p2p_; }
+
+    /// Starts the peer's share of one layer: the entries i with kind[i] == 2 (x: n_tok rows of n_embd floats).
+    /// multi-GPU: with `out` (pinned, portable host rows) the peer writes its rows straight into
+    /// out[entry] with a scatter kernel (no D2H copy + host memcpy in finish); STRATA_PEER_DIRECT=0: the old way.
+    bool 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 = nullptr);
+    /// Waits for the launched share and writes its rows into `out` (row i = entry i, n_embd floats).
+    bool finish(float* out, std::string& err);
+
+    /// Adaptive tier: swap up to `max_swaps` experts routed at least twice that neither the primary (`res0`, slot
+    /// or -1) nor the peer holds into the peer's least-used slots.  Copies run on the peer's refill stream; a
+    /// 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);
+
+    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_; }
+    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
+
+private:
+    int device_ = -1;
+    bool p2p_ = false;
+    int64_t n_layers_ = 0, n_expert_ = 0;
+    ExpertCache cache_;
+    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
+    int64_t resident_ = 0, entries_ = 0, experts_ = 0, swaps_ = 0;
+    cudaStream_t stream_ = nullptr, refill_ = nullptr;
+    cudaEvent_t refill_ev_ = nullptr;
+    float* h_x_ = nullptr;   float* d_x_ = nullptr;   // n_tok x n_embd
+    float* h_out_ = nullptr; float* d_out_ = nullptr; // compact rows
+    void* h_meta_ = nullptr; void* d_meta_ = nullptr; // groups: ptr | start | dst | tok | count
+    uint8_t* d_q8_ = nullptr;
+    void* d_scratch_ = nullptr;
+    std::vector<int32_t> row_of_;                     ///< compact row -> original entry
+    int64_t launched_rows_ = 0;
+    bool launched_direct_ = false;
+};
+
+}  // namespace strata::core

include/strata/kernels/elementwise.hpp

@@ -104,6 +104,8 @@ void doorbell_wait(const uint32_t* d_flag, const uint32_t* d_seq, void* stream);
 void copy_from_mapped(float* dst, const float* src, int64_t n, void* stream);
 /// `rows` rows of `width` floats from mapped memory, except the rows listed in hit_rows[0, *count) (device),
 /// which are written +0.0 instead (a verify window's GPU-computed entries: the pool leaves zeros there).
+/// multi-GPU: dst[rows[r] * width ..] = src[r * width ..] for r < n (float4; dst may be mapped host memory).
+void scatter_rows_f32(const float* src, float* dst, const int32_t* rows, int64_t n, int64_t width, void* stream);
 void copy_rows_from_mapped(float* dst, const float* src, int64_t rows, int64_t width, const int32_t* hit_rows,
                            const int32_t* count, void* stream);
 

include/strata/prefill/prefill.hpp

@@ -21,6 +21,8 @@
 #include <memory>
 #include <string>
 
+namespace strata::core { class PeerExperts; }
+
 namespace strata::prefill {
 
 struct PrefillStats {
@@ -86,6 +88,11 @@ class Prefill {
 
     const PrefillStats& stats() const { return stats_; }
 
+    /// multi-GPU: the experts the peer GPU holds are computed THERE for every prompt chunk (up to `cap_rows` routed
+    /// rows per layer; the rest of the peer's experts are read by this GPU over P2P).  Allocates the peer's buffers for
+    /// chunks of up to init's `chunk` tokens.  Needs P2P between the two cards.
+    bool set_peer(core::PeerExperts* peer, int64_t cap_rows, std::string& err);
+
     /// Plan v0.3 P6: called after every chunk with the chunk's final multi-stream residual rows (device,
     /// T x hc*n_embd, valid until the next chunk) and the chunk's first position; the MTP draft layer builds its
     /// K/V from them.  The prefill stream is synchronized before the call.

src/core/expert_source.cpp

@@ -1,6 +1,7 @@
 // src/core/expert_source.cpp - the adapter.  See the header for the three clauses of the contract.
 #include "strata/core/expert_source.hpp"
 #include "strata/core/remote_experts.hpp"
+#include "strata/core/peer_experts.hpp"
 #include "strata/kernels/cpu/expert_layout.hpp"
 #include "strata/artifact/gguf_reader.hpp"
 
@@ -1659,7 +1660,8 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
             if (first_of[i] == i) {
                 distinct[nd++] = i;
                 const int32_t e = ids[i];
-                if (e >= 0 && e < d.n_expert && d.host_res[(size_t) d.layers * (size_t) d.n_expert + (size_t) e] < 0) ++nmiss;
+                if (e >= 0 && e < d.n_expert && d.host_res[(size_t) d.layers * (size_t) d.n_expert + (size_t) e] < 0 &&
+                    !(d.peer != nullptr && d.peer->has(d.layers, e))) ++nmiss;
             }
         }
         const bool pcie_ok = d.pcie_num > 0 && d.src->pcie_layer(d.layers);
@@ -1679,6 +1681,8 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
                     kd = 0;
                     ptr = (unsigned long long) (d.cache_base + (d.cache_slot_off ? (size_t) d.cache_slot_off[slot]
                                                                                  : (size_t) slot * (size_t) d.cache_blob));
+                } else if (d.peer != nullptr && d.peer->has(d.layers, e)) {
+                    kd = 2;                        // multi-GPU: the second GPU computes it
                 } else {
                     if (miss_rank >= nmiss - m && fetches < P.staging_cap && fetches < 64) {
                         const uint8_t* src = d.src->pinned(d.layers, e) ? d.src->blob(d.layers, e) : nullptr;
@@ -1733,7 +1737,18 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
         for (int64_t i = 0; i < n; ++i) {
             const int32_t e = ids[i];
             kind[i] = (e >= 0 && e < d.n_expert && d.host_res != nullptr &&
-                       d.host_res[(size_t) d.layers * (size_t) d.n_expert + (size_t) e] >= 0) ? 0 : -1;
+                       d.host_res[(size_t) d.layers * (size_t) d.n_expert + (size_t) e] >= 0) ? 0
+                    : (e >= 0 && e < d.n_expert && d.peer != nullptr && d.peer->has(d.layers, e)) ? 2 : -1;
+        }
+    }
+    if (d.peer != nullptr) {                 // multi-GPU: start the second GPU's share before the CPU's own work
+        std::string perr;
+        if (!d.peer->launch(d.layers, x_f, ids, n_tok, k, kind, perr, out)) {
+            std::fprintf(stderr, "strata: %s (layer %lld)\n", perr.c_str(), (long long) d.layers);
+            d.failed = true;
+            d.fail = "the peer GPU's experts could not be launched";
+            d.fail_layer = d.layers;
+            return;
         }
     }
     if (d.remote_count > 0) {
@@ -1776,9 +1791,12 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
                 d.fail_expert = e;
                 return;
             }
-            if (kind[i] >= 0) {             // CUDA0, PCIe, or a remote result staged into this row below
+            if (kind[i] >= 0) {             // CUDA0, PCIe, or a remote/peer result staged into this row below
                 if (kind[i] == 0) ++d.cache_hits;
-                std::memset(row, 0, (size_t) H * sizeof(float));
+                else if (kind[i] == 2 && d.peer != nullptr) ++d.peer_entries;
+                // multi-GPU: a direct peer launch is writing this row right now - zeroing it would race it
+                if (!(kind[i] == 2 && d.peer != nullptr && d.peer->launched_direct()))
+                    std::memset(row, 0, (size_t) H * sizeof(float));
                 continue;
             }
             ++d.cache_refused;
@@ -1816,6 +1834,16 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
                 d.failed = true; d.fail = remote_error.c_str(); d.fail_layer = d.layers; return;
             }
     }
+    if (d.peer != nullptr) {                 // multi-GPU: the second GPU's rows, into the same mapped rows
+        std::string perr;
+        if (!d.peer->finish(out, perr)) {
+            std::fprintf(stderr, "strata: %s (layer %lld)\n", perr.c_str(), (long long) d.layers);
+            d.failed = true;
+            d.fail = "the peer GPU's experts failed";
+            d.fail_layer = d.layers;
+            return;
+        }
+    }
     const auto c4 = std::chrono::steady_clock::now();
     pt("ran");
     auto ms = [](auto a, auto b) { return std::chrono::duration<double, std::milli>(b - a).count(); };

src/core/mtp.cpp

@@ -4,6 +4,7 @@
 #include "strata/core/on_device.hpp"
 
 #include "strata/core/native_head.hpp"
+#include "strata/core/peer_experts.hpp"
 #include "strata/kernels/bf16_gemv.hpp"
 #include "strata/kernels/cpu/expert.hpp"
 #include "strata/kernels/elementwise.hpp"
@@ -55,7 +56,7 @@ struct Bump {
 };
 
 bool mapped(size_t bytes, void** h, void** d) {
-    if (cudaHostAlloc(h, bytes, cudaHostAllocMapped) != cudaSuccess) return false;
+    if (cudaHostAlloc(h, bytes, cudaHostAllocMapped | (peer_portable() ? cudaHostAllocPortable : 0)) != cudaSuccess) return false;
     std::memset(*h, 0, bytes);
     return cudaHostGetDevicePointer(d, *h, 0) == cudaSuccess;
 }