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

Rotation hands the evicted expert's exchange buffer over as its place in the
RAM copy instead of copying the bytes there. It required one block 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 nine.

A swap exchanges two experts of one layer, so the slot the promoted expert
leaves has the size the evicted one needs. ExchangeStorage now keeps its
spares per block size (initialize_sized): an exchange takes a spare of its
layer's size and hands back one of the same size. commit() refuses a spare of
another size and an expert of another layer. initialize() keeps the uniform
geometry and its checks. exchange_buffer() takes the layer; reserve_exchanges()
takes the number of trailing buffers that only stage copies (--adapt-async's
second half), kept in the largest size alone. With --adapt-async 1 a rotated
round flips in the tick that sees the copies in: there is nothing for the job
to move, so the round ends one verify window sooner.

Still opt-in (STRATA_EXCHANGE_ROTATE=1); unset, nothing changes. The exchange
allocation grows to one spare of every size per swap of a round: 1.06 GiB of
page-locked RAM on IQ3_S with --adapt-swaps 80 (1.26 GiB with --adapt-async 1)
against 0.20 and 0.40 GiB.

Measured on a Core i7-13700K, RTX 5090 32 GB, 64 GB, Windows 11, Swift 1.5
IQ3_S, 128K, --resident-experts, --adapt-every 1 --adapt-swaps 80, p-cores:
  one 8K prompt, the same 1,500 forced tokens, blocking tier, rotation off / on:
    161.6-165.1 tok/s (five runs) / 184.9-188.7 (four runs); two runs of each
    arm with the same cache size (10,681 slots) wrote identical --window-hashes
    for all 434 windows
  --serve, 23 coding-agent turns, greedy, three interleaved rounds, blocking tier:
    147.8 / 159.2 / 158.4 (the build before) -> 175.2 / 184.9 / 172.8
  the same turns with --adapt-async 1, one binary, two rounds:
    166.0 / 170.1 -> 185.6 / 182.5
The tier swaps about 30 experts per verify window on this box (44% of the
experts fit in VRAM), so the copy path moved about 50 MiB of RAM per window.

tests/core/exchange_storage_test.cpp gains a three-size case (6,144 exchanges)
and passes built alone with MSVC 14.50. Built for CUDA only (sm_120, CUDA
13.0.88, MSVC 14.50 with nvcc's --allow-unsupported-compiler); HIP and SYCL
were not compiled (expert_source.cpp and generate.cpp are built by both). Not
measured: Linux, AMD cards, a layer split.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
Xiaomi2008andClaude Opus 5.5 committed 2026-10-11 21:43:48 -07:00
1 parent 7695859aa1
commit 8db3750184
7 files changed
+328 -52

No files matched your search

+4 -2
View File
@@ -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.
+71 -4
View File
@@ -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:
+117 -24
View File
@@ -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 @@ public:
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 @@ public:
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 @@ public:
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;
};
+8 -3
View File
@@ -597,15 +597,20 @@ public:
// buffer and calls `stage_exchange`: `out` is then read from that buffer, and `in` still from here (the CPU
// computes both until the swap lands). Once the slot copy has landed, `commit_exchanges` moves `out` into
// `in`'s place, so the copy keeps holding exactly the experts the GPU does not - with no read of the file.
// STRATA_EXCHANGE_ROTATE=1 can instead transfer buffer ownership (uniform, fully pinned/mapped slots).
// STRATA_EXCHANGE_ROTATE=1 can instead transfer buffer ownership (fully pinned/mapped slots; a pack whose layers
// differ in block size keeps spares per size).
/// Whether the compact copy holds `(layer, expert)`.
bool has_resident(int64_t layer, int64_t expert) const;
/// Host room for `n` evicted blobs (page-locked when possible). Idempotent for the same or a smaller `n`.
bool reserve_exchanges(int64_t n, std::string& err);
/// `bounce`: how many of the last buffers are only staged through and never hold an exchange (STRATA_EXCHANGE_ROTATE
/// keeps those in the largest block size alone).
bool reserve_exchanges(int64_t n, std::string& err, int64_t bounce = 0);
int64_t exchange_capacity() const { return xstage_cap_; }
/// The exchange buffers are page-locked (cudaHostAlloc): copies to and from them are asynchronous.
bool exchange_pinned() const { return xstage_pinned_; }
uint8_t* exchange_buffer(int64_t q) const;
/// Buffer `q`. `layer` >= 0: the buffer an exchange of that layer's experts stages into (STRATA_EXCHANGE_ROTATE with
/// several block sizes keeps one per size); without it, one that holds any layer's block.
uint8_t* exchange_buffer(int64_t q, int64_t layer = -1) const;
/// Requires `has_resident(layer, in)`, `!has_resident(layer, out)` and `exchange_buffer(q)` holding out's blob.
bool stage_exchange(int64_t layer, int64_t in, int64_t out, int64_t q);
/// --pipeline-windows: size the exchange table now, so a `stage_exchange` on the adaptive tier's thread never
+25 -14
View File
@@ -3223,7 +3223,7 @@ bool FileExpertSource::has_resident(int64_t layer, int64_t expert) const {
return resident_blob(index) != nullptr;
}
bool FileExpertSource::reserve_exchanges(int64_t n, std::string& err) {
bool FileExpertSource::reserve_exchanges(int64_t n, std::string& err, int64_t bounce) {
err.clear();
if (n <= xstage_cap_) return true;
if (!staged_.empty()) { err = "FileExpertSource: exchange buffers are in use"; return false; }
@@ -3245,11 +3245,13 @@ bool FileExpertSource::reserve_exchanges(int64_t n, std::string& err) {
}
const char* rotate_env = std::getenv("STRATA_EXCHANGE_ROTATE");
const bool requested = rotate_env && std::strcmp(rotate_env, "1") == 0;
const bool uniform = std::all_of(layer_blob_bytes_.begin(), layer_blob_bytes_.end(),
[blob](uint64_t b) { return b == blob; });
const bool eligible = requested && uniform && complement_pinned_ && !complement_partial_ &&
// layers of several block sizes (a GGUF's mixed i-quants): spares of every size, see ExchangeStorage
bounce = std::max<int64_t>(0, std::min(bounce, n));
const uint64_t rotate_total = detail::ExchangeStorage::exchange_bytes_needed(layer_blob_bytes_, (size_t) n, (size_t) bounce);
const bool eligible = requested && rotate_total != 0 && rotate_total <= std::numeric_limits<size_t>::max() &&
complement_pinned_ && !complement_partial_ &&
complement_host_ && complement_device_ && complement_bytes_ > 0;
const size_t total = (size_t) n * (size_t) blob;
const size_t total = eligible ? (size_t) rotate_total : (size_t) n * (size_t) blob;
void* p = nullptr;
if (cudaHostAlloc(&p, total, eligible ? cudaHostAllocMapped : cudaHostAllocDefault) == cudaSuccess && p != nullptr) {
xstage_pinned_ = true;
@@ -3263,14 +3265,21 @@ bool FileExpertSource::reserve_exchanges(int64_t n, std::string& err) {
xstage_cap_ = n;
xstage_blob_ = blob;
if (requested) {
std::string reason = "requires equal-size expert blocks and a fully mapped/pinned RAM complement";
std::string reason = "requires a fully mapped/pinned RAM complement";
void* device = nullptr;
if (eligible && xstage_pinned_ && cudaHostGetDevicePointer(&device, p, 0) == cudaSuccess && device) {
try {
if (exchange_storage_.initialize(complement_offsets_, (uint8_t*)complement_host_, complement_device_,
complement_bytes_, xstage_, (const uint8_t*)device, (size_t)n, (size_t)blob, reason)) {
std::fprintf(stderr, "FileExpertSource: exchange buffer rotation enabled: %lld buffers, %llu bytes each; no host commit memcpy\n",
(long long)n, (unsigned long long)blob);
std::vector<uint64_t> sizes(complement_offsets_.size(), 0);
for (size_t i = 0; i < sizes.size(); ++i) {
const size_t layer = n_expert_ > 0 ? i / (size_t) n_expert_ : 0;
sizes[i] = layer < layer_blob_bytes_.size() ? layer_blob_bytes_[layer] : 0;
}
if (exchange_storage_.initialize_sized(complement_offsets_, sizes, (uint8_t*)complement_host_, complement_device_,
complement_bytes_, xstage_, (const uint8_t*)device, (uint64_t)total, (size_t)n, (size_t)bounce, reason)) {
std::fprintf(stderr, "FileExpertSource: exchange buffer rotation enabled: %lld buffers in %zu block sizes "
"(largest %llu bytes, %.2f GiB page-locked); no host commit memcpy\n",
(long long)n, exchange_storage_.block_sizes(), (unsigned long long)blob,
(double)total / 1073741824.0);
}
} catch (const std::bad_alloc&) { reason = "buffer metadata allocation failed"; }
} else if (eligible) {
@@ -3283,22 +3292,24 @@ bool FileExpertSource::reserve_exchanges(int64_t n, std::string& err) {
return true;
}
uint8_t* FileExpertSource::exchange_buffer(int64_t q) const {
uint8_t* FileExpertSource::exchange_buffer(int64_t q, int64_t layer) const {
if (xstage_ == nullptr || q < 0 || q >= xstage_cap_) return nullptr;
if (exchange_storage_.active()) return exchange_storage_.spare((size_t)q).host;
if (exchange_storage_.active())
return layer >= 0 && layer < n_layers_ ? exchange_storage_.spare((size_t)q, (size_t)layer_blob_bytes_[(size_t)layer]).host
: exchange_storage_.spare((size_t)q).host;
return xstage_ + (size_t) q * (size_t) xstage_blob_;
}
bool FileExpertSource::stage_exchange(int64_t layer, int64_t in, int64_t out, int64_t q) {
if (out < 0 || out >= n_expert_ || !has_resident(layer, in) || has_resident(layer, out) ||
exchange_buffer(q) == nullptr) return false;
exchange_buffer(q, layer) == nullptr) return false;
const size_t i_in = (size_t) layer * (size_t) n_expert_ + (size_t) in;
const size_t i_out = (size_t) layer * (size_t) n_expert_ + (size_t) out;
if (override_.empty()) override_.assign((size_t) blobs_, nullptr);
if (override_[i_out] != nullptr) return false;
for (const Exchange& x : staged_)
if (x.in == i_in || x.q == q) return false;
override_[i_out] = exchange_buffer(q);
override_[i_out] = exchange_buffer(q, layer);
staged_.push_back({i_in, i_out, q, layer_blob_bytes_[(size_t) layer]});
return true;
}
+6 -4
View File
@@ -290,7 +290,7 @@ bool resident_stage_swaps(strata::core::FileExpertSource& src, const std::vector
if (slot < 0) continue;
const SwapHome home = locate(s.layer);
const strata::core::OnDevice on(home.dev);
if (const cudaError_t e = cudaMemcpyAsync(src.exchange_buffer(q), home.cache->device_slot(slot),
if (const cudaError_t e = cudaMemcpyAsync(src.exchange_buffer(q, s.layer), home.cache->device_slot(slot),
(size_t) strata::kernels::cpu::expert_layout().blob_bytes(s.layer), cudaMemcpyDeviceToHost,
home.stream); e != cudaSuccess) {
std::fprintf(stderr, "strata: copy back of layer %d slot %d (device %d) failed: %s\n", (int) s.layer,
@@ -6765,7 +6765,8 @@ int main(int argc, char** argv) {
if (resident_ok) {
// --adapt-async: twice the exchange buffers - the second half bounces the copies in whose source is not
// page-locked (an asynchronous copy instead of the driver's staged, synchronous one)
if (o.adapt_async && !src.reserve_exchanges(2 * std::min<int64_t>(o.adapt_swaps, 96), err))
if (o.adapt_async && !src.reserve_exchanges(2 * std::min<int64_t>(o.adapt_swaps, 96), err,
std::min<int64_t>(o.adapt_swaps, 96)))
adapt_async_off(err.c_str());
// pageable buffers would make every copy of a round synchronous (the driver stages them): that is the
// blocking tier's cost, so the asynchronous tier is not worth running - said, and the blocking tier stays
@@ -8605,7 +8606,7 @@ int main(int argc, char** argv) {
const strata::core::OnDevice on(h.dev);
for (const ASwap& w : aswaps) {
if (w.home != hi || !w.exchange) continue;
if (const cudaError_t e = cudaMemcpyAsync(src.exchange_buffer(w.j), h.cache->device_slot(w.slot),
if (const cudaError_t e = cudaMemcpyAsync(src.exchange_buffer(w.j, w.layer), h.cache->device_slot(w.slot),
(size_t) strata::kernels::cpu::expert_layout().blob_bytes(w.layer),
cudaMemcpyDeviceToHost, h.stream); e != cudaSuccess) {
std::fprintf(stderr, "strata serve: adaptive tier: copy back of layer %d slot %d (CUDA%d) failed: %s\n",
@@ -8798,8 +8799,9 @@ int main(int argc, char** argv) {
for (size_t hi = 0; hi < ahomes.size(); ++hi)
if (a_ram_due[hi] && a_gap_seq[hi] <= a_ram_mark[hi]) return true;
std::fill(a_ram_due.begin(), a_ram_due.end(), (uint8_t) 0);
ajob->post([&] { src.commit_copies(); });
astate = AState::Commit;
if (src.exchange_rotation()) continue; // STRATA_EXCHANGE_ROTATE: no bytes move, the flip is all there is
ajob->post([&] { src.commit_copies(); });
if (!drain) return true;
continue;
case AState::Commit: // step 4
+97 -1
View File
@@ -110,6 +110,101 @@ void exercise(size_t bytes, size_t rounds) {
!storage.exchanges() && !storage.avoided_bytes(), "clear retained state");
}
// Layers of different block sizes: the spares are kept per size, and a swap (two experts of one layer) hands back a
// buffer of the size it took.
void exercise_sized(size_t rounds) {
const size_t layer_bytes[3] = {33, 4097, 64};
constexpr size_t per_layer = 10, resident_per_layer = 6, n = 3 * per_layer, count = 4, bounce = 1, guard = 19;
std::vector<uint64_t> sizes(n), offsets(n, absent);
uint64_t resident_bytes = 0;
for (size_t i = 0; i < n; ++i) {
sizes[i] = layer_bytes[i / per_layer];
if (i % per_layer < resident_per_layer) { offsets[i] = resident_bytes; resident_bytes += sizes[i]; }
}
const uint64_t need = ExchangeStorage::exchange_bytes_needed(sizes, count, bounce);
require(need == (count - bounce) * (33 + 64 + 4097) + bounce * 4097, "wrong exchange allocation size");
std::vector<uint8_t> arena(guard + resident_bytes + guard, 0xcd), spare(guard + need + guard, 0xef);
std::vector<uint8_t> aliases(resident_bytes), spare_aliases(need);
std::vector<std::vector<uint8_t>> reference(n);
std::mt19937 random(603);
for (size_t i = 0; i < n; ++i) {
reference[i].resize(sizes[i]);
for (auto& v : reference[i]) v = (uint8_t)random();
if (offsets[i] != absent) std::memcpy(arena.data() + guard + offsets[i], reference[i].data(), sizes[i]);
}
ExchangeStorage storage;
std::string error;
require(!storage.initialize_sized(offsets, sizes, arena.data() + guard, aliases.data(), resident_bytes,
spare.data() + guard, spare_aliases.data(), need - 1, count, bounce, error), "short exchange allocation accepted");
require(storage.initialize_sized(offsets, sizes, arena.data() + guard, aliases.data(), resident_bytes,
spare.data() + guard, spare_aliases.data(), need, count, bounce, error), error.c_str());
require(storage.block_sizes() == 3, "wrong number of block sizes");
// the bounce ordinal exists in the largest size only; the others in every size
require(storage.spare(count - 1, 4097).host && !storage.spare(count - 1, 33).host && !storage.spare(count - 1, 64).host,
"bounce buffers of a smaller size");
require(storage.spare(count - 1).host == storage.spare(count - 1, 4097).host, "the size-less spare is not the largest");
require(!storage.spare(0, 34).host && !storage.spare(count, 33).host, "a spare of no size or ordinal");
std::set<uint8_t*> original;
for (size_t i = 0; i < n; ++i) if (auto q = storage.resident(i).host) original.insert(q);
for (size_t c = 0; c < 3; ++c)
for (size_t q = 0; q < count; ++q) if (auto b = storage.spare(q, layer_bytes[c]).host) original.insert(b);
uint64_t exchanges = 0, bytes_moved = 0;
for (size_t round = 0; round < rounds; ++round) {
size_t used[3] = {0, 0, 0}; // the ordinal is the swap's number in the round, whatever its layer
size_t ordinal = 0;
std::vector<size_t> order = {0, 1, 2, 0, 1, 2};
std::shuffle(order.begin(), order.end(), random);
for (const size_t layer : order) {
if (ordinal >= count - bounce) break;
const size_t bytes = layer_bytes[layer], q = ordinal++;
std::vector<size_t> ins, outs;
for (size_t i = layer * per_layer; i < (layer + 1) * per_layer; ++i) (storage.resident(i).host ? ins : outs).push_back(i);
std::shuffle(ins.begin(), ins.end(), random);
std::shuffle(outs.begin(), outs.end(), random);
const size_t in = ins[0], out = outs[0];
const ExchangeStorage::View before = storage.resident(in), staged = storage.spare(q, bytes);
require(staged.host != nullptr, "no spare of the layer's size");
std::memcpy(staged.host, reference[out].data(), bytes);
const size_t other = layer_bytes[(layer + 1) % 3];
require(!storage.commit(in, out, q, staged.host, other), "another size's commit accepted");
size_t stranger = ((layer + 1) % 3) * per_layer; // an expert of another layer that RAM does not hold
while (storage.resident(stranger).host) ++stranger;
require(!storage.commit(in, stranger, q, staged.host, bytes), "another layer's expert accepted");
require(storage.commit(in, out, q, staged.host, bytes), "valid sized swap rejected");
require(storage.resident(out).host == staged.host && storage.resident(out).device == staged.device,
"evicted buffer was not adopted with its alias");
require(storage.spare(q, bytes).host == before.host && storage.spare(q, bytes).device == before.device,
"old RAM input was not recycled with its alias");
++used[layer]; ++exchanges; bytes_moved += bytes;
}
std::set<uint8_t*> live;
for (size_t i = 0; i < n; ++i) if (auto q = storage.resident(i).host) {
require(live.insert(q).second, "two experts share storage");
require(!std::memcmp(q, reference[i].data(), sizes[i]), "resident bytes corrupted");
}
for (size_t c = 0; c < 3; ++c)
for (size_t q = 0; q < count; ++q) if (auto b = storage.spare(q, layer_bytes[c]).host)
require(live.insert(b).second, "spare aliases live expert");
require(live == original, "storage leaked or duplicated");
for (size_t layer = 0; layer < 3; ++layer) {
size_t held = 0;
for (size_t i = layer * per_layer; i < (layer + 1) * per_layer; ++i) held += storage.resident(i).host != nullptr;
require(held == resident_per_layer, "a layer's resident count changed");
}
}
require(storage.exchanges() == exchanges && storage.avoided_bytes() == bytes_moved, "wrong accounting");
require(std::all_of(arena.begin(), arena.begin() + guard, [](auto x) { return x == 0xcd; }) &&
std::all_of(arena.end() - guard, arena.end(), [](auto x) { return x == 0xcd; }) &&
std::all_of(spare.begin(), spare.begin() + guard, [](auto x) { return x == 0xef; }) &&
std::all_of(spare.end() - guard, spare.end(), [](auto x) { return x == 0xef; }), "guard overwritten");
// overlapping resident blocks are refused
std::vector<uint64_t> bad = offsets;
bad[1] = offsets[0] + 1;
ExchangeStorage t;
require(!t.initialize_sized(bad, sizes, arena.data() + guard, aliases.data(), resident_bytes,
spare.data() + guard, spare_aliases.data(), need, count, bounce, error) && !t.active(), "overlapping blocks accepted");
}
void invalid_geometry() {
uint8_t h[64]{}, d[64]{}, x[64]{}, xd[64]{};
ExchangeStorage s; std::string e;
@@ -128,6 +223,7 @@ int main() {
invalid_geometry();
for (size_t bytes : {1u, 33u, 4097u}) exercise(bytes, 1024);
exercise(5222400, 4); // actual Q8 expert size
std::cout << "exchange_storage_test: PASS (byte/alias/ownership checks, 12304 exchanges)\n";
exercise_sized(2048);
std::cout << "exchange_storage_test: PASS (byte/alias/ownership checks, 12304 exchanges; several block sizes)\n";
} catch (const std::exception& e) { std::cerr << e.what() << '\n'; return 1; }
}