STRATA_EXCHANGE_ROTATE: rotation for packs whose layers differ in expert block size (spares kept per size)

docs/DETAILS.md

@@ -178,8 +178,10 @@ other ~18 GB), a 32 GB PC with a 12-16 GB GPU the Coder; IQ3_XXS on a 32 GB PC s
   It stays on the blocking tier (said in the log) when the exchange buffers are not page-locked; the stats line
   reports ms per round.
 - `STRATA_EXCHANGE_ROTATE=1` (opt-in): an adaptive swap hands buffer ownership over instead of copying the evicted
-  blob into the RAM copy (equal-size blobs, fully page-locked copy). Same tokens, fewer host copies; it works with
-  `--adapt-async 1` too. Details and the measurement: [EXCHANGE_ROTATION.md](EXCHANGE_ROTATION.md).
+  blob into the RAM copy (fully page-locked copy; layers of different blob sizes keep spare buffers per size). Same
+  tokens, fewer host copies; it works with `--adapt-async 1` too. Swift 1.5 IQ3_S on an RTX 5090 with 64 GB RAM
+  (Windows, about 30 swaps per verify window): 164 -> 188 tok/s on one forced 1,500-token text, 155 -> 178 over 23
+  coding-agent turns. Details and the measurements: [EXCHANGE_ROTATION.md](EXCHANGE_ROTATION.md).
 - The answers are the plain mapped mode's for the same expert placement: the bytes are the file's. With a page-locked
   copy the GPU also takes its usual share of the misses over PCIe (`--pcie-frac`), as with enough RAM; `--pcie-frac 0`
   (or `STRATA_RESIDENT_PIN=0`) gives the mapped mode's exact tokens.

docs/EXCHANGE_ROTATION.md

@@ -21,10 +21,12 @@ $env:STRATA_EXCHANGE_ROTATE = '1'
 
 Only the exact value `1` opts in. Set it to `0` or remove it to restore the copy
 path. This setting does not select resident RAM mode or change any preset.
-It applies only when all layers have equal-size expert blocks, the resident
-complement is fully pinned and GPU-mapped, and the exchange buffers can also be
-pinned and mapped. Other layouts keep the copy path and log why rotation could
-not be enabled. Configurations without a resident complement do not use it.
+It applies only when the resident complement is fully pinned and GPU-mapped
+and the exchange buffers can also be pinned and mapped. Other layouts keep the
+copy path and log why rotation could not be enabled. Configurations without a
+resident complement do not use it. Layers may differ in expert block size (a
+GGUF with mixed quantizations per layer): see
+[Layers of different block sizes](#layers-of-different-block-sizes).
 
 The startup diagnostic confirms activation:
 
@@ -46,6 +48,71 @@ serialized commit; rotation adds no overlapping compute or new synchronization
 policy. Reserve the maximum exchange capacity before rotation starts: growing a
 live exchange arena is rejected because it can now hold resident experts.
 
+## Layers of different block sizes
+
+A swap exchanges two experts of the same layer, so the RAM slot the promoted
+expert leaves has exactly the size of the block the evicted one needs. Rotation
+therefore keeps its spare buffers per block size: an exchange of a layer takes
+a spare of that layer's size and hands back one of the same size, and every
+size keeps its number of spares for good. The first version required one size
+for all layers and fell back to the copy path on any other pack; Swift 1.5
+IQ3_S has seven sizes (1,510,400 to 2,662,400 bytes) and IQ3_XXS has nine
+too, so neither could rotate.
+
+The exchange allocation holds one spare of every size per swap of a round:
+`--adapt-swaps` (at most 96) times the sum of the distinct sizes. With
+`--adapt-async 1` the second half of the buffers only stages copies and is
+kept in the largest size alone. On IQ3_S with `--adapt-swaps 80` that is
+1.06 GiB of page-locked RAM (1.26 GiB with `--adapt-async 1`), against 0.20
+and 0.40 GiB for the copy path's buffers. The startup line says how many sizes
+and how much:
+
+```text
+FileExpertSource: exchange buffer rotation enabled: 80 buffers in 7 block sizes (largest 2662400 bytes, 1.06 GiB page-locked); no host commit memcpy
+```
+
+With `--adapt-async 1` a rotated round also ends one verify window sooner:
+there are no bytes for the background job to move, so the ownership flip runs
+in the tick that sees the copies in.
+
+`tests/core/exchange_storage_test.cpp` covers it without CUDA (three sizes,
+6,144 exchanges: bytes, aliases, a spare of another size or an expert of
+another layer refused, overlapping resident blocks refused).
+
+### Measured: Swift 1.5 IQ3_S on an RTX 5090 with 64 GB (Windows 11)
+
+Core i7-13700K, RTX 5090 32 GB, 64 GB DDR5, Windows 11, driver 617.42, CUDA
+13.0.88, engine 0.1.42 built from source with this change. Swift 1.5 IQ3_S,
+128K context, `--resident-experts` (31.4 GiB of experts page-locked in RAM,
+about 10,800 of 24,576 in VRAM), `--adapt-every 1 --adapt-swaps 80
+--adapt-decay 0.97`, `--spec 4`, `--pool-affinity p-cores`. In the first and third
+rows one binary ran both arms and only `STRATA_EXCHANGE_ROTATE` differs; in the
+second the off arm is the build before this change (the third row's binary also
+carried extra timers).
+
+| Run | Rotation off | Rotation on |
+|---|---|---|
+| One 8K prompt, the same 1,500 forced tokens (`--spec-follow`), blocking tier, decode tok/s | 163.4, 165.1, 163.1, 161.6, 162.6 | 187.4, 188.7, 185.8, 184.9 |
+| `--serve`, 23 coding-agent turns (`tools/agent_turns_bench.py`'s turns, greedy, up to 400 new tokens), blocking tier, tok/s over all turns, three interleaved rounds | 147.8, 159.2, 158.4 | 175.2, 184.9, 172.8 |
+| The same turns with `--adapt-async 1`, one binary, two interleaved rounds | 166.0, 170.1 | 185.6, 182.5 |
+
+In the first row, two runs of each arm that sized the expert cache alike (10,681
+slots) wrote the same `--window-hashes` file for all 434 windows: the same
+arithmetic, window by window. (`--expert-cache auto` sizes the cache from the
+free VRAM at start, and a run with another slot count places experts
+differently, so its hashes differ with or without rotation.) The tier swapped
+about 14,000 experts in such a run, 10,100 to 10,600 of them exchanges. In the serve rows the
+rotated arm was faster on 23 of 23 turns (blocking tier) and 22 of 23
+(`--adapt-async 1`). On this box the tier swaps about 30 experts per verify
+window all through a run (the cache holds 44% of the experts), so the copy
+path moved about 50 MiB of RAM per window on the decode thread; that is the
+time rotation removes. With `--adapt-async 1` the same copy ran on the
+background job, and the decode thread still lost about 1 ms per window (0.4 to 2.5 by request)
+while it ran; rotated, the tier's part of a window is 0.06 ms.
+
+Not measured: Linux, AMD cards, a layer split, a pack with one block size on
+this machine (the code path for it is the one above with a single size).
+
 ## Tests
 
 The ownership test can run without CUDA:

include/strata/core/exchange_storage.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <algorithm>
 #include <atomic>
 #include <cstddef>
 #include <memory>
@@ -13,7 +14,12 @@ namespace strata::core::detail {
 
 // Non-owning views: the complement and exchange allocations retain their original
 // owners until FileExpertSource::close(). Ownership changes via slot IDs; the
-// physical buffer addresses, bytes and CUDA aliases stay in place. This first path requires uniform, fully mapped/pinned expert slots.
+// physical buffer addresses, bytes and CUDA aliases stay in place. Requires fully mapped/pinned expert slots.
+//
+// Block sizes: `initialize` is the uniform geometry (every expert block the same size). `initialize_sized` takes a
+// size per expert (a pack whose layers differ, as the GGUF's mixed i-quants): the spares are kept per size, and
+// since a swap exchanges two experts of one layer, the buffer a commit hands back has the size of the one it took,
+// so every size keeps its number of spares for good.
 class ExchangeStorage {
 public:
     struct View {
@@ -36,34 +42,93 @@ class ExchangeStorage {
         if (resident_slots > std::numeric_limits<size_t>::max() - count) {
             error = "exchange slot count overflow"; return false;
         }
-        auto experts = std::make_unique<std::atomic<size_t>[]>(offsets.size());
-        std::vector<size_t> spares(count);
-        std::vector<View> buffers(resident_slots + count);
-        for (size_t slot = 0; slot < resident_slots; ++slot)
-            buffers[slot] = {resident_host + slot * blob_bytes, resident_device + slot * blob_bytes};
         std::vector<bool> used(resident_slots, false);
         for (size_t i = 0; i < offsets.size(); ++i) {
             const uint64_t at = offsets[i];
-            experts[i].store(kAbsent, std::memory_order_relaxed);
             if (at == ~uint64_t{0}) continue;
             if (at >= resident_bytes || at % blob_bytes || used[(size_t)(at / blob_bytes)]) {
                 error = "invalid or duplicate resident slot"; return false;
             }
             used[(size_t)(at / blob_bytes)] = true;
-            experts[i].store((size_t)(at / blob_bytes), std::memory_order_relaxed);
         }
         for (bool present : used) if (!present) {
             error = "unassigned resident slot"; return false;
         }
-        for (size_t q = 0; q < count; ++q) {
-            spares[q] = resident_slots + q;
-            buffers[spares[q]] = {exchange_host + q * blob_bytes, exchange_device + q * blob_bytes};
+        const std::vector<uint64_t> sizes(offsets.size(), (uint64_t) blob_bytes);
+        return initialize_sized(offsets, sizes, resident_host, resident_device, resident_bytes, exchange_host,
+                                exchange_device, (uint64_t) count * (uint64_t) blob_bytes, count, 0, error);
+    }
+
+    /// The exchange allocation `initialize_sized` lays out: `count - bounce` buffers of every distinct size in
+    /// `sizes` (zeros skipped), then `bounce` buffers of the largest.  0 when the geometry is empty or overflows.
+    static uint64_t exchange_bytes_needed(const std::vector<uint64_t>& sizes, size_t count, size_t bounce) {
+        const std::vector<size_t> classes = classes_of(sizes);
+        if (classes.empty() || count == 0 || bounce > count) return 0;
+        uint64_t sum = 0;
+        for (const size_t c : classes) sum += c;
+        const uint64_t lim = std::numeric_limits<uint64_t>::max();
+        if (sum > lim / count) return 0;
+        return (uint64_t) (count - bounce) * sum + (uint64_t) bounce * classes.back();
+    }
+
+    /// `sizes[i]`: expert i's block size (its layer's).  Spare ordinals `q < count - bounce` exist in every size;
+    /// the last `bounce` ordinals only in the largest (buffers that are staged through and never committed).
+    bool initialize_sized(const std::vector<uint64_t>& offsets, const std::vector<uint64_t>& sizes,
+                          uint8_t* resident_host, const uint8_t* resident_device, uint64_t resident_bytes,
+                          uint8_t* exchange_host, const uint8_t* exchange_device, uint64_t exchange_bytes,
+                          size_t count, size_t bounce, std::string& error) {
+        error.clear();
+        if (active()) { error = "exchange storage already initialized"; return false; }
+        if (!resident_host || !resident_device || !exchange_host || !exchange_device || !count || bounce > count ||
+            !resident_bytes || sizes.size() != offsets.size()) {
+            error = "invalid mapped exchange geometry"; return false;
+        }
+        const std::vector<size_t> classes = classes_of(sizes);
+        const uint64_t need = exchange_bytes_needed(sizes, count, bounce);
+        if (classes.empty() || need == 0 || need > exchange_bytes ||
+            classes.size() > std::numeric_limits<size_t>::max() / count) {
+            error = "the exchange allocation does not hold the spares of every block size"; return false;
+        }
+        auto experts = std::make_unique<std::atomic<size_t>[]>(offsets.size());
+        std::vector<View> buffers;
+        std::vector<size_t> buffer_bytes;
+        std::vector<std::pair<uint64_t, uint64_t>> spans;   // (offset, size) of the resident blocks: none may overlap
+        for (size_t i = 0; i < offsets.size(); ++i) {
+            const uint64_t at = offsets[i];
+            experts[i].store(kAbsent, std::memory_order_relaxed);
+            if (at == ~uint64_t{0}) continue;
+            if (sizes[i] == 0 || at >= resident_bytes || sizes[i] > resident_bytes - at) {
+                error = "invalid resident slot"; return false;
+            }
+            experts[i].store(buffers.size(), std::memory_order_relaxed);
+            buffers.push_back({resident_host + (size_t) at, resident_device + (size_t) at});
+            buffer_bytes.push_back((size_t) sizes[i]);
+            spans.emplace_back(at, sizes[i]);
         }
+        std::sort(spans.begin(), spans.end());
+        for (size_t i = 1; i < spans.size(); ++i)
+            if (spans[i - 1].first + spans[i - 1].second > spans[i].first) {
+                error = "invalid or duplicate resident slot"; return false;
+            }
+        std::vector<size_t> spares(classes.size() * count, kAbsent);
+        uint64_t at = 0;
+        for (size_t c = 0; c < classes.size(); ++c)
+            for (size_t q = 0; q < count; ++q) {
+                if (q >= count - bounce && c + 1 != classes.size()) continue;   // the bounce ordinals: the largest size only
+                spares[c * count + q] = buffers.size();
+                buffers.push_back({exchange_host + (size_t) at, exchange_device + (size_t) at});
+                buffer_bytes.push_back(classes[c]);
+                at += classes[c];
+            }
         experts_ = std::move(experts);
         expert_count_ = offsets.size();
+        expert_bytes_.assign(sizes.begin(), sizes.end());
         spares_.swap(spares);
         buffers_.swap(buffers);
-        blob_bytes_ = blob_bytes;
+        buffer_bytes_.swap(buffer_bytes);
+        classes_ = classes;
+        count_ = count;
+        blob_bytes_ = classes.back();
         return true;
     }
 
@@ -73,21 +138,31 @@ class ExchangeStorage {
         const size_t slot = experts_[expert].load(std::memory_order_acquire);
         return slot != kAbsent ? buffers_[slot] : View{};
     }
-    View spare(size_t q) const { return q < spares_.size() ? buffers_[spares_[q]] : View{}; }
+    /// Spare `q` of the largest block size (any block fits it).
+    View spare(size_t q) const { return active() ? spare(q, blob_bytes_) : View{}; }
+    /// Spare `q` of blocks of exactly `bytes`.
+    View spare(size_t q, size_t bytes) const {
+        const size_t c = class_of(bytes);
+        if (c == kAbsent || q >= count_) return {};
+        const size_t id = spares_[c * count_ + q];
+        return id != kAbsent ? buffers_[id] : View{};
+    }
 
     // Only after every CPU reader and H2D copy of `in` has completed, and `out`
-    // has landed in spare(q). No allocation, memcpy, or CUDA operation here.
+    // has landed in spare(q, bytes). No allocation, memcpy, or CUDA operation here.
     bool commit(size_t in, size_t out, size_t q, const uint8_t* staged, size_t bytes) {
-        if (in >= expert_count_ || out >= expert_count_ || q >= spares_.size() ||
-            !staged || bytes != blob_bytes_) return false;
+        const size_t c = class_of(bytes);
+        if (in >= expert_count_ || out >= expert_count_ || c == kAbsent || q >= count_ || !staged) return false;
+        size_t& spare_id = spares_[c * count_ + q];
         const size_t incoming_slot = experts_[in].load(std::memory_order_acquire);
-        if (incoming_slot == kAbsent || experts_[out].load(std::memory_order_acquire) != kAbsent ||
-            buffers_[spares_[q]].host != staged) return false;
+        if (spare_id == kAbsent || incoming_slot == kAbsent || experts_[out].load(std::memory_order_acquire) != kAbsent ||
+            buffer_bytes_[incoming_slot] != bytes || expert_bytes_[out] != bytes || buffers_[spare_id].host != staged)
+            return false;
         // A background router lookahead can query residency to skip redundant
         // file prefetches. Publish one atomic slot ID; the pointer/alias pair is
         // immutable. Actual compute readers still obey the completion contract.
-        experts_[out].store(spares_[q], std::memory_order_release);
-        spares_[q] = incoming_slot;
+        experts_[out].store(spare_id, std::memory_order_release);
+        spare_id = incoming_slot;
         experts_[in].store(kAbsent, std::memory_order_release);
         ++exchanges_;
         avoided_bytes_ += bytes;
@@ -96,18 +171,36 @@ class ExchangeStorage {
 
     uint64_t exchanges() const { return exchanges_; }
     uint64_t avoided_bytes() const { return avoided_bytes_; } // memcpy payload, not read+write total
+    size_t block_sizes() const { return classes_.size(); }
     void clear() {
-        experts_.reset(); expert_count_ = 0; buffers_.clear(); spares_.clear();
-        blob_bytes_ = 0; exchanges_ = avoided_bytes_ = 0;
+        experts_.reset(); expert_count_ = 0; expert_bytes_.clear(); buffers_.clear(); buffer_bytes_.clear(); spares_.clear();
+        classes_.clear();
+        count_ = 0; blob_bytes_ = 0; exchanges_ = avoided_bytes_ = 0;
     }
 
 private:
     static constexpr size_t kAbsent = std::numeric_limits<size_t>::max();
+    static std::vector<size_t> classes_of(const std::vector<uint64_t>& sizes) {
+        std::vector<size_t> c;
+        for (const uint64_t b : sizes)
+            if (b != 0 && b <= std::numeric_limits<size_t>::max()) c.push_back((size_t) b);
+        std::sort(c.begin(), c.end());
+        c.erase(std::unique(c.begin(), c.end()), c.end());
+        return c;
+    }
+    size_t class_of(size_t bytes) const {
+        const auto it = std::lower_bound(classes_.begin(), classes_.end(), bytes);
+        return it != classes_.end() && *it == bytes ? (size_t) (it - classes_.begin()) : kAbsent;
+    }
     std::unique_ptr<std::atomic<size_t>[]> experts_;
     size_t expert_count_ = 0;
+    std::vector<uint64_t> expert_bytes_; // each expert's block size, immutable
     std::vector<View> buffers_; // immutable physical addresses after initialize
-    std::vector<size_t> spares_; // touched only by the serialized cache manager
-    size_t blob_bytes_ = 0;
+    std::vector<size_t> buffer_bytes_; // each buffer's block size, immutable
+    std::vector<size_t> spares_; // [size class * count_ + q]; touched only by the serialized cache manager
+    std::vector<size_t> classes_; // the distinct block sizes, ascending
+    size_t count_ = 0;
+    size_t blob_bytes_ = 0; // the largest block size; 0 = not initialized
     uint64_t exchanges_ = 0, avoided_bytes_ = 0;
 };