Multi-GPU: a second GPU as an expert tier for decode and prompts (--peer-device), + multi-conversation cache

CMakeLists.txt

@@ -246,7 +246,7 @@ if(STRATA_ENABLE_CUDA)
   # library is the honest shape of the dependency graph: core knows about devices and weights, kernels know
   # 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/expert_cache.cpp src/core/native_head.cpp src/core/native_dense.cpp src/core/verify.cpp src/core/mtp.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/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 CUDA::cudart)
 

include/strata/core/expert_source.hpp

@@ -31,6 +31,8 @@
 
 namespace strata::core {
 
+class PeerExperts;   // the second GPU's expert tier (peer_experts.hpp)
+
 /// Where one routed expert's bytes come from.
 ///
 /// Phase 2 has NO cache (`phase-2-correct-engine.md`: hit rate `h = 0`), so the only implementation is a
@@ -212,6 +214,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
+    /// 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,96 @@
+// 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 {
+
+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).
+    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);
+    /// 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_; }
+    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;
+};
+
+}  // namespace strata::core

include/strata/core/verify.hpp

@@ -32,6 +32,7 @@
 
 #include <cstdint>
 #include <string>
+#include <vector>
 
 namespace strata::core {
 
@@ -104,9 +105,14 @@ class Verifier {
     /// arena directly, 2 = a copy kernel stages it inside the graph (no API calls on the pool's thread; best when
     /// the CPU is RAM-bound, Q2_0).  Set before the first `run`.
     void set_pcie_mode(int mode) { sink_.pcie_mode = mode; }
+    /// the pool never plans a PCIe share (--pcie-frac 0): the window skips that path.  Before the first run.
+    void set_pcie_off(bool off) { pcie_off_ = off; }
 
     double ms_wait = 0, ms_pool = 0, ms_host = 0, ms_commit = 0;
     int64_t windows = 0;
+    /// STRATA_VERIFY_PROFILE=1 - GPU stage times of the windows since the last call (ms per
+    /// window), as one line; empty when off.
+    std::string profile_report();
 
 private:
     bool capture(int T, std::string& err);
@@ -121,6 +127,12 @@ class Verifier {
     bool head_sampling_ = true;          ///< set_head_sampling
     bool capture_commit(std::string& err);
     bool record_window(int T, cudaStream_t cs, std::string& err);
+    static constexpr int kProfPer = 32;              // stamps per layer
+    bool prof_on_ = false;
+    unsigned long long* prof_ = nullptr;              // device: n_layers * kProfPer + 4 stamps
+    std::vector<unsigned long long> prof_h_;
+    double prof_sum_[2][kProfPer] = {};   // [GDN / QSA layers][stage]
+    int64_t prof_windows_ = 0;
 
     const WeightTable* wt_ = nullptr;
     const ModelGeometry* g_ = nullptr;
@@ -150,6 +162,9 @@ class Verifier {
     uint32_t* h_flag_ = nullptr; uint32_t* m_flag_ = nullptr;
     uint32_t* h_flagA_ = nullptr; uint32_t* m_flagA_ = nullptr;  // the GPU plan is in place
     uint32_t* h_flagB_ = nullptr; uint32_t* m_flagB_ = nullptr;  // the PCIe share's DMA copies have landed
+    cudaStream_t aux_[2] = {};                                    // parallel branches of the captured window
+    cudaEvent_t fev_[4] = {};
+    bool pcie_off_ = false;                                       // no PCIe share (--pcie-frac 0)
     cudaStream_t copy_ = nullptr;                                 // the copy engine's stream (DMA of missed experts)
     struct FlagSet { uint32_t* flag; uint32_t value; };
     FlagSet flag_sets_[2 * 64 * 2] = {};                          // host-function arguments, one per (layer, group)

include/strata/kernels/bf16_gemv.hpp

@@ -44,5 +44,9 @@ void bf16_gemv_split(const uint16_t* x, const uint16_t* w, float* y, int64_t n_i
 /// No allocations or synchronization, including when stream is null (the CUDA default stream).
 void bf16_gemv_fp32_mmvf(const float* x, const uint16_t* w, float* y,
                          int64_t n_in, int64_t n_out, void* stream);
+/// bf16_gemv_fp32_mmvf for n_tok (1..8) activation rows x[t * ldx], outputs y[t * ldy + j]; one launch,
+/// the weight read once, every output bit-identical to its own single-row call.
+void bf16_gemv_fp32_mmvf_multi(const float* x, int64_t ldx, const uint16_t* w, float* y, int64_t ldy,
+                               int64_t n_in, int64_t n_out, int n_tok, void* stream);
 
 }  // namespace strata::kernels

include/strata/kernels/elementwise.hpp

@@ -102,6 +102,10 @@ void doorbell_wait(const uint32_t* d_flag, const uint32_t* d_seq, void* stream);
 /// device memory with a kernel, so the handoff stays on the compute queue (a memcpy node is a copy-engine
 /// operation, which WDDM submits separately and which measured 67 flushes per token).
 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).
+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);
 
 /// Plan v0.3 P3: the doorbell's payload and its ring in ONE kernel.  Copies `x` (n floats), `ids` and `weights`
 /// (k each) into the mapped host regions, fences, and increments the mapped sequence number - replacing three

include/strata/kernels/fused_gr.hpp

@@ -46,6 +46,7 @@ void fused_gr_read(const FusedGrArgs& a, void* stream);
 /// weight pointers and eps must be the same for every t); `xn_scratch` is n_tok * hc * n_embd floats.  Every
 /// token's outputs are bitwise `fused_gr_read(a[t])`.
 constexpr int kFusedGrMaxT = 8;
-void fused_gr_read_multi(const FusedGrArgs* a, int n_tok, float* xn_scratch, void* stream);
+void fused_gr_read_multi(const FusedGrArgs* a, int n_tok, float* xn_scratch, void* stream,
+                         unsigned long long* stamp_buf = nullptr, int stamp_i0 = 0);
 
 }  // namespace strata::kernels

include/strata/kernels/native_moe.hpp

@@ -15,4 +15,7 @@ bool native_moe_combine_enabled();
 // Requires a nonnull ordered stream and disjoint output. No allocation or sync.
 void native_moe_combine(const float* parts, const float* weights, const float* shared,
                         float* output, int64_t n_embd, int64_t k, void* stream);
+/// n_tok rows (parts [n,k,N], weights [n,k], shared/output [n,N]) in one launch, each as the single call.
+void native_moe_combine_multi(const float* parts, const float* weights, const float* shared, float* output,
+                              int64_t n_embd, int64_t k, int n_tok, void* stream);
 }

include/strata/kernels/native_qsa_indexer.hpp

@@ -32,4 +32,10 @@ void native_qsa_indexer_append(const float* raw, const int32_t* relative_pos_dev
                                const QsaIndexerBuffers& buffers, const QsaShapes& shapes,
                                int64_t max_cells, float freq_base, void* stream);
 
+// the same appends for T contiguous cells pos0 .. pos0+T-1 (raw [T,128]) in two launches
+// instead of T; the resulting tail / dead / pooled / block_pos are bit-identical to T single appends.
+void native_qsa_indexer_append_chunk(const float* raw, int64_t T, int64_t pos0, int32_t pos_base,
+                                     const float* gamma, float epsilon, const QsaIndexerBuffers& buffers,
+                                     const QsaShapes& shapes, int64_t max_cells, float freq_base, void* stream);
+
 } // namespace strata::kernels

include/strata/kernels/native_router.hpp

@@ -14,4 +14,6 @@ bool native_router_enabled();
 // four-byte aligned and outputs disjoint from each other and the input.
 // Requires a nonnull ordered CUDA stream. No allocation or synchronization.
 void native_router_top10(const float* logits, int32_t* ids, float* weights, void* stream);
+/// n_tok rows at once (logits [n,512], ids/weights [n,10]); each row exactly as the single call.
+void native_router_top10_multi(const float* logits, int32_t* ids, float* weights, int n_tok, void* stream);
 }

include/strata/kernels/ngram.hpp

@@ -110,6 +110,9 @@ struct PleIoOptions {
     uint32_t max_inflight = 64;      ///< outstanding SSD reads (decode needs 16; prefill chunks use more)
     uint64_t cache_rows = 1u << 20;  ///< bounded row cache: 1,048,576 rows x 90 B ~ 95 MB; 0 disables
     bool io_thread = true;           ///< reads submitted by a worker thread, not the caller
+    /// --ple-io ram: mmap mode with the whole table locked in RAM at open (a box with RAM to spare:
+    /// no SSD read ever sits on the prompt or token path)
+    bool lock = false;
 };
 
 /// The PLE table.  Held by pointer-to-impl so this header does not drag `<windows.h>` into every

include/strata/kernels/verify_kernels.hpp

@@ -38,6 +38,8 @@ void gdn_step_norm_multi(float* state, const float* h, int conv_channels, const
 /// Spin until *flag >= value (a mapped host flag).  The value is fixed at capture, so several rings can be
 /// outstanding at once (the split verify window keeps two).
 void wait_flag_ge(const uint32_t* flag, uint32_t value, void* stream);
+/// the GPU's %globaltimer (ns) into buf[i] (a one-thread kernel: the verify window's stage profiler).
+void gpu_stamp(unsigned long long* buf, int i, void* stream);
 
 /// Rows of the S2/S4/S8 embedding for T token ids read from DEVICE memory; out (T, n).  Bitwise `embedding_gather`.
 void embedding_gather_dev(const uint8_t* codes, const float* scales, const float* offsets, const int32_t* tokens,

include/strata/prefill/kernels.hpp

@@ -29,11 +29,14 @@ void gr_broadcast(const float* e, float* R, int64_t T, void* stream);
 void gdn_gates(const float* ab, const float* dt, const float* ssm_a, float* gate, float* beta, int64_t T, void* stream);
 /// The 4-tap causal conv + SiLU over the chunk (history [C][3] in, updated to the chunk's last three inputs), then
 /// the L2 norm of the q and k heads of every token.  h: [T, C].
-void gdn_conv(float* history, const float* qkv, const float* conv_w, float* h, int64_t T, float eps, void* stream);
+// mode: -1 = STRATA_PF_CONV_PAR (default parallel over tokens), 0 = the per-channel walk, 1 = parallel
+void gdn_conv(float* history, const float* qkv, const float* conv_w, float* h, int64_t T, float eps, void* stream,
+              int mode = -1);
 /// The recurrence over the chunk, block per value head, state in registers; y[t] = rmsnorm(o) * gamma * sigmoid(z)
 /// (FP32 and FP16 bits: the out projection is quantized).
 void gdn_recurrence(float* state, const float* h, const float* gate, const float* beta, const float* z,
-                    const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, void* stream);
+                    const float* gamma, float eps, float* y, uint16_t* y16, int64_t T, void* stream,
+                    int mode = -1);   // -1: STRATA_PF_REC_SPLIT (default: the column-split kernel), 0: pipelined, 1: split
 
 // ---- MoE
 /// softmax over 512, top-10 (ties to the lower id), weights renormalised over the ten (the native router).

include/strata/prefill/prefill.hpp

@@ -21,6 +21,8 @@
 #include <memory>
 #include <string>
 
+namespace strata::core { class PeerExperts; }
+
 namespace strata::prefill {
 
 struct PrefillStats {
@@ -67,6 +69,11 @@ class Prefill {
 
     const PrefillStats& stats() const { return stats_; }
 
+    /// 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.

serve/server.py

@@ -375,6 +375,10 @@ def __init__(self, cfg: dict, log=None, env: dict | None = None):
             args += ["--threads", str(cfg["threads"])]
         if cfg.get("max_tokens"):
             args += ["--max-tokens", str(cfg["max_tokens"])]
+        if cfg.get("env"):
+            # e.g. {"CUDA_VISIBLE_DEVICES": "0"} puts the encoder on the helper card, not the primary
+            env = dict(env if env is not None else os.environ)
+            env.update({str(k): str(v) for k, v in cfg["env"].items()})
         self.dir = Path(tempfile.mkdtemp(prefix="strata-vision-"))
         self.proc = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=log or subprocess.DEVNULL,
                                      text=True, encoding="utf-8", bufsize=1, env=env)