Prompt path: the experts the main GPU's cache lacks run on the second GPU

A prompt chunk uses nearly every expert, and the 3090 streamed all it did
not hold over its PCIe 4.0 x4 link.  With --second-gpu the prompt path now
hands them to the second GPU (x16): per layer the 3090 copies the chunk's
MoE input (FP16) to pinned host memory while its router and shared expert
run; the second GPU computes its experts from its own cache or streamed
from the pinned arena, and returns each token's weighted sum (FP32), which
the 3090's combine adds to its cached experts' rows (prefill/offload.cpp).

Experts go through in batches of consecutive ids (up to 16 experts and 4096
rows): one staging copy per run of adjacent arena blobs, one dequant launch
(iq_dequant_experts_f16), one gather, per-expert GEMMs (cuBLAS does not group
FP16 -> FP32), one SwiGLU and one ordered per-token add (moe_gather_add,
the same sums as one moe_scatter_add per expert).  --no-prompt-offload
streams them to the main GPU as before.

16K-token prompt, UD-Q4_K_XL: 21.6 -> 11.9 s.

CMakeLists.txt

@@ -249,7 +249,8 @@ if(STRATA_ENABLE_CUDA)
   # what composes them into a token.  No tokenizer - the prompt arrives as ids - because Gate C1 compares logits
   # against llama.cpp on the SAME ids and a tokenizer on one side of that comparison is a second variable.
   # plan v0.3 P5: batched prompt processing (tensor-core GEMMs through cuBLAS).
-  add_library(strata_prefill STATIC src/prefill/gemm.cu src/prefill/kernels.cu src/prefill/prefill.cpp)
+  add_library(strata_prefill STATIC src/prefill/gemm.cu src/prefill/kernels.cu src/prefill/offload.cpp
+                                    src/prefill/prefill.cpp)
   target_include_directories(strata_prefill PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
   target_link_libraries(strata_prefill PUBLIC strata_engine CUDA::cublas CUDA::cudart)
 

include/strata/kernels/iq_kernels.hpp

@@ -34,6 +34,10 @@ void iq_embed_rows(int ggml_type, const void* table, size_t row_bytes, const int
 /// uses: row 2r = gate row r, row 2r+1 = up row r.
 void iq_dequant_gu_f16(int ggml_type, const void* gate, const void* up, int64_t n_ff, int64_t n_embd, uint16_t* dst,
                        void* stream);
+/// iq_dequant_gu_f16 and iq_dequant_f16 (the down matrix) of n experts in one launch: expert i's blob at blobs[i]
+/// (a DEVICE array), its gate/up into gu + i * 2 * n_ff * n_embd and its down into dn + i * n_embd * n_ff.
+void iq_dequant_experts_f16(int gu_type, int d_type, const uint8_t* const* blobs, int n, size_t up_off, size_t down_off,
+                            int64_t n_ff, int64_t n_embd, uint16_t* gu, uint16_t* dn, void* stream);
 
 /// The layout of one native expert blob: [gate rows | up rows | down rows], raw GGUF blocks.
 struct NativeExpertLayout {

include/strata/prefill/gemm.hpp

@@ -32,6 +32,12 @@ public:
     void f16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_t N, int64_t K, int64_t ldy = 0,
              float beta = 0.0f);
 
+    /// 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.
+    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);
+
     /// W given as native GGUF blocks of `ggml_type`, dequantized to FP16 in the scratch, X in FP16.
     void native(const uint16_t* X, int ggml_type, const void* W_blocks, float* Y, int64_t T, int64_t N, int64_t K,
                 int64_t ldy = 0, float beta = 0.0f);
@@ -47,6 +53,7 @@ private:
     int64_t scratch_elems_ = 0;
     void* workspace_ = nullptr;
     bool external_ = false;
+    bool grouped_ = true;   // cuBLAS groups FP16 -> FP32 GEMMs (cleared at the first refusal)
 };
 
 

include/strata/prefill/kernels.hpp

@@ -46,9 +46,16 @@ 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);
-/// bo[t, :] = shared[t, :] * sigmoid(sg[t]) + sum_k w[t, k] * D[slot[t, k], :]
+/// bo[t, :] = sum_k w[t, k] * D[slot[t, k], :] + shared[t, :] * sigmoid(sg[t]) (+ partial[t, :]); a slot of -1 is
+/// an entry the second GPU computed, whose sum `partial` holds.
 void moe_combine(const float* D, const int32_t* slot, const float* w, const float* shared, const float* sg, float* bo,
-                 int64_t T, void* stream);
+                 int64_t T, void* stream, const float* partial = nullptr);
+/// sum[src[r], :] += w[r] * rows[r, :] for n rows (one expert's rows: no token twice).
+void moe_scatter_add(float* sum, const float* rows, const float* w, const int32_t* src, int64_t n, void* stream);
+/// The same for several experts' rows (row r at rows[r - r0]), where a token may recur: token tok[b] adds its rows
+/// list[start[b] .. start[b + 1]) in that order, so it matches moe_scatter_add per expert in the list's order.
+void moe_gather_add(float* sum, const float* rows, int64_t r0, const float* w, const int32_t* tok, const int32_t* start,
+                    const int32_t* list, int64_t n_tok, void* stream);
 
 // ---- QSA helpers
 /// In place: x[r, :] = x[r, :] * rsqrt(mean x^2 + eps) * w  over rows of `cols` (row stride `ld`).

include/strata/prefill/offload.hpp

@@ -0,0 +1,61 @@
+// include/strata/prefill/offload.hpp - the prompt path's experts on a second GPU.
+//
+// A prompt chunk uses nearly every expert of every layer, and the first GPU (the 3090, on a PCIe 4.0 x4 link) has
+// to stream all the ones its cache does not hold: ~0.65 GB (IQ3_XXS) to ~1.4 GB (UD-Q4_K_XL) a layer at 6.4 GB/s.
+// The second GPU sits on a x16 link.  Per layer the first GPU copies the chunk's MoE input (FP16) to pinned host
+// memory while its router and shared expert run; the second GPU copies it in, computes each of its experts -
+// from its own cache when it holds them, else streamed from the pinned arena through a ring of staging slots -
+// and adds each expert's weighted rows into a per-token FP32 sum, in ascending expert order.  The sum returns
+// through pinned host memory to the first GPU, which computes the experts its own cache holds meanwhile and adds
+// the sum in its combine.  Neither GPU reads the other's memory: Windows gives GeForce cards no peer access.
+#pragma once
+
+#include "strata/core/expert_cache.hpp"
+#include "strata/core/expert_source.hpp"
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace strata::prefill {
+
+class Offload {
+public:
+    Offload();
+    ~Offload();
+    Offload(const Offload&) = delete;
+    Offload& operator=(const Offload&) = delete;
+
+    /// Buffers for chunks of up to `chunk` tokens on `device`; `main_device` is made current again after every
+    /// call.  `cache`/`res`: the device's expert cache and its residency (n_layers x n_expert slots, -1 when not
+    /// held), or null.  `borrow`/`borrow_bytes`: device memory for the buffers (lent cache slots), else allocated.
+    bool init(int device, int main_device, core::ExpertSource* src, const core::ExpertCache* cache, const int32_t* res,
+              int64_t n_expert, int64_t chunk, std::string& err, void* borrow = nullptr, uint64_t borrow_bytes = 0);
+    /// Device bytes `init` needs for chunks of `chunk` tokens (what a borrowed region must hold).
+    static uint64_t bytes_needed(int64_t chunk, int64_t n_expert);
+
+    /// Pinned [chunk, n_embd]: the first GPU writes the MoE input (FP16) here, and reads the sums (FP32) here.
+    uint16_t* host_input() const;
+    float* 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].  Queues the work behind `input_ready` (an event of the first GPU, recorded
+    /// once host_input holds the input) and returns; done() is recorded once host_sum holds the sums.
+    bool run_layer(int64_t layer, int64_t T, const std::vector<int32_t>& experts, const std::vector<int32_t>& off,
+                   const std::vector<int32_t>& src, const std::vector<float>& w, cudaEvent_t input_ready, std::string& err);
+    cudaEvent_t done() const;
+
+    int64_t experts_resident = 0;   ///< expert-layer groups served from its cache
+    int64_t experts_streamed = 0;   ///< expert blobs copied from the arena
+    double ms_host = 0;             ///< host time queuing its work
+    double ms_gpu = 0;              ///< its GPU time from the input's arrival (all layers but the last)
+
+private:
+    struct Impl;
+    std::unique_ptr<Impl> impl_;
+};
+
+}  // namespace strata::prefill

include/strata/prefill/prefill.hpp

@@ -23,6 +23,8 @@
 
 namespace strata::prefill {
 
+class Offload;
+
 /// The sections `Prefill::profile` times, in the order a layer runs them.
 enum PrefillSection {
     kPsPle, kPsHcRead, kPsGdn, kPsGdnConv, kPsGdnScan, kPsGdnNorm, kPsGdnOut, kPsQsaProj, kPsQsaIndexer, kPsQsaScores, kPsQsaAttn, kPsQsaOut, kPsHcFfn, kPsRouter,
@@ -78,6 +80,9 @@ public:
     /// one after every chunk).
     bool profile = false;
 
+    /// The experts this GPU's cache does not hold, computed on a second GPU instead of streamed here; null: here.
+    Offload* offload = nullptr;
+
 private:
     struct Impl;
     std::unique_ptr<Impl> impl_;

src/kernels/cuda/iq_kernels.cu

@@ -937,6 +937,24 @@ __global__ void dequant_gu_kernel(int ty, const void* __restrict__ gate, const v
     dq_dispatch<__half>(ty, parity ? up : gate, i, y + ((2 * r + parity) * per_row + c) * QK_K, threadIdx.x);
 }
 
+// several experts at once, expert z's blob at blobs[z]: first its gate/up superblocks (as dequant_gu_kernel), then
+// its down superblocks (as dequant_flat_kernel)
+__global__ void dequant_experts_kernel(int gu_ty, int d_ty, const uint8_t* const* __restrict__ blobs, size_t up_off,
+                                       size_t down_off, int64_t per_row, int64_t gu_blocks, __half* __restrict__ gu,
+                                       __half* __restrict__ dn, int64_t gu_elems, int64_t dn_elems) {
+    const uint8_t* b = blobs[blockIdx.z];
+    const int64_t i = blockIdx.x;
+    if (i < 2 * gu_blocks) {
+        const int parity = (int) (i / gu_blocks);
+        const int64_t j = i % gu_blocks, r = j / per_row, c = j % per_row;
+        dq_dispatch<__half>(gu_ty, parity ? b + up_off : b, j,
+                            gu + blockIdx.z * gu_elems + ((2 * r + parity) * per_row + c) * QK_K, threadIdx.x);
+    } else {
+        const int64_t j = i - 2 * gu_blocks;
+        dq_dispatch<__half>(d_ty, b + down_off, j, dn + blockIdx.z * dn_elems + j * QK_K, threadIdx.x);
+    }
+}
+
 // the formats dq_dispatch dequantizes
 bool is_iq(int t) {
     return t == 16 || t == 17 || t == 18 || t == 20 || t == 21 || t == 22 || t == 23 || t == 29 || t == 42 || t == 11 ||
@@ -1029,6 +1047,20 @@ void iq_dequant_gu_f16(int t, const void* gate, const void* up, int64_t n_ff, in
     check("iq_dequant_gu_f16");
 }
 
+void iq_dequant_experts_f16(int gu_type, int d_type, const uint8_t* const* blobs, int n, size_t up_off, size_t down_off,
+                            int64_t n_ff, int64_t n_embd, uint16_t* gu, uint16_t* dn, void* stream) {
+    if (n <= 0) return;
+    if (n_embd % 256 != 0 || (n_embd * n_ff) % 256 != 0 || !is_iq(gu_type) || !is_iq(d_type) || n > 65535) {
+        std::fprintf(stderr, "iq_dequant_experts_f16: bad arguments\n");
+        std::exit(1);
+    }
+    const int64_t per_row = n_embd / 256, gu_blocks = n_ff * per_row, d_blocks = n_embd * n_ff / 256;
+    dequant_experts_kernel<<<dim3((unsigned) (2 * gu_blocks + d_blocks), 1, (unsigned) n), 32, 0, (cudaStream_t) stream>>>(
+        gu_type, d_type, blobs, up_off, down_off, per_row, gu_blocks, (__half*) gu, (__half*) dn, 2 * n_ff * n_embd,
+        n_embd * n_ff);
+    check("iq_dequant_experts_f16");
+}
+
 bool native_expert_supported(int gu_type, int d_type, int64_t n_embd, int64_t n_ff) noexcept {
     const int qg = fmt_qk(gu_type), qd = fmt_qk(d_type);
     return qg > 0 && qd > 0 && is_iq(gu_type) && is_iq(d_type) && n_embd % qg == 0 && n_ff % qd == 0;

src/prefill/gemm.cu

@@ -7,6 +7,7 @@
 
 #include <cstdio>
 #include <cstdlib>
+#include <vector>
 
 namespace strata::prefill {
 namespace {
@@ -86,6 +87,27 @@ void Gemm::f16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_
        "cublasGemmEx f16");
 }
 
+void Gemm::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) {
+    if (n <= 0) return;
+    if (grouped_) {
+        // column-major, as f16(): Y_i^T[N, rows] = W_i (K x N col-major, transposed) . X_i^T[K, rows]; a group per problem
+        const std::vector<cublasOperation_t> ta((size_t) n, CUBLAS_OP_T), tb((size_t) n, CUBLAS_OP_N);
+        const std::vector<int> m((size_t) n, (int) N), k((size_t) n, (int) K), ld_wx((size_t) n, (int) K),
+            ld_y((size_t) n, (int) N), size((size_t) n, 1);
+        const std::vector<float> alpha((size_t) n, 1.0f), beta((size_t) n, 0.0f);
+        const cublasStatus_t st = cublasGemmGroupedBatchedEx(
+            (cublasHandle_t) handle_, ta.data(), tb.data(), m.data(), rows, k.data(), alpha.data(), (const void* const*) W,
+            CUDA_R_16F, ld_wx.data(), (const void* const*) X, CUDA_R_16F, ld_wx.data(), beta.data(), (void* const*) Y,
+            CUDA_R_32F, ld_y.data(), n, size.data(), CUBLAS_COMPUTE_32F);
+        if (st == CUBLAS_STATUS_SUCCESS) return;
+        if (st != CUBLAS_STATUS_NOT_SUPPORTED && st != CUBLAS_STATUS_INVALID_VALUE) ck(st, "cublasGemmGroupedBatchedEx");
+        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);
+}
+
 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,
                   int64_t ldy, float beta) {
     if (N * K > scratch_elems_) {

src/prefill/kernels.cu

@@ -348,14 +348,27 @@ __global__ void gather_rows16_kernel(const uint16_t* __restrict__ x, const int32
 }
 __global__ void moe_combine_kernel(const float* __restrict__ Dm, const int32_t* __restrict__ slot,
                                    const float* __restrict__ w, const float* __restrict__ shared,
-                                   const float* __restrict__ sg, float* __restrict__ bo, int64_t T) {
+                                   const float* __restrict__ sg, const float* __restrict__ partial,
+                                   float* __restrict__ bo, int64_t T) {
     const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
     if (i >= T * N) return;
     const int64_t t = i / N, d = i % N;
     float s = 0.0f;
 #pragma unroll
-    for (int k = 0; k < 10; ++k) s = fmaf(w[t * 10 + k], Dm[(int64_t) slot[t * 10 + k] * N + d], s);
-    bo[i] = s + shared[i] * sigm(sg[t]);
+    for (int k = 0; k < 10; ++k) {
+        const int32_t r = slot[t * 10 + k];
+        if (r >= 0) s = fmaf(w[t * 10 + k], Dm[(int64_t) r * N + d], s);
+    }
+    const float v = s + shared[i] * sigm(sg[t]);
+    bo[i] = partial ? v + partial[i] : v;
+}
+__global__ void moe_scatter_add_kernel(float* __restrict__ sum, const float* __restrict__ rows,
+                                       const float* __restrict__ w, const int32_t* __restrict__ src, int64_t n) {
+    const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n * N) return;
+    const int64_t r = i / N, d = i % N;
+    float* o = sum + (int64_t) src[r] * N + d;
+    *o = fmaf(w[r], rows[i], *o);
 }
 
 // ---------------------------------------------------------------- QSA helpers
@@ -526,10 +539,36 @@ void gather_rows16(const uint16_t* x16, const int32_t* src, uint16_t* dst16, int
     check("gather_rows16");
 }
 void moe_combine(const float* Dm, const int32_t* slot, const float* w, const float* shared, const float* sg, float* bo,
-                 int64_t T, void* stream) {
-    moe_combine_kernel<<<blocks_for(T * N), 256, 0, (cudaStream_t) stream>>>(Dm, slot, w, shared, sg, bo, T);
+                 int64_t T, void* stream, const float* partial) {
+    moe_combine_kernel<<<blocks_for(T * N), 256, 0, (cudaStream_t) stream>>>(Dm, slot, w, shared, sg, partial, bo, T);
     check("moe_combine");
 }
+__global__ void moe_gather_add_kernel(float* __restrict__ sum, const float* __restrict__ rows, int64_t r0,
+                                      const float* __restrict__ w, const int32_t* __restrict__ tok,
+                                      const int32_t* __restrict__ start, const int32_t* __restrict__ list) {
+    const int64_t d = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+    if (d >= N) return;
+    const int b = blockIdx.y;
+    float* o = sum + (int64_t) tok[b] * N + d;
+    float s = *o;
+    for (int i = start[b]; i < start[b + 1]; ++i) {
+        const int r = list[i];
+        s = fmaf(w[r], rows[(r - r0) * N + d], s);
+    }
+    *o = s;
+}
+void moe_gather_add(float* sum, const float* rows, int64_t r0, const float* w, const int32_t* tok, const int32_t* start,
+                    const int32_t* list, int64_t n_tok, void* stream) {
+    if (n_tok <= 0) return;
+    moe_gather_add_kernel<<<dim3((unsigned) ((N + 255) / 256), (unsigned) n_tok), 256, 0, (cudaStream_t) stream>>>(
+        sum, rows, r0, w, tok, start, list);
+    check("moe_gather_add");
+}
+void moe_scatter_add(float* sum, const float* rows, const float* w, const int32_t* src, int64_t n, void* stream) {
+    if (n <= 0) return;
+    moe_scatter_add_kernel<<<blocks_for(n * N), 256, 0, (cudaStream_t) stream>>>(sum, rows, w, src, n);
+    check("moe_scatter_add");
+}
 void rms_rows(float* x, const float* w, int64_t rows, int64_t cols, int64_t ld, float eps, void* stream) {
     if (rows <= 0) return;
     rms_rows_kernel<<<(unsigned) rows, 256, 0, (cudaStream_t) stream>>>(x, w, cols, ld, eps);

src/prefill/offload.cpp

@@ -0,0 +1,360 @@
+// src/prefill/offload.cpp - see include/strata/prefill/offload.hpp.
+#include "strata/prefill/offload.hpp"
+
+#include "strata/kernels/cpu/expert_layout.hpp"
+#include "strata/kernels/iq_kernels.hpp"
+#include "strata/prefill/gemm.hpp"
+#include "strata/prefill/kernels.hpp"
+
+#include <algorithm>
+#include <chrono>
+#include <cstring>
+
+namespace strata::prefill {
+namespace {
+
+constexpr int64_t N = 2560, FF2 = 1280, FF = 640;
+// Experts go through in batches of consecutive ids, at most G experts and ROWS rows (an expert with more rows is a
+// batch of its own, in pieces of ROWS): one staging copy per run of adjacent arena blobs, one dequant launch, one
+// gather, two grouped GEMMs, one SwiGLU and one scatter-add per batch.  Per-expert launches made the host, not the
+// GPU, the limit (~10 calls for each of ~500 experts a layer).
+constexpr int G = 16;
+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
+
+struct DeviceScope {   // the second GPU current for one call, the main one again afterwards
+    int main;
+    DeviceScope(int dev, int main_dev) : main(main_dev) { cudaSetDevice(dev); }
+    ~DeviceScope() { cudaSetDevice(main); }
+};
+
+// Bump allocation from a borrowed region, else cudaMalloc; with neither it only counts (bytes_needed).
+struct Alloc {
+    uint8_t* base = nullptr;
+    uint64_t cap = 0, used = 0;
+    bool count_only = false;
+    std::vector<void*>* owned = nullptr;
+    template <typename T> T* take(size_t n, bool& ok) {
+        const uint64_t bytes = ((uint64_t) n * sizeof(T) + 255) & ~255ull;
+        if (count_only) { used += bytes; return nullptr; }
+        if (base != nullptr) {
+            if (used + bytes > cap) { ok = false; return nullptr; }
+            T* p = (T*) (base + used);
+            used += bytes;
+            return p;
+        }
+        void* p = nullptr;
+        if (cudaMalloc(&p, bytes) != cudaSuccess) { ok = false; return nullptr; }
+        owned->push_back(p);
+        used += bytes;
+        return (T*) p;
+    }
+};
+
+}  // namespace
+
+struct Offload::Impl {
+    int dev = -1, main = 0;
+    core::ExpertSource* src = nullptr;
+    const core::ExpertCache* cache = nullptr;
+    const int32_t* res = nullptr;
+    int64_t n_expert = 0, chunk = 0;
+    size_t blob_cap = 0;   // staging bytes per expert
+    cudaStream_t s = nullptr, copy = nullptr;
+    cudaEvent_t done = nullptr, copied[SLOTS] = {}, used[SLOTS] = {};
+    cudaEvent_t t0 = nullptr, t1 = nullptr;   // the previous layer's GPU time, read at the next one
+    bool timed = false, live[SLOTS] = {};
+    Gemm gemm;
+    std::vector<void*> owned;
+    uint16_t *input = nullptr, *xs = nullptr, *hh = nullptr, *dq_gu = nullptr, *dq_d = nullptr;
+    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;
+    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;
+
+    // the same sequence counted (bytes_needed) or carved
+    bool carve(Alloc& a, int64_t T) {
+        bool ok = true;
+        const size_t t = (size_t) T;
+        input = a.take<uint16_t>(t * N, ok);
+        sum = a.take<float>(t * N, ok);
+        rows_src = a.take<int32_t>(t * 10, ok);
+        w = a.take<float>(t * 10, ok);
+        adds = a.take<int32_t>(t * 10 * 3 + (size_t) n_expert * 2, ok);
+        xs = a.take<uint16_t>((size_t) ROWS * N, ok);
+        gu = a.take<float>((size_t) ROWS * FF2, ok);
+        hh = a.take<uint16_t>((size_t) ROWS * FF, ok);
+        d = a.take<float>((size_t) ROWS * N, ok);
+        dq_gu = a.take<uint16_t>((size_t) G * FF2 * N, ok);
+        dq_d = a.take<uint16_t>((size_t) G * N * FF, ok);
+        for (auto& p : stage) p = a.take<uint8_t>((size_t) G * blob_cap, ok);
+        ptrs = a.take<void*>((size_t) NPTR * (size_t) n_expert, ok);
+        void* ws = a.take<uint8_t>(GEMM_WS, ok);
+        if (ok && !a.count_only) {
+            std::string err;
+            ok = gemm.init_external(s, nullptr, 0, ws, GEMM_WS, err);
+        }
+        return ok;
+    }
+};
+
+Offload::Offload() : impl_(new Impl) {}
+
+Offload::~Offload() {
+    Impl& m = *impl_;
+    if (m.dev < 0) return;
+    DeviceScope scope(m.dev, m.main);
+    if (m.s) cudaStreamSynchronize(m.s);
+    if (m.copy) cudaStreamSynchronize(m.copy);
+    for (int i = 0; i < SLOTS; ++i) {
+        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);
+    if (m.h_input) cudaFreeHost(m.h_input);
+    if (m.h_sum) cudaFreeHost(m.h_sum);
+    for (void* p : m.owned) cudaFree(p);
+    if (m.copy) cudaStreamDestroy(m.copy);
+    if (m.s) cudaStreamDestroy(m.s);
+}
+
+uint64_t Offload::bytes_needed(int64_t chunk, int64_t n_expert) {
+    Impl m;
+    m.n_expert = n_expert;
+    m.blob_cap = ((size_t) strata::kernels::cpu::expert_layout().max_blob + 255) & ~(size_t) 255;
+    Alloc a;
+    a.count_only = true;
+    m.carve(a, chunk);
+    return a.used + (1u << 20);
+}
+
+bool Offload::init(int device, int main_device, core::ExpertSource* src, const core::ExpertCache* cache,
+                   const int32_t* res, int64_t n_expert, int64_t chunk, std::string& err, void* borrow,
+                   uint64_t borrow_bytes) {
+    Impl& m = *impl_;
+    if (!strata::kernels::cpu::expert_layout().native) {
+        err = "prompt offload: needs a native pack";
+        return false;
+    }
+    m.dev = device;
+    m.main = main_device;
+    m.src = src;
+    m.cache = cache;
+    m.res = res;
+    m.n_expert = n_expert;
+    m.chunk = chunk;
+    m.blob_cap = ((size_t) strata::kernels::cpu::expert_layout().max_blob + 255) & ~(size_t) 255;
+    m.ptrs_host.resize((size_t) NPTR * (size_t) n_expert);
+    m.adds_host.reserve((size_t) chunk * 30);
+    m.mark.assign((size_t) chunk, -1);
+    m.first.resize((size_t) chunk);
+    m.count.resize((size_t) chunk);
+    DeviceScope scope(device, main_device);
+    bool ok = cudaStreamCreateWithFlags(&m.s, cudaStreamNonBlocking) == cudaSuccess &&
+              cudaStreamCreateWithFlags(&m.copy, cudaStreamNonBlocking) == cudaSuccess &&
+              cudaEventCreateWithFlags(&m.done, cudaEventDisableTiming) == cudaSuccess &&
+              cudaEventCreate(&m.t0) == cudaSuccess && cudaEventCreate(&m.t1) == cudaSuccess;
+    for (int i = 0; ok && i < SLOTS; ++i)
+        ok = cudaEventCreateWithFlags(&m.copied[i], cudaEventDisableTiming) == cudaSuccess &&
+             cudaEventCreateWithFlags(&m.used[i], cudaEventDisableTiming) == cudaSuccess;
+    if (!ok) { err = "prompt offload: streams or events"; return false; }
+    // the input and the sums: pinned and portable, both GPUs copy them
+    if (cudaHostAlloc((void**) &m.h_input, (size_t) chunk * N * 2, cudaHostAllocPortable) != cudaSuccess ||
+        cudaHostAlloc((void**) &m.h_sum, (size_t) chunk * N * 4, cudaHostAllocPortable) != cudaSuccess) {
+        err = "prompt offload: pinned input and sum buffers";
+        return false;
+    }
+    Alloc a;
+    a.base = (uint8_t*) borrow;
+    a.cap = borrow_bytes;
+    a.owned = &m.owned;
+    if (!m.carve(a, chunk)) {
+        err = "prompt offload: device buffers for a chunk of " + std::to_string(chunk) + " tokens do not fit";
+        return false;
+    }
+    return true;
+}
+
+uint16_t* Offload::host_input() const { return impl_->h_input; }
+float* Offload::host_sum() const { return impl_->h_sum; }
+cudaEvent_t Offload::done() const { return impl_->done; }
+
+bool Offload::run_layer(int64_t layer, int64_t T, const std::vector<int32_t>& experts, const std::vector<int32_t>& off,
+                        const std::vector<int32_t>& src, const std::vector<float>& w, cudaEvent_t input_ready,
+                        std::string& err) {
+    Impl& m = *impl_;
+    const size_t n_exp = experts.size(), n_rows = src.size();
+    if (T > m.chunk || off.size() != n_exp + 1 || w.size() != n_rows || (size_t) off.back() != n_rows ||
+        n_rows > (size_t) T * 10 || n_exp > (size_t) m.n_expert) {
+        err = "prompt offload: a layer's share is out of range";
+        return false;
+    }
+    DeviceScope scope(m.dev, m.main);
+    const auto h0 = std::chrono::steady_clock::now();
+    if (m.timed) {   // the previous layer is done by now (the main GPU waited for it)
+        float ms = 0;
+        cudaEventElapsedTime(&ms, m.t0, m.t1);
+        ms_gpu += ms;
+    }
+    const auto& f = strata::kernels::cpu::expert_layout().fmt[(size_t) layer];
+    const size_t bb = (size_t) strata::kernels::cpu::expert_layout().blob_bytes(layer);
+    const size_t stride = (bb + 255) & ~(size_t) 255;   // staging: a run of arena blobs lands as one copy when equal
+    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;
+    };
+    // ---- 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;) {
+        size_t j1 = j0 + 1;
+        int64_t rows = off[j0 + 1] - off[j0];
+        if (rows <= ROWS)
+            while (j1 < n_exp && j1 - j0 < (size_t) G && rows + (off[j1 + 1] - off[j1]) <= ROWS) rows += off[j1 + 1] - off[j1++];
+        batches.emplace_back(j0, j1);
+        j0 = j1;
+    }
+    void** P = m.ptrs_host.data();
+    const size_t E = (size_t) m.n_expert;   // array stride
+    std::vector<int> rows_of(n_exp);
+    for (size_t b = 0; b < batches.size(); ++b) {
+        const auto [j0, j1] = batches[b];
+        size_t streamed = 0;
+        for (size_t j = j0; j < j1; ++j) {
+            const int64_t r = off[j] - off[j0];
+            const int32_t slot = held(experts[j]);
+            P[0 * E + j] = slot >= 0 ? (void*) m.cache->device_slot(slot) : (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;
+            P[4 * E + j] = m.hh + r * FF;
+            P[5 * E + j] = m.dq_d + (j - j0) * (size_t) (N * FF);
+            P[6 * E + j] = m.d + r * N;
+            rows_of[j] = off[j + 1] - off[j];
+        }
+    }
+    // per batch, its tokens and each one's rows in order (a token may be routed to several of the batch's experts):
+    // [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();
+    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;
+        for (int32_t r = r0; r < r1; ++r) {
+            const int32_t t = src[(size_t) r];
+            if (m.mark[(size_t) t] != (int32_t) b) {
+                m.mark[(size_t) t] = (int32_t) b;
+                m.first[(size_t) t] = (int32_t) toks.size();
+                m.count[(size_t) toks.size()] = 0;
+                toks.push_back(t);
+            }
+            ++m.count[(size_t) m.first[(size_t) t]];
+        }
+        const size_t nt = toks.size(), at = m.adds_host.size();
+        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;
+        int32_t* h_start = h_tok + nt;
+        int32_t* h_list = h_start + nt + 1;
+        h_start[0] = 0;
+        for (size_t i = 0; i < nt; ++i) {
+            h_tok[i] = toks[i];
+            h_start[i + 1] = h_start[i] + m.count[i];
+            m.count[i] = h_start[i];   // the fill cursor
+        }
+        for (int32_t r = r0; r < r1; ++r) h_list[m.count[(size_t) m.first[(size_t) src[(size_t) r]]]++] = r;
+        for (int32_t t : toks) m.mark[(size_t) t] = -1;
+    }
+    cudaStreamWaitEvent(m.s, input_ready, 0);
+    cudaEventRecord(m.t0, m.s);
+    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
+    cudaMemcpyAsync(m.ptrs, P, NPTR * E * sizeof(void*), cudaMemcpyHostToDevice, m.s);
+    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);
+    }
+    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) {
+        const auto [j0, j1] = batches[b];
+        const int sl = (int) (b % SLOTS);
+        bool any = false;
+        for (size_t j = j0; j < j1;) {
+            const int32_t e = experts[j];
+            if (held(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]) &&
+                       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
+            experts_streamed += (int64_t) (k - j);
+            j = k;
+        }
+        if (any) {
+            cudaEventRecord(m.copied[sl], m.copy);
+            m.live[sl] = true;
+        }
+        return any;
+    };
+    std::vector<char> has_stream(batches.size(), 0);
+    if (!batches.empty()) has_stream[0] = stage_batch(0);
+    auto** dp = (const uint16_t**) m.ptrs;   // device arrays by index
+    for (size_t b = 0; b < batches.size(); ++b) {
+        if (b + 1 < batches.size()) has_stream[b + 1] = stage_batch(b + 1);   // the next batch copies meanwhile
+        const auto [j0, j1] = batches[b];
+        const int n = (int) (j1 - j0);
+        const int sl = (int) (b % SLOTS);
+        if (has_stream[b]) cudaStreamWaitEvent(m.s, m.copied[sl], 0);
+        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);
+        if (has_stream[b]) cudaEventRecord(m.used[sl], m.s);
+        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
+            for (int64_t p = 0; p < rows; p += ROWS) {
+                const int64_t nr = std::min<int64_t>(ROWS, rows - p);
+                gather_rows16(m.input, m.rows_src + r0 + p, m.xs, nr, N, m.s);
+                m.gemm.f16(m.xs, m.dq_gu, m.gu, nr, FF2, N);
+                swiglu_interleaved(m.gu, m.hh, nr, m.s);
+                m.gemm.f16(m.hh, m.dq_d, m.d, nr, N, FF);
+                moe_scatter_add(m.sum, m.d, m.w + r0 + p, m.rows_src + r0 + p, nr, m.s);
+            }
+            continue;
+        }
+        gather_rows16(m.input, m.rows_src + r0, m.xs, rows, N, m.s);
+        m.gemm.f16_grouped(dp + 1 * E + j0, dp + 2 * E + j0, (float* const*) (m.ptrs + 3 * E + j0),
+                           (const uint16_t* const*) (P + 1 * E + j0), (const uint16_t* const*) (P + 2 * E + j0),
+                           (float* const*) (P + 3 * E + j0), rows_of.data() + j0, n, FF2, N);
+        swiglu_interleaved(m.gu, m.hh, rows, m.s);
+        m.gemm.f16_grouped(dp + 4 * E + j0, dp + 5 * E + j0, (float* const*) (m.ptrs + 6 * E + j0),
+                           (const uint16_t* const*) (P + 4 * E + j0), (const uint16_t* const*) (P + 5 * E + j0),
+                           (float* const*) (P + 6 * E + j0), rows_of.data() + j0, n, N, FF);
+        const int32_t* a = m.adds + adds_at[b];
+        moe_gather_add(m.sum, m.d, r0, m.w, a, a + ntok[b], a + 2 * ntok[b] + 1, ntok[b], m.s);
+    }
+    cudaMemcpyAsync(m.h_sum, m.sum, (size_t) T * N * 4, cudaMemcpyDeviceToHost, m.s);
+    cudaEventRecord(m.done, m.s);
+    cudaEventRecord(m.t1, m.s);
+    m.timed = true;
+    ms_host += std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - h0).count();
+    const cudaError_t e = cudaGetLastError();
+    if (e != cudaSuccess) { err = std::string("prompt offload: ") + cudaGetErrorString(e); return false; }
+    return true;
+}
+
+}  // namespace strata::prefill

src/prefill/prefill.cpp

@@ -13,6 +13,7 @@
 #include "strata/kernels/qsa_select.hpp"
 #include "strata/prefill/gemm.hpp"
 #include "strata/prefill/kernels.hpp"
+#include "strata/prefill/offload.hpp"
 
 #include <cuda_runtime.h>
 
@@ -98,6 +99,11 @@ struct Prefill::Impl {
     int32_t *ids = nullptr, *slot_dev = nullptr, *src_dev = nullptr;
     uint16_t *Xs = nullptr, *Hh = nullptr, *sh_h = nullptr;
     std::vector<int32_t> ids_host, slot_host, src_host, cnt, off;
+    // the second GPU's share (Prefill::offload): the input goes out and the sums come back on `xfer`
+    cudaStream_t xfer = nullptr;
+    cudaEvent_t ev_mixed = nullptr, ev_input = nullptr, ev_sum = nullptr;
+    std::vector<int32_t> entry_of, experts2, off2, src2;
+    std::vector<float> w_host, w2;
     uint16_t* dq_gu[DQ] = {};
     uint16_t* dq_d[DQ] = {};
     uint8_t* stage_dev[STAGE] = {};
@@ -121,6 +127,9 @@ Prefill::~Prefill() {
         if (impl_->stage_host[i]) cudaFreeHost(impl_->stage_host[i]);
     }
     if (impl_->copy) cudaStreamDestroy(impl_->copy);
+    if (impl_->xfer) cudaStreamSynchronize(impl_->xfer);
+    for (cudaEvent_t e : {impl_->ev_mixed, impl_->ev_input, impl_->ev_sum}) if (e) cudaEventDestroy(e);
+    if (impl_->xfer) cudaStreamDestroy(impl_->xfer);
     for (void* p : impl_->owned) cudaFree(p);
 }
 
@@ -138,7 +147,14 @@ bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, co
     if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_expert != NE || ss.k != K) {
         err = "prefill: geometry differs from the artifact's"; return false;
     }
-    if (cudaStreamCreateWithFlags(&m.copy, cudaStreamNonBlocking) != cudaSuccess) { err = "prefill: copy stream"; return false; }
+    if (cudaStreamCreateWithFlags(&m.copy, cudaStreamNonBlocking) != cudaSuccess ||
+        cudaStreamCreateWithFlags(&m.xfer, cudaStreamNonBlocking) != cudaSuccess ||
+        cudaEventCreateWithFlags(&m.ev_mixed, cudaEventDisableTiming) != cudaSuccess ||
+        cudaEventCreateWithFlags(&m.ev_input, cudaEventDisableTiming) != cudaSuccess ||
+        cudaEventCreateWithFlags(&m.ev_sum, cudaEventDisableTiming) != cudaSuccess) {
+        err = "prefill: streams and events";
+        return false;
+    }
     const size_t T = (size_t) chunk;
     bool ok = true;
     Alloc o;
@@ -186,6 +202,7 @@ bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, co
     }
     m.ple_emb = o.take<float>(T * N, ok);
     m.ids_host.resize(T * K); m.slot_host.resize(T * K); m.src_host.resize(T * K); m.cnt.resize(NE); m.off.resize(NE + 1);
+    m.entry_of.resize(T * K); m.w_host.resize(T * K);
     m.ple_emb_host.resize(T * N); m.ple_rows.resize(T * strata::kernels::PLE_N_HEADS);
     if (!ok) { err = "prefill: device buffers for a chunk of " + std::to_string(chunk) + " tokens do not fit"; return false; }
     return true;
@@ -430,6 +447,12 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                                           *wsu = need(v, "ffn_up_shexp.weight", err),
                                           *wsd = need(v, "ffn_down_shexp.weight", err);
                     if (!wr || !wgi || !wsg || !wsu || !wsd) return false;
+                    if (offload) {   // the input's trip to the second GPU starts while the router and shared expert run
+                        cudaEventRecord(m.ev_mixed, m.cs);
+                        cudaStreamWaitEvent(m.xfer, m.ev_mixed, 0);
+                        cudaMemcpyAsync(offload->host_input(), m.mixed_h, (size_t) T * N * 2, cudaMemcpyDeviceToHost, m.xfer);
+                        cudaEventRecord(m.ev_input, m.xfer);
+                    }
                     mark(kPsRouter);
                     if (!bf16_proj(m.gemm, wr, m.mixed_bf, m.logits, T, v.name("ffn_gate_inp.weight"), err)) return false;
                     route(m.logits, m.ids, m.w, T, m.cs);
@@ -442,6 +465,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                     m.gemm.bf16(m.mixed_bf, (const uint16_t*) wgi->data, m.sg, T, 1, N);
                     // group the (token, k) pairs by expert on the host
                     cudaMemcpyAsync(m.ids_host.data(), m.ids, (size_t) T * K * 4, cudaMemcpyDeviceToHost, m.cs);
+                    if (offload) cudaMemcpyAsync(m.w_host.data(), m.w, (size_t) T * K * 4, cudaMemcpyDeviceToHost, m.cs);
                     cudaStreamSynchronize(m.cs);
                     std::fill(m.cnt.begin(), m.cnt.end(), 0);
                     for (int64_t i = 0; i < T * K; ++i) {
@@ -457,14 +481,38 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                         const int32_t p = fill[(size_t) e]++;
                         m.slot_host[(size_t) i] = p;
                         m.src_host[(size_t) p] = (int32_t) (i / K);
+                        m.entry_of[(size_t) p] = (int32_t) i;
+                    }
+                    auto resident = [&](int32_t e) {
+                        return m.host_res && m.cache && m.host_res[(size_t) l * NE + e] >= 0;
+                    };
+                    if (offload) {   // the experts this GPU's cache does not hold go to the second GPU
+                        m.experts2.clear();
+                        m.off2.assign(1, 0);
+                        m.src2.clear();
+                        m.w2.clear();
+                        for (int32_t e = 0; e < NE; ++e) {
+                            if (m.cnt[(size_t) e] == 0 || resident(e)) continue;
+                            m.experts2.push_back(e);
+                            for (int32_t p = m.off[(size_t) e]; p < m.off[(size_t) e + 1]; ++p) {
+                                const int32_t i = m.entry_of[(size_t) p];
+                                m.src2.push_back(i / (int32_t) K);
+                                m.w2.push_back(m.w_host[(size_t) i]);
+                                m.slot_host[(size_t) i] = -1;   // the combine takes its row from the returned sum
+                            }
+                            m.off2.push_back((int32_t) m.src2.size());
+                        }
+                        if (!offload->run_layer(l, T, m.experts2, m.off2, m.src2, m.w2, m.ev_input, err)) return false;
                     }
                     cudaMemcpyAsync(m.slot_dev, m.slot_host.data(), (size_t) T * K * 4, cudaMemcpyHostToDevice, m.cs);
                     cudaMemcpyAsync(m.src_dev, m.src_host.data(), (size_t) T * K * 4, cudaMemcpyHostToDevice, m.cs);
                     gather_rows16(m.mixed_h, m.src_dev, m.Xs, T * K, N, m.cs);
                     mark(kPsExperts);
-                    // the experts, in id order: resident ones from VRAM, the others through the staging ring
+                    // the experts, in id order: resident ones from VRAM, the others through the staging ring (or on
+                    // the second GPU)
                     std::vector<int32_t> order;
-                    for (int32_t e = 0; e < NE; ++e) if (m.cnt[(size_t) e] > 0) order.push_back(e);
+                    for (int32_t e = 0; e < NE; ++e)
+                        if (m.cnt[(size_t) e] > 0 && (!offload || resident(e))) order.push_back(e);
                     const strata::kernels::cpu::ExpertLayout& lay = strata::kernels::cpu::expert_layout();
                     // Stage ahead: the copy stream moves blobs host -> device while the compute stream works.
                     int stage_next = 0;
@@ -527,8 +575,15 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                         swiglu_interleaved(m.GU + o0 * 1280, m.Hh + o0 * 640, ne, m.cs);
                         m.gemm.f16(m.Hh + o0 * 640, m.dq_d[q], m.Dm + o0 * N, ne, N, 640);
                     }
+                    float* sums = m.emb;   // the chunk's embeddings are in R by now
+                    if (offload) {
+                        cudaStreamWaitEvent(m.xfer, offload->done(), 0);
+                        cudaMemcpyAsync(sums, offload->host_sum(), (size_t) T * N * 4, cudaMemcpyHostToDevice, m.xfer);
+                        cudaEventRecord(m.ev_sum, m.xfer);
+                        cudaStreamWaitEvent(m.cs, m.ev_sum, 0);
+                    }
                     mark(kPsCombine);
-                    moe_combine(m.Dm, m.slot_dev, m.w, m.shared, m.sg, m.bo, T, m.cs);
+                    moe_combine(m.Dm, m.slot_dev, m.w, m.shared, m.sg, m.bo, T, m.cs, offload ? sums : nullptr);
                 }
                 // ---- the hyper-connection write of this half
                 if (half == 0) mark(kPsHcFfn);

src/program/generate.cpp

@@ -44,6 +44,7 @@
 #include "strata/core/verify.hpp"
 #include "strata/core/mtp.hpp"
 #include "strata/core/prompt_cache.hpp"
+#include "strata/prefill/offload.hpp"
 #include "strata/prefill/prefill.hpp"
 #include "strata/core/native_dense.hpp"
 #include "strata/program/logits_selection.hpp"
@@ -203,6 +204,9 @@ struct Options {
     /// Plan v0.3 P5: the prompt path borrows the top expert-cache slots for its buffers and refills them after
     /// the prompt (default); `--no-prefill-borrow` reserves the buffers' VRAM for the whole session instead.
     bool no_prefill_borrow = false;
+    /// With --second-gpu, the prompt path's experts the main GPU's cache does not hold run on the second GPU
+    /// (default); `--no-prompt-offload` streams them to the main GPU instead.
+    bool no_prompt_offload = false;
     /// Plan v0.3 P5 validation: batch only positions [0, P) and run the rest of the prompt through the token path
     /// (teacher-forced), so the logits of positions >= P - which depend on the batched state - can be scored
     /// against the oracle at many positions.  0 = the whole prompt but the last position.
@@ -344,6 +348,8 @@ void usage() {
                  "  --second-gpu-gib G   its VRAM for experts (default: all but --second-gpu-reserve-mib, 2048)\n"
                  "  --second-gpu-min-mb M  a layer's missed experts from which it takes its share (default 4; below\n"
                  "                       it the CPU is quicker than its ~100 us round trip)\n"
+                 "  --no-prompt-offload  with --second-gpu: the prompt path streams the experts the main GPU's cache\n"
+                 "                       lacks to the main GPU instead of computing them on the second\n"
                  "  --spec-follow PATH   benchmarks: emit this continuation (token ids) instead of the argmax and\n"
                  "                       accept the drafts that match it, so speculation settings compare on the\n"
                  "                       same text (not --serve)\n"
@@ -584,6 +590,7 @@ int main(int argc, char** argv) {
         else if (a == "--prefill") o.prefill_chunk = std::atoll(next("--prefill"));
         else if (a == "--no-split-rows") o.no_split_rows = true;
         else if (a == "--no-prefill-borrow") o.no_prefill_borrow = true;
+        else if (a == "--no-prompt-offload") o.no_prompt_offload = true;
         else if (a == "--prefill-until") o.prefill_until = std::atoll(next("--prefill-until"));
         else if (a == "--dump-final-r") o.dump_final_r = next("--dump-final-r");
         else if (a == "--spec") o.spec = std::atoi(next("--spec"));
@@ -1995,6 +2002,15 @@ int main(int argc, char** argv) {
             std::fprintf(stderr, "strata serve: %s\n", err.c_str());
             return 1;
         }
+        strata::prefill::Offload offload;
+        if (o.second_gpu >= 0 && !o.no_prompt_offload) {
+            if (!offload.init(o.second_gpu, o.main_gpu, srcp, &gpu2.cache(), host_res2.empty() ? nullptr : host_res2.data(),
+                              g.n_expert, o.prefill_chunk, err)) {
+                std::fprintf(stderr, "strata serve: %s\n", err.c_str());
+                return 1;
+            }
+            prefill.offload = &offload;
+        }
         strata::core::Verifier ver;
         strata::core::VerifyHits vh;
         vh.d_res = thits.d_res;
@@ -2327,6 +2343,7 @@ int main(int argc, char** argv) {
     int64_t pos_start = 0;
     int64_t spec_pos = 0;   // plan v0.3 P6: where the speculative loop starts (0 = not used)
     strata::prefill::Prefill prefill;
+    strata::prefill::Offload offload;
     double prefill_batched_ms = 0;
     std::FILE* final_r = o.dump_final_r.empty() ? nullptr : std::fopen(o.dump_final_r.c_str(), "wb");
     std::vector<float> final_r_host(final_r ? (size_t) (g.hc * g.n_embd) : 0);
@@ -2366,6 +2383,14 @@ int main(int argc, char** argv) {
             std::fprintf(stderr, "strata generate: %s\n", err.c_str());
             return 1;
         }
+        if (o.second_gpu >= 0 && !o.no_prompt_offload) {
+            if (!offload.init(o.second_gpu, o.main_gpu, srcp, &gpu2.cache(), host_res2.empty() ? nullptr : host_res2.data(),
+                              g.n_expert, o.prefill_chunk, err)) {
+                std::fprintf(stderr, "strata generate: %s\n", err.c_str());
+                return 1;
+            }
+            prefill.offload = &offload;
+        }
         if (!o.mtp.empty()) {
             if (!mtp.bind(wt, &native_head, nullptr, err)) {
                 std::fprintf(stderr, "strata generate: %s\n", err.c_str());
@@ -2414,6 +2439,10 @@ int main(int argc, char** argv) {
                      (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,
                      (long long) ps.experts_dma, ps.ms_experts_host, (long long) ps.experts_resident, ps.ms_ple);
+        if (prefill.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);
         if (o.prefill_profile) {
             std::fprintf(stderr, "strata generate: prefill GPU ms by section:");
             for (int i = 0; i < strata::prefill::kPsCount; ++i)