MoE streaming: expert reads in 256 KiB chunks by 16 threads (pin llama.cpp-vulkan cecf3ee)

docs/moe-streaming.md

@@ -586,6 +586,30 @@ Perplexity of the exact runs: 11.32.
 - **The drive is the limit.** At 4,608 slots, 4.0 ms per layer goes to waiting for reads
   against 1.7 ms of everything else.
 
+**Prefetch makes it worse.** On Flash-Next, the gate-ahead predictor (router l+1 applied to
+layer l's 2560-wide FFN input) is right about two times in three, but every wrong guess costs a
+read from a drive that is already the limit. At 4,608 slots, depth 1 took 247 s per pass
+against 133 s without (240 against 157 GiB read), and depth 2 took 253 s. Keep
+`ONEBIT_MOE_PREFETCH=0` on this model.
+
+**Reads in flight matter more than prefetch.** `ONEBIT_MOE_STATS` times every read. With whole
+expert parts (about 1 MiB) and 8 reader threads, a layer's misses make only 6-7 reads, and each
+1 MiB `pread` took 4.5-4.8 ms: about 1.8 GB/s, from a drive that gives 3.4 GB/s at that depth
+when the box is idle (`rdtest`: 3.2-4.1 GB/s idle, 1.3-2.9 under the other tenants' load).
+Reading parts in 256 KiB chunks with 16 threads keeps the queue deep (the streamer's default
+since fork PR #23; `ONEBIT_MOE_CHUNK_KB`, `ONEBIT_MOE_IO`):
+
+| Flash-Next, 4,608 slots | s / pass (two rounds) | pread per read |
+|---|---|---|
+| 8 threads, whole parts | 174, 177 | 4.5-4.8 ms (1 MiB) |
+| 16 threads, 256 KiB | 143, 151 | 2.0 ms |
+| 16 threads, 512 KiB | 152, 149 | 3.2 ms |
+| 32 threads, 256 KiB | 157, 205 | 3.4-4.9 ms |
+
+That is 14-18% faster, with perplexity unchanged in all 8 runs. The rounds ran at load 7-22,
+which is why the baseline here is slower than the 133-145 s above. On Coder-30B at 1,536 slots
+the change is within the noise (four rounds, 6.7 against 7.0 tok/s).
+
 ## Next
 
 1. **Fewer host round trips.** The remap per layer caps streamed decode at about 57 tok/s on

moe/cache_bench.cpp

@@ -90,6 +90,7 @@ void usage() {
         "  --slots N           experts held in RAM (default 1024)\n"
         "  --per-layer         slots / layers per layer instead of one shared LRU\n"
         "  --io N              reads in flight (default 8)\n"
+        "  --chunk-kb N        read expert parts in N KiB chunks (default 0: whole parts)\n"
         "  --prefetch none|k|2k|oracle  gate-ahead prefetch of the trace's predictions (default k);\n"
         "                      oracle queues every layer's real experts when a token starts: the\n"
         "                      bound for a perfect predictor with every read in flight at once\n"
@@ -120,6 +121,7 @@ int run_moe_cache(int argc, char** argv) {
         else if (a == "--slots") opt.slots = std::stoi(next());
         else if (a == "--per-layer") opt.per_layer = true;
         else if (a == "--io") opt.io_threads = std::stoi(next());
+        else if (a == "--chunk-kb") opt.read_chunk = size_t(std::stoi(next())) << 10;
         else if (a == "--prefetch") prefetch = next();
         else if (a == "--lookahead") lookahead = std::stoi(next());
         else if (a == "--compute-ms") compute_ms = std::stod(next());

moe/expert_cache.cpp

@@ -158,14 +158,23 @@ int ExpertCache::start_load(int layer, int e, bool demand) {
     auto& lru = lru_[s.lru_list];
     lru.splice(lru.begin(), lru, s.lru_it);
     const auto& parts = index_.experts.at(layer);
-    s.pending = (int) parts.size();
+    s.pending = 0;
+    const uint64_t chunk = opt_.read_chunk ? up(opt_.read_chunk) : 0;
     for (size_t k = 0; k < parts.size(); ++k) {
         const auto& p = parts[k];
         const uint64_t off = p.base + uint64_t(e) * p.stride;
         const uint64_t a = down(off), b = up(off + p.bytes);
         uint8_t* dst = lay.base + lay.part_off[k] + size_t(s.index) * lay.part_bytes[k];
-        Read r{si, fds_[p.file], a, b - a, dst, off - a, p.bytes};
-        (demand ? demand_q_ : prefetch_q_).push_back(r);
+        // [a, b) aligned for O_DIRECT, in chunks; each copies its overlap with [off, off + bytes)
+        const uint64_t step = chunk ? chunk : b - a;
+        for (uint64_t c0 = a; c0 < b; c0 += step) {
+            const uint64_t c1 = std::min(b, c0 + step);
+            const uint64_t lo = std::max(c0, off), hi = std::min(c1, off + p.bytes);
+            if (lo >= hi) continue;
+            Read r{si, fds_[p.file], c0, c1 - c0, dst + (lo - off), lo - c0, hi - lo};
+            (demand ? demand_q_ : prefetch_q_).push_back(r);
+            s.pending++;
+        }
         st_.bytes_read += b - a;
     }
     work_cv_.notify_all();
@@ -190,6 +199,7 @@ void ExpertCache::reader() {
             bounce_bytes = r.len;
             if (posix_memalign((void**) &bounce, kAlign, bounce_bytes) != 0) throw std::runtime_error("out of memory");
         }
+        const auto t0 = std::chrono::steady_clock::now();
         uint64_t done = 0;
         while (done < r.len) {
             const ssize_t n = ::pread(r.fd, bounce + done, r.len - done, (off_t) (r.off + done));
@@ -201,9 +211,14 @@ void ExpertCache::reader() {
             }
             done += (uint64_t) n;
         }
+        const auto t1 = std::chrono::steady_clock::now();
         std::memcpy(r.dst, bounce + r.lead, r.bytes);
+        const auto t2 = std::chrono::steady_clock::now();
         {
             std::lock_guard<std::mutex> lk(mu_);
+            st_.read_ms += std::chrono::duration<double, std::milli>(t1 - t0).count();
+            st_.copy_ms += std::chrono::duration<double, std::milli>(t2 - t1).count();
+            st_.reads++;
             Slot& s = slots_[r.slot];
             if (--s.pending == 0) s.state = State::Ready;
         }

moe/expert_cache.h

@@ -40,6 +40,9 @@ struct CacheOptions {
     int slots = 1024;         // experts held in RAM, all layers together (each layer keeps the same share)
     bool per_layer = false;   // false: one LRU over all layers (it hit more in replays); true: slots / layers each
     int io_threads = 8;       // reads in flight
+    // Parts larger than this are read in chunks of this size (a multiple of 4 KiB), so one
+    // expert's miss keeps several reads in flight; 0 reads each part whole.
+    size_t read_chunk = 0;
     bool pin = true;          // mlock the slots
     // Memory for each list's slots, owned by the caller (for example a GPU buffer the host can
     // map); null: the cache maps and pins its own. Called once per list with its byte size.
@@ -54,6 +57,9 @@ struct CacheStats {
     uint64_t prefetch_wasted = 0;   // prefetched, then evicted before any use
     uint64_t bytes_read = 0;
     double stall_ms = 0;            // time acquire() waited for reads
+    double read_ms = 0;             // reader threads: time in pread, summed over threads
+    double copy_ms = 0;             // reader threads: time copying out of the bounce buffer
+    uint64_t reads = 0;             // part reads done
 };
 
 class ExpertCache {

registry/architectures.json

@@ -1,7 +1,7 @@
 {
  "about": "HF architecture -> GGUF architecture and the backends whose code accepts it. Generated by tools/registry_build.py from the pinned sources; do not edit.",
  "sources": {
-  "llama.cpp (vulkan)": "546800c1d285845a74d8108d8a382e2912fa1bc7",
+  "llama.cpp (vulkan)": "cecf3ee01d9d99378e98bfea95f51cac714b04b8",
   "llama.cpp (hrx)": "8dd75eb48f282198abdb0bd748aa0e4594745f1d",
   "zinc": "29bc350ac4cbf9f110ec628b9e177ea04ac816fb"
  },