Prompt path: expert blobs of long chunks copied ahead while the dense steps run; per-expert GEMMs over side streams; the second GPU's sums return as FP16; pinned uploads

include/strata/prefill/experts.hpp

@@ -4,14 +4,16 @@
 // batch one staging copy for each run of adjacent arena blobs among the experts the GPU's cache does not hold, one
 // dequantization launch, one gather of the batch's rows, the two GEMMs of each expert with a SwiGLU between them, and
 // one ordered add of every row, times its routing weight, into its token's sum.  A token's experts add up in ascending
-// id order, whatever the batches.
+// id order, whatever the batches.  A chunk of a few hundred tokens or more uses nearly every expert, so the blobs a
+// layer will stream can be copied ahead (`prefetch`), while the prompt path's dense work runs and before the routing
+// is known; experts that do not fit there stream batch by batch.
 //
 // Local: on the prompt path's own GPU and stream, reading the chunk's MoE input and adding into its MoE output there.
 // Remote: on a second GPU.  The first GPU (the 3090) sits on a PCIe 4.0 x4 link, over which an expert its cache lacks
 // streams at 6.4 GB/s (~0.65 GB (IQ3_XXS) to ~1.4 GB (UD-Q4_K_XL) a layer); the second GPU sits on x16.  The first GPU
 // copies the chunk's MoE input (FP16) to pinned host memory; the second copies it in, computes the experts the first
-// one's cache lacks - from its own cache when it holds them - and returns their sums (FP32) the same way, while the
-// first computes the experts its cache holds.  Neither GPU reads the other's memory: Windows gives GeForce cards no
+// one's cache lacks - from its own cache when it holds them - and returns their sums (FP16: they stay below ~20) the
+// same way, while the first computes the experts its cache holds.  Neither GPU reads the other's memory: Windows gives GeForce cards no
 // peer access.
 #pragma once
 
@@ -40,16 +42,30 @@ public:
     /// whether experts its cache does not hold come its way (staging buffers).
     bool init(int device, int main_device, void* stream, core::ExpertSource* src, const core::ExpertCache* cache,
               const int32_t* res, int64_t n_expert, int64_t max_chunk, bool streams, std::string& err);
+    /// Chunks of at least this many tokens prefetch: they use nearly every expert, and their dense steps last long
+    /// enough to hide the copies (a 521-token chunk prefetching was 2.2 s slower on IQ3_XXS: routing is skewed, so
+    /// many of the experts no cache holds go unused in a short chunk).
+    static constexpr int64_t kPrefetchMin = 4096;
+
     /// Device bytes of the buffers for chunks of `chunk` tokens.
     static uint64_t bytes_needed(int64_t chunk, int64_t n_expert, bool remote, bool streams);
+    /// The same with its prefetch area, for the residency `skip` (see prefetch_need) and its own as they are now.
+    uint64_t bytes_for(int64_t chunk, const int32_t* skip) const;
     /// The buffers for chunks of up to `chunk` tokens, carved from `region` (`bytes` long, e.g. lent cache slots), or
     /// with a null region allocated for `max_chunk` tokens (once).  Before the first layer of every prompt that
-    /// uses a region.
-    bool bind(void* region, uint64_t bytes, int64_t chunk, std::string& err);
+    /// uses a region.  `area`: device memory for the prefetched blobs; null: what the region has left.
+    bool bind(void* region, uint64_t bytes, int64_t chunk, std::string& err, void* area = nullptr,
+              uint64_t area_bytes = 0);
+    /// Bytes of prefetch area that hold a layer's streamed experts: the largest layer's experts that neither its
+    /// cache nor `skip` (another GPU's residency, n_layers x n_expert, or null) holds, and a margin.
+    uint64_t prefetch_need(const int32_t* skip) const;
+    /// Before a layer's routing: copies the blobs of `experts` (ascending ids) its cache does not hold into the
+    /// prefetch area, as many as fit, on its copy stream; run_layer of that layer reads them there.
+    bool prefetch(int64_t layer, const std::vector<int32_t>& experts, std::string& err);
 
-    /// Remote: pinned [max_chunk, n_embd]: the MoE input (FP16) goes in here, the sums (FP32) come out here.
+    /// Remote: pinned [max_chunk, n_embd]: the MoE input goes in here, the sums come out here (both FP16).
     uint16_t* host_input() const;
-    float* host_sum() const;
+    uint16_t* host_sum() const;
 
     /// Layer `layer` of a chunk of T tokens: expert experts[j] serves rows [off[j], off[j + 1]), row r being token
     /// src[r] with routing weight w[r]; each row times its weight adds to its token's sum.  Local: reads `input`
@@ -63,6 +79,7 @@ public:
 
     int64_t experts_resident = 0;   ///< expert-layer groups served from its cache
     int64_t experts_streamed = 0;   ///< expert blobs copied from the arena
+    int64_t experts_prefetched = 0; ///< ...of which ahead of the routing
     double ms_host = 0;             ///< host time queuing its work
     double ms_gpu = 0;              ///< remote: its GPU time from the input's arrival (all layers but the last)
 

include/strata/prefill/gemm.hpp

@@ -18,8 +18,9 @@ public:
     Gemm(const Gemm&) = delete;
     Gemm& operator=(const Gemm&) = delete;
 
-    /// The cuBLAS handle on `stream`; set_buffers() before the first product.
-    bool init(void* stream, std::string& err);
+    /// The cuBLAS handle on `stream`, and with `side_streams` the side streams f16_grouped() spreads over;
+    /// set_buffers() before the first product.
+    bool init(void* stream, std::string& err, bool side_streams = false);
     /// Caller-owned device buffers (e.g. lent expert-cache slots): the dequantization scratch (`scratch_elems` FP16
     /// elements: the largest weight dequantized at once; may be null for no native() calls) and the cuBLAS workspace.
     void set_buffers(uint16_t* scratch, int64_t scratch_elems, void* workspace, size_t ws_bytes);
@@ -34,7 +35,9 @@ public:
 
     /// Y_i[rows[i], N] = X_i[rows[i], K] . W_i[N, K]^T (FP16 in, FP32 out) for n problems of one shape but their
     /// own row counts, in one cuBLAS grouped call; the pointer arrays are DEVICE memory, `rows` host memory.  Where
-    /// cuBLAS does not group these types it runs f16() per problem, from the host copies of the pointers.
+    /// cuBLAS does not group these types it runs f16() per problem, from the host copies of the pointers, spread
+    /// over the side streams when the problems average 128 rows or more (a GEMM of a few hundred rows leaves most
+    /// of the GPU idle; below that the launches, not the GPU, set the pace).
     void f16_grouped(const uint16_t* const* X, const uint16_t* const* W, float* const* Y, const uint16_t* const* X_host,
                      const uint16_t* const* W_host, float* const* Y_host, const int* rows, int n, int64_t N, int64_t K);
 
@@ -47,8 +50,13 @@ public:
     void* stream() const { return stream_; }
 
 private:
+    static constexpr int AUX = 3;   // side streams for independent products
     void* handle_ = nullptr;
     void* stream_ = nullptr;
+    void* aux_handle_[AUX] = {};
+    void* aux_stream_[AUX] = {};
+    void* fork_ = nullptr;
+    void* join_[AUX] = {};
     uint16_t* scratch_ = nullptr;
     int64_t scratch_elems_ = 0;
     bool grouped_ = true;   // cuBLAS groups FP16 -> FP32 GEMMs (cleared at the first refusal)

include/strata/prefill/kernels.hpp

@@ -17,9 +17,9 @@ void gr_silu(const float* lo, uint16_t* lo16, int64_t T, void* stream);
 /// mixed[t, d] = mean_c xn[t, c, d] * sigmoid(gated[t, c, d]); FP32, BF16 and FP16 (any of them may be null).
 void gr_mix(const float* xn, const float* gated, float* mixed, uint16_t* mixed16, int64_t T, void* stream,
             uint16_t* mixed_h = nullptr);
-/// R[t, c, d] += (bo[t, d] (+ partial[t, d])) * 2 sigmoid(inj[t, c] / hc)   (inj has row stride inj_ld)
+/// R[t, c, d] += (bo[t, d] (+ partial[t, d], FP16)) * 2 sigmoid(inj[t, c] / hc)   (inj has row stride inj_ld)
 void gr_write(float* R, const float* bo, const float* inj, int64_t inj_ld, int64_t T, void* stream,
-              const float* partial = nullptr);
+              const uint16_t* partial = nullptr);
 /// R[t, c, :] = e[t, :] for all four streams (the embedding broadcast).
 void gr_broadcast(const float* e, float* R, int64_t T, void* stream);
 
@@ -47,6 +47,8 @@ void swiglu_interleaved(const float* gu, uint16_t* h16, int64_t n, void* stream)
 void swiglu_pair(const float* g, const float* u, uint16_t* h16, int64_t n, void* stream);
 /// Gather rows: dst16[i, :] = x16[src[i], :] (n rows of `width` BF16).
 void gather_rows16(const uint16_t* x16, const int32_t* src, uint16_t* dst16, int64_t n, int64_t width, void* stream);
+/// y = FP16 bits of x, saturated at +-65504 (the second GPU's expert sums on their way back: below ~20 measured)
+void sums_to_f16(const float* x, uint16_t* y, int64_t n, void* stream);
 /// In place: x[t, :] *= sigmoid(sg[t])   (the shared expert's scalar gate)
 void moe_shared_gate(float* x, const float* sg, int64_t T, void* stream);
 /// sum[src[r], :] += w[r] * rows[r, :] for n rows (one expert's rows: no token twice).

include/strata/prefill/prefill.hpp

@@ -38,6 +38,7 @@ struct PrefillStats {
     double ms_total = 0;
     double ms_experts_host = 0;     ///< host time queuing this GPU's experts
     int64_t experts_streamed = 0;   ///< expert blobs copied host -> this GPU
+    int64_t experts_prefetched = 0; ///< ...of which ahead of the routing
     int64_t experts_resident = 0;   ///< expert-layer groups served from its VRAM tier
     double ms_ple = 0;
     double ms_section[kPsCount] = {};   ///< GPU ms per section, with `Prefill::profile`
@@ -60,9 +61,11 @@ public:
     /// Device bytes of the buffers for chunks of `chunk` tokens.
     static uint64_t bytes_needed(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk,
                                  bool offload);
+    /// The same with the area its experts are prefetched into (without a second GPU), for the residency as it is now.
+    uint64_t bytes_for(int64_t chunk) const;
     /// The buffers for chunks of up to `chunk` tokens, carved from `region` (`bytes` long: lent expert-cache
-    /// slots), or with a null region allocated for the longest chunk (once).  Before the first chunk of every prompt
-    /// that uses a region.
+    /// slots; what is left over takes prefetched experts), or with a null region allocated for the longest chunk
+    /// (once).  Before the first chunk of every prompt that uses a region.
     bool bind(void* region, uint64_t bytes, int64_t chunk, std::string& err);
 
     /// Positions [pos0, pos0 + n) holding `tokens`, in chunks of the bound length; `ss.ple_prev` must be the two

src/prefill/experts.cpp

@@ -21,6 +21,7 @@ constexpr int64_t ROWS = 4096;
 constexpr int SLOTS = 2;                 // batches in staging: the copies of one overlap the work on the other
 constexpr size_t GEMM_WS = 32u << 20;    // cuBLAS workspace
 constexpr int NPTR = 7;                  // per-expert device pointers: blob, xs, gate/up, gu, h, down, d
+constexpr int PRE_GROUP = 16;            // prefetched blobs per copy event
 
 struct OnDevice {   // a remote runner's GPU current for one call, the main one again afterwards
     int main;
@@ -44,6 +45,12 @@ struct Alloc {
 
 size_t blob_cap() { return ((size_t) strata::kernels::cpu::expert_layout().max_blob + 255) & ~(size_t) 255; }
 
+// the pinned upload block: NPTR x n_expert pointers, then the adds (at most 3 per row, 2 per expert), rows, weights
+size_t adds_cap(int64_t chunk, int64_t n_expert) { return (size_t) (chunk * K * 3 + n_expert * 2); }
+size_t up_bytes(int64_t chunk, int64_t n_expert) {
+    return (size_t) NPTR * (size_t) n_expert * sizeof(void*) + (adds_cap(chunk, n_expert) + (size_t) chunk * K * 2) * 4;
+}
+
 }  // namespace
 
 struct ExpertRunner::Impl {
@@ -63,12 +70,24 @@ struct ExpertRunner::Impl {
     float *sum = nullptr, *w = nullptr, *gu = nullptr, *d = nullptr;
     int32_t* rows_src = nullptr;
     int32_t* adds = nullptr;   // per batch: its tokens, their row lists' starts, the row lists (moe_gather_add)
-    std::vector<int32_t> adds_host, mark, first, count;
+    std::vector<int32_t> mark, first, count;
     uint8_t* stage[SLOTS] = {};
     void** ptrs = nullptr;                    // NPTR arrays of n_expert device pointers
-    std::vector<void*> ptrs_host;
-    uint16_t* h_input = nullptr;
-    float* h_sum = nullptr;
+    // a layer's uploads, staged in pinned memory (a pageable copy blocks the host behind queued expert copies):
+    // [pointers | adds | rows | weights], for chunks of up to max_chunk tokens
+    uint8_t* h_up = nullptr;
+    cudaEvent_t up_done = nullptr;   // the previous layer's uploads have been read
+    bool up_live = false;
+    uint16_t *h_input = nullptr, *h_sum = nullptr;
+    // a layer's blobs copied ahead: at pre_off[e] in `pre` (-1: not there), covered by event pre_ev[pre_grp[e]]
+    uint8_t* pre = nullptr;
+    uint64_t pre_cap = 0;
+    int64_t pre_layer = -1;
+    std::vector<int64_t> pre_off;
+    std::vector<int32_t> pre_grp;
+    std::vector<cudaEvent_t> pre_ev;
+    cudaEvent_t pre_used = nullptr;   // the last read of the area
+    bool pre_live = false;
 
     // the same sequence counted (bytes_needed) or carved
     static void carve(Alloc& a, int64_t T, int64_t n_expert, bool remote, bool streams, Impl* m, void** ws, bool& ok) {
@@ -109,9 +128,11 @@ ExpertRunner::~ExpertRunner() {
         if (m.copied[i]) cudaEventDestroy(m.copied[i]);
         if (m.used[i]) cudaEventDestroy(m.used[i]);
     }
-    for (cudaEvent_t e : {m.done, m.t0, m.t1}) if (e) cudaEventDestroy(e);
+    for (cudaEvent_t e : {m.done, m.t0, m.t1, m.pre_used, m.up_done}) if (e) cudaEventDestroy(e);
+    for (cudaEvent_t e : m.pre_ev) if (e) cudaEventDestroy(e);
     if (m.h_input) cudaFreeHost(m.h_input);
     if (m.h_sum) cudaFreeHost(m.h_sum);
+    if (m.h_up) cudaFreeHost(m.h_up);
     if (m.owned) cudaFree(m.owned);
     if (m.copy) cudaStreamDestroy(m.copy);
     if (m.remote && m.s) cudaStreamDestroy(m.s);
@@ -142,14 +163,14 @@ bool ExpertRunner::init(int device, int main_device, void* stream, core::ExpertS
     m.res = res;
     m.n_expert = n_expert;
     m.max_chunk = max_chunk;
-    m.ptrs_host.resize((size_t) NPTR * (size_t) n_expert);
-    m.adds_host.reserve((size_t) max_chunk * 30);
     m.mark.assign((size_t) max_chunk, -1);
     m.first.resize((size_t) max_chunk);
     m.count.resize((size_t) max_chunk);
     OnDevice on(device, main_device, m.remote);
     bool ok = cudaEventCreateWithFlags(&m.done, cudaEventDisableTiming) == cudaSuccess &&
-              cudaEventCreate(&m.t0) == cudaSuccess && cudaEventCreate(&m.t1) == cudaSuccess;
+              cudaEventCreateWithFlags(&m.up_done, cudaEventDisableTiming) == cudaSuccess &&
+              cudaEventCreate(&m.t0) == cudaSuccess && cudaEventCreate(&m.t1) == cudaSuccess &&
+              cudaHostAlloc((void**) &m.h_up, up_bytes(max_chunk, n_expert), cudaHostAllocDefault) == cudaSuccess;
     if (m.remote) ok = ok && cudaStreamCreateWithFlags(&m.s, cudaStreamNonBlocking) == cudaSuccess;
     else m.s = (cudaStream_t) stream;
     if (streams) {
@@ -157,23 +178,29 @@ bool ExpertRunner::init(int device, int main_device, void* stream, core::ExpertS
         for (int i = 0; ok && i < SLOTS; ++i)
             ok = cudaEventCreateWithFlags(&m.copied[i], cudaEventDisableTiming) == cudaSuccess &&
                  cudaEventCreateWithFlags(&m.used[i], cudaEventDisableTiming) == cudaSuccess;
+        m.pre_ev.assign((size_t) ((n_expert + PRE_GROUP - 1) / PRE_GROUP), nullptr);
+        for (auto& e : m.pre_ev) ok = ok && cudaEventCreateWithFlags(&e, cudaEventDisableTiming) == cudaSuccess;
+        ok = ok && cudaEventCreateWithFlags(&m.pre_used, cudaEventDisableTiming) == cudaSuccess;
+        m.pre_off.assign((size_t) n_expert, -1);
+        m.pre_grp.assign((size_t) n_expert, 0);
     }
-    if (!ok || !m.gemm.init(m.s, err)) {
+    if (!ok || !m.gemm.init(m.s, err, true)) {
         if (err.empty()) err = "prompt experts: streams or events";
         return false;
     }
     // the input and the sums: pinned and portable, both GPUs copy them
     if (m.remote && (cudaHostAlloc((void**) &m.h_input, (size_t) max_chunk * N * 2, cudaHostAllocPortable) != cudaSuccess ||
-                     cudaHostAlloc((void**) &m.h_sum, (size_t) max_chunk * N * 4, cudaHostAllocPortable) != cudaSuccess)) {
+                     cudaHostAlloc((void**) &m.h_sum, (size_t) max_chunk * N * 2, cudaHostAllocPortable) != cudaSuccess)) {
         err = "prompt experts: pinned input and sum buffers";
         return false;
     }
     return true;
 }
 
-bool ExpertRunner::bind(void* region, uint64_t bytes, int64_t chunk, std::string& err) {
+bool ExpertRunner::bind(void* region, uint64_t bytes, int64_t chunk, std::string& err, void* area, uint64_t area_bytes) {
     Impl& m = *impl_;
     OnDevice on(m.dev, m.main, m.remote);
+    m.pre_layer = -1;
     if (region == nullptr) {
         if (m.owned != nullptr) return true;
         chunk = m.max_chunk;
@@ -201,11 +228,87 @@ bool ExpertRunner::bind(void* region, uint64_t bytes, int64_t chunk, std::string
     }
     m.gemm.set_buffers(nullptr, 0, ws, GEMM_WS);
     m.T = chunk;
+    if (area == nullptr && region != m.owned && bytes > a.used) {
+        area = (uint8_t*) region + a.used;
+        area_bytes = bytes - a.used;
+    }
+    m.pre = m.streams ? (uint8_t*) area : nullptr;
+    m.pre_cap = m.pre != nullptr ? area_bytes : 0;
+    return true;
+}
+
+uint64_t ExpertRunner::bytes_for(int64_t chunk, const int32_t* skip) const {
+    const Impl& m = *impl_;
+    return bytes_needed(chunk, m.n_expert, m.remote, m.streams) + (chunk >= kPrefetchMin ? prefetch_need(skip) : 0);
+}
+
+uint64_t ExpertRunner::prefetch_need(const int32_t* skip) const {
+    const Impl& m = *impl_;
+    if (!m.streams) return 0;
+    const auto& lay = strata::kernels::cpu::expert_layout();
+    uint64_t most = 0;
+    for (int64_t l = 0; l < lay.n_layers; ++l) {
+        const size_t stride = ((size_t) lay.blob_bytes(l) + 255) & ~(size_t) 255;
+        uint64_t sum = 0;
+        for (int64_t e = 0; e < m.n_expert; ++e) {
+            const size_t i = (size_t) (l * m.n_expert + e);
+            if ((m.res == nullptr || m.res[i] < 0) && (skip == nullptr || skip[i] < 0)) sum += stride;
+        }
+        most = std::max(most, sum);
+    }
+    return most + most / 10;   // the slots lent for the prompt make a few more experts stream
+}
+
+bool ExpertRunner::prefetch(int64_t layer, const std::vector<int32_t>& experts, std::string& err) {
+    Impl& m = *impl_;
+    if (m.pre_cap == 0) return true;
+    OnDevice on(m.dev, m.main, m.remote);
+    const auto& lay = strata::kernels::cpu::expert_layout();
+    const size_t bb = (size_t) lay.blob_bytes(layer);
+    const size_t stride = (bb + 255) & ~(size_t) 255;
+    std::fill(m.pre_off.begin(), m.pre_off.end(), -1);
+    m.pre_layer = layer;
+    if (m.pre_live) cudaStreamWaitEvent(m.copy, m.pre_used, 0);   // the previous layer has read the area
+    uint64_t at = 0;
+    int32_t grp = 0, in_grp = 0;
+    const size_t n = experts.size();
+    for (size_t j = 0; j < n;) {
+        const int32_t e = experts[j];
+        const bool held = m.res && m.cache && m.res[(size_t) layer * (size_t) m.n_expert + (size_t) e] >= 0;
+        if (held) { ++j; continue; }
+        if (at + stride > m.pre_cap) break;
+        // a run: the next experts' blobs follow this one in the arena, and fit
+        const uint8_t* b0 = m.src->blob(layer, e);
+        if (b0 == nullptr) { err = "prompt experts: the expert source has no blob"; return false; }
+        size_t k = j + 1;
+        if (m.src->pinned(layer, e) && stride == bb)
+            while (k < n && experts[(size_t) k] == e + (int32_t) (k - j) && at + (k - j + 1) * stride <= m.pre_cap &&
+                   !(m.res && m.cache && m.res[(size_t) layer * (size_t) m.n_expert + (size_t) experts[k]] >= 0) &&
+                   m.src->pinned(layer, experts[k]) && m.src->blob(layer, experts[k]) == b0 + (k - j) * bb &&
+                   in_grp + (int32_t) (k - j) < PRE_GROUP) ++k;
+        cudaMemcpyAsync(m.pre + at, b0, (k - j) * bb, cudaMemcpyHostToDevice, m.copy);   // pageable blobs: staged
+        for (size_t i = j; i < k; ++i) {
+            m.pre_off[(size_t) experts[i]] = (int64_t) (at + (i - j) * stride);
+            m.pre_grp[(size_t) experts[i]] = grp;
+        }
+        at += (k - j) * stride;
+        experts_prefetched += (int64_t) (k - j);
+        experts_streamed += (int64_t) (k - j);
+        in_grp += (int32_t) (k - j);
+        j = k;
+        if (in_grp >= PRE_GROUP) {
+            cudaEventRecord(m.pre_ev[(size_t) grp++], m.copy);
+            in_grp = 0;
+        }
+    }
+    if (in_grp > 0) cudaEventRecord(m.pre_ev[(size_t) grp], m.copy);
+    const cudaError_t e = cudaGetLastError();
+    if (e != cudaSuccess) { err = std::string("prompt experts: prefetch: ") + cudaGetErrorString(e); return false; }
     return true;
 }
 
 uint16_t* ExpertRunner::host_input() const { return impl_->h_input; }
-float* ExpertRunner::host_sum() const { return impl_->h_sum; }
+uint16_t* ExpertRunner::host_sum() const { return impl_->h_sum; }
 cudaEvent_t ExpertRunner::done() const { return impl_->done; }
 
 bool ExpertRunner::run_layer(int64_t layer, int64_t T, const std::vector<int32_t>& experts, const std::vector<int32_t>& off,
@@ -233,6 +336,7 @@ bool ExpertRunner::run_layer(int64_t layer, int64_t T, const std::vector<int32_t
     auto held = [&](int32_t e) -> int32_t {
         return m.res && m.cache ? m.res[(size_t) layer * (size_t) m.n_expert + (size_t) e] : -1;
     };
+    auto pre = [&](int32_t e) -> int64_t { return m.pre_layer == layer ? m.pre_off[(size_t) e] : -1; };
     // ---- the batches, and every expert's device pointers (one upload)
     std::vector<std::pair<size_t, size_t>> batches;
     for (size_t j0 = 0; j0 < n_exp;) {
@@ -243,8 +347,12 @@ bool ExpertRunner::run_layer(int64_t layer, int64_t T, const std::vector<int32_t
         batches.emplace_back(j0, j1);
         j0 = j1;
     }
-    void** P = m.ptrs_host.data();
+    if (m.up_live) cudaEventSynchronize(m.up_done);   // the host block is free again
     const size_t E = (size_t) m.n_expert;   // array stride
+    void** P = (void**) m.h_up;
+    int32_t* h_adds = (int32_t*) (P + NPTR * E);
+    int32_t* h_rows = h_adds + adds_cap(m.max_chunk, m.n_expert);
+    float* h_w = (float*) (h_rows + (size_t) m.max_chunk * K);
     std::vector<int> rows_of(n_exp);
     for (size_t b = 0; b < batches.size(); ++b) {
         const auto [j0, j1] = batches[b];
@@ -256,7 +364,9 @@ bool ExpertRunner::run_layer(int64_t layer, int64_t T, const std::vector<int32_t
                 err = "prompt experts: an expert its cache does not hold, and no staging";
                 return false;
             }
-            P[0 * E + j] = slot >= 0 ? (void*) m.cache->device_slot(slot) : (void*) (m.stage[b % SLOTS] + streamed++ * stride);
+            P[0 * E + j] = slot >= 0 ? (void*) m.cache->device_slot(slot)
+                         : pre(experts[j]) >= 0 ? (void*) (m.pre + pre(experts[j]))
+                         : (void*) (m.stage[b % SLOTS] + streamed++ * stride);
             P[1 * E + j] = m.xs + r * N;
             P[2 * E + j] = m.dq_gu + (j - j0) * (size_t) (FF2 * N);
             P[3 * E + j] = m.gu + r * FF2;
@@ -270,7 +380,7 @@ bool ExpertRunner::run_layer(int64_t layer, int64_t T, const std::vector<int32_t
     // [tokens | starts (n + 1) | rows] at adds_at[b]
     std::vector<size_t> adds_at(batches.size());
     std::vector<int32_t> ntok(batches.size());
-    m.adds_host.clear();
+    size_t n_adds = 0;
     for (size_t b = 0; b < batches.size(); ++b) {
         const int32_t r0 = off[batches[b].first], r1 = off[batches[b].second];
         std::vector<int32_t> toks;
@@ -284,11 +394,11 @@ bool ExpertRunner::run_layer(int64_t layer, int64_t T, const std::vector<int32_t
             }
             ++m.count[(size_t) m.first[(size_t) t]];
         }
-        const size_t nt = toks.size(), at = m.adds_host.size();
+        const size_t nt = toks.size(), at = n_adds;
         adds_at[b] = at;
         ntok[b] = (int32_t) nt;
-        m.adds_host.resize(at + nt + (nt + 1) + (size_t) (r1 - r0));
-        int32_t* h_tok = m.adds_host.data() + at;
+        n_adds = at + nt + (nt + 1) + (size_t) (r1 - r0);
+        int32_t* h_tok = h_adds + at;
         int32_t* h_start = h_tok + nt;
         int32_t* h_list = h_start + nt + 1;
         h_start[0] = 0;
@@ -306,15 +416,17 @@ bool ExpertRunner::run_layer(int64_t layer, int64_t T, const std::vector<int32_t
     }
     const uint16_t* in = m.remote ? m.input : input;
     float* out = m.remote ? m.sum : sum;
-    if (!m.adds_host.empty())
-        cudaMemcpyAsync(m.adds, m.adds_host.data(), m.adds_host.size() * 4, cudaMemcpyHostToDevice, m.s);
-    // pageable sources: the copies are staged before they return, so the caller may reuse the vectors
+    std::copy(src.begin(), src.end(), h_rows);
+    std::copy(w.begin(), w.end(), h_w);
+    if (n_adds > 0) cudaMemcpyAsync(m.adds, h_adds, n_adds * 4, cudaMemcpyHostToDevice, m.s);
     cudaMemcpyAsync(m.ptrs, P, NPTR * E * sizeof(void*), cudaMemcpyHostToDevice, m.s);
     if (m.remote) cudaMemcpyAsync(m.input, m.h_input, (size_t) T * N * 2, cudaMemcpyHostToDevice, m.s);
     if (n_rows > 0) {
-        cudaMemcpyAsync(m.rows_src, src.data(), n_rows * 4, cudaMemcpyHostToDevice, m.s);
-        cudaMemcpyAsync(m.w, w.data(), n_rows * 4, cudaMemcpyHostToDevice, m.s);
+        cudaMemcpyAsync(m.rows_src, h_rows, n_rows * 4, cudaMemcpyHostToDevice, m.s);
+        cudaMemcpyAsync(m.w, h_w, n_rows * 4, cudaMemcpyHostToDevice, m.s);
     }
+    cudaEventRecord(m.up_done, m.s);
+    m.up_live = true;
     if (m.remote) cudaMemsetAsync(m.sum, 0, (size_t) T * N * 4, m.s);
     // ---- staging: the streamed experts of batch b into slot b % SLOTS, a copy per run of adjacent arena blobs
     auto stage_batch = [&](size_t b) {
@@ -323,14 +435,14 @@ bool ExpertRunner::run_layer(int64_t layer, int64_t T, const std::vector<int32_t
         bool any = false;
         for (size_t j = j0; j < j1;) {
             const int32_t e = experts[j];
-            if (held(e) >= 0) { ++j; continue; }
+            if (held(e) >= 0 || pre(e) >= 0) { ++j; continue; }
             if (!any && m.live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0);   // its previous batch is dequantized
             any = true;
             const uint8_t* b0 = m.src->blob(layer, e);
             uint8_t* dst = (uint8_t*) P[0 * E + j];
             size_t k = j + 1;   // the run: the next experts' blobs follow this one in the arena and in staging
             if (m.src->pinned(layer, e) && stride == bb)
-                while (k < j1 && held(experts[k]) < 0 && m.src->pinned(layer, experts[k]) &&
+                while (k < j1 && held(experts[k]) < 0 && pre(experts[k]) < 0 && m.src->pinned(layer, experts[k]) &&
                        m.src->blob(layer, experts[k]) == b0 + (k - j) * bb &&
                        (uint8_t*) P[0 * E + k] == dst + (k - j) * bb) ++k;
             cudaMemcpyAsync(dst, b0, (k - j) * bb, cudaMemcpyHostToDevice, m.copy);   // pageable blobs: staged
@@ -352,6 +464,10 @@ bool ExpertRunner::run_layer(int64_t layer, int64_t T, const std::vector<int32_t
         const int n = (int) (j1 - j0);
         const int sl = (int) (b % SLOTS);
         if (has_stream[b]) cudaStreamWaitEvent(m.s, m.copied[sl], 0);
+        int32_t last_grp = -1;   // the batch's prefetched blobs have landed with their last group
+        for (size_t j = j0; j < j1; ++j)
+            if (held(experts[j]) < 0 && pre(experts[j]) >= 0) last_grp = std::max(last_grp, m.pre_grp[(size_t) experts[j]]);
+        if (last_grp >= 0) cudaStreamWaitEvent(m.s, m.pre_ev[(size_t) last_grp], 0);
         if (lay.native)
             strata::kernels::iq_dequant_experts_f16(f.gu_type, f.d_type, (const uint8_t* const*) (m.ptrs + 0 * E + j0), n,
                                                     f.up_off, f.down_off, f.n_ff, f.n_embd, m.dq_gu, m.dq_d, m.s);
@@ -359,6 +475,10 @@ bool ExpertRunner::run_layer(int64_t layer, int64_t T, const std::vector<int32_t
             for (size_t j = j0; j < j1; ++j)
                 blob_dequant_f16((const uint8_t*) P[0 * E + j], (uint16_t*) P[2 * E + j], (uint16_t*) P[5 * E + j], m.s);
         if (has_stream[b]) cudaEventRecord(m.used[sl], m.s);
+        if (last_grp >= 0) {
+            cudaEventRecord(m.pre_used, m.s);
+            m.pre_live = true;
+        }
         for (size_t j = j0; j < j1; ++j) if (held(experts[j]) >= 0) ++experts_resident;
         const int64_t r0 = off[j0], rows = off[j1] - r0;
         if (rows > ROWS) {   // one expert: its rows in pieces
@@ -383,8 +503,9 @@ bool ExpertRunner::run_layer(int64_t layer, int64_t T, const std::vector<int32_t
         const int32_t* a = m.adds + adds_at[b];
         moe_gather_add(out, m.d, r0, m.w, a, a + ntok[b], a + 2 * ntok[b] + 1, ntok[b], m.s);
     }
-    if (m.remote) {
-        cudaMemcpyAsync(m.h_sum, m.sum, (size_t) T * N * 4, cudaMemcpyDeviceToHost, m.s);
+    if (m.remote) {   // the sums as FP16 in the input's place: half the bytes over the first GPU's x4 link
+        sums_to_f16(m.sum, m.input, T * N, m.s);
+        cudaMemcpyAsync(m.h_sum, m.input, (size_t) T * N * 2, cudaMemcpyDeviceToHost, m.s);
         cudaEventRecord(m.done, m.s);
         cudaEventRecord(m.t1, m.s);
         m.timed = true;

src/prefill/gemm.cu

@@ -5,6 +5,7 @@
 #include <cublas_v2.h>
 #include <cuda_runtime.h>
 
+#include <algorithm>
 #include <cstdio>
 #include <cstdlib>
 #include <vector>
@@ -23,21 +24,51 @@ void ck(cublasStatus_t s, const char* what) {
 
 Gemm::~Gemm() {
     if (handle_) cublasDestroy((cublasHandle_t) handle_);
+    for (int i = 0; i < AUX; ++i) {
+        if (aux_handle_[i]) cublasDestroy((cublasHandle_t) aux_handle_[i]);
+        if (join_[i]) cudaEventDestroy((cudaEvent_t) join_[i]);
+        if (aux_stream_[i]) {
+            cudaStreamSynchronize((cudaStream_t) aux_stream_[i]);
+            cudaStreamDestroy((cudaStream_t) aux_stream_[i]);
+        }
+    }
+    if (fork_) cudaEventDestroy((cudaEvent_t) fork_);
 }
 
-bool Gemm::init(void* stream, std::string& err) {
+bool Gemm::init(void* stream, std::string& err, bool side_streams) {
     cublasHandle_t h = nullptr;
     if (cublasCreate(&h) != CUBLAS_STATUS_SUCCESS) { err = "prefill gemm: cublasCreate failed"; return false; }
     handle_ = h;
     stream_ = stream;
     cublasSetStream(h, (cudaStream_t) stream);
     cublasSetMathMode(h, CUBLAS_DEFAULT_MATH);
+    if (!side_streams) return true;
+    if (cudaEventCreateWithFlags((cudaEvent_t*) &fork_, cudaEventDisableTiming) != cudaSuccess) {
+        err = "prefill gemm: events";
+        return false;
+    }
+    for (int i = 0; i < AUX; ++i) {
+        cublasHandle_t a = nullptr;
+        if (cudaStreamCreateWithFlags((cudaStream_t*) &aux_stream_[i], cudaStreamNonBlocking) != cudaSuccess ||
+            cudaEventCreateWithFlags((cudaEvent_t*) &join_[i], cudaEventDisableTiming) != cudaSuccess ||
+            cublasCreate(&a) != CUBLAS_STATUS_SUCCESS) {
+            err = "prefill gemm: side streams";
+            return false;
+        }
+        aux_handle_[i] = a;
+        cublasSetStream(a, (cudaStream_t) aux_stream_[i]);
+        cublasSetMathMode(a, CUBLAS_DEFAULT_MATH);
+    }
     return true;
 }
 
 void Gemm::set_buffers(uint16_t* scratch, int64_t scratch_elems, void* workspace, size_t ws_bytes) {
-    // a fixed workspace, so the handle never allocates on the way
-    cublasSetWorkspace((cublasHandle_t) handle_, workspace, ws_bytes);
+    // a fixed workspace, so the handles never allocate on the way: a slice for each stream
+    const int ns = aux_handle_[0] ? AUX + 1 : 1;
+    const size_t slice = ws_bytes / (size_t) ns / 256 * 256;
+    cublasSetWorkspace((cublasHandle_t) handle_, workspace, slice);
+    for (int i = 0; i + 1 < ns; ++i)
+        cublasSetWorkspace((cublasHandle_t) aux_handle_[i], (uint8_t*) workspace + (size_t) (i + 1) * slice, slice);
     scratch_ = scratch;
     scratch_elems_ = scratch_elems;
 }
@@ -83,7 +114,24 @@ void Gemm::f16_grouped(const uint16_t* const* X, const uint16_t* const* W, float
         grouped_ = false;
         std::fprintf(stderr, "prefill gemm: cuBLAS does not group FP16 -> FP32 GEMMs here; one call per expert\n");
     }
-    for (int i = 0; i < n; ++i) f16(X_host[i], W_host[i], Y_host[i], rows[i], N, K);
+    // round robin over the stream and its side streams, which join it again
+    int64_t sum_rows = 0;
+    for (int i = 0; i < n; ++i) sum_rows += rows[i];
+    const int ns = aux_handle_[0] && sum_rows >= 128 * (int64_t) n ? std::min(n, AUX + 1) : 1;
+    if (ns > 1) {
+        cudaEventRecord((cudaEvent_t) fork_, (cudaStream_t) stream_);
+        for (int s = 1; s < ns; ++s) cudaStreamWaitEvent((cudaStream_t) aux_stream_[s - 1], (cudaEvent_t) fork_, 0);
+    }
+    void* const main = handle_;
+    for (int i = 0; i < n; ++i) {
+        handle_ = i % ns == 0 ? main : aux_handle_[i % ns - 1];
+        f16(X_host[i], W_host[i], Y_host[i], rows[i], N, K);
+    }
+    handle_ = main;
+    for (int s = 1; s < ns; ++s) {
+        cudaEventRecord((cudaEvent_t) join_[s - 1], (cudaStream_t) aux_stream_[s - 1]);
+        cudaStreamWaitEvent((cudaStream_t) stream_, (cudaEvent_t) join_[s - 1], 0);
+    }
 }
 
 void Gemm::native(const uint16_t* X, int ggml_type, const void* W_blocks, float* Y, int64_t T, int64_t N, int64_t K,

src/prefill/kernels.cu

@@ -94,11 +94,11 @@ __global__ void gr_mix_kernel(const float* __restrict__ xn, const float* __restr
     if (mixed_h) mixed_h[i] = hf(s);
 }
 __global__ void gr_write_kernel(float* __restrict__ R, const float* __restrict__ bo, const float* __restrict__ inj,
-                                int64_t inj_ld, int64_t T, const float* __restrict__ partial) {
+                                int64_t inj_ld, int64_t T, const uint16_t* __restrict__ partial) {
     const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
     if (i >= T * D) return;
     const int64_t t = i / D, c = (i % D) / N, d = i % N;
-    const float b = partial ? bo[t * N + d] + partial[t * N + d] : bo[t * N + d];
+    const float b = partial ? bo[t * N + d] + __half2float(__ushort_as_half(partial[t * N + d])) : bo[t * N + d];
     R[i] = fmaf(b, 2.0f * sigm(inj[t * inj_ld + c] / (float) HC), R[i]);
 }
 __global__ void gr_broadcast_kernel(const float* __restrict__ e, float* __restrict__ R, int64_t T) {
@@ -469,7 +469,8 @@ void gr_mix(const float* xn, const float* gated, float* mixed, uint16_t* mixed16
     gr_mix_kernel<<<blocks_for(T * N), 256, 0, (cudaStream_t) stream>>>(xn, gated, mixed, mixed16, T, mixed_h);
     check("gr_mix");
 }
-void gr_write(float* R, const float* bo, const float* inj, int64_t inj_ld, int64_t T, void* stream, const float* partial) {
+void gr_write(float* R, const float* bo, const float* inj, int64_t inj_ld, int64_t T, void* stream,
+              const uint16_t* partial) {
     gr_write_kernel<<<blocks_for(T * D), 256, 0, (cudaStream_t) stream>>>(R, bo, inj, inj_ld, T, partial);
     check("gr_write");
 }
@@ -528,6 +529,14 @@ void gather_rows16(const uint16_t* x16, const int32_t* src, uint16_t* dst16, int
     gather_rows16_kernel<<<blocks_for(n * (width / 8)), 256, 0, (cudaStream_t) stream>>>(x16, src, dst16, n, width);
     check("gather_rows16");
 }
+__global__ void sums_to_f16_kernel(const float* __restrict__ x, uint16_t* __restrict__ y, int64_t n) {
+    const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) y[i] = hf(fminf(fmaxf(x[i], -65504.0f), 65504.0f));
+}
+void sums_to_f16(const float* x, uint16_t* y, int64_t n, void* stream) {
+    sums_to_f16_kernel<<<blocks_for(n), 256, 0, (cudaStream_t) stream>>>(x, y, n);
+    check("sums_to_f16");
+}
 void moe_shared_gate(float* x, const float* sg, int64_t T, void* stream) {
     moe_shared_gate_kernel<<<blocks_for(T * N), 256, 0, (cudaStream_t) stream>>>(x, sg, T);
     check("moe_shared_gate");

src/prefill/prefill.cpp

@@ -53,8 +53,8 @@ struct Alloc {
 struct Bufs {
     uint16_t* w16 = nullptr;
     void* gemm_ws = nullptr;
-    // the chunk's: residual stream, the MoE's input (FP16) and output, the FFN half's injection, routing, steps
-    float *R = nullptr, *bo_moe = nullptr, *inj_f = nullptr, *w = nullptr;
+    // the chunk's: residual stream, the MoE's input (FP16) and output, the FFN half's injection, routing, steps, PLE rows
+    float *R = nullptr, *bo_moe = nullptr, *inj_f = nullptr, *w = nullptr, *ple_emb = nullptr;
     uint16_t* mixed_h = nullptr;
     int32_t *ids = nullptr, *steps = nullptr;
     // a sub-chunk's: the mixer's BF16 input, its output, the mixer half's injection
@@ -64,7 +64,7 @@ struct Bufs {
     float* emb = nullptr;
     float *xn = nullptr, *lo = nullptr, *gated = nullptr;
     uint16_t *xn16 = nullptr, *lo16 = nullptr;
-    float *ple_emb = nullptr, *ple_kp = nullptr, *ple_key = nullptr, *ple_q = nullptr, *ple_gated = nullptr;
+    float *ple_kp = nullptr, *ple_key = nullptr, *ple_q = nullptr, *ple_gated = nullptr;
     float *qkv = nullptr, *z = nullptr, *ab = nullptr, *gate = nullptr, *beta = nullptr, *hbuf = nullptr, *y = nullptr;
     uint16_t* y_h = nullptr;
     float *Kc = nullptr, *Vc = nullptr, *Qf = nullptr, *q = nullptr, *idx_raw = nullptr, *q_idx = nullptr,
@@ -73,7 +73,7 @@ struct Bufs {
     int32_t* sel_ids = nullptr;
     float *logits = nullptr, *sgate = nullptr, *sup = nullptr, *sg = nullptr;
     uint16_t* sh_h = nullptr;
-    float* sums = nullptr;        // the second GPU's sums
+    uint16_t* sums = nullptr;     // the second GPU's sums (FP16)
     uint8_t* local = nullptr;     // this GPU's expert runner
     uint64_t local_bytes = 0;
 };
@@ -93,6 +93,7 @@ Bufs layout(Alloc& a, int64_t T, int64_t cap, int64_t max_blocks, bool offload)
     b.ids = a.take<int32_t>(t * K);
     b.w = a.take<float>(t * K);
     b.steps = a.take<int32_t>(t * strata::kernels::kStepCount);
+    b.ple_emb = a.take<float>(t * N);
     b.mixed_bf = a.take<uint16_t>(p * N);
     b.bo = a.take<float>(p * N);
     b.inj_a = a.take<float>(p * HC);
@@ -104,7 +105,7 @@ Bufs layout(Alloc& a, int64_t T, int64_t cap, int64_t max_blocks, bool offload)
     b.xn = a.take<float>(p * D); b.xn16 = a.take<uint16_t>(p * D); b.lo = a.take<float>(p * LR);
     b.lo16 = a.take<uint16_t>(p * LR); b.gated = a.take<float>(p * D);
     next();   // the PLE block
-    b.ple_emb = a.take<float>(p * N); b.ple_kp = a.take<float>(p * D); b.ple_key = a.take<float>(p * D);
+    b.ple_kp = a.take<float>(p * D); b.ple_key = a.take<float>(p * D);
     b.ple_q = a.take<float>(p * D); b.ple_gated = a.take<float>(p * D);
     next();   // GDN
     b.qkv = a.take<float>(p * C); b.z = a.take<float>(p * ZV); b.ab = a.take<float>(p * 2 * HV);
@@ -119,7 +120,7 @@ Bufs layout(Alloc& a, int64_t T, int64_t cap, int64_t max_blocks, bool offload)
     b.logits = a.take<float>(p * NE); b.sgate = a.take<float>(p * 640); b.sup = a.take<float>(p * 640);
     b.sh_h = a.take<uint16_t>(p * 640); b.sg = a.take<float>(p);
     next();   // the routed experts
-    b.sums = offload ? a.take<float>(t * N) : nullptr;
+    b.sums = offload ? a.take<uint16_t>(t * N) : nullptr;
     b.local_bytes = ExpertRunner::bytes_needed(T, NE, false, !offload);
     b.local = a.take<uint8_t>(b.local_bytes);
     next();
@@ -220,6 +221,12 @@ uint64_t Prefill::bytes_needed(const core::ModelGeometry& g, const core::Session
     return a.used;
 }
 
+uint64_t Prefill::bytes_for(int64_t chunk) const {
+    const Impl& m = *impl_;
+    return bytes_needed(*m.g, *m.ss, chunk, m.offload != nullptr) +
+           (m.offload == nullptr && chunk >= ExpertRunner::kPrefetchMin ? m.local.prefetch_need(nullptr) : 0);
+}
+
 bool Prefill::bind(void* region, uint64_t bytes, int64_t chunk, std::string& err) {
     Impl& m = *impl_;
     if (region == nullptr) {
@@ -246,7 +253,11 @@ bool Prefill::bind(void* region, uint64_t bytes, int64_t chunk, std::string& err
         return false;
     }
     m.gemm.set_buffers(nullptr, 0, m.b.gemm_ws, GEMM_WS);
-    if (!m.local.bind(m.b.local, m.b.local_bytes, chunk, err)) return false;
+    // the rest of a lent region takes this GPU's prefetched experts (outside the union: copies run beside the steps)
+    const bool area = region != m.owned && bytes > a.used;
+    if (!m.local.bind(m.b.local, m.b.local_bytes, chunk, err, area ? (uint8_t*) region + a.used : nullptr,
+                      area ? bytes - a.used : 0))
+        return false;
     m.T = chunk;
     return true;
 }
@@ -345,6 +356,8 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                 prev[1] = tok;
             }
             if (!ss.ple.table->gather_batch(m.ple_rows.data(), (size_t) T, m.ple_emb_host.data(), err)) return false;
+            // before any expert copy of the chunk takes the link
+            cudaMemcpyAsync(b.ple_emb, m.ple_emb_host.data(), (size_t) T * N * 4, cudaMemcpyHostToDevice, m.cs);
             stats_.ms_ple += ms_since(tp);
         } else {
             for (int64_t t = 0; t < T; ++t) { prev[0] = prev[1]; prev[1] = (int32_t) tokens[c0 + t]; }
@@ -358,6 +371,14 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
             const core::LayerView v(*m.wt, l);
             const bool qsa = core::is_qsa_layer(g, l);
             new_layer();
+            // a long chunk uses nearly every expert: the ones this GPU's cache lacks are copied ahead, here while
+            // the dense steps run, or on the second GPU
+            if (T >= ExpertRunner::kPrefetchMin) {
+                m.ex1.clear();
+                for (int32_t e = 0; e < NE; ++e)
+                    if (!(m.host_res && m.cache && m.host_res[(size_t) l * NE + e] >= 0)) m.ex1.push_back(e);
+                if (!(m.offload ? m.offload : &m.local)->prefetch(l, m.ex1, err)) return false;
+            }
             mark(kPsPle);
             // ---- the PLE block at layer 1: the key and value projections, then the decode path's per-token
             // arithmetic (the conv reads the previous tokens' normalized rows)
@@ -365,13 +386,13 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                 const strata::kernels::PleWeights& pw = ss.ple.w;
                 for (int64_t t0 = 0; t0 < T; t0 += SUB) {
                     const int64_t P = std::min(SUB, T - t0);
-                    cudaMemcpyAsync(b.ple_emb, m.ple_emb_host.data() + t0 * N, (size_t) P * N * 4, cudaMemcpyHostToDevice, m.cs);
-                    to_bf16(b.ple_emb, b.mixed_bf, P * N, m.cs);
+                    const float* ple_emb = b.ple_emb + t0 * N;
+                    to_bf16(ple_emb, b.mixed_bf, P * N, m.cs);
                     if (pw.key_bf16 != nullptr) {
                         m.gemm.bf16(b.mixed_bf, pw.key_bf16, b.ple_kp, P, D, N);
                     } else if (const uint16_t* W = pw.key_native_data ? w16(pw.key_native_type, pw.key_native_data, D, N)
                                                                       : nullptr) {
-                        to_f16(b.ple_emb, b.mixed_h, P * N, m.cs);
+                        to_f16(ple_emb, b.mixed_h, P * N, m.cs);
                         m.gemm.f16(b.mixed_h, W, b.ple_kp, P, D, N);
                     } else {
                         err = "prefill: the PLE key has neither a BF16 nor a native GGUF form (run with --native)";
@@ -568,7 +589,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
             mark(kPsCombine);
             if (m.offload) {
                 cudaStreamWaitEvent(m.xfer, m.offload->done(), 0);
-                cudaMemcpyAsync(b.sums, m.offload->host_sum(), (size_t) T * N * 4, cudaMemcpyHostToDevice, m.xfer);
+                cudaMemcpyAsync(b.sums, m.offload->host_sum(), (size_t) T * N * 2, cudaMemcpyHostToDevice, m.xfer);
                 cudaEventRecord(m.ev_sum, m.xfer);
                 cudaStreamWaitEvent(m.cs, m.ev_sum, 0);
             }
@@ -601,6 +622,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
         return false;
     }
     stats_.experts_streamed = m.local.experts_streamed;
+    stats_.experts_prefetched = m.local.experts_prefetched;
     stats_.experts_resident = m.local.experts_resident;
     stats_.ms_experts_host = m.local.ms_host;
     stats_.ms_total += ms_since(t_start);

src/program/generate.cpp

@@ -564,13 +564,16 @@ struct PromptBuffers {
     }
     /// Before a prompt whose chunks hold up to `chunk` tokens.
     bool take(int64_t chunk, std::string& err) {
+        // sized before either loan, from the residency both GPUs' prefetch areas are sized for
+        const uint64_t need = lend ? prefill->bytes_for(chunk) : 0;
+        const uint64_t need2 = lend2 ? offload->bytes_for(chunk, loan.res ? loan.res->data() : nullptr) : 0;
         if (lend) {
-            loan.plan(strata::prefill::Prefill::bytes_needed(*g, *ss, chunk, offload != nullptr));
+            loan.plan(need);
             loan.lend();
             if (!prefill->bind(loan.base(), loan.bytes(), chunk, err)) return false;
         }
         if (lend2) {
-            loan2.plan(strata::prefill::ExpertRunner::bytes_needed(chunk, g->n_expert, true, true));
+            loan2.plan(need2);
             loan2.lend();
             if (!offload->bind(loan2.base(), loan2.bytes(), chunk, err)) return false;
         }
@@ -2496,14 +2499,15 @@ int main(int argc, char** argv) {
         ss.ple_prev[1] = pos_start >= 1 ? (int32_t) o.tokens[(size_t) (pos_start - 1)] : -1;
         const strata::prefill::PrefillStats& ps = prefill.stats();
         std::fprintf(stderr, "strata generate: prefill %lld tokens in %lld chunks, %.1f ms (%.1f tok/s); experts "
-                             "streamed %lld (host %.1f ms), resident %lld; PLE %.1f ms\n",
+                             "streamed %lld (%lld ahead), host %.1f ms, resident %lld; PLE %.1f ms\n",
                      (long long) ps.tokens, (long long) ps.chunks, ps.ms_total,
                      ps.ms_total > 0 ? 1000.0 * (double) ps.tokens / ps.ms_total : 0.0, (long long) ps.experts_streamed,
-                     ps.ms_experts_host, (long long) ps.experts_resident, ps.ms_ple);
+                     (long long) ps.experts_prefetched, ps.ms_experts_host, (long long) ps.experts_resident, ps.ms_ple);
         if (pbuf.offload != nullptr)
-            std::fprintf(stderr, "strata generate: prefill on the second GPU: experts streamed %lld, resident %lld; "
-                                 "%.1f ms queuing, %.1f ms GPU\n", (long long) offload.experts_streamed,
-                         (long long) offload.experts_resident, offload.ms_host, offload.ms_gpu);
+            std::fprintf(stderr, "strata generate: prefill on the second GPU: experts streamed %lld (%lld ahead), "
+                                 "resident %lld; %.1f ms queuing, %.1f ms GPU\n", (long long) offload.experts_streamed,
+                         (long long) offload.experts_prefetched, (long long) offload.experts_resident, offload.ms_host,
+                         offload.ms_gpu);
         if (o.prefill_profile) {
             std::fprintf(stderr, "strata generate: prefill GPU ms by section:");
             for (int i = 0; i < strata::prefill::kPsCount; ++i)