feat(wp): ensure_batch — concurrent QD=N expert page-in (Colibri pattern)

Decode paged the MoE active-expert set through a per-expert prefetch+ensure
loop whose reserved slots stayed UNPINNED until harvest, so under decode
eviction a later expert's read could evict an earlier one's not-yet-loaded
slot, forcing a serial page_in_sync_ re-read — collapsing effective io_uring
queue depth to ~1 (raising queue depth did nothing; the eviction window, not
the depth, was the bottleneck).

ensure_batch (opt-in WP_ENSURE_BATCH=1) reserves AND pins every cold-miss slot
up front (alloc_slot skips pinned, so no sibling read can evict it), then on the
P2P/direct-to-device path submits all misses in one io_uring batch (true QD=N)
straight into VRAM, harvests, and on any read failure syncs into the SAME pinned
slot. Returns the pinned pages for the caller to release on the next callback
(existing per-op/per-range pin lifecycle). page_in_sync_ gains a reuse_slot arg
so the failure fallback reads into the pinned slot without releasing it.

Default off; existing prefetch+ensure path preserved in the else branch for A/B
+ rollback. Build-gated (HIP), pending R9700 decode A/B validation.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014ZRfPpL8XFzk1hep9MMg9P

Author: kmbandy

src/weight-pager/wp-eval-cb.cpp

@@ -595,6 +595,56 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
                                         active.insert((int) idx);
                                     }
 
+                                    // WP_ENSURE_BATCH (opt-in): Colibri-style
+                                    // concurrent batch page-in. Reserve+PIN every
+                                    // active-expert slot up front, then issue all
+                                    // cold-miss reads in ONE io_uring batch (true
+                                    // QD=N). Closes the eviction window that let a
+                                    // later expert's read evict an earlier one's
+                                    // not-yet-harvested slot, collapsing effective
+                                    // queue depth to ~1 under decode. Default off;
+                                    // the else branch keeps the current path (A/B).
+                                    static int s_ensure_batch_env = -1;
+                                    if (s_ensure_batch_env < 0) {
+                                        const char * eb = std::getenv("WP_ENSURE_BATCH");
+                                        s_ensure_batch_env = (eb != nullptr && eb[0] == '1') ? 1 : 0;
+                                    }
+                                    int    n_ensures = 0;
+                                    void * first_active_slot = nullptr;
+                                    if (s_ensure_batch_env == 1) {
+                                        std::vector<int> active_pages;
+                                        active_pages.reserve(active.size());
+                                        for (int e : active) {
+                                            active_pages.push_back(weight_page + 1 + e);
+                                        }
+                                        std::vector<void *> active_ptrs;
+                                        std::vector<int>    active_pinned;
+                                        pager->ensure_batch(active_pages, active_ptrs, active_pinned);
+                                        std::size_t ap = 0;
+                                        for (int e : active) {
+                                            void * slot = active_ptrs[ap++];
+                                            if (slot != nullptr) {
+                                                slot = capture_ptr_for_page(weight_page + 1 + e, slot);
+                                                host_ptrs[(size_t) e] = slot;
+                                                if (first_active_slot == nullptr) {
+                                                    first_active_slot = slot;
+                                                }
+                                                ++n_ensures;
+#if defined(GGML_USE_HIP)
+                                                enqueue_async_wait_for_page(weight_page + 1 + e, s_async_events_prev_op);
+#endif
+                                            }
+                                        }
+                                        // Record the pins ensure_batch took so the
+                                        // per-op / per-range lifecycle releases them.
+                                        for (int p : active_pinned) {
+                                            (paged_batch ? s_range_pins : s_pinned_pages_prev_op).push_back(p);
+                                            if (paged_batch) { s_range_pinned_bytes += pager->page_meta(p).size; }
+                                        }
+#if defined(GGML_USE_HIP)
+                                        if (!active_pinned.empty()) { s_prev_op_pager = pager; }
+#endif
+                                    } else {
                                     // Pass 1: fire async prefetch for every
                                     // active expert. With io_uring (depth 4),
                                     // multiple preads can be in flight at
@@ -643,8 +693,6 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
                                     // hit; for in-flight pages it waits on
                                     // the async completion; for unloaded
                                     // pages it falls back to sync.
-                                    int n_ensures = 0;
-                                    void * first_active_slot = nullptr;
                                     for (int e : active) {
                                         const int sub_page_idx = weight_page + 1 + e;
                                         void * slot = pager->ensure(sub_page_idx);
@@ -669,6 +717,8 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
 #endif
                                         }
                                     }
+                                    }  // end WP_ENSURE_BATCH else (existing prefetch+ensure path)
+
                                     // Safety: fill INACTIVE expert slots with a non-null
                                     // sentinel (first active slot) so a kernel that reads
                                     // expert_ptrs[inactive_idx] gets a valid (wrong) pointer

src/weight-pager/wp-pager.cpp

@@ -873,6 +873,114 @@ void * WeightPager::ensure(int page_idx) {
     return slot_ptr_(slot);
 }
 
+void WeightPager::ensure_batch(const std::vector<int> & page_indices,
+                               std::vector<void *>     & out_ptrs,
+                               std::vector<int>        & out_pinned) {
+    out_ptrs.assign(page_indices.size(), nullptr);
+    out_pinned.clear();
+    if (!initialized_) return;
+
+    // Pass 1 — resolve hits/pinned inline; for every cold miss reserve AND PIN
+    // a slot up front. Pinning before any sibling read means alloc_slot for a
+    // later miss can never evict an earlier miss's in-flight slot (alloc_slot
+    // skips pinned), which is what collapsed effective QD to ~1 under decode
+    // eviction. Each pinned page is reported in out_pinned for the caller to
+    // release next callback.
+    struct Miss { int page; int slot; std::size_t out_i; };
+    std::vector<Miss> misses;
+    misses.reserve(page_indices.size());
+    for (std::size_t i = 0; i < page_indices.size(); ++i) {
+        const int p = page_indices[i];
+        if (p < 0 || p >= catalog_.size()) continue;
+        const PageMeta & m = catalog_.at(p);
+        if (m.is_pinned) {                       // MAD-236: always-resident, no slot
+            out_ptrs[i] = m.resident_ptr;
+            continue;
+        }
+        if (page_loaded_[p]) {                   // hit — bump LRU, pin, harvest
+            const int s = page_to_slot_[p];
+            pool_.mark_used(s);
+            pool_.pin_slot(s);
+            out_pinned.push_back(p);
+            out_ptrs[i] = slot_ptr_(s);
+            continue;
+        }
+        if (page_to_slot_[p] >= 0) {
+            // Reserved by an in-flight (e.g. cross-layer) prefetch. Harvest via
+            // the tested ensure() path (waits on the prefetch, or syncs), then
+            // pin the result. Not part of the concurrent batch.
+            void * ptr = ensure(p);
+            if (ptr != nullptr) {
+                pool_.pin_slot(page_to_slot_[p]);
+                out_pinned.push_back(p);
+                out_ptrs[i] = ptr;
+            }
+            continue;
+        }
+        const int s = pool_.alloc_slot(m.size);
+        if (s < 0) { ++stats_.sync_fallbacks; continue; }  // pool exhausted (rare)
+        ensure_slot_map_(s);
+        page_to_slot_[p]  = s;
+        slot_to_page_[s]  = p;
+        pool_.pin_slot(s);                       // PIN before the next alloc
+        out_pinned.push_back(p);
+        misses.push_back({ p, s, i });
+    }
+    if (misses.empty()) return;
+
+    // Pass 2 — issue all cold-miss reads. On the P2P/direct-to-device path the
+    // reads land straight in the VRAM slots, so one io_uring batch keeps N reads
+    // in flight at once (true QD=N). Off P2P the shared staging buffer can't
+    // hold N pages, so fall back to serial sync into each pinned slot (still
+    // correct; the throughput case that matters for decode is P2P).
+    if (file_io_->direct_to_device()) {
+        std::vector<FileIOBatchRequest> reqs;
+        reqs.reserve(misses.size());
+        for (std::size_t k = 0; k < misses.size(); ++k) {
+            const PageMeta & m = catalog_.at(misses[k].page);
+            reqs.push_back({ (uint64_t) k, (int) m.file_idx, m.file_offset,
+                             m.size, slot_ptr_(misses[k].slot) });
+        }
+        const int n_sub = file_io_->submit_batch(reqs);
+        file_io_->flush();
+        std::vector<bool> ok(misses.size(), false);
+        int reaped = 0;
+        while (reaped < n_sub) {
+            IoResult r = file_io_->wait_any(/*timeout_ms=*/-1);
+            if (r.req_id < misses.size()) {
+                ok[(std::size_t) r.req_id] =
+                    (r.status == IoStatus::Ok &&
+                     r.bytes_read == (int) catalog_.at(misses[(std::size_t) r.req_id].page).size);
+                ++reaped;
+            }
+            // Unknown req_id (stale prefetch completion): drop, keep waiting.
+        }
+        const auto io_t0 = std::chrono::steady_clock::now();
+        for (std::size_t k = 0; k < misses.size(); ++k) {
+            const Miss & mm = misses[k];
+            const PageMeta & m = catalog_.at(mm.page);
+            if (ok[k] && zero_device_padding(slot_ptr_(mm.slot), m.size, pool_.slot_size(mm.slot))) {
+                page_to_slot_[mm.page] = mm.slot;
+                page_loaded_[mm.page]  = true;
+                slot_to_page_[mm.slot] = mm.page;
+                pool_.mark_used(mm.slot);
+                record_page_in_(m.size, seconds_since(io_t0));
+                out_ptrs[mm.out_i] = slot_ptr_(mm.slot);
+            } else {
+                // read (or padding) failed — sync-fallback into the SAME pinned
+                // slot so the up-front pin/out_pinned bookkeeping stays valid.
+                const int s = page_in_sync_(mm.page, /*reuse_slot=*/mm.slot);
+                out_ptrs[mm.out_i] = (s < 0) ? nullptr : slot_ptr_(s);
+            }
+        }
+    } else {
+        for (const Miss & mm : misses) {
+            const int s = page_in_sync_(mm.page, /*reuse_slot=*/mm.slot);
+            out_ptrs[mm.out_i] = (s < 0) ? nullptr : slot_ptr_(s);
+        }
+    }
+}
+
 int WeightPager::take_async_transfer_event(int page_idx) {
     if (!initialized_) return -1;
     if (page_idx < 0 || page_idx >= (int) page_async_event_.size()) return -1;
@@ -1119,10 +1227,13 @@ void WeightPager::mark_cross_layer_prefetch_candidates(const std::vector<int> &
     }
 }
 
-int WeightPager::page_in_sync_(int page_idx) {
-    // Synchronous read into a fresh slot, bypassing the prefetch pipeline.
-    // Used by ensure() on miss. Host transports read into staging and then
-    // hand off to GpuTransport; P2P transports read into the VRAM slot.
+int WeightPager::page_in_sync_(int page_idx, int reuse_slot) {
+    // Synchronous read into a slot, bypassing the prefetch pipeline. Used by
+    // ensure() on miss (reuse_slot < 0: alloc a fresh slot) and by ensure_batch
+    // as its per-page failure fallback (reuse_slot >= 0: read into that already-
+    // reserved+pinned slot; never release it here). Host transports read into
+    // staging and then hand off to GpuTransport; P2P transports read into VRAM.
+    const bool owns_slot = (reuse_slot < 0);
 
     static int s_diag_count = 0;
     const bool diag = (s_diag_count < 5);
@@ -1137,9 +1248,12 @@ int WeightPager::page_in_sync_(int page_idx) {
 
     const PageMeta & m = catalog_.at(page_idx);
 
-    const int slot = pool_.alloc_slot(m.size);
-    if (slot < 0) return -1;
-    ensure_slot_map_(slot);
+    int slot = reuse_slot;
+    if (slot < 0) {
+        slot = pool_.alloc_slot(m.size);
+        if (slot < 0) return -1;
+        ensure_slot_map_(slot);
+    }
     void * dst = slot_ptr_(slot);
     if (diag) LLAMA_LOG_ERROR("[DIAG] page_in_sync_[%d]: alloc_slot ok, slot=%d dst=%p\n", s_diag_count, slot, dst);
 
@@ -1151,7 +1265,7 @@ int WeightPager::page_in_sync_(int page_idx) {
     if (staging == nullptr || m.size > sync_staging_size_) {
         LLAMA_LOG_ERROR("wp::WeightPager::page_in_sync_: page %d size %zu exceeds shared staging size %zu\n",
                         page_idx, m.size, sync_staging_size_);
-        pool_.release_slot(slot);
+        if (owns_slot) pool_.release_slot(slot);
         return -1;
     }
 
@@ -1163,7 +1277,7 @@ int WeightPager::page_in_sync_(int page_idx) {
             if (evt < 0) {
                 LLAMA_LOG_WARN("wp::WeightPager::page_in_sync_: host tier gpu stage_in failed for page %d\n",
                                page_idx);
-                pool_.release_slot(slot);
+                if (owns_slot) pool_.release_slot(slot);
                 return -1;
             }
             transport_.release_event(evt);
@@ -1212,13 +1326,13 @@ int WeightPager::page_in_sync_(int page_idx) {
     if (diag) LLAMA_LOG_ERROR("[DIAG] page_in_sync_[%d]: stage1 done ok=%d\n", s_diag_count, (int)ok);
     if (!ok) {
         LLAMA_LOG_WARN("wp::WeightPager::page_in_sync_: file IO failed for page %d\n", page_idx);
-        pool_.release_slot(slot);
+        if (owns_slot) pool_.release_slot(slot);
         return -1;
     }
 
     if (direct_to_device) {
         if (!zero_device_padding(dst, m.size, pool_.slot_size(slot))) {
-            pool_.release_slot(slot);
+            if (owns_slot) pool_.release_slot(slot);
             return -1;
         }
         page_to_slot_[page_idx] = slot;
@@ -1241,7 +1355,7 @@ int WeightPager::page_in_sync_(int page_idx) {
     if (diag) LLAMA_LOG_ERROR("[DIAG] page_in_sync_[%d]: stage_in returned evt=%d\n", s_diag_count, evt);
     if (evt < 0) {
         LLAMA_LOG_WARN("wp::WeightPager::page_in_sync_: gpu stage_in failed for page %d\n", page_idx);
-        pool_.release_slot(slot);
+        if (owns_slot) pool_.release_slot(slot);
         return -1;
     }
     transport_.release_event(evt);

src/weight-pager/wp-pager.h

@@ -168,6 +168,20 @@ class WeightPager {
     // page_idx is out of range or any underlying op fails.
     void * ensure(int page_idx);
 
+    // Batch-ensure a set of pages with all cold-miss reads issued CONCURRENTLY
+    // (Colibri pattern). Each miss's slot is reserved AND pinned up front so no
+    // read in the batch can evict a sibling's not-yet-loaded slot — the eviction
+    // window that collapsed effective io_uring queue depth to ~1 under decode
+    // pressure. On the P2P (direct-to-device) IO path the misses are submitted in
+    // one io_uring batch (true QD=N) reading straight into the VRAM slots.
+    // Fills out_ptrs[i] with the slot pointer for pages[i] (nullptr on failure /
+    // pool exhaustion), and out_pinned with every page this call pinned — the
+    // CALLER must record those and unpin them in the next eval callback (matches
+    // the per-op pin lifecycle used by ensure()+pin_page).
+    void ensure_batch(const std::vector<int> & page_indices,
+                      std::vector<void *>     & out_ptrs,
+                      std::vector<int>        & out_pinned);
+
     // WP_ASYNC_ENSURE handoff. ensure() stashes the transfer event here when
     // it returns before stage 2 has completed; the eval callback takes it,
     // queues a compute-stream wait, and releases it after that op completes.
@@ -254,7 +268,11 @@ class WeightPager {
     // Internal helper: synchronous page-in (used by ensure() on miss).
     // Reads the page's bytes via FileIOLayer (sync path), copies to VRAM,
     // and zeros the padding. Returns the slot index or -1 on failure.
-    int  page_in_sync_(int page_idx);
+    // reuse_slot >= 0: read into that caller-owned (typically pinned) slot
+    // instead of allocating a fresh one, and do NOT release it on failure — the
+    // caller owns its lifecycle. reuse_slot < 0 keeps the original behavior
+    // (alloc a slot, release it on any error). Returns the slot index or -1.
+    int  page_in_sync_(int page_idx, int reuse_slot = -1);
 
     // Resolve a slot index to a VRAM pointer.
     void * slot_ptr_(int slot_idx) const { return pool_.slot_ptr(slot_idx); }