moe: expert slots in size classes (Flash-Next: 34.4 -> 27.0 GiB pinned for the same slots)

docs/moe-streaming.md

@@ -172,6 +172,10 @@ co-occurrence with the previous layer) cover under 30%.
   - `acquire(layer, experts)` blocks until the experts are resident and pins them until
     `release`. `prefetch(layer, experts)` only queues reads.
   - Eviction is LRU, over one budget shared by all layers (the replays favour it) or per layer.
+  - Slots come in size classes. Unsloth's UD quants use bigger types for some layers' experts,
+    so layers are grouped by the bytes one expert needs, and each group gets its own slots
+    (same share of experts per layer). Flash-Next's 9,216 slots pin 27.0 GiB this way,
+    instead of 34.4 GiB with one slot size.
 - **`1bit moe-cache`** replays a trace through the cache against the real model file, with
   real reads. Compute is simulated per layer:
   1. at the layer's start, the prefetch is issued (lookahead 1: this layer's gate-ahead
@@ -332,7 +336,7 @@ resident, as in a fleet), the cache sets the speed:
 | Qwen3-Coder-30B-A3B | Q4_K_M, all on Vulkan0 | 92-93 | UD-Q4_K_XL: 0.027 at a similar size | 13.2 GiB of 17.5 GiB | 32-48 (code 34, chat 48, long 32) |
 | Qwen3.6-35B-A3B | **UD-Q5_K_XL**, all on Vulkan0 (faster than Q8_0 at nearly its quality) | 52-53 | 0.009 | ≈16.5 GiB of ≈21.9 GiB | 32-39 |
 | GLM-4.7-Flash | Q4_K_M, all on Vulkan0 | 71 | no other quant measured | ≈12.6 GiB of ≈16.8 GiB | 33-40 |
-| Qwen3.8-Flash-Next | UD-Q4_K_XL + MTP n = 3, resident on Vulkan (about 77 GiB) | 40-49 | | 27.5 GiB of 71.7 GiB (34.4 GiB pinned) | 7-12 on this drive (ceiling 24-34 at 3.7 GB/s) |
+| Qwen3.8-Flash-Next | UD-Q4_K_XL + MTP n = 3, resident on Vulkan (about 77 GiB) | 40-49 | | 27.5 GiB of 71.7 GiB (27.0 GiB pinned) | 7-12 on this drive (ceiling 24-34 at 3.7 GB/s) |
 
 MTP is available only for Flash-Next, and it is worth it both ways: 1.6-1.9 times resident,
 and it leaves the drive's reads per token unchanged when streamed.
@@ -378,8 +382,8 @@ resident) and 22 ms with it (45.5 tok/s):
 | 384 (55.1 GiB) | 37 / 43 MB | 26.0 / 26.0 | 31.3 / 29.8 | 45.5 / 45.5 |
 
 **Measured on the drive** (`1bit moe-cache`, chat and long traces, 192 per layer = 9,216
-slots: 27.5 GiB of experts in 34.4 GiB of pinned slots. Each slot is sized for the largest
-layer's expert, and the UD quant uses bigger types in some layers):
+slots: 27.5 GiB of experts in 34.4 GiB of pinned slots. At the time, every slot was sized for
+the largest layer's expert. Size classes (below) now pin 27.0 GiB for the same slots):
 
 | mode | compute per token | decode tok/s | stall per token |
 |---|---|---|---|
@@ -420,7 +424,5 @@ below the floor.
    caches. Candidates: frequency with decay per layer, and hints from the router's scores.
 3. **Flash-Next's gate-ahead prediction.** Its layers keep four 2560-wide residual streams
    (hyper-connections), so the predictor needs that model's mixing step before the router.
-4. **Slots sized per layer.** Mixed-quant (UD) models waste up to 25% of the pinned memory
-   with one slot size. Flash-Next's 9,216 slots pin 34.4 GiB for 27.5 GiB of experts.
-5. **The quants not on the box** (UD-Q2/Q3/Q4_K_XL): their decode speed, once there is room
+4. **The quants not on the box** (UD-Q2/Q3/Q4_K_XL): their decode speed, once there is room
    on the drive.

moe/cache_bench.cpp

@@ -139,10 +139,11 @@ int run_moe_cache(int argc, char** argv) {
     ExpertCache cache(index, opt);
     const int n_moe = (int) index.experts.size();
     const double layer_ms = compute_ms / n_moe;
-    std::printf("%s: %d MoE layers, %d experts, %.2f MiB per expert (layer %d); %d slots of %.2f MiB = %.1f GiB%s\n",
-                model.c_str(), n_moe, index.n_expert, index.expert_bytes(index.experts.begin()->first) / 1048576.0,
-                index.experts.begin()->first, cache.slots(), cache.slot_bytes() / 1048576.0,
-                cache.slots() * double(cache.slot_bytes()) / (1u << 30), opt.per_layer ? ", per layer" : ", shared");
+    double expert_bytes = 0;
+    for (const auto& [l, parts] : index.experts) expert_bytes += index.expert_bytes(l);
+    std::printf("%s: %d MoE layers, %d experts, %.2f MiB per expert on average; %d slots, %.1f GiB pinned (largest slot %.2f MiB)%s\n",
+                model.c_str(), n_moe, index.n_expert, expert_bytes / n_moe / 1048576.0, cache.slots(),
+                cache.pinned_bytes() / double(1u << 30), cache.slot_bytes() / 1048576.0, opt.per_layer ? ", per layer" : ", shared");
 
     auto sleep_ms = [](double ms) {  // compute stand-in: spin for sub-millisecond precision
         const auto end = std::chrono::steady_clock::now() + std::chrono::duration<double, std::milli>(ms);

moe/expert_cache.cpp

@@ -19,6 +19,7 @@
 #include <chrono>
 #include <cstring>
 #include <fcntl.h>
+#include <map>
 #include <stdexcept>
 #include <string>
 #include <sys/mman.h>
@@ -34,21 +35,33 @@ uint64_t up(uint64_t v) { return (v + kAlign - 1) / kAlign * kAlign; }
 
 ExpertCache::ExpertCache(const GgufIndex& index, const CacheOptions& opt) : index_(index), opt_(opt) {
     if (index.experts.empty()) throw std::runtime_error("the model has no routed expert tensors");
-    // a slot holds every part's aligned read window: [down(off), up(off + bytes))
+    // A slot holds every part's aligned read window, [down(off), up(off + bytes)). Mixed-quant
+    // models (Unsloth's UD) use bigger types in some layers, so layers are grouped by that size:
+    // each size class has its own slots, and LRU (shared mode) or one per layer.
+    std::map<size_t, std::vector<int>> classes;  // slot size -> layers
     for (const auto& [l, parts] : index.experts) {
         size_t n = 0;
         for (const auto& p : parts) n += up(p.bytes) + kAlign;
+        classes[n].push_back(l);
         slot_bytes_ = std::max(slot_bytes_, n);
     }
-    int n_lists = 1;
-    if (opt_.per_layer) {
-        n_lists = (int) index.experts.size();
-        int i = 0;
-        for (const auto& [l, parts] : index.experts) layer_lru_[l] = i++;
+    // the budget is in experts; every layer keeps the same share of its experts
+    const double share = std::min(1.0, double(opt_.slots) / (double(index.experts.size()) * index.n_expert));
+    struct List { size_t bytes; int slots; };
+    std::vector<List> lists;
+    for (const auto& [bytes, layers] : classes) {
+        if (opt_.per_layer) {
+            for (int l : layers) {
+                layer_lru_[l] = (int) lists.size();
+                lists.push_back({bytes, std::max(1, (int) (share * index.n_expert))});
+            }
+        } else {
+            for (int l : layers) layer_lru_[l] = (int) lists.size();
+            lists.push_back({bytes, std::max(1, (int) (share * index.n_expert * layers.size()))});
+        }
     }
-    const int per_list = std::max(1, opt_.slots / n_lists);
-    const int n_slots = per_list * n_lists;
-    mem_bytes_ = size_t(n_slots) * slot_bytes_;
+    int n_slots = 0;
+    for (const auto& li : lists) { n_slots += li.slots; mem_bytes_ += size_t(li.slots) * li.bytes; }
     void* p = ::mmap(nullptr, mem_bytes_, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
     if (p == MAP_FAILED) throw std::runtime_error("cannot map " + std::to_string(mem_bytes_ >> 20) + " MiB of expert slots");
     mem_ = static_cast<uint8_t*>(p);
@@ -61,12 +74,17 @@ ExpertCache::ExpertCache(const GgufIndex& index, const CacheOptions& opt) : inde
         fds_.push_back(fd);
     }
     slots_.resize(n_slots);
-    lru_.resize(n_lists);
-    for (int s = 0; s < n_slots; ++s) {
-        slots_[s].lru_list = s / per_list;
-        lru_[slots_[s].lru_list].push_back(s);  // empty slots at the back: taken first
-        slots_[s].lru_it = std::prev(lru_[slots_[s].lru_list].end());
-    }
+    lru_.resize(lists.size());
+    size_t off = 0;
+    int si = 0;
+    for (int li = 0; li < (int) lists.size(); ++li)
+        for (int k = 0; k < lists[li].slots; ++k, ++si) {
+            slots_[si].lru_list = li;
+            slots_[si].off = off;
+            off += lists[li].bytes;
+            lru_[li].push_back(si);  // empty slots at the back: taken first
+            slots_[si].lru_it = std::prev(lru_[li].end());
+        }
     for (int i = 0; i < std::max(1, opt_.io_threads); ++i) threads_.emplace_back([this] { reader(); });
 }
 
@@ -85,7 +103,6 @@ ExpertCache::~ExpertCache() {
 }
 
 int ExpertCache::lru_of(int layer) const {
-    if (!opt_.per_layer) return 0;
     auto it = layer_lru_.find(layer);
     if (it == layer_lru_.end()) throw std::runtime_error("layer " + std::to_string(layer) + " has no routed experts");
     return it->second;
@@ -117,7 +134,7 @@ int ExpertCache::start_load(int layer, int e, bool demand) {
     where_[s.key] = si;
     auto& lru = lru_[s.lru_list];
     lru.splice(lru.begin(), lru, s.lru_it);
-    uint8_t* base = mem_ + size_t(si) * slot_bytes_;
+    uint8_t* base = mem_ + s.off;
     size_t at = 0;
     const auto& parts = index_.experts.at(layer);
     s.pending = (int) parts.size();
@@ -193,7 +210,7 @@ std::vector<ExpertCache::Resident> ExpertCache::acquire(int layer, const std::ve
     std::vector<Resident> out(experts.size());
     for (size_t i = 0; i < experts.size(); ++i) {
         const Slot& s = slots_[si[i]];
-        const uint8_t* base = mem_ + size_t(si[i]) * slot_bytes_;
+        const uint8_t* base = mem_ + s.off;
         for (size_t off : s.part_off) out[i].parts.push_back(base + off);
     }
     return out;

moe/expert_cache.h

@@ -36,7 +36,7 @@
 namespace onebit::moe {
 
 struct CacheOptions {
-    int slots = 1024;         // experts held in RAM, all layers together
+    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
     bool pin = true;          // mlock the slots
@@ -70,7 +70,8 @@ class ExpertCache {
 
     CacheStats stats() const;
     void reset_stats();
-    size_t slot_bytes() const { return slot_bytes_; }
+    size_t slot_bytes() const { return slot_bytes_; }  // the largest slot
+    size_t pinned_bytes() const { return mem_bytes_; }
     int slots() const { return (int) slots_.size(); }
 
 private:
@@ -82,6 +83,7 @@ class ExpertCache {
         int pending = 0;          // reads not finished
         bool prefetched = false;  // loaded by prefetch() and not acquired since
         std::vector<size_t> part_off;  // byte offset of each part's data inside the slot
+        size_t off = 0;           // byte offset of the slot in the pinned region
         std::list<int>::iterator lru_it;
         int lru_list = 0;
     };