prefill: STRATA_PREFILL_ROUTED_STREAM=1 (opt-in) does not copy the streamed experts no token of the chunk routes to; test

A prompt chunk of 1,024 tokens or more streams every non-resident expert of every layer through the ring, and the issuer
thread enqueues each copy as soon as the ring has room. The ring is longer than a layer's entries, so a layer's copies are
all on the copy stream before its router has run: an expert no token routes to is copied anyway.

With the switch on, for a streamed chunk on one GPU with the issuer thread:
- the issuer keeps at most STRATA_PREFILL_ROUTED_STREAM_DEPTH copies (default 16) queued ahead of the link: it waits until
  the copy issued that many copies earlier has landed (cudaEventQuery on its slot's copied event). The ring's own rule is
  unchanged, so the stream still runs as far ahead of the compute as before;
- after each layer's host grouping the compute thread publishes which experts at least one token routes to;
- the issuer drops an entry it has not issued yet whose layer is published and whose expert has no token. Entries already
  issued stay. An entry staged through the Stager's pinned buffers is never dropped (its job queue is in order).

Every routed entry is still copied, in the same order, into the same ring slot (k % ring), with the same slot waits and
copied event; the consumer's walk is unchanged and never reads an unrouted entry's slot, and the copy stream stays in order,
so a group's single wait on its last copy still covers the group. The prompt's products are the same bits. After the read
the borrowed slots are refilled from RAM blob by blob into the same slots, independent of what the ring held.

On v0.2.6 (4b9e05a2) at the release defaults (RTX 3090, PCIe 3.0 x16, 250 W; the same binary with the switch off and on,
4 pairs):
a 1,790-token follow-up read on an 18.8K-token conversation took 17.9% less time, fresh reads of 2K, 4K and 32K tokens
ran 20%, 8% and 5% faster, and the streamed experts fell by 23%, 16% and 4%. In a deterministic lane the replies are
byte-identical to that base at 1K to 32K with the switch off and on, and the switch dropped 42% of the streamed experts at
a 1,046-token read and 3% at 31K tokens.

Unset, nothing new runs: no gate, no publication, and the default issuer loop is not edited. The fused path, the peer
prompt ring, the inline issuer (STRATA_PREFILL_ISSUER=0) and the staged walk are untouched. With STRATA_TRACE=1 each chunk
prints the ring entries it copied and dropped.

The issuer walk and the route gate are header-only (include/strata/prefill/routed_stream.hpp); routed_stream_test runs
the walk itself against a simulated in-order link and a consumer that walks the layers as prefill.cpp does (every routed
entry issued into its own slot and read with its own bytes, only unrouted entries of published layers dropped, never more
than the depth queued, never past the ring), and on three threads with real atomics.

STRATA_PREFILL_COPY_EVENT_EVERY > 1 and STRATA_PREFILL_COPY_WAIT_DEDUPE=1 (#4005) keep the default issuer: the routed
walk records a copied event for every copy, but it does not publish the records a sparse consumer waits on (a_rec,
rec_of) and does not deduplicate release waits, so with either set the consumer would wait forever or the wait policy
would be bypassed. routed_stream::eligible holds the rule and routed_stream_test checks it. Those two settings are now
read before the switch's own, which moves their two declarations up; nothing else in the default issuer loop changes.
This commit is contained in:
Stefano Sinden committed 2026-10-31 13:17:30 +11:00
1 parent 870aedd9e1
commit 3956b94a75
5 files changed
+631 -8

No files matched your search

+7
View File
@@ -391,6 +391,13 @@ if(STRATA_BUILD_TESTS)
target_include_directories(share_rules_test PRIVATE include)
add_test(NAME share_rules_test COMMAND share_rules_test)
# STRATA_PREFILL_ROUTED_STREAM: the prompt stream's issuer walk and route gate, on a simulated link and consumer
# (header-only, no GPU)
add_executable(routed_stream_test tests/prefill/routed_stream_test.cpp)
target_include_directories(routed_stream_test PRIVATE include)
target_link_libraries(routed_stream_test PRIVATE Threads::Threads)
add_test(NAME routed_stream_test COMMAND routed_stream_test)
# the Windows HIP auto prefill chunk cap (#1630) and the commit-limit verdict (#1607): header-only, no GPU
add_executable(host_limits_test tests/platform/host_limits_test.cpp)
target_include_directories(host_limits_test PRIVATE include)
+12
View File
@@ -1066,6 +1066,18 @@ measured. "Exact" means the 10-prompt and 4K/20K identity gate with the switch o
`STRATA_DEQUANT_WIDE` (wide dequantization of Q4_K, Q5_K, Q6_K and IQ4_XS for the prompt's dense weights) is on by default on
sm_86 and sm_61, bit-exact. `STRATA_DEQUANT_WIDE=0` turns it off.
`STRATA_PREFILL_ROUTED_STREAM=1` (opt-in; the streamed prompt walk of a chunk of 1,024 tokens or more, one GPU, the issuer
thread): the issuer keeps at most `STRATA_PREFILL_ROUTED_STREAM_DEPTH` copies (default 16) queued ahead of the PCIe link,
each layer's routing is handed to it as soon as the layer's router has run, and an expert not yet copied that no token of
the chunk routes to is not copied at all. The routed experts go into the same ring slots as without the switch, so the
prompt's results are the same bits, and the slots the prompt borrowed are refilled from RAM as before. With `STRATA_TRACE=1`
each chunk prints how many ring entries it copied and how many it dropped. With `STRATA_PREFILL_COPY_EVENT_EVERY` above 1 or
`STRATA_PREFILL_COPY_WAIT_DEDUPE=1` the switch stays off: those keep the default issuer, whose copy records the routed walk
does not write. On an RTX 3090 (PCIe 3.0 x16, 250 W) with the Swift IQ3_XXS pack at the v0.2.6 release defaults, 23% fewer
experts were streamed for a 2K prompt and 4% fewer at 32K (8,192-token chunks); a 1,790-token follow-up read on an
18.8K-token conversation took 18% less time and fresh 2K, 4K and 32K reads ran 20%, 8% and 5% faster (4 pairs), with
byte-identical replies in a deterministic lane at 1K to 32K.
### RTX 20 (Turing, sm_75) kernels (opt-in)
Kernels by rafatxf (#2069-#2072, #2093, #2220-#2222), imanu86 (#2081) and xrip (#1821) for compute capability 7.5 are bitwise
+150
View File
@@ -0,0 +1,150 @@
#pragma once
// STRATA_PREFILL_ROUTED_STREAM (opt-in, off by default): the streamed prompt walk drops the ring entries no token of
// the chunk routes to, when their layer's routing is known before the entry is issued.
//
// A chunk of stream_all_min() tokens or more streams every non-resident expert of every layer through the ring
// (prefill.cpp, step 3), and its issuer thread enqueues each copy as soon as the ring has room: with a ring longer
// than a layer's entries, a layer's copies are all on the copy stream before its router has run, so an expert no
// token routes to is copied anyway (measured on an RTX 3090, IQ3_XXS: 47% of the streamed experts at 1,045 tokens,
// 35% at 2K, 15% at 8,192-token chunks). With the switch on:
// - the issuer keeps at most `depth` copies queued ahead of the link (it waits until the copy issued `depth`
// copies earlier has landed), so the copies run back to back but are not committed far ahead of the link;
// how far the stream runs ahead of the compute (the ring) is unchanged;
// - after a layer's router, the compute thread publishes which experts at least one token routes to (RouteGate);
// - the issuer drops an entry not yet issued whose layer is published and whose expert no token routes to.
// Entries already issued stay; every routed entry is issued, into the ring slot it has today (k % ring), and
// the consumer's walk is unchanged (it never reads an unrouted entry's slot), so the prompt's products are the
// same bits. Only entries copied straight from the pinned arena are dropped: an entry staged through the
// Stager's pinned buffers keeps its copy (its job queue is consumed in order).
// Header-only so the walk and its bookkeeping are tested on the CPU (tests/prefill/routed_stream_test.cpp).
#include <algorithm>
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <vector>
namespace strata::prefill::routed_stream {
constexpr int kDefaultDepth = 16; // about 28 MB of 1.7 MB blobs, 2.4 ms of a PCIe 3.0 x16 link
constexpr int kMaxDepth = 256;
// STRATA_PREFILL_ROUTED_STREAM: unset, empty or 0 is off; any other number is on.
inline bool enabled_from(const char* v) { return v != nullptr && v[0] != '\0' && std::atoi(v) != 0; }
// STRATA_PREFILL_ROUTED_STREAM_DEPTH: the copies kept queued ahead of the link (default 16, clamped to 1..256).
inline int depth_from(const char* v) {
if (v == nullptr || v[0] == '\0') return kDefaultDepth;
return std::clamp(std::atoi(v), 1, kMaxDepth);
}
// Whether a chunk takes the routed walk. The routed walk records a copied event for every copy it issues, but it does
// not publish the entries a sparse consumer waits on (STRATA_PREFILL_COPY_EVENT_EVERY > 1: a_rec / rec_of) and does not
// deduplicate release waits (STRATA_PREFILL_COPY_WAIT_DEDUPE=1). With either set the chunk keeps the default issuer:
// otherwise the consumer would wait on a record that never comes, or the requested wait policy would be bypassed.
inline bool eligible(bool stream_all, bool issuer_on, bool peer, int copy_event_every, bool copy_wait_dedupe) {
return stream_all && issuer_on && !peer && copy_event_every == 1 && !copy_wait_dedupe;
}
// Which experts each layer's tokens route to, published by the compute thread once the layer's router has run and
// read by the issuer thread. One writer per layer per chunk: the routed bytes are written before the layer's flag
// is set (release), and the reader reads them only after it sees the flag (acquire).
class RouteGate {
public:
void reset(int64_t n_layers, int32_t n_expert) {
n_layers_ = n_layers > 0 ? n_layers : 0;
n_expert_ = n_expert > 0 ? n_expert : 0;
routed_.assign((size_t) n_layers_ * (size_t) n_expert_, 0);
pub_.reset(new std::atomic<uint8_t>[(size_t) (n_layers_ > 0 ? n_layers_ : 1)]);
for (int64_t l = 0; l < n_layers_; ++l) pub_[(size_t) l].store(0, std::memory_order_relaxed);
}
// `cnt[e]`: the tokens of the chunk that route to expert e in layer l (the host grouping's counts)
void publish(int64_t l, const int32_t* cnt) {
if (l < 0 || l >= n_layers_ || cnt == nullptr) return;
uint8_t* row = routed_.data() + (size_t) l * (size_t) n_expert_;
for (int32_t e = 0; e < n_expert_; ++e) row[e] = cnt[e] > 0 ? 1 : 0;
pub_[(size_t) l].store(1, std::memory_order_release);
}
bool published(int64_t l) const {
return l >= 0 && l < n_layers_ && pub_[(size_t) l].load(std::memory_order_acquire) != 0;
}
// true only when layer l is published and no token of the chunk routes to expert e
bool unrouted(int64_t l, int32_t e) const {
if (e < 0 || e >= n_expert_ || !published(l)) return false;
return routed_[(size_t) l * (size_t) n_expert_ + (size_t) e] == 0;
}
private:
int64_t n_layers_ = 0;
int32_t n_expert_ = 0;
std::vector<uint8_t> routed_;
std::unique_ptr<std::atomic<uint8_t>[]> pub_;
};
// The walk's operations as six callables (prefill.cpp passes lambdas; a test may pass any type with these members).
template <class C, class S, class L, class D, class I, class Y>
struct FnOps {
C consumed;
S stop;
L landed;
D droppable;
I issue;
Y idle;
};
template <class C, class S, class L, class D, class I, class Y>
FnOps<C, S, L, D, I, Y> make_ops(C c, S s, L l, D d, I i, Y y) {
return FnOps<C, S, L, D, I, Y>{c, s, l, d, i, y};
}
struct WalkResult {
bool stopped = false; // the walk was told to stop before its end
size_t issued = 0; // entries copied
size_t dropped = 0; // entries not copied: published, unrouted and droppable
};
// The issuer's walk over a chunk's entries `seq` (each with members `l` and `e`), entry k going to ring slot
// k % ring as without the switch. `a_issued` is advanced past every entry, copied or dropped (the consumer's
// wait_issued reads it). Ops:
// size_t consumed() the consumer's give-back count (acquire load of a_consumed)
// bool stop() the chunk is ending early
// bool landed(int slot) the copy last issued into `slot` has completed (a query, never a wait)
// bool droppable(const E&) the entry may be dropped (its blob is copied straight from the pinned arena)
// void issue(size_t k, int sl) enqueue entry k's copy into slot sl exactly as the default walk does
// void idle() nothing to do yet (yield)
// An entry is copied only when (1) the ring has room for it (k < consumed() + ring, as without the switch) and
// (2) fewer than `depth` copies are queued (the copy issued `depth` copies before it has landed). While it waits for
// either, its layer's routing may be published; then an unrouted entry is dropped instead.
template <class Entry, class Ops>
WalkResult issue_walk(const std::vector<Entry>& seq, size_t ring, int depth, const RouteGate& gate,
std::atomic<size_t>& a_issued, Ops& ops) {
WalkResult r;
if (ring == 0) return r;
const size_t q = (size_t) std::clamp(depth, 1, kMaxDepth);
int win[kMaxDepth]; // the slots of the last q copies, the n-th copy at n % q
size_t n_out = 0;
for (size_t k = 0; k < seq.size(); ++k) {
const Entry& en = seq[k];
const bool may_drop = ops.droppable(en);
bool drop = false;
for (;;) {
if (may_drop && gate.unrouted(en.l, en.e)) { drop = true; break; }
if (k < ops.consumed() + ring && (n_out < q || ops.landed(win[n_out % q]))) break;
if (ops.stop()) { r.stopped = true; return r; }
ops.idle();
}
if (drop) {
++r.dropped;
} else {
const int sl = (int) (k % ring);
ops.issue(k, sl);
win[n_out % q] = sl;
++n_out;
++r.issued;
}
a_issued.store(k + 1, std::memory_order_release);
}
return r;
}
} // namespace strata::prefill::routed_stream
+68 -8
View File
@@ -47,6 +47,7 @@
#include "strata/prefill/moe_ws.hpp"
#include "strata/prefill/moe_fused_iq.hpp"
#include "strata/prefill/moe_mmq.hpp"
#include "strata/prefill/routed_stream.hpp"
#include "strata/core/peer_experts.hpp"
#include "strata/core/expert_nibble.hpp"
#include "strata/core/sm120_series.hpp"
@@ -3459,6 +3460,36 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
const char* v = std::getenv("STRATA_PREFILL_ISSUER");
return v == nullptr || std::atoi(v) != 0;
}();
static const int copy_event_every = [] {
const char* v = std::getenv("STRATA_PREFILL_COPY_EVENT_EVERY");
return strata::prefill::copy_event_every_from(v ? std::atoi(v) : 1);
}();
static const bool copy_wait_dedupe = [] {
const char* v = std::getenv("STRATA_PREFILL_COPY_WAIT_DEDUPE");
return v != nullptr && std::atoi(v) != 0;
}();
// STRATA_PREFILL_ROUTED_STREAM=1 (opt-in, routed_stream.hpp): the issuer keeps at most a few copies queued
// ahead of the link (STRATA_PREFILL_ROUTED_STREAM_DEPTH, default 16), each layer's routing is published once
// its router has run (host grouping below), and an entry not yet issued that no token of the chunk routes to
// is dropped. Routed entries are copied as below, into the same slots: the same bits. The threaded issuer
// on one GPU only (no peer prompt ring), and not with STRATA_PREFILL_COPY_EVENT_EVERY > 1 or
// STRATA_PREFILL_COPY_WAIT_DEDUPE=1, whose copy records the routed walk does not write (routed_stream::eligible);
// unset, nothing here runs and the loop below is the walk. Declared
// before the issuer thread, so they outlive it on every return (IssuerJoin joins it first).
static const bool rs_env = routed_stream::enabled_from(std::getenv("STRATA_PREFILL_ROUTED_STREAM"));
static const int rs_depth = routed_stream::depth_from(std::getenv("STRATA_PREFILL_ROUTED_STREAM_DEPTH"));
const bool rs_on = rs_env && routed_stream::eligible(
stream_all, issuer_on, m.pp != nullptr, copy_event_every, copy_wait_dedupe);
routed_stream::RouteGate rs_gate;
size_t rs_dropped = 0;
if (rs_on) {
rs_gate.reset(g.n_layers, (int32_t) m.g->n_expert);
if (static bool said = false; !said) {
said = true;
std::fprintf(stderr, "strata: STRATA_PREFILL_ROUTED_STREAM=1: a streamed prompt chunk drops the experts no token "
"routes to before they are copied (at most %d copies queued ahead of the link)\n", rs_depth);
}
}
std::atomic<size_t> a_issued{0}, a_consumed{0};
std::atomic<bool> a_stop{false};
double iss_ms = 0;
@@ -3478,14 +3509,6 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
// after the one it needs (the copy stream is in order). STRATA_PREFILL_COPY_WAIT_DEDUPE=1 skips the copy
// stream's wait on a slot's release event when the previous slot was released by the same event (one flush).
// The same copies into the same slots: the same bytes.
static const int copy_event_every = [] {
const char* v = std::getenv("STRATA_PREFILL_COPY_EVENT_EVERY");
return strata::prefill::copy_event_every_from(v ? std::atoi(v) : 1);
}();
static const bool copy_wait_dedupe = [] {
const char* v = std::getenv("STRATA_PREFILL_COPY_WAIT_DEDUPE");
return v != nullptr && std::atoi(v) != 0;
}();
const bool sparse_copied = threaded_issue && copy_event_every > 1;
if (threaded_issue && (sparse_copied || copy_wait_dedupe)) {
static bool said = false;
@@ -3505,6 +3528,34 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
strata::platform::name_this_thread("strata-issue");
strata::aux_cpus::pin_current_thread();
const core::OnDevice od(m.device);
if (rs_on) {
auto ops = routed_stream::make_ops(
[&] { return a_consumed.load(std::memory_order_acquire); },
[&] { return a_stop.load(std::memory_order_acquire); },
[&](int sl) { return cudaEventQuery(m.copied[sl]) != cudaErrorNotReady; },
[](const StreamEntry& en) { return en.job < 0; },
[&](size_t idx, int sl) { // the copy of the loop below, unchanged
const StreamEntry& en = seq[idx];
const auto th = Clock::now();
const size_t bytes = (size_t) lay0.blob_bytes(en.l);
if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[m.used_of[sl]], 0);
if (en.job < 0) {
cudaMemcpyAsync(m.stage_dev[sl], en.blob, bytes, cudaMemcpyHostToDevice, m.copy);
++iss_dma;
} else {
const uint8_t* hb = m.stager->wait(en.job);
cudaMemcpyAsync(m.stage_dev[sl], hb, bytes, cudaMemcpyHostToDevice, m.copy);
m.stager->issued_one(en.job, m.copy);
}
cudaEventRecord(m.copied[sl], m.copy);
m.stage_live[sl] = true;
iss_ms += ms_since(th);
++iss_streamed;
},
[] { std::this_thread::yield(); });
rs_dropped = routed_stream::issue_walk(seq, (size_t) m.ring, rs_depth, rs_gate, a_issued, ops).dropped;
return;
}
int64_t last_rec = -1; // sparse_copied: the last entry whose copied event is recorded
int prev_rel = -1; // dedupe: the release event the copy stream waited on for the previous entry
auto record = [&](size_t i) {
@@ -4447,6 +4498,9 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
if (e < 0 || e >= m.g->n_expert) { err = "prefill: routed id out of range"; return false; }
++m.cnt[(size_t) e];
}
// STRATA_PREFILL_ROUTED_STREAM: this layer's routing to the issuer (it drops the unrouted entries
// it has not issued yet; the walk below gives every unrouted entry back without reading it)
if (rs_on) rs_gate.publish(l, m.cnt.data());
if (cpu_maybe) {
std::vector<std::pair<int32_t, int32_t>> cand; // (tokens, expert)
int64_t nstream = 0;
@@ -5611,6 +5665,12 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
stats_.ms_experts_host += iss_ms;
stats_.experts_streamed += iss_streamed;
stats_.experts_dma += iss_dma;
if (rs_on && std::getenv("STRATA_TRACE")) {
std::fprintf(stderr, "strata trace: routed stream: chunk from token %lld (%lld tokens): %lld of %zu ring entries "
"copied, %zu dropped (no token routes to them), depth %d\n", (long long) p0, (long long) T,
(long long) iss_streamed, seq.size(), rs_dropped, rs_depth);
std::fflush(stderr);
}
if (copy_wait_dedupe && pt.on)
std::fprintf(stderr, "strata prefill copies: %lld slot waits, %lld skipped as the same release event\n",
(long long) iss_waits, (long long) iss_dedup);
+394
View File
@@ -0,0 +1,394 @@
// routed_stream_test: STRATA_PREFILL_ROUTED_STREAM's issuer walk and route gate (include/strata/prefill/routed_stream.hpp).
// CPU only. A simulated chunk: an in-order link that lands one copy every few ticks, a consumer that walks the layers as
// prefill.cpp's streamed walk does (publish a layer's routing, give back unrouted entries, read a routed entry's slot
// once its copy has landed), and the issuer walk itself, each idle() of which advances the simulation by one tick.
#include "strata/prefill/routed_stream.hpp"
#include <atomic>
#include <cstdio>
#include <deque>
#include <mutex>
#include <random>
#include <thread>
#include <vector>
namespace rs = strata::prefill::routed_stream;
static int g_fail = 0;
static void expect(bool ok, const char* what) {
if (!ok) {
++g_fail;
std::fprintf(stderr, "FAIL: %s\n", what);
}
}
struct Entry { int32_t l, e; int job; };
struct Case {
int layers = 6, n_expert = 64, resident_every = 3, ring = 48, depth = 4;
double unrouted = 0.4; // share of the streamed experts no token routes to
int link_ticks = 3; // ticks per copy
int publish_delay = 5; // ticks from the end of layer l-1 to layer l's routing
int compute_ticks = 1; // ticks per routed entry
bool publish = true; // false: no layer is ever published (the fused path, or the switch off)
int staged_every = 0; // every n-th entry has a stager job (never dropped); 0: none
unsigned seed = 1;
};
struct Sim {
Case c;
std::vector<Entry> seq;
std::vector<size_t> seq_start;
std::vector<std::vector<int32_t>> cnt; // [layer][expert] tokens
rs::RouteGate gate;
std::atomic<size_t> a_issued{0}, a_consumed{0};
// the link: copies in issue order, each lands link_ticks after it starts
struct Copy { size_t k; int slot; };
std::deque<Copy> queue;
int link_left = 0;
std::vector<long> slot_content; // the entry whose bytes a ring slot holds (-1: none)
std::vector<long> slot_last_issue; // the entry last issued into a slot (its `copied` event)
std::vector<char> landed_k; // per entry: its copy has landed
std::vector<char> issued_k, dropped_k, pub_at_drop;
size_t max_out = 0, outstanding = 0;
bool ring_violation = false, depth_violation = false, wrong_bytes = false, read_before_issue = false;
// the consumer
int64_t cl = 0; // layer being walked
size_t ck = 0; // next entry of the walk
int phase = 0; // 0: waiting to publish, 1: walking the layer
int wait = 0;
int64_t tick = 0;
explicit Sim(const Case& cc) : c(cc) {
std::mt19937 rng(c.seed);
std::uniform_real_distribution<double> u(0.0, 1.0);
seq_start.assign((size_t) c.layers + 1, 0);
cnt.assign((size_t) c.layers, std::vector<int32_t>((size_t) c.n_expert, 0));
for (int l = 0; l < c.layers; ++l) {
seq_start[(size_t) l] = seq.size();
for (int e = 0; e < c.n_expert; ++e) {
cnt[(size_t) l][(size_t) e] = u(rng) < c.unrouted ? 0 : 1 + (int) (u(rng) * 7);
if (c.resident_every > 0 && e % c.resident_every == 0) continue; // resident: not in the stream
const int job = c.staged_every > 0 && seq.size() % (size_t) c.staged_every == 0 ? (int) seq.size() : -1;
seq.push_back({l, e, job});
}
}
seq_start[(size_t) c.layers] = seq.size();
gate.reset(c.layers, c.n_expert);
slot_content.assign((size_t) c.ring, -1);
slot_last_issue.assign((size_t) c.ring, -1);
landed_k.assign(seq.size(), 0);
issued_k.assign(seq.size(), 0);
dropped_k.assign(seq.size(), 0);
pub_at_drop.assign(seq.size(), 0);
wait = c.publish_delay;
}
void step_link() {
if (queue.empty()) return;
if (--link_left > 0) return;
const Copy cp = queue.front();
queue.pop_front();
slot_content[(size_t) cp.slot] = (long) cp.k; // the bytes land
landed_k[cp.k] = 1;
--outstanding;
link_left = c.link_ticks;
}
void give_back(size_t upto) { a_consumed.store(upto, std::memory_order_release); }
void step_consumer() {
if (cl >= c.layers) return;
const size_t kend = seq_start[(size_t) cl + 1];
if (phase == 0) { // the layer's attention and router: then the routing is known
if (--wait > 0) return;
if (c.publish) gate.publish(cl, cnt[(size_t) cl].data());
ck = seq_start[(size_t) cl];
phase = 1;
wait = 0;
return;
}
if (wait > 0) { --wait; return; }
if (ck < kend) {
const Entry& en = seq[ck];
if (cnt[(size_t) cl][(size_t) en.e] == 0) { // release_to: an unrouted entry only gives its slot back
give_back(++ck);
return;
}
// a routed entry: wait_issued, then the copy's event, then compute from the slot
if (a_issued.load(std::memory_order_acquire) <= ck) return;
if (!issued_k[ck]) { read_before_issue = true; give_back(++ck); return; }
if (!landed_k[ck]) return;
if (slot_content[ck % (size_t) c.ring] != (long) ck) wrong_bytes = true;
wait = c.compute_ticks;
give_back(++ck);
return;
}
++cl; // the next layer's attention and router
phase = 0;
wait = c.publish_delay;
}
void tick_world() {
++tick;
step_link();
step_consumer();
}
};
struct SimOps {
Sim& s;
size_t consumed() { return s.a_consumed.load(std::memory_order_acquire); }
bool stop() { return s.tick > 50000000; }
bool landed(int slot) {
const long k = s.slot_last_issue[(size_t) slot];
return k < 0 || s.landed_k[(size_t) k];
}
bool droppable(const Entry& en) { return en.job < 0; }
void issue(size_t k, int sl) {
if (k >= consumed() + (size_t) s.c.ring) s.ring_violation = true;
// the copy waits for the slot's previous entry to be released (its `used` event): the ring rule above
s.slot_last_issue[(size_t) sl] = (long) k;
s.queue.push_back({k, sl});
if (s.queue.size() == 1) s.link_left = s.c.link_ticks;
++s.outstanding;
s.max_out = std::max(s.max_out, s.outstanding);
if (s.outstanding > (size_t) s.c.depth) s.depth_violation = true;
s.issued_k[k] = 1;
}
void idle() { s.tick_world(); }
};
// drops are recorded through the gate as seen at the walk's end; check each dropped entry was unrouted and published
static void run_case(const Case& c, const char* name, size_t* dropped_out = nullptr, size_t* unrouted_out = nullptr) {
Sim s(c);
// the walk ticks the world only while it waits; ticks between its steps (the issuer's own work) advance it too
struct TickingOps : SimOps {
void issue(size_t k, int sl) { SimOps::issue(k, sl); s.tick_world(); }
} tops{{s}};
const rs::WalkResult r = rs::issue_walk(s.seq, (size_t) c.ring, c.depth, s.gate, s.a_issued, tops);
for (size_t k = 0; k < s.seq.size(); ++k)
if (!s.issued_k[k]) s.dropped_k[k] = 1;
// let the consumer finish the chunk
while (s.cl < c.layers && s.tick < 50000000) s.tick_world();
char buf[256];
std::snprintf(buf, sizeof buf, "%s: walk not stopped", name);
expect(!r.stopped, buf);
std::snprintf(buf, sizeof buf, "%s: issued + dropped == entries (%zu + %zu vs %zu)", name, r.issued, r.dropped, s.seq.size());
expect(r.issued + r.dropped == s.seq.size(), buf);
std::snprintf(buf, sizeof buf, "%s: a_issued covers every entry", name);
expect(s.a_issued.load() == s.seq.size(), buf);
std::snprintf(buf, sizeof buf, "%s: the consumer walked every layer", name);
expect(s.cl == c.layers, buf);
std::snprintf(buf, sizeof buf, "%s: no copy beyond the ring", name);
expect(!s.ring_violation, buf);
std::snprintf(buf, sizeof buf, "%s: at most depth copies queued (max %zu, depth %d)", name, s.max_out, c.depth);
expect(!s.depth_violation, buf);
std::snprintf(buf, sizeof buf, "%s: every routed entry read its own bytes from slot k %% ring", name);
expect(!s.wrong_bytes && !s.read_before_issue, buf);
size_t unrouted = 0, bad_drop = 0;
for (size_t k = 0; k < s.seq.size(); ++k) {
const Entry& en = s.seq[k];
const bool unr = s.cnt[(size_t) en.l][(size_t) en.e] == 0;
if (unr && en.job < 0) ++unrouted;
if (s.dropped_k[k] && (!unr || en.job >= 0 || !c.publish)) ++bad_drop;
}
std::snprintf(buf, sizeof buf, "%s: only unrouted, arena-copied entries of published layers dropped (%zu bad)", name, bad_drop);
expect(bad_drop == 0, buf);
std::printf(" %-34s entries %5zu unrouted %5zu dropped %5zu (%.0f%% of unrouted) max queued %zu ticks %lld\n", name,
s.seq.size(), unrouted, r.dropped, unrouted ? 100.0 * (double) r.dropped / (double) unrouted : 0.0,
s.max_out, (long long) s.tick);
if (dropped_out) *dropped_out = r.dropped;
if (unrouted_out) *unrouted_out = unrouted;
}
// The same walk on three threads (issuer, consumer, link), real atomics, no simulated clock: the bookkeeping holds
// whatever the interleaving.
static void threaded_case(unsigned seed, int depth) {
Case c;
c.layers = 8; c.n_expert = 128; c.ring = 96; c.depth = depth; c.seed = seed; c.unrouted = 0.45;
Sim s(c);
std::atomic<bool> done{false};
std::vector<std::atomic<int>> landed(s.seq.size());
for (auto& x : landed) x.store(0);
std::vector<std::atomic<long>> content((size_t) c.ring);
for (auto& x : content) x.store(-1);
std::vector<std::atomic<long>> last((size_t) c.ring);
for (auto& x : last) x.store(-1);
std::vector<std::atomic<int>> issued(s.seq.size());
for (auto& x : issued) x.store(0);
std::mutex qm;
std::deque<std::pair<size_t, int>> q;
std::atomic<size_t> out{0};
std::atomic<bool> depth_bad{false}, ring_bad{false}, bytes_bad{false};
struct ThreadOps {
Sim& s; std::mutex& qm; std::deque<std::pair<size_t, int>>& q; std::vector<std::atomic<long>>& last;
std::vector<std::atomic<int>>& landed_v; std::vector<std::atomic<int>>& issued; std::atomic<size_t>& out;
std::atomic<bool>& depth_bad; std::atomic<bool>& ring_bad; int depth;
size_t consumed() { return s.a_consumed.load(std::memory_order_acquire); }
bool stop() { return false; }
bool landed(int slot) { const long k = last[(size_t) slot].load(); return k < 0 || landed_v[(size_t) k].load(); }
bool droppable(const Entry& en) { return en.job < 0; }
void issue(size_t k, int sl) {
if (k >= consumed() + (size_t) s.c.ring) ring_bad = true;
last[(size_t) sl].store((long) k);
issued[k].store(1);
if (out.fetch_add(1) + 1 > (size_t) depth) depth_bad = true;
std::lock_guard<std::mutex> g(qm);
q.emplace_back(k, sl);
}
void idle() { std::this_thread::yield(); }
} ops{s, qm, q, last, landed, issued, out, depth_bad, ring_bad, depth};
std::thread link([&] {
while (!done.load()) {
std::pair<size_t, int> cp;
{
std::lock_guard<std::mutex> g(qm);
if (q.empty()) { cp.second = -1; } else { cp = q.front(); q.pop_front(); }
}
if (cp.second < 0) { std::this_thread::yield(); continue; }
for (int i = 0; i < 50; ++i) std::this_thread::yield(); // the copy's time
content[(size_t) cp.second].store((long) cp.first);
out.fetch_sub(1); // before `landed`: the issuer counts a queued copy until it sees this one land
landed[cp.first].store(1);
}
});
std::thread consumer([&] {
for (int l = 0; l < c.layers; ++l) {
for (int i = 0; i < 200; ++i) std::this_thread::yield(); // attention and router
s.gate.publish(l, s.cnt[(size_t) l].data());
for (size_t k = s.seq_start[(size_t) l]; k < s.seq_start[(size_t) l + 1]; ++k) {
if (s.cnt[(size_t) l][(size_t) s.seq[k].e] == 0) { s.a_consumed.store(k + 1, std::memory_order_release); continue; }
while (s.a_issued.load(std::memory_order_acquire) <= k) std::this_thread::yield();
while (!landed[k].load()) std::this_thread::yield();
if (!issued[k].load() || content[k % (size_t) c.ring].load() != (long) k) bytes_bad = true;
s.a_consumed.store(k + 1, std::memory_order_release);
}
}
});
const rs::WalkResult r = rs::issue_walk(s.seq, (size_t) c.ring, depth, s.gate, s.a_issued, ops);
consumer.join();
done = true;
link.join();
char buf[200];
std::snprintf(buf, sizeof buf, "threaded seed %u depth %d: bookkeeping (issued %zu dropped %zu of %zu)", seed, depth,
r.issued, r.dropped, s.seq.size());
expect(r.issued + r.dropped == s.seq.size() && !depth_bad && !ring_bad && !bytes_bad, buf);
size_t bad = 0;
for (size_t k = 0; k < s.seq.size(); ++k)
if (!issued[k].load() && s.cnt[(size_t) s.seq[k].l][(size_t) s.seq[k].e] != 0) ++bad;
std::snprintf(buf, sizeof buf, "threaded seed %u depth %d: every routed entry issued", seed, depth);
expect(bad == 0, buf);
std::printf(" threaded seed %u depth %-3d entries %5zu dropped %5zu\n", seed, depth, s.seq.size(), r.dropped);
}
int main() {
// the switch and its depth
expect(!rs::enabled_from(nullptr), "unset: off");
expect(!rs::enabled_from(""), "empty: off");
expect(!rs::enabled_from("0"), "0: off");
expect(rs::enabled_from("1"), "1: on");
expect(rs::depth_from(nullptr) == rs::kDefaultDepth, "depth unset: default");
expect(rs::depth_from("8") == 8, "depth 8");
expect(rs::depth_from("0") == 1 && rs::depth_from("-4") == 1, "depth clamps to 1");
expect(rs::depth_from("100000") == rs::kMaxDepth, "depth clamps to the maximum");
// Which chunks take the routed walk: the copy policies of STRATA_PREFILL_COPY_EVENT_EVERY > 1 and
// STRATA_PREFILL_COPY_WAIT_DEDUPE=1 keep the default issuer. Otherwise a sparse consumer's wait_copied would spin
// on an a_rec the routed walk never publishes, and the requested release-wait deduplication would be bypassed.
expect(rs::eligible(true, true, false, 1, false), "streamed threaded single GPU: routed walk eligible");
expect(!rs::eligible(false, true, false, 1, false), "short prompt: default walk");
expect(!rs::eligible(true, false, false, 1, false), "inline issuer: default walk");
expect(!rs::eligible(true, true, true, 1, false), "peer prompt ring: default walk");
for (int every : {2, 16, 64}) {
expect(!rs::eligible(true, true, false, every, false), "sparse copied events: default walk");
expect(!rs::eligible(true, true, false, every, true), "sparse events and dedupe: default walk");
}
expect(!rs::eligible(true, true, false, 1, true), "release-wait dedupe: default walk");
// the gate
{
rs::RouteGate g;
g.reset(3, 8);
const int32_t c1[8] = {0, 2, 0, 1, 0, 0, 5, 0};
expect(!g.published(1) && !g.unrouted(1, 0), "nothing published: nothing unrouted");
g.publish(1, c1);
expect(g.published(1) && !g.published(0) && !g.published(2), "one layer published");
bool ok = true;
for (int e = 0; e < 8; ++e) ok = ok && g.unrouted(1, e) == (c1[e] == 0);
expect(ok, "unrouted exactly where the count is 0");
expect(!g.unrouted(0, 0) && !g.unrouted(2, 0), "other layers keep every entry");
expect(!g.unrouted(1, -1) && !g.unrouted(1, 8) && !g.unrouted(-1, 0) && !g.unrouted(3, 0), "out of range: kept");
g.reset(3, 8);
expect(!g.published(1) && !g.unrouted(1, 0), "reset: nothing published");
}
std::printf("routed_stream_test: simulated chunks\n");
{
Case c; // routing never published (the default walk's behaviour): every entry copied, in order
c.publish = false;
size_t d = 1;
run_case(c, "never published", &d);
expect(d == 0, "never published: nothing dropped");
}
{
Case c; // the routing of every layer known before the walk could reach it: every unrouted entry dropped
c.publish_delay = 1; c.link_ticks = 40; c.ring = 20;
size_t d = 0, u = 0;
run_case(c, "routing long before the link", &d, &u);
expect(d > 0 && d >= u * 8 / 10, "early routing: most unrouted entries dropped");
}
{
Case c; // a link-bound chunk: some of each layer is issued before its routing lands
c.publish_delay = 30; c.link_ticks = 4;
size_t d = 0;
run_case(c, "link-bound", &d);
expect(d > 0, "link-bound: some entries dropped");
}
{
Case c; // stager jobs are never dropped
c.staged_every = 3; c.publish_delay = 1; c.link_ticks = 20;
run_case(c, "with stager jobs");
}
{
Case c; // a ring smaller than a layer, depth larger than the ring
c.ring = 10; c.depth = 64; c.link_ticks = 2;
run_case(c, "ring < layer, depth > ring");
}
{
Case c; // depth 1: one copy in flight at a time
c.depth = 1; c.link_ticks = 2;
run_case(c, "depth 1");
}
{
Case c; // every expert routed: nothing to drop
c.unrouted = 0.0;
size_t d = 1;
run_case(c, "all routed", &d);
expect(d == 0, "all routed: nothing dropped");
}
{
Case c; // nothing routed in the stream (all unrouted), fast compute
c.unrouted = 1.0; c.link_ticks = 5;
run_case(c, "none routed");
}
for (unsigned seed = 2; seed < 40; ++seed) {
Case c;
c.seed = seed;
c.layers = 3 + (int) (seed % 5);
c.n_expert = 32 + (int) (seed * 7 % 97);
c.ring = 4 + (int) (seed * 13 % 120);
c.depth = 1 + (int) (seed * 5 % 40);
c.unrouted = (seed % 10) / 10.0;
c.link_ticks = 1 + (int) (seed % 6);
c.publish_delay = 1 + (int) (seed * 3 % 60);
c.compute_ticks = (int) (seed % 4);
c.staged_every = seed % 3 == 0 ? 4 : 0;
char name[64];
std::snprintf(name, sizeof name, "random seed %u", seed);
run_case(c, name);
}
std::printf("routed_stream_test: threaded walks\n");
for (unsigned seed = 1; seed <= 6; ++seed) threaded_case(seed, seed % 2 ? 4 : 16);
if (g_fail) { std::fprintf(stderr, "routed_stream_test: %d failure(s)\n", g_fail); return 1; }
std::printf("routed_stream_test: OK\n");
return 0;
}