From 47a1047750cae87d9e50d6cfa2c40e13a5fa7154 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 02:49:04 +0100 Subject: [PATCH 01/17] WIP multi-prompt layer-major prefill (Prefill::run_multi) and BCO co-reads in the batch admission --- include/strata/prefill/prefill.hpp | 17 +++ src/prefill/prefill.cpp | 190 ++++++++++++++++++++++------- src/program/generate.cpp | 90 +++++++++++++- 3 files changed, 254 insertions(+), 43 deletions(-) diff --git a/include/strata/prefill/prefill.hpp b/include/strata/prefill/prefill.hpp index 549aff38..5009f8fd 100644 --- a/include/strata/prefill/prefill.hpp +++ b/include/strata/prefill/prefill.hpp @@ -10,6 +10,8 @@ // Requires the native weights (`--native`): every quantized projection must carry its GGUF blocks. #pragma once +#include + #include "strata/core/expert_cache.hpp" #include "strata/core/expert_source.hpp" #include "strata/core/layer.hpp" @@ -265,6 +267,19 @@ public: /// Layer-major prefill (STRATA_PREFILL_LAYER_MAJOR=1): true while `on_chunk` reports a sub-chunk before the end of /// its super-chunk. The session state is already at the super-chunk's end then, so no checkpoint may be taken. /// On a layer split it is the last stage's (the one that reports the rows), asked through the first. + /// One prompt of a multi-prompt read (run_multi): its tokens, the position it starts at and, per stage of a layer split + /// (index 0 = the first stage), the session it reads into (nullptr / missing: that stage's own main session). + struct Seg { + const int64_t* tokens = nullptr; + int64_t n = 0; + int64_t pos0 = 0; + std::vector stage_ss; + }; + /// Several prompts in ONE layer-major pass (the streamed experts of a layer are read once for all of them). Each + /// prompt fills its own sessions; `on_chunk` / `on_stage_chunk` report the FIRST prompt only. Returns false with `err` + /// when the layer-major walk does not apply (then nothing was read; the caller reads the prompts one by one). + bool run_multi(const std::vector& segs, std::string& err); + bool layer_major_mid() const { const Prefill* p = this; while (p->next_ != nullptr) p = p->next_; @@ -282,6 +297,8 @@ private: int64_t lm_prepare_chain(int64_t n, std::string& why); void lm_release_chain(); bool run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err); + struct LmSub { int seg = 0; int64_t c0 = 0, T = 0, t0 = 0; }; ///< a sub-chunk: prompt, first token, tokens, row in the super-chunk + bool lm_run(const std::vector& segs, int depth, const std::vector* given_super, std::string& err); void lm_release(); bool lm_mid_ = false; bool lm_by_default_ = false; ///< the walk is on because nothing was set (a wider VRAM margin, the measured super-chunk size) diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index 16a609a5..24ba3277 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -6148,15 +6148,55 @@ void Prefill::lm_release_chain() { } bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err) { + Seg one; + one.tokens = tokens; + one.n = n; + one.pos0 = pos0; + return lm_run(std::vector{one}, 0, nullptr, err); +} + +// Several prompts in ONE layer-major pass (PAR-H, 2026-10-30): the super-chunk is a list of sub-chunks, each from one +// prompt ("segment") with its own sessions; a layer's attention halves run per sub-chunk against that segment's +// session (GDN / QSA / PLE state), and the layer's experts are staged ONCE for every sub-chunk of the super-chunk - +// the streamed experts cost is shared by the prompts. A single segment with no sessions given is exactly the usual run. +bool Prefill::run_multi(const std::vector& segs, std::string& err) { + err.clear(); + if (segs.empty()) { err = "prefill: no prompt"; return false; } + int64_t total = 0; + for (const Seg& sg : segs) total += sg.n; + std::string why; + if (hand_in_ != nullptr || !impl_) { err = "prefill: multi-prompt read declined (chunk hand-off in use)"; return false; } + if (lm_prepare_chain(total, why) <= 0) { err = "prefill: multi-prompt read declined: " + why; return false; } + bool ok = lm_run(segs, 0, nullptr, err); + std::string drain_err; + if (!drain_pipeline(drain_err) && ok) { + ok = false; + err = drain_err; + } + for (Prefill* q = this; q != nullptr; q = q->next_) q->lm_mid_ = false; + lm_release_chain(); + return ok; +} + +bool Prefill::lm_run(const std::vector& segs, int depth, const std::vector* given_super, std::string& err) { err.clear(); Impl& m = *impl_; const core::OnDevice on_device(m.device); const core::ModelGeometry& g = *m.g; - core::SessionState& ss = *m.ss; + // the session of segment i on THIS stage: the segment's own (stage depth `depth`), else this stage's main session + auto ss_of = [&](size_t i) -> core::SessionState* { + const Seg& sg = segs[i]; + return (size_t) depth < sg.stage_ss.size() && sg.stage_ss[(size_t) depth] != nullptr ? sg.stage_ss[(size_t) depth] : m.ss; + }; + core::SessionState* ssp = ss_of(0); + const bool multi = segs.size() > 1; + int64_t n = 0; // all the prompts' tokens + for (const Seg& sg : segs) n += sg.n; + const int64_t pos0 = segs[0].pos0; // (only the progress line's) const auto t_start = Clock::now(); const cudaStream_t cs = (cudaStream_t) m.cs; Impl::LmBufs& lb = m.lm; - const int64_t S = lb.S, L = g.n_layers; + const int64_t L = g.n_layers; // a layer split's stage (2026-10-24): its layers [LB, LE); a later stage's rows arrive in R_host, written by the // stage before it, and a stage before the last writes its final rows into the next stage's R_host, then runs // the next stage on that super-chunk on a thread while it reads its own next super-chunk (one hand-off buffer: @@ -6166,38 +6206,85 @@ bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, st strata::kernels::QsaShapes s = core::qsa_shapes(g); // as 0.2.2's run_impl const uint64_t gdn_floats = (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size + (uint64_t) g.ssm_conv_channels * (g.ssm_d_conv - 1); - int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]}; + std::vector> prevs(segs.size()); + for (size_t i = 0; i < segs.size(); ++i) prevs[i] = {ss_of(i)->ple_prev[0], ss_of(i)->ple_prev[1]}; PfTimer pt; // the MMQ row table lives in the borrowed cache slots when the prompt path borrows: written for every prompt // (none without an MMQ layer: the FP16 path has no such table) if (m.ids_identity) mmq::iota(m.ids_identity, m.T * K, m.cs); // ---- the PLE rows, read per chunk on a thread as run_impl does (the same state machine, in chunk order) - const bool ple_on = ss.ple.ready(); + const bool ple_on = ssp->ple.ready(); static const bool ple_batch_env = [] { const char* v = std::getenv("STRATA_PLE_BATCH"); return v == nullptr || std::atoi(v) != 0; }(); const bool ple_batch = ple_on && ple_batch_env && strata::kernels::ple_native_postops_enabled() && - (ss.ple.w.key_bf16 != nullptr || ss.ple.w.key_native_data != nullptr) && + (ssp->ple.w.key_bf16 != nullptr || ssp->ple.w.key_native_data != nullptr) && m.region_bytes / ((uint64_t) (3 * strata::kernels::NG_HC_DIM + N + 4) * 4 + (uint64_t) N * 2 + 4096) >= 64; - const int32_t prev0[2] = {prev[0], prev[1]}; - auto ple_gather = [&m, &ss, tokens, n, prev0](int64_t c0, int buf, std::string& e) -> bool { - const int64_t T = std::min(m.T, n - c0); - auto at = [&](int64_t i) { return i < 2 ? prev0[i] : (int32_t) tokens[i - 2]; }; + // ---- the plan. A segment is cut into pieces of at most S tokens (the usual super-chunks), each piece into + // sub-chunks of m.T tokens from the piece's start; pieces of several segments are packed into one super-chunk while + // their rows fit S. One segment: exactly the usual super-chunks. + const int64_t S = m.lm.S; + std::vector> supers_own; + const std::vector>* supers_p = &supers_own; + std::vector> one_given; + if (given_super != nullptr) { + one_given.push_back(*given_super); + supers_p = &one_given; + } else { + std::vector cur; + int64_t cur_rows = 0; + for (size_t si = 0; si < segs.size(); ++si) + for (int64_t p0s = 0; p0s < segs[si].n; p0s += S) { + const int64_t piece = std::min(S, segs[si].n - p0s); + if (!cur.empty() && cur_rows + piece > S) { + supers_own.push_back(std::move(cur)); + cur.clear(); + cur_rows = 0; + } + for (int64_t c0s = 0; c0s < piece; c0s += m.T) { + LmSub sb; + sb.seg = (int) si; + sb.c0 = p0s + c0s; + sb.T = std::min(m.T, piece - c0s); + sb.t0 = cur_rows + c0s; + cur.push_back(sb); + } + cur_rows += piece; + } + if (!cur.empty()) supers_own.push_back(std::move(cur)); + } + const std::vector>& supers = *supers_p; + std::vector> flat; // every sub-chunk in order: (super-chunk, sub-chunk) + for (size_t si = 0; si < supers.size(); ++si) + for (size_t c = 0; c < supers[si].size(); ++c) flat.push_back({(int) si, (int) c}); + std::vector flat_base(supers.size(), 0); + for (size_t si = 1; si < supers.size(); ++si) flat_base[si] = flat_base[si - 1] + (int64_t) supers[si - 1].size(); + auto sub_at = [&](int64_t gi) -> const LmSub& { return supers[(size_t) flat[(size_t) gi].first][(size_t) flat[(size_t) gi].second]; }; + const std::array prev_none{-1, -1}; + (void) prev_none; + std::vector> prev0s = prevs; + auto ple_gather = [&m, &segs, &ss_of, &prev0s, &sub_at](int64_t gi, int buf, std::string& e) -> bool { + const LmSub& sb = sub_at(gi); + const int64_t* tokens = segs[(size_t) sb.seg].tokens; + core::SessionState* ssq = ss_of((size_t) sb.seg); + const int64_t c0 = sb.c0, T = sb.T; + const std::array& prev0 = prev0s[(size_t) sb.seg]; + auto at = [&](int64_t i) { return i < 2 ? prev0[(size_t) i] : (int32_t) tokens[i - 2]; }; int32_t pv[2] = {at(c0), at(c0 + 1)}; for (int64_t t = 0; t < T; ++t) { const int32_t tok = (int32_t) tokens[c0 + t]; - strata::kernels::ngram_rows(&tok, pv, 1, ss.ple.consts, + strata::kernels::ngram_rows(&tok, pv, 1, ssq->ple.consts, m.ple_rows[buf].data() + t * strata::kernels::PLE_N_HEADS); pv[0] = pv[1]; pv[1] = tok; } - return ss.ple.table->gather_batch(m.ple_rows[buf].data(), (size_t) T, m.ple_emb_host[buf], e); + return ssq->ple.table->gather_batch(m.ple_rows[buf].data(), (size_t) T, m.ple_emb_host[buf], e); }; std::string ple_next_err; std::future ple_next; int ple_buf = 0; - auto ple_land = [&](int64_t c0, int64_t T) -> bool { + auto ple_land = [&](int64_t gi, int64_t T) -> bool { const auto tp = Clock::now(); if (!ple_next.get()) { err = ple_next_err; return false; } // #1996: the pageable fallback (no RAM to pin) goes up through a pinned slot; the pinned buffers as before @@ -6208,12 +6295,12 @@ bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, st err = std::string("prefill: the PLE rows' upload failed: ") + cudaGetErrorString(cudaGetLastError()); return false; } - if (c0 + m.T < n) { + if (gi + 1 < (int64_t) flat.size()) { if (cudaEventSynchronize(m.ple_copied[ple_buf ^ 1]) != cudaSuccess) { err = std::string("prefill: the PLE rows' upload failed: ") + cudaGetErrorString(cudaGetLastError()); return false; } - ple_next = std::async(std::launch::async, [&ple_gather, &ple_next_err, c1 = c0 + m.T, b = ple_buf ^ 1] { + ple_next = std::async(std::launch::async, [&ple_gather, &ple_next_err, c1 = gi + 1, b = ple_buf ^ 1] { strata::aux_cpus::pin_current_thread(); return ple_gather(c1, b, ple_next_err); }); @@ -6556,9 +6643,13 @@ bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, st return true; }; - for (int64_t sc0 = 0; sc0 < n; sc0 += S) { - const int64_t Sn = std::min(S, n - sc0); - const int64_t nc = (Sn + m.T - 1) / m.T; + for (size_t sci = 0; sci < supers.size(); ++sci) { + const std::vector& subs = supers[sci]; + int64_t Sn = 0; + for (const LmSub& sb : subs) Sn = std::max(Sn, sb.t0 + sb.T); + const int64_t nc = (int64_t) subs.size(); + int64_t last_seg0 = -1; + for (int64_t c = 0; c < nc; ++c) if (subs[(size_t) c].seg == 0) last_seg0 = c; ++lm_supers; lm_mid_ = true; for (int64_t l = LB; l <= LE; ++l) { @@ -6567,8 +6658,8 @@ bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, st // host mode: sub-chunk c's rows go up on the copy stream into buffer c % 2 (issued one sub-chunk ahead) auto r_buf = [&](int64_t c) { return (c & 1) ? lb.R2 : R_chunk; }; auto issue_up = [&](int64_t c) { - const int64_t cc0 = sc0 + c * m.T, TT = std::min(m.T, n - cc0); - cudaMemcpyAsync(r_buf(c), lb.R_host + (size_t) (c * m.T) * D, (size_t) TT * D * 4, cudaMemcpyHostToDevice, + const int64_t TT = subs[(size_t) c].T; + cudaMemcpyAsync(r_buf(c), lb.R_host + (size_t) subs[(size_t) c].t0 * D, (size_t) TT * D * 4, cudaMemcpyHostToDevice, m.copy); cudaEventRecord(lb.ev_in[c & 1], m.copy); }; @@ -6579,12 +6670,17 @@ bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, st for (int64_t c = 0; c < nc; ++c) { if (should_stop && should_stop()) { err = "cancelled"; return false; } core::progress_beat(); - core::progress_at("reading the prompt (layer-major): layer", l, pos0 + sc0 + c * m.T); + const LmSub& sub = subs[(size_t) c]; + const Seg& sgm = segs[(size_t) sub.seg]; + const int64_t* const tokens = sgm.tokens; + ssp = ss_of((size_t) sub.seg); + core::progress_at("reading the prompt (layer-major): layer", l, sgm.pos0 + sub.c0); qsa_index = qsa_index0; gdn_index = gdn_index0; - const int64_t c0 = sc0 + c * m.T; // in this run's tokens - const int64_t T = std::min(m.T, n - c0), p0 = pos0 + c0; - const int64_t t0 = c * m.T; // in the super-chunk + const int64_t c0 = sub.c0; // in this prompt's tokens + const int64_t T = sub.T, p0 = sgm.pos0 + c0; + const int64_t t0 = sub.t0; // in the super-chunk + const int64_t gi = flat_base[sci] + c; // in the whole run m.R = lb.r_dev ? lb.R + t0 * D : r_buf(c); if (!lb.r_dev) { if (l > 0) { @@ -6635,11 +6731,14 @@ bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, st if (!core::embed_row(*m.wt, g, tokens[c0 + t], m.emb + t * N, cs, err)) return false; gr_broadcast(m.emb, m.R, T, cs); if (ple_on && !ple_next.valid()) - ple_next = std::async(std::launch::async, [&ple_gather, &ple_next_err, c0, b = ple_buf] { + ple_next = std::async(std::launch::async, [&ple_gather, &ple_next_err, gi, b = ple_buf] { strata::aux_cpus::pin_current_thread(); - return ple_gather(c0, b, ple_next_err); + return ple_gather(gi, b, ple_next_err); }); - for (int64_t t = 0; t < T; ++t) { prev[0] = prev[1]; prev[1] = (int32_t) tokens[c0 + t]; } + { + std::array& pvv = prevs[(size_t) sub.seg]; + for (int64_t t = 0; t < T; ++t) { pvv[0] = pvv[1]; pvv[1] = (int32_t) tokens[c0 + t]; } + } } else if (l > LB) { // (a later stage's first layer: the stage before wrote layer LB - 1) // the previous layer's experts: the hyper-connection write of its FFN half gr_write(m.R, acc, m.inj, HC, T, cs); @@ -6684,19 +6783,21 @@ bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, st err = std::string("prefill: ") + cudaGetErrorString(cudaGetLastError()); return false; } - lm_mid_ = c + 1 < nc; - if (on_stage_chunk && !lm_mid_ && !on_stage_chunk(p0 + T, err)) return false; - if (on_chunk && !on_chunk(m.R, T, p0, err)) return false; + lm_mid_ = c != last_seg0 && last_seg0 >= 0; + if (sub.seg == 0) { // (the reports are the first prompt's; the others only fill their sessions) + if (on_stage_chunk && !lm_mid_ && !on_stage_chunk(p0 + T, err)) return false; + if (on_chunk && !on_chunk(m.R, T, p0, err)) return false; + } } continue; } const core::LayerView v(*m.wt, l); - if (l == 1 && ple_on && !ple_land(c0, T)) return false; + if (l == 1 && ple_on && !ple_land(gi, T)) return false; // ---- the PLE block at layer 1 (run_impl's, batched or token by token) if (l == 1 && ple_on && ple_batch) { pt.mark(kPfPle, cs); const auto tp = Clock::now(); - const strata::kernels::PleWeights& pw = ss.ple.w; + const strata::kernels::PleWeights& pw = ssp->ple.w; constexpr int64_t HD = strata::kernels::NG_HC_DIM; const uint64_t per_token = (uint64_t) (3 * HD + N + 4) * 4 + (uint64_t) N * (bf16x2(m.f16_io) ? 4 : 2) + 4096; const int64_t SB = std::min(T, (int64_t) (m.region_bytes / per_token)); @@ -6724,7 +6825,7 @@ bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, st m.gemm.bf16(e16, pw.value_bf16, val, nb, N, N); if (e16_lo) m.gemm.bf16(e16_lo, pw.value_bf16, val, nb, N, N, 0, 1.0f); try { - strata::kernels::native_ple_postops_batch(key, m.R + s0 * D, val, ss.ple.hist, pw, qn, gated, + strata::kernels::native_ple_postops_batch(key, m.R + s0 * D, val, ssp->ple.hist, pw, qn, gated, gate, (int) nb, cs); } catch (const std::exception& e) { err = std::string("prefill PLE: ") + e.what(); return false; } } @@ -6737,10 +6838,10 @@ bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, st po.normalized = m.ple_norm; po.result = m.R + t * D; try { - strata::kernels::ple_block(m.ple_emb + t * N, m.R + t * D, ss.ple.hist, ss.ple.w, po, - ss.ple.scratch, cs); + strata::kernels::ple_block(m.ple_emb + t * N, m.R + t * D, ssp->ple.hist, ssp->ple.w, po, + ssp->ple.scratch, cs); } catch (const std::exception& e) { err = std::string("prefill PLE: ") + e.what(); return false; } - strata::kernels::ple_history_advance(ss.ple.hist, m.ple_norm, cs); + strata::kernels::ple_history_advance(ssp->ple.hist, m.ple_norm, cs); } stats_.ms_ple += ms_since(tp); } @@ -6798,7 +6899,7 @@ bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, st *wsa = need(v, "ssm_a", err); if (!wqkv || !wg || !wo || !wa || !wb || !wc || !wnm || !wdt || !wsa) return false; pt.mark(kPfGdn, cs); - float* state = ss.gdn_state + (size_t) (gdn_index - ss.gdn_ord0) * gdn_floats; + float* state = ssp->gdn_state + (size_t) (gdn_index - ssp->gdn_ord0) * gdn_floats; float* conv = state + (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size; if (!native_proj(m.gemm, wqkv, m.mixed_h, m.qkv, T, v.name("attn_qkv.weight"), err)) return false; if (!native_proj(m.gemm, wg, m.mixed_h, m.z, T, v.name("attn_gate.weight"), err)) return false; @@ -6814,7 +6915,7 @@ bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, st ++gdn_index; } else if (half == 0) { // ======================= QSA (run_impl's, without its debug dumps) ======================= - const core::QsaState& st = ss.qsa_states[qsa_index]; + const core::QsaState& st = ssp->qsa_states[qsa_index]; const core::WeightRef *wq = need(v, "attn_q.weight", err), *wk = need(v, "attn_k.weight", err), *wv = need(v, "attn_v.weight", err), *wo = need(v, "attn_output.weight", err), *wik = need(v, "indexer.k_proj.weight", err), @@ -6986,17 +7087,22 @@ bool Prefill::run_layer_major(const int64_t* tokens, int64_t n, int64_t pos0, st cudaGetErrorString(cudaGetLastError()); return false; } - if (on_stage_chunk && !on_stage_chunk(pos0 + sc0 + Sn, err)) return false; + if (on_stage_chunk && last_seg0 >= 0 && + !on_stage_chunk(segs[0].pos0 + subs[(size_t) last_seg0].c0 + subs[(size_t) last_seg0].T, err)) + return false; next_err_.clear(); - next_run_ = std::async(std::launch::async, [this, tokens, sc0, Sn, pos0] { + next_run_ = std::async(std::launch::async, [this, segs_c = std::make_shared>(segs), depth, + sp_super = std::make_shared>(supers[sci])] { strata::aux_cpus::pin_current_thread(); - return next_->run_layer_major(tokens + sc0, Sn, pos0 + sc0, next_err_); + return next_->lm_run(*segs_c, depth + 1, sp_super.get(), next_err_); }); } } lm_mid_ = false; - ss.ple_prev[0] = prev[0]; - ss.ple_prev[1] = prev[1]; + for (size_t i = 0; i < segs.size(); ++i) { + ss_of(i)->ple_prev[0] = prevs[i][0]; + ss_of(i)->ple_prev[1] = prevs[i][1]; + } if (ple_next.valid()) ple_next.wait(); if (cudaStreamSynchronize(cs) != cudaSuccess) { err = std::string("prefill: ") + cudaGetErrorString(cudaGetLastError()); diff --git a/src/program/generate.cpp b/src/program/generate.cpp index 956d2c56..475b1e52 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -12491,6 +12491,12 @@ int main(int argc, char** argv) { }; int admit_slot = -1; ///< the slot the request being read will continue in (BGEN) long long admit_max_new = 0; + // BCO (research, PAR-H; STRATA_COALESCE_PREFILL on the server): prompts of waiting requests, read in + // the SAME layer-major pass as the next BGEN's prompt into the named idle slots' sessions (the experts a layer + // streams are shared), so their own admissions find the prompt in the slot and read nothing. BCOK + // says a slot holds its prompt; a declined or failed co-read says nothing and costs nothing. + struct CoReq { int slot = -1; std::vector ids; }; + std::vector co_pending, co_set; auto batch_on = [&] { for (const BSlot& b : bs) if (b.active) return true; return false; }; auto try_next_line = [&](std::string& out) -> bool { std::lock_guard lk(in_mu); @@ -13430,6 +13436,21 @@ int main(int argc, char** argv) { if (b >= 0 && b < (int) bs.size()) bs[(size_t) b].stop = true; continue; } + if (line.rfind("BCO ", 0) == 0) { // BCO : kept for the next BGEN + CoReq cr; + char* e1 = nullptr; + cr.slot = (int) std::strtol(line.c_str() + 4, &e1, 10); + for (const char* q = e1; q != nullptr && *q != '\0';) { + while (*q == ' ' || *q == ',') ++q; + if (*q == '\0') break; + char* e2 = nullptr; + cr.ids.push_back((int64_t) std::strtoll(q, &e2, 10)); + if (e2 == q) { cr.ids.clear(); break; } + q = e2; + } + if (cr.slot >= 0 && cr.slot < (int) bs.size() && cr.ids.size() >= 3) co_pending.push_back(std::move(cr)); + continue; + } if (line.rfind("BYIELD", 0) == 0) { // for a prompt read that has ended meanwhile if (trace && piped) std::fprintf(stderr, "strata trace: BYIELD dropped, %d groups in flight\n", @@ -13465,6 +13486,8 @@ int main(int argc, char** argv) { admit_slot = (int) b; admit_max_new = mn; adm_drain_ms = line_drain_ms; + co_set = std::move(co_pending); + co_pending.clear(); line = std::string(bimg ? "GENI 1" : "GEN 1") + e2; } // #477: every --expert-profile-save-every minutes, before the next request (at QUIT: after the loop) @@ -15061,9 +15084,74 @@ int main(int argc, char** argv) { ConvCheckpoint::PromptPlan& p; ~SegPlanScope() { p = {}; } } seg_plan_scope{seg_plan}; + // the first run of a BGEN with co-reads: the waiting prompts are read in the same layer-major pass + bool co_tried = false; + auto co_run = [&](int64_t q, int64_t r, std::string& ce) -> bool { + using PSeg = strata::prefill::Prefill::Seg; + std::vector segs; + PSeg lead; + lead.tokens = ids.data() + q; + lead.n = r - q; + lead.pos0 = q; + segs.push_back(lead); + std::vector used; + for (const CoReq& c : co_set) { + if (c.slot < 0 || c.slot >= (int) bs.size() || c.slot == admit_slot || bs[(size_t) c.slot].active || + bs[(size_t) c.slot].partial) + continue; + bool dup = false; + for (const CoReq* u : used) dup = dup || u->slot == c.slot; + if (dup) continue; + PSeg sg; + sg.tokens = c.ids.data(); + sg.n = (int64_t) c.ids.size() - 1; // (the last token goes through the slot's own first window) + sg.pos0 = 0; + for (size_t k = 0; k < bslot_ss.size(); ++k) sg.stage_ss.push_back(bslot_ss[k][(size_t) c.slot].get()); + segs.push_back(sg); + used.push_back(&c); + } + if (segs.size() < 2) return sp.run(lead.tokens, lead.n, lead.pos0, ce); + for (const CoReq* c : used) { + park_slot(c->slot); // the conversation the slot holds, before this read overwrites it + for (size_t k = 0; k < bslot_ss.size(); ++k) { + const strata::core::OnDevice on_k(k == 0 ? 0 : stages[k - 1]->dev); + strata::core::session_zero(*bslot_ss[k][(size_t) c->slot], g, nullptr, nullptr); + cudaDeviceSynchronize(); + } + bs[(size_t) c->slot] = BSlot{}; // (it holds nothing until the read is done) + } + const auto tco = Clock::now(); + if (!sp.run_multi(segs, ce)) { + if (ce.rfind("prefill: multi-prompt read declined", 0) == 0) { // nothing was read: the usual run + std::fprintf(stderr, "strata batch: co-read of %zu prompt(s) declined (%s)\n", used.size(), ce.c_str()); + ce.clear(); + return sp.run(lead.tokens, lead.n, lead.pos0, ce); + } + return false; + } + for (const CoReq* c : used) { + BSlot& sl = bs[(size_t) c->slot]; + sl = BSlot{}; + sl.ids.assign(c->ids.begin(), c->ids.end() - 1); + sl.cached = true; + sl.cvec = want_cvec; + std::printf("BCOK %d %lld\n", c->slot, (long long) sl.ids.size()); + } + std::fflush(stdout); + std::fprintf(stderr, "strata batch: co-read: %zu more prompt(s) read in the same pass as this request's %lld tokens " + "(%.0f ms)\n", used.size(), (long long) lead.n, + std::chrono::duration(Clock::now() - tco).count()); + return true; + }; auto run = [&](int64_t q, int64_t r) -> bool { if (planned) sp.force_chunk(plan.chunk, plan_single, plan.chunked); - const bool run_ok = sp.run(ids.data() + q, r - q, q, e); + bool run_ok; + if (!co_set.empty() && !co_tried && admit_slot >= 0 && req_imgs.empty()) { + co_tried = true; + run_ok = co_run(q, r, e); + } else { + run_ok = sp.run(ids.data() + q, r - q, q, e); + } if (run_ok && q == 0 && sp.kva_floor() > 0) { // KVA approximated the start of this prompt: remember which tokens kva_floors.note(ids.data(), sp.kva_floor()); if (use_mtp && mtp.cell_floor_enabled()) mtp.set_cell_floor(sp.kva_floor()); From faa1be268ceb91a3f760ef7ad9711a192d59b744 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 02:51:33 +0100 Subject: [PATCH 02/17] co-read: server side (STRATA_COALESCE_PREFILL), par_load --keep-text, cmp_texts --- serve/server.py | 64 ++++- tools/par/cmp_texts.py | 19 ++ tools/par_load.py | 544 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 624 insertions(+), 3 deletions(-) create mode 100644 tools/par/cmp_texts.py create mode 100644 tools/par_load.py diff --git a/serve/server.py b/serve/server.py index 5614dbf1..a2f5d66f 100644 --- a/serve/server.py +++ b/serve/server.py @@ -336,6 +336,12 @@ try: except ValueError: ADMIT_SHORT_FIRST = 0 ADMIT_MAX_WAIT_S = _timeout_env("STRATA_ADMIT_MAX_WAIT_S", 30.0) or 0.0 +# STRATA_COALESCE_PREFILL=N (research, PAR-H; needs the engine's BCO / BCOK lines): a request that takes the control lines +# for its admission also sends the prompts of up to N-1 requests still waiting for them (BCO , into free slots +# it reserves for them); the engine reads all the prompts in ONE layer-major pass, so the experts a layer streams are +# read once for all of them, and each waiting request then finds its prompt in its slot. 0 / unset: off. +COALESCE_PREFILL = max(0, int(os.environ.get("STRATA_COALESCE_PREFILL") or 0)) +COALESCE_TOKENS = max(256, int(os.environ.get("STRATA_COALESCE_TOKENS") or 16384)) # prompt tokens read together at most STALL_CPU_EPS_S = 0.5 # CPU seconds the engine may use in the whole stall window and still count as frozen STALL_IO_EPS_B = 1 << 20 # bytes read or written in that window @@ -1575,6 +1581,8 @@ class StrataEngine: f = line.split() self._ctl_result = ("badm", len(f) >= 3 and f[2] == "1") return + elif line.startswith("BCOK "): # STRATA_COALESCE_PREFILL: : the slot holds a waiting prompt + pass elif line.startswith("YIELDED "): # the read gave way (BYIELD): f = line.split() if len(f) >= 3 and f[1].lstrip("-").isdigit() and f[2].isdigit(): @@ -1726,11 +1734,11 @@ class StrataEngine: YIELDS_MAX = 2 # #656: how often one request's prompt read gives way to a shorter waiting one - def _take_control(self, cancel, plen: int, after_epoch: int | None = None, new: bool = False): + def _take_control(self, cancel, plen: int, after_epoch: int | None = None, new: bool = False, prompt=None): """Waits for the control lines (one prompt read at a time), yielding None heartbeats; False when cancelled. `after_epoch`: a request whose read gave way lets the requests waiting then go first. `new`: the request's first admission (STRATA_ADMIT_SHORT_FIRST orders only these).""" - entry = [plen, time.monotonic(), new] # (length, arrival, ordered: STRATA_ADMIT_SHORT_FIRST) + entry = [plen, time.monotonic(), new, prompt, cancel] # (length, arrival, ordered: STRATA_ADMIT_SHORT_FIRST; the prompt and request, for STRATA_COALESCE_PREFILL) with self.slot_cv: self.waiting += 1 self.wait_lens.append(entry) @@ -1784,6 +1792,49 @@ class StrataEngine: self.slot_cv.notify_all() return True + def _release_preslot(self, cancel): + """STRATA_COALESCE_PREFILL: the request that was given a slot ahead of its turn left before using it.""" + slot = getattr(cancel, "_strata_preslot", None) + if slot is None: + return + cancel._strata_preslot = None + with self.slot_cv: + self.slot_busy[slot] = False + self.slot_claims.discard(slot) + self.slot_cv.notify_all() + + def _coalesce(self, cancel, lead_prompt): + """STRATA_COALESCE_PREFILL (the caller holds the control lines and sends its BGEN right after): reserve free slots + for requests still waiting for the control lines and send their prompts as BCO lines; the engine reads them in the + same layer-major pass as the BGEN's prompt. Each follower finds its slot in cancel._strata_preslot.""" + picked = [] + total = len(lead_prompt) + with self.slot_cv: + for e in list(self.wait_lens): + if len(picked) >= COALESCE_PREFILL - 1: + break + prompt = e[3] if len(e) > 4 else None + rq = e[4] if len(e) > 4 else None + if prompt is None or rq is None or rq is cancel or getattr(rq, "_strata_preslot", None) is not None: + continue + if getattr(rq, "_strata_co", False) or len(prompt) < 32 or total + len(prompt) > COALESCE_TOKENS: + continue + slot = self.pick_slot(prompt) + if slot is None: + break + self.slot_busy[slot] = True + self.slot_claims.add(slot) + self.slot_held[slot] = [] + rq._strata_preslot = slot + rq._strata_co = True + total += len(prompt) + picked.append((slot, prompt)) + for slot, prompt in picked: + self._send(f"BCO {slot} {','.join(str(int(t)) for t in prompt)}") + if picked: + print(f"[strata] coalesce: {len(picked)} waiting prompt(s) ({total - len(lead_prompt)} tokens) read with this " + f"request's {len(lead_prompt)}", flush=True) + def _admit_next(self): """STRATA_ADMIT_SHORT_FIRST= (the caller holds slot_cv): the new request that takes the control lines next - one whose prompt is at most that long (it reads little) before a longer one, in arrival order within @@ -1839,12 +1890,17 @@ class StrataEngine: prompt, left = list(ids), int(max_new) reused0 = None # the first read's reused tokens: a continued leg's count includes this request's own output carry = {} # #1869: the figures of the legs that ended in a hand-back to the solo path (the final DONE adds to them) - ok = yield from self._take_control(cancel, len(prompt), new=True) + ok = yield from self._take_control(cancel, len(prompt), new=True, + prompt=prompt if (COALESCE_PREFILL > 1 and not embeddings) else None) if not ok: + self._release_preslot(cancel) return holding, born = True, self.gen # (the engine it now has the control lines of) btrace("ctl acquired") slot, gen0, reserved = None, None, None + pre = getattr(cancel, "_strata_preslot", None) + if pre is not None: # STRATA_COALESCE_PREFILL: a request before this one read the prompt into this slot + slot, cancel._strata_preslot = pre, None phase = "none" # solo -> (admit -> slot) ; "done" once the engine has finished with this request stop_sent = False yields, solo_again, refused = 0, 0, 0 @@ -1993,6 +2049,8 @@ class StrataEngine: raise QueueWaitExpired("the inference queue wait deadline expired; retry later") self.slot_live[slot] = live self._last_done = None # (#879: this admission's DONE decides below) + if COALESCE_PREFILL > 1 and not embeddings: + self._coalesce(cancel, prompt) self._send(f"{head} {','.join(str(int(t)) for t in prompt)}") if deadline is not None: cancel._strata_queue_deadline = None diff --git a/tools/par/cmp_texts.py b/tools/par/cmp_texts.py new file mode 100644 index 00000000..19a01e19 --- /dev/null +++ b/tools/par/cmp_texts.py @@ -0,0 +1,19 @@ +#!/usr/bin/env python3 +"""cmp_texts.py A.json B.json - compare the answers (par_load --keep-text) of two ledgers request by request (same nonce).""" +import json, sys +a, b = (json.load(open(x)) for x in sys.argv[1:3]) +ra = {(r["workload"], r.get("level"), r["idx"]): r for r in a["requests"]} +rb = {(r["workload"], r.get("level"), r["idx"]): r for r in b["requests"]} +same = diff = 0 +for k, r in sorted(ra.items(), key=str): + o = rb.get(k) + if o is None or "text" not in r or "text" not in o: + continue + if r["text"] == o["text"]: + same += 1 + else: + diff += 1 + n = next((i for i, (x, y) in enumerate(zip(r["text"], o["text"])) if x != y), min(len(r["text"]), len(o["text"]))) + print("DIFF", k, "first difference at char", n, "|", repr(r["text"][max(0, n - 20):n + 30]), "vs", repr(o["text"][max(0, n - 20):n + 30])) +print("identical", same, "different", diff) +sys.exit(1 if diff else 0) diff --git a/tools/par_load.py b/tools/par_load.py new file mode 100644 index 00000000..1b68567a --- /dev/null +++ b/tools/par_load.py @@ -0,0 +1,544 @@ +#!/usr/bin/env python3 +"""Shared load harness for the parallelism program (research tool, standard library only). + +Wraps the OpenAI-compatible server the way tools/serve_load.py does, and adds what an SLO-capacity harness needs: + + * closed-loop concurrency staircases (fixed work, shared wall window incl. the final drain), + * open-loop seeded Poisson arrivals (offered rate independent of completions; late/unfinished requests are failures), + * token counts from the server's `usage` (completion_tokens / prompt_tokens), NEVER from SSE chunks; a request with + no usage is marked `tokens_source: "sse"` and the run is flagged invalid, + * per request: TTFT, TPOT, max gap between content chunks (a chunk can hold several tokens when the engine flushes + speculative windows, so "gap" is a chunk gap), client scheduling lag, finish reason, cached prompt tokens, + * aggregate tok/s = total committed completion tokens / ONE shared wall-clock window (first submit to last completion), + * the workloads W1 solo, W2 short concurrent, W4 long joiner, W6 shared-prefix cohorts, + * one JSON ledger per run (args, server status before/after, per-request records, per-level blocks). + + python3 tools/par_load.py http://127.0.0.1:8080 --workload w2 --conc 1,2,4,8 --rounds 3 --out run.json + python3 tools/par_load.py URL --workload w1 --prompt-tokens 512 --max-tokens 512 --rounds 3 --out w1.json + python3 tools/par_load.py URL --workload w4 --conc 4 --joiner-tokens 32768 --out w4.json + python3 tools/par_load.py URL --workload w6 --conc 4 --share 0,0.5,0.9 --out w6.json + python3 tools/par_load.py URL --workload w2 --open-rates 0.25,0.5,1 --duration 120 --slo-ttft 10 --out open.json + +Labelling for the paired analysis (tools/par_analyze.py): --arm A|B --pair N --label "free text". The ledger's +`blocks` list holds one entry per (workload, level) with the headline metrics. See Desktop 30thOctober results/par/HARNESS.md. +Prompts are deterministic from --seed and carry a unique nonce, so the prompt cache cannot help unless W6 asks for a shared prefix. +""" +import argparse, http.client, json, math, os, platform, random, statistics, sys, threading, time, urllib.parse, urllib.request + +HARNESS_VERSION = "1.0" +WORDS = ("time year people way day man thing woman life child world school state family student group country problem hand " + "part place case week company system program question work government number night point home water room mother area " + "money story fact month lot right study book eye job word business issue side kind head house service friend father " + "power hour game line end member law car city community name president team minute idea kid body information back " + "parent face others level office door health person art war history party result change morning reason research girl " + "guy moment air teacher force education foot boy age policy music market sense nation plan college interest death " + "experience effect use class control care field development role effort rate heart drug show leader light voice wife " + "police mind price report decision son view relationship town road arm difference value building action model season " + "society tax director position player record paper space ground form event official matter center couple site project " + "activity star table need court oil situation cost industry figure street image phone data picture practice piece " + "land product doctor wall patient worker news test movie north love support technology step baby computer type " + "attention film tree source organization hair window culture chance brother energy period course summer plant").split() +TOPICS = ["the history of lighthouses", "how a mechanical watch keeps time", "why bridges are built as arches", + "the water cycle and droughts", "how bread rises", "the design of a medieval castle", + "how a transistor switches", "the life of a honey bee colony"] + + +# ----------------------------------------------------------------------------------------------------------- prompts +class Prompter: + """Deterministic filler text of a target token count. The words-per-token ratio is calibrated once against the + server's own usage.prompt_tokens (a max_tokens=1 probe), so the target counts tokens after the real template.""" + + def __init__(self, url, key, model, seed, think, calib=None): + self.url, self.key, self.model, self.seed, self.think = url, key, model, seed, think + self.tok_per_word = (calib or {}).get("tok_per_word") + self.overhead = (calib or {}).get("overhead", 40) + self.calib = calib or {} + + def words(self, n, rng): + return " ".join(rng.choice(WORDS) for _ in range(n)) + + def body_text(self, key, n_words): + return self.words(n_words, random.Random(f"{self.seed}/{key}")) + + def calibrate(self): + if self.tok_per_word: + return + n = 1500 + text = "[calibration] " + self.body_text("calib", n) + "\nSay ok." + u = chat_once(self.url, self.key, self.model, text, 1, self.think) + pt = u.get("usage", {}).get("prompt_tokens") + if not pt: + self.tok_per_word, self.overhead = 1.3, 40 # server gave no usage: a guess, the ledger says so + self.calib = {"tok_per_word": 1.3, "overhead": 40, "guess": True} + return + # prompt = template overhead (about 25-40 tokens) + n words; assume the overhead is the same as the tail + self.overhead = 40 + self.tok_per_word = max(0.5, (pt - self.overhead) / n) + self.calib = {"tok_per_word": self.tok_per_word, "overhead": self.overhead, "probe_prompt_tokens": pt, "probe_words": n} + + def make(self, key, target_tokens, prefix=None, share=0.0, tail="Now write a long, detailed essay about "): + """A prompt of ~target_tokens. With prefix (a stable string) and share: the first share*target tokens are + the prefix's words (the same text for the whole cohort), the rest is unique to `key`.""" + self.calibrate() + nw = max(8, int((target_tokens - self.overhead) / self.tok_per_word)) + topic = TOPICS[hash_int(key) % len(TOPICS)] + nonce = f"[req {key}] " + if prefix is not None and share > 0: + ns = int(nw * share) + pre = self.body_text("prefix/" + prefix, ns) + rest = self.body_text(key, nw - ns) + return pre + "\n" + nonce + rest + f"\n{tail}{topic}." + return nonce + self.body_text(key, nw) + f"\n{tail}{topic}." + + +def hash_int(s): + h = 1469598103934665603 + for ch in s.encode(): + h = ((h ^ ch) * 1099511628211) & 0xFFFFFFFFFFFFFFFF + return h + + +# ------------------------------------------------------------------------------------------------------------ request +def _conn(url): + u = urllib.parse.urlparse(url) + cls = http.client.HTTPSConnection if u.scheme == "https" else http.client.HTTPConnection + return cls(u.hostname, u.port or (443 if u.scheme == "https" else 80), timeout=1800), u.path.rstrip("/") + + +def _body(model, text, max_tokens, think, stream, extra=None): + b = {"model": model, "stream": stream, "max_tokens": max_tokens, "temperature": 0, + "messages": [{"role": "user", "content": text}]} + if stream: + b["stream_options"] = {"include_usage": True} + if not think: + b["reasoning_effort"] = "none" + b["chat_template_kwargs"] = {"enable_thinking": False} + if extra: + b.update(extra) + return b + + +def chat_once(url, key, model, text, max_tokens, think): + c, base = _conn(url) + h = {"Content-Type": "application/json", **({"Authorization": "Bearer " + key} if key else {})} + c.request("POST", base + "/v1/chat/completions", json.dumps(_body(model, text, max_tokens, think, False)), h) + r = c.getresponse() + d = json.loads(r.read() or b"{}") + c.close() + return d + + +def stream_request(url, key, model, text, max_tokens, think, sched_t0, extra=None, keep_times=False, keep_text=False): + """One streamed chat request. Times are seconds on the perf_counter clock of this process; the record also + carries them relative to the run start (t_rel_*) by the caller.""" + rec = {"t_sched": sched_t0, "t_send": None, "t_first": None, "t_last": None, "t_end": None, "completion_tokens": None, + "prompt_tokens": None, "cached_tokens": 0, "finish": None, "tokens_source": None, "chunks": 0, "max_gap": 0.0, + "gaps_ge1s": 0, "err": None, "timings": None, "reasoning_chars": 0, "content_chars": 0} + times = [] if keep_times else None + texts = [] if keep_text else None + rec["t_send"] = time.perf_counter() + last = None + sse_tokens = 0 + try: + c, base = _conn(url) + h = {"Content-Type": "application/json", **({"Authorization": "Bearer " + key} if key else {})} + c.request("POST", base + "/v1/chat/completions", json.dumps(_body(model, text, max_tokens, think, True, extra)), h) + r = c.getresponse() + if r.status != 200: + rec["err"] = f"http {r.status}: {r.read(300).decode('utf-8', 'replace')}" + else: + while True: + raw = r.readline() + if not raw: + break + line = raw.decode("utf-8", "replace").strip() + if not line.startswith("data:"): + continue + d = line[5:].strip() + if d == "[DONE]": + break + try: + j = json.loads(d) + except ValueError: + continue + now = time.perf_counter() + if j.get("error"): + rec["err"] = str(j["error"])[:300] + for ch in j.get("choices", []): + dl = ch.get("delta", {}) + a, b = dl.get("content") or "", dl.get("reasoning_content") or "" + if a or b: + if rec["t_first"] is None: + rec["t_first"] = now + elif last is not None: + g = now - last + if g > rec["max_gap"]: + rec["max_gap"] = g + if g >= 1.0: + rec["gaps_ge1s"] += 1 + last = now + rec["t_last"] = now + rec["chunks"] += 1 + sse_tokens += 1 + rec["content_chars"] += len(a) + if texts is not None: + texts.append(a + b) + rec["reasoning_chars"] += len(b) + if times is not None: + times.append(now) + if ch.get("finish_reason"): + rec["finish"] = ch["finish_reason"] + u = j.get("usage") + if u: + rec["completion_tokens"] = u.get("completion_tokens") + rec["prompt_tokens"] = u.get("prompt_tokens") + rec["cached_tokens"] = (u.get("prompt_tokens_details") or {}).get("cached_tokens", 0) or 0 + rec["tokens_source"] = "usage" + if j.get("timings"): + rec["timings"] = j["timings"] + c.close() + except Exception as e: # noqa: BLE001 + rec["err"] = repr(e)[:300] + rec["t_end"] = time.perf_counter() + if rec["completion_tokens"] is None: + rec["completion_tokens"] = sse_tokens + rec["tokens_source"] = "sse" # NOT a valid token count; the ledger marks the run invalid + if times is not None: + rec["chunk_times"] = times + if texts is not None: + rec["text"] = "".join(texts) + return rec + + +def finish_record(rec, t0): + """Derived per-request metrics, and times relative to the run start.""" + ttft = rec["t_first"] - rec["t_send"] if rec["t_first"] is not None else None + n = rec["completion_tokens"] or 0 + tpot = (rec["t_last"] - rec["t_first"]) / (n - 1) if (ttft is not None and n > 1 and rec["t_last"] is not None) else None + out = {k: rec[k] for k in ("completion_tokens", "prompt_tokens", "cached_tokens", "finish", "tokens_source", "chunks", "max_gap", + "gaps_ge1s", "err", "timings", "reasoning_chars", "content_chars")} + out.update({"t_sched": rec["t_sched"] - t0, "t_send": rec["t_send"] - t0, "t_end": rec["t_end"] - t0, + "t_first": (rec["t_first"] - t0) if rec["t_first"] is not None else None, + "lag": rec["t_send"] - rec["t_sched"], "ttft": ttft, "tpot": tpot, "latency": rec["t_end"] - rec["t_send"]}) + if "text" in rec: + out["text"] = rec["text"] + if "chunk_times" in rec: + out["chunk_t"] = [round(x - t0, 4) for x in rec["chunk_times"]] + return out + + +# --------------------------------------------------------------------------------------------------------- aggregates +def pct(v, p): + v = sorted(x for x in v if x is not None) + if not v: + return None + k = (len(v) - 1) * p + lo = int(k) + hi = min(lo + 1, len(v) - 1) + return v[lo] + (v[hi] - v[lo]) * (k - lo) + + +def aggregate(recs, window_s, slo=None): + """Level metrics. window_s = one shared wall window for all requests (first submit to last completion).""" + ok = [r for r in recs if not r["err"]] + toks = sum(r["completion_tokens"] or 0 for r in ok) + agg = {"requests": len(recs), "errors": len(recs) - len(ok), "window_s": window_s, "tokens": toks, + "agg_tps": toks / window_s if window_s > 0 else None, + "prompt_tokens": sum(r["prompt_tokens"] or 0 for r in ok), + "invalid_token_source": sum(1 for r in recs if r["tokens_source"] != "usage"), + "length_stops": sum(1 for r in ok if r["finish"] == "length"), "early_stops": sum(1 for r in ok if r["finish"] == "stop")} + for k in ("ttft", "tpot", "max_gap", "latency", "lag"): + v = [r[k] for r in ok] + agg[k + "_p50"], agg[k + "_p95"], agg[k + "_max"] = pct(v, .5), pct(v, .95), (max(x for x in v if x is not None) if any(x is not None for x in v) else None) + per = [(r["completion_tokens"] - 1) / (r["t_end"] - r["t_first"]) for r in ok + if r["t_first"] is not None and (r["completion_tokens"] or 0) > 1 and r["t_end"] > r["t_first"]] + agg["per_req_tps_p50"] = pct(per, .5) + if slo: + good = [r for r in ok if (slo.get("ttft") is None or (r["ttft"] is not None and r["ttft"] <= slo["ttft"])) + and (slo.get("gap") is None or r["max_gap"] <= slo["gap"])] + agg["slo"] = slo + agg["slo_pass"] = len(good) + agg["goodput_tps"] = sum(r["completion_tokens"] or 0 for r in good) / window_s if window_s > 0 else None + return agg + + +# --------------------------------------------------------------------------------------------------------- runners +class Ctx: + def __init__(self, a, prompter): + self.a, self.p = a, prompter + + +def closed_loop(ctx, level, n_requests, make_job, label): + """level concurrent clients pull jobs 0..n_requests-1 from one queue; the window is first submit -> last end.""" + a = ctx.a + jobs = list(range(n_requests)) + lock = threading.Lock() + recs = [None] * n_requests + t0 = time.perf_counter() + + def worker(): + while True: + with lock: + if not jobs: + return + i = jobs.pop(0) + text, mt, extra = make_job(i) + recs[i] = stream_request(a.url, a.key, a.model, text, mt, a.think, time.perf_counter(), extra, a.keep_times, a.keep_text) + + th = [threading.Thread(target=worker) for _ in range(level)] + for t in th: + t.start() + for t in th: + t.join() + t1 = max(r["t_end"] for r in recs) + fr = [finish_record(r, t0) for r in recs] + return fr, t1 - t0 + + +def run_closed(ctx, wl, levels): + a = ctx.a + blocks, allreq = [], [] + for c in levels: + n = max(c * a.rounds, a.min_requests) + tag = f"{a.nonce}-{wl}-c{c}" + + def make_job(i, c=c, tag=tag): + return ctx.p.make(f"{tag}-{i}", a.prompt_tokens), a.max_tokens, None + + if a.warm_requests: # warmup outside the measured window + for w in range(a.warm_requests): + stream_request(a.url, a.key, a.model, ctx.p.make(f"{tag}-warm{w}", a.prompt_tokens), min(a.max_tokens, 64), a.think, time.perf_counter()) + fr, wall = closed_loop(ctx, c, n, make_job, tag) + agg = aggregate(fr, wall) + agg.update({"workload": wl, "level": c, "mode": "closed", "prompt_target": a.prompt_tokens, "max_tokens": a.max_tokens}) + blocks.append(agg) + for i, r in enumerate(fr): + r.update({"workload": wl, "level": c, "idx": i}) + allreq += fr + print(f"{wl} c={c}: {agg['tokens']} tokens / {wall:.1f}s = {agg['agg_tps']:.1f} tok/s ttft p50/p95 " + f"{fmt(agg['ttft_p50'])}/{fmt(agg['ttft_p95'])} tpot p50 {fmt(agg['tpot_p50'], 4)} max_gap p95 {fmt(agg['max_gap_p95'])}" + f" err {agg['errors']}", flush=True) + return blocks, allreq + + +def fmt(x, d=2): + return "n/a" if x is None else f"{x:.{d}f}" + + +def run_w4(ctx, levels): + """W4: c W2-style decoders running, then one long joiner (joiner_tokens) arrives joiner_at seconds after they start.""" + a = ctx.a + blocks, allreq = [], [] + for c in levels: + for rd in range(a.rounds): + tag = f"{a.nonce}-w4-c{c}-r{rd}" + res = {} + t0 = time.perf_counter() + + def dec(i): + res[("d", i)] = stream_request(a.url, a.key, a.model, ctx.p.make(f"{tag}-d{i}", a.prompt_tokens), a.decoder_tokens, a.think, time.perf_counter(), None, a.keep_times or True) + + def joiner(): + time.sleep(a.joiner_at) + res[("j", 0)] = stream_request(a.url, a.key, a.model, ctx.p.make(f"{tag}-j", a.joiner_tokens), a.joiner_out, a.think, time.perf_counter(), None, True) + + th = [threading.Thread(target=dec, args=(i,)) for i in range(c)] + [threading.Thread(target=joiner)] + for t in th: + t.start() + for t in th: + t.join() + wall = max(r["t_end"] for r in res.values()) - t0 + fr = {k: finish_record(r, t0) for k, r in res.items()} + j = fr[("j", 0)] + ts_j = j["t_send"] + te_j = j["t_first"] if j["t_first"] is not None else j["t_end"] + # incumbent gaps between the joiner's submit and its first token (the interference window) + gaps = [] + for (k, i), r in fr.items(): + if k != "d": + continue + ct = [x for x in r.get("chunk_t", [])] + for x, y in zip(ct, ct[1:]): + if y >= ts_j and x <= te_j + 1.0: + gaps.append(y - x) + agg = aggregate(list(fr.values()), wall) + allreq += [dict(r, workload="w4", level=c, round=rd, role=k[0]) for k, r in fr.items()] + for r in allreq[-len(fr):]: + r.pop("chunk_t", None) if not a.keep_times else None + decs = [r for k, r in fr.items() if k[0] == "d"] + agg.update({"workload": "w4", "level": c, "mode": "closed", "round": rd, "joiner_tokens": a.joiner_tokens, + "joiner_ttft": j["ttft"], "joiner_prompt_tokens": j["prompt_tokens"], + "incumbent_gap_max_during_join": max(gaps) if gaps else None, + "incumbent_gap_p99_during_join": pct(gaps, .99), + "incumbent_tps_p50": pct([(d["completion_tokens"] - 1) / max(d["t_end"] - d["t_first"], 1e-9) for d in decs if d["t_first"] is not None and d["completion_tokens"] and d["completion_tokens"] > 1], .5), + "makespan_s": wall}) + blocks.append(agg) + print(f"w4 c={c} r={rd}: makespan {wall:.1f}s agg {fmt(agg['agg_tps'], 1)} tok/s, joiner ttft {fmt(j['ttft'])}s ({j['prompt_tokens']} prompt tok), " + f"incumbent max gap during join {fmt(agg['incumbent_gap_max_during_join'])}s", flush=True) + return blocks, allreq + + +def run_w6(ctx, levels): + a = ctx.a + blocks, allreq = [], [] + for c in levels: + for share in [float(x) for x in a.share.split(",")]: + tag = f"{a.nonce}-w6-c{c}-s{share}" + prefix = f"{tag}" + if a.prime and share > 0: # prefix-warm: one priming request, not measured + stream_request(a.url, a.key, a.model, ctx.p.make(f"{tag}-prime", a.prompt_tokens, prefix, share), 8, a.think, time.perf_counter()) + n = max(c * a.rounds, a.min_requests) + + def make_job(i, tag=tag, prefix=prefix, share=share): + return ctx.p.make(f"{tag}-{i}", a.prompt_tokens, prefix, share), a.max_tokens, None + + fr, wall = closed_loop(ctx, c, n, make_job, tag) + agg = aggregate(fr, wall) + agg.update({"workload": "w6", "level": c, "share": share, "primed": bool(a.prime), "mode": "closed", + "cached_tokens_mean": statistics.mean(r["cached_tokens"] for r in fr) if fr else None}) + blocks.append(agg) + for i, r in enumerate(fr): + r.update({"workload": "w6", "level": c, "share": share, "idx": i}) + allreq += fr + print(f"w6 c={c} share={share}: {fmt(agg['agg_tps'], 1)} tok/s ttft p50 {fmt(agg['ttft_p50'])} cached mean {agg['cached_tokens_mean']:.0f}", flush=True) + return blocks, allreq + + +def run_open(ctx, wl, rates): + """Open loop: seeded exponential inter-arrival times, independent of completions. Requests still running at + duration + drain_s are failures (kept in the ledger as unfinished).""" + a = ctx.a + blocks, allreq = [], [] + for rate in rates: + rng = random.Random(f"{a.seed}/open/{rate}") + sched, t = [], 0.0 + while True: + t += rng.expovariate(rate) + if t >= a.duration: + break + sched.append(t) + tag = f"{a.nonce}-{wl}-open{rate}" + t0 = time.perf_counter() + res = [None] * len(sched) + th = [] + + def one(i, at): + res[i] = stream_request(a.url, a.key, a.model, ctx.p.make(f"{tag}-{i}", a.prompt_tokens), a.max_tokens, a.think, t0 + at, None, False) + + def launcher(): + for i, at in enumerate(sched): + d = t0 + at - time.perf_counter() + if d > 0: + time.sleep(d) + t_ = threading.Thread(target=one, args=(i, at), daemon=True) + t_.start() + th.append(t_) + + lt = threading.Thread(target=launcher) + lt.start() + lt.join() + deadline = t0 + a.duration + a.drain_s + for t_ in th: + t_.join(max(0.0, deadline - time.perf_counter())) + done = [r for r in res if r is not None and r["t_end"] is not None] + unfinished = len(sched) - len(done) + window = max([r["t_end"] for r in done], default=t0) - t0 + fr = [finish_record(r, t0) for r in done] + slo = {"ttft": a.slo_ttft, "gap": a.slo_gap} + # tokens finished inside the arrival window over that window; the drain is reported separately + inwin = [r for r in fr if r["t_end"] <= a.duration] + agg = aggregate(fr, window, slo) + agg["tps_in_window"] = sum(r["completion_tokens"] or 0 for r in inwin if not r["err"]) / a.duration + agg.update({"workload": wl, "mode": "open", "offered_rate": rate, "offered_requests": len(sched), "unfinished": unfinished, + "duration": a.duration, "drain_s": a.drain_s, "achieved_rate": len(inwin) / a.duration, + "slo_pass_frac": (agg.get("slo_pass", 0) / len(sched)) if sched else None}) + blocks.append(agg) + allreq += [dict(r, workload=wl, offered_rate=rate, idx=i) for i, r in enumerate(fr)] + print(f"{wl} open rate={rate}/s: offered {len(sched)}, done {len(done)}, unfinished {unfinished}, slo pass {agg.get('slo_pass')}, " + f"goodput {fmt(agg.get('goodput_tps'), 1)} tok/s, ttft p95 {fmt(agg['ttft_p95'])}, lag max {fmt(agg['lag_max'], 3)}", flush=True) + return blocks, allreq + + +# ------------------------------------------------------------------------------------------------------------- main +def get_json(url): + try: + with urllib.request.urlopen(url, timeout=10) as r: + return json.loads(r.read()) + except Exception as e: # noqa: BLE001 + return {"error": repr(e)[:200]} + + +def main(argv=None): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("url") + ap.add_argument("--workload", default="w2", choices=["w1", "w2", "w3", "w4", "w6"]) + ap.add_argument("--conc", default="1,2,4,8", help="closed-loop concurrency levels (w1 forces 1)") + ap.add_argument("--rounds", type=int, default=3, help="requests per client per level (a level runs conc*rounds requests)") + ap.add_argument("--min-requests", type=int, default=0) + ap.add_argument("--warm-requests", type=int, default=0, help="unmeasured short requests before each level") + ap.add_argument("--prompt-tokens", type=int, default=None, help="default: w1 512, w2 1024, w3 8192, w6 2048") + ap.add_argument("--max-tokens", type=int, default=None, help="default: w1 512, w2 256, w3 128, w6 128") + ap.add_argument("--joiner-tokens", type=int, default=28000) + ap.add_argument("--joiner-out", type=int, default=64) + ap.add_argument("--joiner-at", type=float, default=6.0, help="seconds after the decoders start") + ap.add_argument("--decoder-tokens", type=int, default=512) + ap.add_argument("--share", default="0,0.5,0.9", help="w6 shared-prefix fractions") + ap.add_argument("--prime", action="store_true", help="w6: one unmeasured request primes the prefix (prefix-warm)") + ap.add_argument("--open-rates", default=None, help="open-loop offered rates (req/s), e.g. 0.25,0.5,1") + ap.add_argument("--duration", type=float, default=120.0, help="open loop: arrival window (s)") + ap.add_argument("--drain-s", type=float, default=120.0, help="open loop: extra time to finish before counting failures") + ap.add_argument("--slo-ttft", type=float, default=None) + ap.add_argument("--slo-gap", type=float, default=None) + ap.add_argument("--seed", default="par29") + ap.add_argument("--nonce", default=str(int(time.time()))) + ap.add_argument("--model", default="strata") + ap.add_argument("--key", default=os.environ.get("STRATA_KEY", "")) + ap.add_argument("--think", action="store_true") + ap.add_argument("--keep-times", action="store_true", help="keep every chunk timestamp in the ledger (large)") + ap.add_argument("--keep-text", action="store_true", help="keep every answer's text in the ledger (identity checks)") + ap.add_argument("--calib", default=None, help="a ledger/calibration JSON to reuse the words-per-token calibration from") + ap.add_argument("--arm", default="") + ap.add_argument("--pair", default="") + ap.add_argument("--label", default="") + ap.add_argument("--settings", default=None, help="file with the effective engine settings/provenance to embed (launcher writes it)") + ap.add_argument("--out", required=True) + a = ap.parse_args(argv) + a.url = a.url.rstrip("/") + d = {"w1": (512, 512), "w2": (1024, 256), "w3": (8192, 128), "w4": (1024, 256), "w6": (2048, 128)}[a.workload] + a.prompt_tokens = a.prompt_tokens or d[0] + a.max_tokens = a.max_tokens or d[1] + levels = [1] if a.workload == "w1" else [int(x) for x in a.conc.split(",")] + calib = None + if a.calib and os.path.exists(a.calib): + calib = json.load(open(a.calib)).get("calibration") + p = Prompter(a.url, a.key, a.model, a.seed, a.think, calib) + ctx = Ctx(a, p) + status0 = get_json(a.url + "/v1/status") + started = time.time() + if a.open_rates: + blocks, reqs = run_open(ctx, a.workload, [float(x) for x in a.open_rates.split(",")]) + elif a.workload == "w4": + blocks, reqs = run_w4(ctx, levels) + elif a.workload == "w6": + blocks, reqs = run_w6(ctx, levels) + else: + blocks, reqs = run_closed(ctx, a.workload, levels) + status1 = get_json(a.url + "/v1/status") + settings = None + if a.settings and os.path.exists(a.settings): + settings = open(a.settings, errors="replace").read()[:20000] + for b in blocks: + b.update({"arm": a.arm, "pair": a.pair, "label": a.label}) + invalid = sum(b.get("invalid_token_source", 0) for b in blocks) + ledger = {"harness": "par_load", "version": HARNESS_VERSION, "started": started, "host": platform.node(), "argv": sys.argv[1:], + "args": {k: v for k, v in vars(a).items() if k != "key"}, "calibration": p.calib, "valid": invalid == 0, + "invalid_reason": None if invalid == 0 else f"{invalid} requests without server usage: token counts are SSE chunks", + "status_before": status0, "status_after": status1, "settings": settings, "blocks": blocks, "requests": reqs} + os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True) + json.dump(ledger, open(a.out, "w"), indent=1) + print("ledger:", a.out, "valid" if invalid == 0 else "INVALID (no usage)") + return 0 if invalid == 0 and all(b["errors"] == 0 for b in blocks) else 1 + + +if __name__ == "__main__": + sys.exit(main()) From a9f154e82b2a4c321ef6e9ccd27b3eeff547f020 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 03:01:14 +0100 Subject: [PATCH 03/17] run_multi: layer-major even for a short total --- include/strata/prefill/prefill.hpp | 1 + src/prefill/prefill.cpp | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/include/strata/prefill/prefill.hpp b/include/strata/prefill/prefill.hpp index 5009f8fd..3e2e13c0 100644 --- a/include/strata/prefill/prefill.hpp +++ b/include/strata/prefill/prefill.hpp @@ -301,6 +301,7 @@ private: bool lm_run(const std::vector& segs, int depth, const std::vector* given_super, std::string& err); void lm_release(); bool lm_mid_ = false; + bool lm_force_ = false; ///< run_multi: the walk even for a short total bool lm_by_default_ = false; ///< the walk is on because nothing was set (a wider VRAM margin, the measured super-chunk size) // Stage-1 pipeline: intermediate stages return after handing their chunk to // the direct successor. The public run() drains the chain once at prompt end. diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index 24ba3277..5d049e20 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -5996,7 +5996,7 @@ int64_t Prefill::lm_prepare(int64_t n, std::string& why, int64_t cap) { if (m.fused_bufs) { why = "the fused expert layout"; return 0; } if (!m.GU || !m.Dm || (n_mmq > 0 && !m.ids_identity)) { why = "the MoE buffers are missing"; return 0; } // (an explicit STRATA_PREFILL_LM_TOKENS may be one chunk: the A/B against the default path, experts aside) - const int64_t min_s = lm_tokens_env() > 0 ? m.T : 2 * m.T; + const int64_t min_s = lm_force_ ? 1 : (lm_tokens_env() > 0 ? m.T : 2 * m.T); // (lm_force_: run_multi - several prompts, however short) if (n < min_s) { why = "a short prompt"; return 0; } const int64_t want = lm_tokens_env() > 0 ? lm_tokens_env() : (lm_by_default_ ? kLmDefaultOnTokens : kLmDefaultTokens); int64_t S = std::max(min_s, want / m.T * m.T); @@ -6166,6 +6166,8 @@ bool Prefill::run_multi(const std::vector& segs, std::string& err) { for (const Seg& sg : segs) total += sg.n; std::string why; if (hand_in_ != nullptr || !impl_) { err = "prefill: multi-prompt read declined (chunk hand-off in use)"; return false; } + for (Prefill* q = this; q != nullptr; q = q->next_) q->lm_force_ = true; + struct ForceOff { Prefill* h; ~ForceOff() { for (Prefill* q = h; q != nullptr; q = q->next_) q->lm_force_ = false; } } force_off{this}; if (lm_prepare_chain(total, why) <= 0) { err = "prefill: multi-prompt read declined: " + why; return false; } bool ok = lm_run(segs, 0, nullptr, err); std::string drain_err; From 1d9b4c48e9f0385eae651b63f922412b9ab12e8b Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 03:21:45 +0100 Subject: [PATCH 04/17] prefill: STRATA_LM_HIP=1 opts the layer-major (and multi-prompt) read in on HIP builds --- src/prefill/prefill.cpp | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index 5d049e20..58531af2 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -5936,9 +5936,15 @@ int64_t Prefill::lm_prepare(int64_t n, std::string& why, int64_t cap) { Impl& m = *impl_; if (m.g == nullptr) { why = "not initialized"; return 0; } #if defined(STRATA_USE_HIP) || defined(STRATA_HIP_GFX906) - // only NVIDIA CUDA has run this path; on an AMD build the prompt stays on the chunk-major path - why = "layer-major prefill has only been run on NVIDIA CUDA"; - return 0; + // only NVIDIA CUDA has run this path by default; on an AMD build the prompt stays on the chunk-major path unless + // STRATA_LM_HIP=1 (opt-in, PAR-A 2026-10-30: gfx1201 is being qualified) + { + static const bool lm_hip = [] { const char* v = std::getenv("STRATA_LM_HIP"); return v != nullptr && v[0] == '1'; }(); + if (!lm_hip) { + why = "layer-major prefill has only been run on NVIDIA CUDA"; + return 0; + } + } #endif const core::ModelGeometry& g = *m.g; const core::OnDevice on_device(m.device); // a layer split's later stage is sized on its own GPU From 3ccfd5aed3f51fe64875c7a1594ea9be7676177a Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 03:50:58 +0100 Subject: [PATCH 05/17] replicas: opt-in shared expert arena (one /dev/shm file for all replicas, stale-file sweep, private fallback); the engine reserves a tmpfs arena up front --- serve/replicas.py | 105 +++++++++++++++++++++++++++++++++++++++++ serve/server.py | 17 ++++++- serve/test_replicas.py | 83 +++++++++++++++++++++++++++++++- src/core/pinned.cu | 15 ++++++ 4 files changed, 218 insertions(+), 2 deletions(-) diff --git a/serve/replicas.py b/serve/replicas.py index 5a139efa..9229a64d 100644 --- a/serve/replicas.py +++ b/serve/replicas.py @@ -18,6 +18,15 @@ conversation's next turn therefore goes back to its replica, until that replica reading the prompt again is cheaper than waiting. What each replica holds is tracked as hashes of 64-token blocks of the prompts and answers it served (the engine itself checks the real prefix, as it always does). +Host RAM. Every engine keeps its own copy of the experts in RAM (about 47 GiB for Qwen3.8-Flash-Next IQ3_S), so two +replicas need two copies and four need four. "replica_shared_arena": true (or STRATA_REPLICA_SHARED_ARENA=1; opt-in) +backs the experts' arena of all replicas with ONE file in /dev/shm (the engine's --shared-expert-arena, MAP_SHARED): +the experts are in RAM once, whatever the number of replicas. Every engine still reads the experts into the file +itself (the same bytes, so a replica that starts, restarts or dies at any time never leaves another one with half an +arena), the pack hash in the file's header refuses another model, the file is unlinked when the server ends (a +running engine keeps its mapping: the kernel frees the pages with the last one) and a file whose server is gone is +removed at the next start. See shared_arena_plan. + CPU cores. On Linux the replicas get disjoint physical cores (serve/cpu_affinity.py, from PR #1436 by agorevski), so the host threads of the engines (the expert workers) do not compete. "replica_cpus": false turns it off. """ @@ -26,6 +35,7 @@ from __future__ import annotations import collections import copy import os +import sys import threading import time @@ -81,6 +91,100 @@ def _cards(g) -> list[int]: raise ReplicaConfigError(f"cards must be numbers, not {g!r}") from None +ARENA_PREFIX = "strata-replicas-" +ARENA_MIN_FREE = 8 << 30 # /dev/shm must have at least this much free to try (the engine reserves the real size) +ARENA_INCOMPATIBLE = ("--mmap-experts", "--resident-experts", "--shared-expert-arena", "--resident-budget-gib") + + +def _pid_alive(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def sweep_stale_arenas(shm: str = "/dev/shm", uid: int | None = None) -> list[str]: + """Remove the arena files of servers that no longer run (a killed server cannot unlink its own). A file of a living + server, or of another user, is never touched; a mapping an engine still holds stays valid either way.""" + if uid is None: + uid = os.getuid() if hasattr(os, "getuid") else 0 + gone = [] + try: + names = os.listdir(shm) + except OSError: + return gone + head = f"{ARENA_PREFIX}{uid}-" + for n in names: + if not (n.startswith(head) and n.endswith(".bin")): + continue + try: + pid = int(n[len(head):-4]) + except ValueError: + continue + if _pid_alive(pid): + continue + try: + os.unlink(os.path.join(shm, n)) + gone.append(n) + except OSError: + pass + return gone + + +def _free(path: str) -> int: + try: + st = os.statvfs(path) + return st.f_bavail * st.f_frsize + except OSError: + return 0 + + +def shared_arena_plan(cfg: dict, n_replicas: int, *, shm: str = "/dev/shm", platform: str | None = None, + environ=None, free_bytes=None) -> tuple[str | None, str]: + """(the file for --shared-expert-arena or None, a line for the log). Opt-in: "replica_shared_arena": true (or + STRATA_REPLICA_SHARED_ARENA=1). Only on Linux (the engine's MAP_SHARED backing), for Qwen's arena (not DeepSeek, + not the mapped or budgeted RAM modes, which have no arena), and with room in /dev/shm. The engine checks the real + size and reserves it before it loads (a tmpfs that is too small is an error at start, not a bus error later).""" + environ = os.environ if environ is None else environ + platform = sys.platform if platform is None else platform + want = cfg.get("replica_shared_arena") + if want is None: + want = str(environ.get("STRATA_REPLICA_SHARED_ARENA", "")).strip().lower() in ("1", "true", "on", "yes") + if not want or n_replicas < 2: + return None, "" + if not platform.startswith("linux"): + return None, "shared expert arena: Linux only; each replica keeps its own copy of the experts in RAM" + args = [str(a) for a in cfg.get("args") or []] + if isinstance(cfg.get("replicas"), list): + for e in cfg["replicas"]: + if isinstance(e, dict): + args += [str(a) for a in e.get("args") or []] + if cfg.get("arch") == "deepseek41" or any(a in args for a in ARENA_INCOMPATIBLE): + return None, "shared expert arena: this model/config has no private expert arena to share; each replica as before" + if not (os.path.isdir(shm) and os.access(shm, os.W_OK)): + return None, f"shared expert arena: {shm} is not usable; each replica keeps its own copy of the experts in RAM" + free = free_bytes if free_bytes is not None else _free(shm) + if free < ARENA_MIN_FREE: + return None, (f"shared expert arena: {shm} has only {free / 2**30:.1f} GiB free (a tmpfs is half the RAM by " + "default); each replica keeps its own copy of the experts in RAM") + sweep_stale_arenas(shm) + uid = os.getuid() if hasattr(os, "getuid") else 0 + path = os.path.join(shm, f"{ARENA_PREFIX}{uid}-{os.getpid()}.bin") + return path, f"one shared expert arena for {n_replicas} replicas ({path}): the experts are in RAM once" + + +def remove_arena(path: str | None) -> None: + """Unlink the arena file (the server's end, or a start that fell back): engines that map it keep their pages.""" + if path: + try: + os.unlink(path) + except OSError: + pass + + def _suffixed(path, i: int): if not isinstance(path, str) or not path: return path @@ -337,6 +441,7 @@ class ReplicaEngine: e.close() except BaseException as x: # noqa: BLE001 err = err or x + remove_arena(self.__dict__.get("arena_file")) # (an engine that still maps it keeps its pages) if err: raise err diff --git a/serve/server.py b/serve/server.py index a2f5d66f..5a77af9a 100644 --- a/serve/server.py +++ b/serve/server.py @@ -8559,9 +8559,14 @@ def start_replicas(cfg: dict, groups: list[list[int]], exe: str, effort_end: lis flush=True) engines: list = [None] * n errors: list = [] + arena, arena_note = rp.shared_arena_plan(cfg, n) # opt-in: one copy of the experts in RAM for all replicas + if arena_note: + print(f"[strata] {arena_note}", flush=True) def one(i: int): rc = rp.replica_cfg(cfg, i, groups[i]) + if arena: + rc["args"] = list(rc["args"]) + ["--shared-expert-arena", arena] try: engines[i] = StrataEngine(exe, engine_args(rc) + effort_end, cwd=rc.get("cwd"), log=rc.get("log"), env=child_env(rc, offline=offline), lazy=lazy, cpus=cpus[i] if cpus else None) @@ -8585,10 +8590,20 @@ def start_replicas(cfg: dict, groups: list[list[int]], exe: str, effort_end: lis e.close() except Exception: # noqa: BLE001 pass + rp.remove_arena(arena) i, e = errors[0] + if arena: # a shared arena that could not be made (no room in /dev/shm, a file of another pack): as before + print(f"[strata] replica {i} did not start with the shared expert arena ({e}); starting every replica with " + "its own copy of the experts instead", flush=True) + return start_replicas({**cfg, "replica_shared_arena": False}, groups, exe, effort_end, lazy, offline=offline) raise SystemExit(f"[strata] replica {i} did not start: {e}") load = int(cfg.get("replica_load_tokens") or rp.LOAD_TOKENS) - return rp.ReplicaEngine(engines, load_tokens=load) + eng = rp.ReplicaEngine(engines, load_tokens=load) + if arena: + eng.__dict__["arena_file"] = arena + import atexit + atexit.register(rp.remove_arena, arena) # (a killed server leaves the file: the next start sweeps it) + return eng def start_split(cfg: dict, spec: dict, exe: str, effort_end: list[str], lazy: bool, *, offline: bool = False): diff --git a/serve/test_replicas.py b/serve/test_replicas.py index 95f1a12e..7791adbd 100644 --- a/serve/test_replicas.py +++ b/serve/test_replicas.py @@ -11,7 +11,8 @@ from pathlib import Path from unittest import mock from serve.frontend import ChatTemplate -from serve.replicas import BLOCK, ReplicaConfigError, ReplicaEngine, Router, cpu_sets, gpu_groups, replica_cfg +from serve.replicas import (BLOCK, ReplicaConfigError, ReplicaEngine, Router, cpu_sets, gpu_groups, remove_arena, + replica_cfg, shared_arena_plan, sweep_stale_arenas) from serve.server import ByteTokenizer, Service, StrataEngine, serve from serve.test_parallel import FAKE_BATCH @@ -309,5 +310,85 @@ class ReplicaService(unittest.TestCase): self.assertEqual(self.served()[1], 3) +class SharedArena(unittest.TestCase): + """The opt-in shared expert arena of the replicas (shared_arena_plan): when it is used, when it is not, what a stale + file means.""" + CFG = {"replica_shared_arena": True, "args": ["--pack", "/p"], "gpu": [0, 1]} + + def plan(self, cfg=None, n=2, tmp=None, **kw): + return shared_arena_plan(cfg or self.CFG, n, shm=tmp, platform=kw.pop("platform", "linux"), + environ=kw.pop("environ", {}), free_bytes=kw.pop("free_bytes", 100 << 30), **kw) + + def test_off_by_default_and_for_one_engine(self): + with tempfile.TemporaryDirectory() as d: + self.assertEqual(self.plan({"args": []}, tmp=d), (None, "")) + self.assertEqual(self.plan(n=1, tmp=d), (None, "")) + self.assertEqual(self.plan({"args": []}, tmp=d, environ={"STRATA_REPLICA_SHARED_ARENA": "0"}), (None, "")) + + def test_on_by_config_or_environment(self): + with tempfile.TemporaryDirectory() as d: + path, note = self.plan(tmp=d) + self.assertTrue(path and path.startswith(d) and path.endswith(".bin"), path) + self.assertIn("once", note) + path2, _ = self.plan({"args": []}, tmp=d, environ={"STRATA_REPLICA_SHARED_ARENA": "1"}) + self.assertTrue(path2) + self.assertEqual(path, path2) # (one server: one file) + + def test_only_where_there_is_an_arena_to_share(self): + with tempfile.TemporaryDirectory() as d: + self.assertIsNone(self.plan(tmp=d, platform="win32")[0]) + self.assertIsNone(self.plan(tmp=d, platform="darwin")[0]) + for flag in ("--mmap-experts", "--resident-experts", "--resident-budget-gib", "--shared-expert-arena"): + self.assertIsNone(self.plan({**self.CFG, "args": ["--pack", "/p", flag]}, tmp=d)[0], flag) + self.assertIsNone(self.plan({**self.CFG, "replicas": [{"gpus": [0], "args": ["--mmap-experts"]}, {"gpus": [1]}]}, tmp=d)[0]) + self.assertIsNone(self.plan({**self.CFG, "arch": "deepseek41"}, tmp=d)[0]) + + def test_no_room_or_no_shm_says_so_and_falls_back(self): + with tempfile.TemporaryDirectory() as d: + path, note = self.plan(tmp=d, free_bytes=2 << 30) + self.assertIsNone(path) + self.assertIn("2.0 GiB free", note) + path, note = self.plan(tmp=d + "/nope") + self.assertIsNone(path) + self.assertIn("not usable", note) + + def test_a_stale_file_goes_a_live_one_stays(self): + with tempfile.TemporaryDirectory() as d: + uid = __import__("os").getuid() if hasattr(__import__("os"), "getuid") else 0 + dead = Path(d) / f"strata-replicas-{uid}-2999999.bin" + live = Path(d) / f"strata-replicas-{uid}-{__import__('os').getpid()}.bin" + other = Path(d) / f"strata-replicas-{uid + 1}-2999999.bin" + junk = Path(d) / "strata-replicas-x.bin" + for f in (dead, live, other, junk): + f.write_bytes(b"x") + with mock.patch("serve.replicas._pid_alive", side_effect=lambda pid: pid != 2999999): + gone = sweep_stale_arenas(d, uid) + self.assertEqual(gone, [dead.name]) + self.assertFalse(dead.exists()) + self.assertTrue(live.exists() and other.exists() and junk.exists()) + + def test_removing_twice_or_nothing_is_harmless(self): + with tempfile.TemporaryDirectory() as d: + f = Path(d) / "a.bin" + f.write_bytes(b"x") + remove_arena(str(f)) + remove_arena(str(f)) + remove_arena(None) + self.assertFalse(f.exists()) + + def test_the_replica_engine_unlinks_the_file_at_close(self): + with tempfile.TemporaryDirectory() as d: + f = Path(d) / "a.bin" + f.write_bytes(b"x") + + class E: + def close(self): + pass + r = ReplicaEngine([E(), E()]) + r.__dict__["arena_file"] = str(f) + r.close() + self.assertFalse(f.exists()) + + if __name__ == "__main__": unittest.main() diff --git a/src/core/pinned.cu b/src/core/pinned.cu index 0c45c051..53cb4842 100644 --- a/src/core/pinned.cu +++ b/src/core/pinned.cu @@ -32,6 +32,7 @@ #include #include #include +#include #include #include #ifndef MAP_HUGE_2MB @@ -177,6 +178,20 @@ void* reserve(uint64_t bytes, PageBacking& got, std::string& note, const std::st note = "cannot size shared arena " + shared_file + ": " + std::strerror(e); return nullptr; } + // tmpfs: take the pages now. ftruncate only sets the length; with /dev/shm too small the load would die + // with a bus error part-way, whereas here a full tmpfs is a plain error at start (the server then starts + // its replicas with private copies). A file on a disk is left sparse, as before. + struct statfs sf{}; + if (fstatfs(fd, &sf) == 0 && (unsigned long) sf.f_type == 0x01021994ul) { + const int fe = posix_fallocate(fd, 0, (off_t) file_bytes); + if (fe != 0 && fe != EOPNOTSUPP && fe != EINVAL) { + close(fd); + unlink(shared_file.c_str()); // (created by this call, never filled: nobody maps it yet) + note = "not enough room for the shared arena " + shared_file + " (" + + std::to_string((unsigned long long) (file_bytes >> 20)) + " MiB): " + std::strerror(fe); + return nullptr; + } + } SharedArenaHeader hdr{}; std::memcpy(hdr.magic, kSharedArenaMagic, sizeof hdr.magic); hdr.version = 1; From 0c94851835de3577c23d8a070ef4d2ec8a3602db Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 04:13:24 +0100 Subject: [PATCH 06/17] shared arena: the header is written before any engine reserves pages, an importer waits for a header still being written, and every engine reserves the tmpfs pages --- src/core/pinned.cu | 51 ++++++++++++++++++++++++++++++++-------------- 1 file changed, 36 insertions(+), 15 deletions(-) diff --git a/src/core/pinned.cu b/src/core/pinned.cu index 53cb4842..ae6b9bbd 100644 --- a/src/core/pinned.cu +++ b/src/core/pinned.cu @@ -178,20 +178,6 @@ void* reserve(uint64_t bytes, PageBacking& got, std::string& note, const std::st note = "cannot size shared arena " + shared_file + ": " + std::strerror(e); return nullptr; } - // tmpfs: take the pages now. ftruncate only sets the length; with /dev/shm too small the load would die - // with a bus error part-way, whereas here a full tmpfs is a plain error at start (the server then starts - // its replicas with private copies). A file on a disk is left sparse, as before. - struct statfs sf{}; - if (fstatfs(fd, &sf) == 0 && (unsigned long) sf.f_type == 0x01021994ul) { - const int fe = posix_fallocate(fd, 0, (off_t) file_bytes); - if (fe != 0 && fe != EOPNOTSUPP && fe != EINVAL) { - close(fd); - unlink(shared_file.c_str()); // (created by this call, never filled: nobody maps it yet) - note = "not enough room for the shared arena " + shared_file + " (" + - std::to_string((unsigned long long) (file_bytes >> 20)) + " MiB): " + std::strerror(fe); - return nullptr; - } - } SharedArenaHeader hdr{}; std::memcpy(hdr.magic, kSharedArenaMagic, sizeof hdr.magic); hdr.version = 1; @@ -215,7 +201,27 @@ void* reserve(uint64_t bytes, PageBacking& got, std::string& note, const std::st return nullptr; } SharedArenaHeader hdr{}; - const ssize_t got_header = pread(fd, &hdr, sizeof hdr, 0); + ssize_t got_header = pread(fd, &hdr, sizeof hdr, 0); + // the creator sizes the file first and writes the header right after: an engine that starts in that gap waits + // for it (up to 10 s) instead of calling the file broken + for (int w = 0; w < 100 && got_header == (ssize_t) sizeof hdr && + std::memcmp(hdr.magic, kSharedArenaMagic, sizeof hdr.magic) != 0 && + hdr.magic[0] == 0 && hdr.magic[1] == 0; ++w) { + usleep(100 * 1000); + got_header = pread(fd, &hdr, sizeof hdr, 0); + } + if (got_header == (ssize_t) sizeof hdr && std::memcmp(hdr.magic, kSharedArenaMagic, sizeof hdr.magic) != 0 && + hdr.magic[0] == 0 && hdr.magic[1] == 0) { + // still no header after 10 s: the engine that sized the file died before writing it. Every engine + // fills the arena itself, so nothing in the file is worth keeping: write the header as the creator would + std::memset(&hdr, 0, sizeof hdr); + std::memcpy(hdr.magic, kSharedArenaMagic, sizeof hdr.magic); + hdr.version = 1; + hdr.header_bytes = (uint32_t) kSharedArenaHeaderBytes; + hdr.arena_bytes = bytes; + hdr.pack_hash = shared_pack_hash; + got_header = pwrite(fd, &hdr, sizeof hdr, 0); + } if (got_header != (ssize_t) sizeof hdr || std::memcmp(hdr.magic, kSharedArenaMagic, sizeof hdr.magic) != 0 || hdr.version != 1 || hdr.header_bytes != kSharedArenaHeaderBytes || hdr.arena_bytes != bytes) { @@ -235,6 +241,21 @@ void* reserve(uint64_t bytes, PageBacking& got, std::string& note, const std::st } } + // tmpfs: take the pages now (ftruncate only sets the length). With /dev/shm too small the load would die with a bus + // error part-way; here a full tmpfs is a plain error at start, for every engine that maps the file (pages already + // taken cost nothing), and the server then starts its replicas with private copies. A file on a disk stays sparse. + { + struct statfs sf{}; + if (fstatfs(fd, &sf) == 0 && (unsigned long) sf.f_type == 0x01021994ul) { + const int fe = posix_fallocate(fd, 0, (off_t) file_bytes); + if (fe != 0 && fe != EOPNOTSUPP && fe != EINVAL) { + close(fd); + note = "not enough room for the shared arena " + shared_file + " (" + + std::to_string((unsigned long long) (file_bytes >> 20)) + " MiB): " + std::strerror(fe); + return nullptr; + } + } + } void* map = mmap(nullptr, (size_t) file_bytes, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); const int e = errno; close(fd); From 3782008944a5648a15242ded95fc20128a3fcf8e Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 04:46:51 +0100 Subject: [PATCH 07/17] shared arena: memfd backing (not limited by the /dev/shm size), one engine reads the experts and the others skip the read, RAM check before the pages are taken --- include/strata/core/pinned.hpp | 3 ++ serve/replicas.py | 61 ++++++++++++++++++++++++++++------ serve/test_replicas.py | 24 ++++++++++++- src/core/expert_source.cpp | 53 ++++++++++++++++++++++++++++- src/core/pinned.cu | 61 ++++++++++++++++++++++++++++++++++ 5 files changed, 190 insertions(+), 12 deletions(-) diff --git a/include/strata/core/pinned.hpp b/include/strata/core/pinned.hpp index c9d7da07..f97f3c52 100644 --- a/include/strata/core/pinned.hpp +++ b/include/strata/core/pinned.hpp @@ -90,6 +90,9 @@ struct PinnedArena { std::vector slice_starts; void* mapping_base = nullptr; ///< actual mapping start; differs from base when a shared-file header exists uint64_t mapping_bytes = 0; ///< bytes to release from mapping_base + /// shared-file arena only: the load state word in the file's header (0: nobody has loaded it, a pid: that engine is + /// loading it, ~0: it is complete). Shared by every engine that maps the file; null for a private arena. + std::atomic* shared_state = nullptr; ~PinnedArena(); PinnedArena(const PinnedArena&) = delete; PinnedArena& operator=(const PinnedArena&) = delete; diff --git a/serve/replicas.py b/serve/replicas.py index 9229a64d..93309875 100644 --- a/serve/replicas.py +++ b/serve/replicas.py @@ -20,12 +20,15 @@ the prompts and answers it served (the engine itself checks the real prefix, as Host RAM. Every engine keeps its own copy of the experts in RAM (about 47 GiB for Qwen3.8-Flash-Next IQ3_S), so two replicas need two copies and four need four. "replica_shared_arena": true (or STRATA_REPLICA_SHARED_ARENA=1; opt-in) -backs the experts' arena of all replicas with ONE file in /dev/shm (the engine's --shared-expert-arena, MAP_SHARED): -the experts are in RAM once, whatever the number of replicas. Every engine still reads the experts into the file -itself (the same bytes, so a replica that starts, restarts or dies at any time never leaves another one with half an -arena), the pack hash in the file's header refuses another model, the file is unlinked when the server ends (a -running engine keeps its mapping: the kernel frees the pages with the last one) and a file whose server is gone is -removed at the next start. See shared_arena_plan. +backs the experts' arena of all replicas with ONE shared memory object (the engine's --shared-expert-arena, MAP_SHARED): +the experts are in RAM once, whatever the number of replicas. On Linux the object is a memfd the server creates +(os.memfd_create): it is not a file on a mount, so the size of /dev/shm does not limit it; the engines open it as +/proc//fd/. Without memfd (an old kernel, a sandbox that refuses it) a file in /dev/shm is used instead. +The pack hash in the header refuses another model; the first engine to start reads the experts into it and the others +wait for it and skip the read (a state word in the header: nobody / a pid reading / filled; an engine whose reader died +takes the read over); the server closes the memfd (or unlinks the file) when it ends - an engine that still maps it +keeps its pages, the kernel frees them with the last mapping - and a /dev/shm file whose server is gone is removed at +the next start. See shared_arena_plan. CPU cores. On Linux the replicas get disjoint physical cores (serve/cpu_affinity.py, from PR #1436 by agorevski), so the host threads of the engines (the expert workers) do not compete. "replica_cpus": false turns it off. @@ -92,6 +95,8 @@ def _cards(g) -> list[int]: ARENA_PREFIX = "strata-replicas-" +_MEMFDS: dict = {} # the memfds this server created: path (/proc//fd/) -> fd +ARENA_MIN_RAM = 8 << 30 # a memfd arena: at least this much free RAM to try (the engine reserves the real size) ARENA_MIN_FREE = 8 << 30 # /dev/shm must have at least this much free to try (the engine reserves the real size) ARENA_INCOMPATIBLE = ("--mmap-experts", "--resident-experts", "--shared-expert-arena", "--resident-budget-gib") @@ -142,8 +147,19 @@ def _free(path: str) -> int: return 0 +def _mem_available() -> int: + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) * 1024 + except (OSError, ValueError, IndexError): + pass + return 0 + + def shared_arena_plan(cfg: dict, n_replicas: int, *, shm: str = "/dev/shm", platform: str | None = None, - environ=None, free_bytes=None) -> tuple[str | None, str]: + environ=None, free_bytes=None, use_memfd: bool | None = None) -> tuple[str | None, str]: """(the file for --shared-expert-arena or None, a line for the log). Opt-in: "replica_shared_arena": true (or STRATA_REPLICA_SHARED_ARENA=1). Only on Linux (the engine's MAP_SHARED backing), for Qwen's arena (not DeepSeek, not the mapped or budgeted RAM modes, which have no arena), and with room in /dev/shm. The engine checks the real @@ -164,6 +180,22 @@ def shared_arena_plan(cfg: dict, n_replicas: int, *, shm: str = "/dev/shm", plat args += [str(a) for a in e.get("args") or []] if cfg.get("arch") == "deepseek41" or any(a in args for a in ARENA_INCOMPATIBLE): return None, "shared expert arena: this model/config has no private expert arena to share; each replica as before" + if use_memfd is None: + use_memfd = hasattr(os, "memfd_create") and str(cfg.get("replica_shared_arena")).lower() != "devshm" + if use_memfd: + avail = free_bytes if free_bytes is not None else _mem_available() + if avail < ARENA_MIN_RAM: + return None, (f"shared expert arena: only {avail / 2**30:.1f} GiB of RAM is free; each replica keeps its own " + "copy of the experts in RAM") + try: + fd = os.memfd_create("strata-replicas-arena", os.MFD_CLOEXEC) + except OSError: + fd = -1 + if fd >= 0: + path = f"/proc/{os.getpid()}/fd/{fd}" + _MEMFDS[path] = fd + return path, (f"one shared expert arena for {n_replicas} replicas (memory object of the server, not limited by " + "the size of /dev/shm): the experts are in RAM once") if not (os.path.isdir(shm) and os.access(shm, os.W_OK)): return None, f"shared expert arena: {shm} is not usable; each replica keeps its own copy of the experts in RAM" free = free_bytes if free_bytes is not None else _free(shm) @@ -177,12 +209,21 @@ def shared_arena_plan(cfg: dict, n_replicas: int, *, shm: str = "/dev/shm", plat def remove_arena(path: str | None) -> None: - """Unlink the arena file (the server's end, or a start that fell back): engines that map it keep their pages.""" - if path: + """Close the memfd / unlink the arena file (the server's end, or a start that fell back): engines that map it keep + their pages.""" + if not path: + return + fd = _MEMFDS.pop(path, None) + if fd is not None: try: - os.unlink(path) + os.close(fd) except OSError: pass + return + try: + os.unlink(path) + except OSError: + pass def _suffixed(path, i: int): diff --git a/serve/test_replicas.py b/serve/test_replicas.py index 7791adbd..ae1d3e0c 100644 --- a/serve/test_replicas.py +++ b/serve/test_replicas.py @@ -317,7 +317,8 @@ class SharedArena(unittest.TestCase): def plan(self, cfg=None, n=2, tmp=None, **kw): return shared_arena_plan(cfg or self.CFG, n, shm=tmp, platform=kw.pop("platform", "linux"), - environ=kw.pop("environ", {}), free_bytes=kw.pop("free_bytes", 100 << 30), **kw) + environ=kw.pop("environ", {}), free_bytes=kw.pop("free_bytes", 100 << 30), + use_memfd=kw.pop("use_memfd", False), **kw) def test_off_by_default_and_for_one_engine(self): with tempfile.TemporaryDirectory() as d: @@ -367,6 +368,27 @@ class SharedArena(unittest.TestCase): self.assertFalse(dead.exists()) self.assertTrue(live.exists() and other.exists() and junk.exists()) + @unittest.skipUnless(hasattr(__import__("os"), "memfd_create") and sys.platform.startswith("linux"), "Linux memfd") + def test_a_memfd_is_not_limited_by_dev_shm_and_is_closed_at_the_end(self): + import os + path, note = shared_arena_plan(self.CFG, 2, platform="linux", environ={}, free_bytes=100 << 30, use_memfd=True) + self.assertTrue(path and path.startswith(f"/proc/{os.getpid()}/fd/"), path) + fd = int(path.rsplit("/", 1)[1]) + os.ftruncate(fd, 1 << 20) + self.assertEqual(os.fstat(fd).st_size, 1 << 20) + with open(path, "r+b") as other: # (how an engine opens it) + other.seek(100) + other.write(b"x") + remove_arena(path) + with self.assertRaises(OSError): + os.fstat(fd) + remove_arena(path) # (twice: harmless) + + def test_a_memfd_needs_free_ram(self): + path, note = shared_arena_plan(self.CFG, 2, platform="linux", environ={}, free_bytes=2 << 30, use_memfd=True) + self.assertIsNone(path) + self.assertIn("2.0 GiB of RAM", note) + def test_removing_twice_or_nothing_is_harmless(self): with tempfile.TemporaryDirectory() as d: f = Path(d) / "a.bin" diff --git a/src/core/expert_source.cpp b/src/core/expert_source.cpp index 583ee07f..3eee2116 100644 --- a/src/core/expert_source.cpp +++ b/src/core/expert_source.cpp @@ -20,6 +20,10 @@ #include "strata/kernels/cpu/pool.hpp" #include +#ifndef _WIN32 +#include +#include +#endif #include #include @@ -6000,6 +6004,45 @@ bool ArenaExpertSource::open(const std::string& pack_dir, int64_t n_layers, int6 const bool unbuffered = experts_unbuffered(files, want + (uint64_t) blob, why); const int readers = unbuffered ? std::max(threads, 16) : threads; // 16 keep a PCIe 5 drive's queue full LoadStats st; + // A shared arena (--shared-expert-arena, the replicas): ONE engine reads the experts into it, the others wait and skip the + // read. The state word in the file's header says who: 0 nobody, a pid loading, ~0 filled. An engine that finds a + // pid that is gone takes the load over (the previous one died mid-read: the arena is read again from the start), so a + // dead loader never leaves the others waiting, and an engine that waits too long loads it itself. + constexpr uint64_t kArenaFilled = ~uint64_t{0}; + std::atomic* arena_state = padded_ ? nullptr : a->shared_state; + bool arena_leader = false, arena_skip = false; +#ifndef _WIN32 + if (arena_state != nullptr) { + const uint64_t me = (uint64_t) getpid(); + const auto t_begin = std::chrono::steady_clock::now(); + bool said = false; + for (;;) { + uint64_t cur = arena_state->load(std::memory_order_acquire); + if (cur == kArenaFilled) { arena_skip = true; break; } + if (cur == 0 || cur == me) { + if (arena_state->compare_exchange_strong(cur, me, std::memory_order_acq_rel)) { arena_leader = true; break; } + continue; + } + const bool alive = !(kill((pid_t) cur, 0) != 0 && errno == ESRCH); + if (!alive) { + if (arena_state->compare_exchange_strong(cur, me, std::memory_order_acq_rel)) { + std::fprintf(stderr, "strata generate: shared expert arena: engine %llu died while reading it; reading it again\n", + (unsigned long long) cur); + arena_leader = true; + break; + } + continue; + } + if (std::chrono::steady_clock::now() - t_begin > std::chrono::minutes(45)) break; // (loads it on its own) + if (!said) { + said = true; + std::fprintf(stderr, "strata generate: shared expert arena: engine %llu is reading the experts; waiting for it " + "instead of reading them again\n", (unsigned long long) cur); + } + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + } +#endif std::atomic ready{0}; std::thread reg; if (deferred) { @@ -6011,7 +6054,12 @@ bool ArenaExpertSource::open(const std::string& pack_dir, int64_t n_layers, int6 }); } const std::atomic* rp = deferred ? &ready : nullptr; - if (from_gguf) { + if (arena_skip) { + st.ok = true; + st.bytes = want; + st.seconds = 0.0; + std::fprintf(stderr, "strata generate: shared expert arena: already filled by another engine; the read is skipped\n"); + } else if (from_gguf) { st = load_experts_gguf(gguf_, a->data(), lay, readers, unbuffered, rp, padded_ ? &pad_ : nullptr); } else { if (unbuffered && !padded_) st = load_experts_direct(path, a->data(), loff, lbytes, readers, /*chunk=*/8u << 20, rp); @@ -6022,15 +6070,18 @@ bool ArenaExpertSource::open(const std::string& pack_dir, int64_t n_layers, int6 std::fprintf(stderr, "strata generate: expert arena read %s (%s)\n", unbuffered ? "unbuffered" : "through the file cache", why.c_str()); if (!st.ok) { + if (arena_leader) arena_state->store(0, std::memory_order_release); // (somebody else may try) delete a; err = "ArenaExpertSource: the expert load was refused: " + (st.error.empty() ? std::string("unknown") : st.error); return false; } if (st.bytes != want) { + if (arena_leader) arena_state->store(0, std::memory_order_release); delete a; err = "ArenaExpertSource: the load read " + std::to_string(st.bytes) + " B of " + std::to_string(want); return false; } + if (arena_leader) arena_state->store(kArenaFilled, std::memory_order_release); // the others may use (and skip reading) it if (std::getenv("STRATA_VERIFY_ARENA") != nullptr) { // tests: the loaded arena's checksum, any loader std::fprintf(stderr, "strata generate: expert arena checksum %016llx (%.2f GiB)\n", (unsigned long long) fnv1a64(a->data(), want), (double) want / (1ull << 30)); diff --git a/src/core/pinned.cu b/src/core/pinned.cu index ae6b9bbd..468971c4 100644 --- a/src/core/pinned.cu +++ b/src/core/pinned.cu @@ -2,6 +2,7 @@ #include "strata/core/pinned.hpp" #include "strata/platform/memory.hpp" #include "strata/core/pin_budget.hpp" +#include #include @@ -57,6 +58,45 @@ struct SharedArenaHeader { }; static_assert(sizeof(SharedArenaHeader) <= kSharedArenaHeaderBytes); +#ifndef _WIN32 +// MemAvailable and, in a cgroup v2 group with a limit, the room left under it; ~0 when neither can be read +uint64_t free_ram_room() { + uint64_t room = ~uint64_t{0}; + if (std::FILE* f = std::fopen("/proc/meminfo", "r")) { + char line[256]; + while (std::fgets(line, sizeof line, f) != nullptr) { + unsigned long long kb = 0; + if (std::sscanf(line, "MemAvailable: %llu kB", &kb) == 1) { room = (uint64_t) kb * 1024u; break; } + } + std::fclose(f); + } + std::string cg; + if (std::FILE* f = std::fopen("/proc/self/cgroup", "r")) { + char line[512]; + while (std::fgets(line, sizeof line, f) != nullptr) + if (std::strncmp(line, "0::", 3) == 0) { cg = line + 3; break; } + std::fclose(f); + } + while (!cg.empty() && (cg.back() == '\n' || cg.back() == '\r')) cg.pop_back(); + if (!cg.empty()) { + unsigned long long lim = 0, cur = 0; + bool have_lim = false; + if (std::FILE* f = std::fopen(("/sys/fs/cgroup" + cg + "/memory.max").c_str(), "r")) { + have_lim = std::fscanf(f, "%llu", &lim) == 1; // ("max" does not parse: no limit) + std::fclose(f); + } + if (have_lim) { + if (std::FILE* f = std::fopen(("/sys/fs/cgroup" + cg + "/memory.current").c_str(), "r")) { + if (std::fscanf(f, "%llu", &cur) != 1) cur = 0; + std::fclose(f); + } + room = std::min(room, lim > cur ? (uint64_t) (lim - cur) : 0); + } + } + return room; +} +#endif + #ifndef _WIN32 // bytes in free blocks of 2 MiB or more (/proc/buddyinfo, order >= 9): what MADV_HUGEPAGE can get without compaction uint64_t free_huge_bytes() { @@ -247,6 +287,21 @@ void* reserve(uint64_t bytes, PageBacking& got, std::string& note, const std::st { struct statfs sf{}; if (fstatfs(fd, &sf) == 0 && (unsigned long) sf.f_type == 0x01021994ul) { + // a memfd (or any tmpfs file not capped by a mount size) would take the RAM and then the OOM killer: check the + // free RAM and the cgroup's room for the pages that are not there yet before taking them + struct stat st2{}; + const uint64_t have = fstat(fd, &st2) == 0 ? (uint64_t) st2.st_blocks * 512u : 0; + const uint64_t need = file_bytes > have ? file_bytes - have : 0; + if (need > 0) { + const uint64_t room = free_ram_room(); + if (room != ~uint64_t{0} && need + (2ull << 30) > room) { + close(fd); + note = "not enough free RAM for the shared arena " + shared_file + ": it needs " + + std::to_string((unsigned long long) (need >> 20)) + " MiB more and " + + std::to_string((unsigned long long) (room >> 20)) + " MiB are free (2 GiB are kept back)"; + return nullptr; + } + } const int fe = posix_fallocate(fd, 0, (off_t) file_bytes); if (fe != 0 && fe != EOPNOTSUPP && fe != EINVAL) { close(fd); @@ -446,6 +501,9 @@ PinnedArena::PinnedArena(uint64_t bytes, const std::vector& bounds, uint64_t shared_pack_hash) : capacity(bytes) { if (bytes == 0) return; base = reserve(bytes, backing, note, shared_file, shared_pack_hash, mapping_base, mapping_bytes); + if (!shared_file.empty() && base != nullptr && mapping_base != nullptr) + shared_state = reinterpret_cast*>(static_cast(mapping_base) + + offsetof(SharedArenaHeader, reserved)); if (base != nullptr && mapping_base == nullptr) { mapping_base = base; mapping_bytes = bytes; @@ -565,6 +623,9 @@ PinnedArena::PinnedArena(uint64_t bytes, ByNeed, const std::vector*>(static_cast(mapping_base) + + offsetof(SharedArenaHeader, reserved)); if (base != nullptr && mapping_base == nullptr) { mapping_base = base; mapping_bytes = bytes; From 2f0d6a1c15167e8d320dd319fb6d042c1d1fcc82 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 04:51:14 +0100 Subject: [PATCH 08/17] replica profile (opt-in): setup --replica-profile picks one engine per card + shared expert arena + the host stack when the RAM and VRAM check passes --- setup.py | 22 ++++- strata_control/replica_profile.py | 131 +++++++++++++++++++++++++ tests/terminal/test_replica_profile.py | 124 +++++++++++++++++++++++ 3 files changed, 276 insertions(+), 1 deletion(-) create mode 100644 strata_control/replica_profile.py create mode 100644 tests/terminal/test_replica_profile.py diff --git a/setup.py b/setup.py index 0c4b1b92..233cb9a2 100644 --- a/setup.py +++ b/setup.py @@ -8953,6 +8953,11 @@ def main() -> int: help="several GPUs: run N independent engines, each on its own equal group of the chosen cards (2 " "replicas on 4 cards = two 2-card engines), requests spread over them (opt-in; more total " "throughput for a model that fits fewer cards; each replica holds its own copy of the model)") + ap.add_argument("--replica-profile", action="store_true", default=False, + help="two or more GPUs, Qwen3.8 IQ3_S/IQ3_XXS/Q2_0: one engine per card behind the server instead of one " + "layer split, the experts shared in RAM, the batch switches on (opt-in; about 1.8 times the " + "throughput at 8-16 requests at once on 4 Radeon AI PRO R9700; needs about 61 GB of RAM for 2 " + "engines and 66 GB for 4; a layer split you chose, replicas or --parallel always win)") ap.add_argument("--kv-streaming", choices=["auto", "on", "off"], default="auto", help="from a 64K context: keep the KV cache in RAM and only the attention's window in VRAM (more " "experts fit on the GPU); auto: when the RAM has room for it") @@ -10095,7 +10100,22 @@ def main() -> int: cfg["gpu_order"] = "as_given" ok(f"layer split across GPUs {multi} ({cfg['layer_split']})") recommend_remote_expert_opt(cfg, off=a.no_remote_expert_opt) - if hip and a.layer_split in (None, "auto"): # four Radeon AI PRO R9700 cards and Qwen3.8 IQ3_S: the measured serving profile + replica_done = False + if a.replicas is None and a.layer_split in (None, "auto") and (getattr(a, "replica_profile", False) or + os.environ.get("STRATA_REPLICA_PROFILE")): + try: # the opt-in replica profile (strata_control/replica_profile.py) + from strata_control import replica_profile + if getattr(a, "replica_profile", False): + cfg["replica_profile"] = True + lines = replica_profile.apply([{"name": c.get("name"), "vram_gb": c.get("vram_gb"), "vendor": c.get("vendor"), + "index": c.get("index")} for c in chosen], cfg, float(ram), + parallel_given=a.parallel is not None) + replica_done = bool(cfg.get("replicas")) + for line in lines: + (ok if replica_done else say)(line) + except Exception as e: # noqa: BLE001 - a profile is never a reason to stop setup + say(f" (the replica profile was not applied: {e})") + if hip and a.layer_split in (None, "auto") and not replica_done: # four Radeon AI PRO R9700 cards and Qwen3.8 IQ3_S: the measured serving profile try: # (strata_control/amd_four_card.py); what the person chose always wins from strata_control import amd_four_card for line in amd_four_card.apply([{"name": c.get("name"), "vram_gb": c.get("vram_gb"), "vendor": "amd"} for c in chosen], diff --git a/strata_control/replica_profile.py b/strata_control/replica_profile.py new file mode 100644 index 00000000..f7a3b516 --- /dev/null +++ b/strata_control/replica_profile.py @@ -0,0 +1,131 @@ +"""The replica serving profile (opt-in): on two or more GPUs, one engine per card behind the server instead of one layer split. + +Measured on 4 x Radeon AI PRO R9700 (Qwen3.8-Flash-Next IQ3_S, 1k-token prompts, 128 tokens out, total tokens per second; +docs/BATCHING.md, "Replicas as a serving profile"): + +* 2 cards: 2 engines of one card each reach 1.8 to 1.9 times what a layer split of the two cards does with `parallel` 8 or + 16 (116.6 against 63.8 tok/s at 8 requests, 126.9 against 65.7 at 16); +* 4 cards: 4 engines reach 1.8 to 1.9 times the shipped four-card profile (190 against 107 at 8, 220 against 116 at 16). + +Nothing is lost for one user: a request alone runs on one card at the same speed (13 ms a token on one card, 13 ms on four). + +The price is host RAM: every engine keeps its own copy of the experts. The profile therefore turns on the shared expert +arena (serve/replicas.py: one memory object holds the experts for all engines, measured 61 GB for 2 engines and 66 GB for 4 +instead of 109 and 211 GB) and chooses the number of engines the RAM allows. Without the shared arena (Windows) it needs the +private copies to fit. + +Recommend, never force: this profile is opt-in (`setup.py --replica-profile`, STRATA_REPLICA_PROFILE=1, or +"replica_profile": true in the run config; "replica_profile": false or STRATA_REPLICA_PROFILE=0 turn it off), and what the +person chose (the config's `replicas`, `--parallel`, `--layer-split`, `--peer-device`, an environment variable) always wins. +It only touches the measured model class; every other computer gets exactly what it got before. + +Pure functions, no I/O (the callers pass the cards, the RAM and the run config). +""" + +from __future__ import annotations + +import os +import sys + +# the experts' size in RAM per quant (setup.py MODELS "arena_gb"): the shared arena holds this once +ARENA_GB = {"iq3_s": 50.3, "iq3_xxs": 42.9, "q2_0": 34.0} +PER_REPLICA_GB = 8.0 # measured: about 7 GB of private RAM per engine beside the arena (2 engines: 61 GB in all, 47 of them shared) +HEADROOM_GB = 12.0 # what the OS, the server and the page cache keep +MIN_CARD_GB = 24.0 # measured on 32 GB cards; a smaller card is not claimed (an A4000 16 GB is measured separately) +MIN_REPLICAS = 2 +SWITCHES = {"STRATA_QSA_SLOTS": "1", "STRATA_BATCH_SELF_COMMIT": "1", "STRATA_SLOT_COPY_DIRECT": "1", + "STRATA_CKPT_POOL": "1", "STRATA_CKPT_POOL_READY": "4", "STRATA_SLOT_CKPT_MOVE": "1", + "STRATA_SLOT_COPY_BATCHSYNC": "1"} +OFF_VALUES = ("0", "off", "no", "false") +ON_VALUES = ("1", "on", "yes", "true") + + +def requested(cfg: dict | None = None, environ=None) -> bool: + """Opt-in, and an explicit no wins: config "replica_profile" true/false, else STRATA_REPLICA_PROFILE.""" + environ = os.environ if environ is None else environ + if cfg and cfg.get("replica_profile") is not None: + return cfg.get("replica_profile") is True + return str(environ.get("STRATA_REPLICA_PROFILE", "")).strip().lower() in ON_VALUES + + +def quant_of(cfg: dict) -> str | None: + """The measured model: Qwen3.8-Flash-Next at a quant whose experts size is known.""" + text = " ".join([str(cfg.get("model_name") or "")] + [str(a) for a in cfg.get("args") or []]).lower().replace("-", "_").replace(".", "_") + if "qwen3_8" not in text: + return None + for q in ARENA_GB: + if q in text: + return q + return None + + +def shared_arena_possible(platform: str | None = None) -> bool: + return (sys.platform if platform is None else platform).startswith("linux") + + +def replicas_for(n_cards: int, arena_gb: float, ram_gb: float, shared: bool) -> int: + """How many one-card engines the RAM allows (at most one per card).""" + if shared: + room = ram_gb - HEADROOM_GB - arena_gb + k = int(room // PER_REPLICA_GB) if room > 0 else 0 + else: + k = int((ram_gb - HEADROOM_GB) // (arena_gb + PER_REPLICA_GB)) + return max(0, min(n_cards, k)) + + +def applies(cards: list[dict], cfg: dict, ram_gb: float, *, platform: str | None = None, environ=None) -> tuple[bool, str]: + """(True, "") when the profile would be applied; else (False, why not) - the reason is for the log, '' when nothing was asked.""" + if not requested(cfg, environ): + return False, "" + args = [str(a) for a in cfg.get("args") or []] + if cfg.get("replicas") or cfg.get("arch") == "deepseek41" or "--peer-device" in args: + return False, "replica profile: the config already names replicas, a peer device or another model family" + if str(cfg.get("layer_split") or "auto") != "auto" or "--layer-split" in args: + return False, "replica profile: you chose a layer split; kept" + if len(cards) < MIN_REPLICAS: + return False, "replica profile: needs two or more GPUs" + q = quant_of(cfg) + if q is None: + return False, "replica profile: measured only for Qwen3.8-Flash-Next IQ3_S, IQ3_XXS and Q2_0" + if any(float(c.get("vram_gb") or 0.0) < MIN_CARD_GB for c in cards): + return False, f"replica profile: a card has under {MIN_CARD_GB:.0f} GB of VRAM (measured on 32 GB cards)" + shared = shared_arena_possible(platform) + k = replicas_for(len(cards), ARENA_GB[q], ram_gb, shared) + if k < MIN_REPLICAS: + need = HEADROOM_GB + ARENA_GB[q] + MIN_REPLICAS * PER_REPLICA_GB if shared else HEADROOM_GB + MIN_REPLICAS * (ARENA_GB[q] + PER_REPLICA_GB) + return False, (f"replica profile: {ram_gb:.0f} GB of RAM; two engines need about {need:.0f} GB" + f"{'' if shared else ' (this OS cannot share the experts between engines)'}; the layer split stays") + return True, "" + + +def apply(cards: list[dict], cfg: dict, ram_gb: float, *, parallel_given: bool = False, platform: str | None = None, + environ=None) -> list[str]: + """Fill the profile into `cfg` (in place) where the person has not chosen; returns the plain lines to show.""" + ok, why = applies(cards, cfg, ram_gb, platform=platform, environ=environ) + if not ok: + return [why] if why else [] + environ = os.environ if environ is None else environ + q = quant_of(cfg) + shared = shared_arena_possible(platform) + k = replicas_for(len(cards), ARENA_GB[q], ram_gb, shared) + gpus = cfg.get("gpu") if isinstance(cfg.get("gpu"), list) else [c.get("index", i) for i, c in enumerate(cards)] + chosen = list(gpus)[:k] + cfg["gpu"] = chosen + cfg["replicas"] = [{"gpus": [g]} for g in chosen] + if shared: + cfg.setdefault("replica_shared_arena", True) + env = dict(cfg.get("env") or {}) + for key, val in SWITCHES.items(): + if key not in environ: # an environment variable of the person's own wins + env.setdefault(key, val) + cfg["env"] = env + notes = [f"{k} engines of one card each instead of one layer split of {len(cards)}: requests spread over them " + f"(about 1.8 times the throughput at 8 to 16 requests at once; one request alone is as fast as before)"] + if k < len(cards): + notes.append(f"{len(cards) - k} card(s) left out: the RAM ({ram_gb:.0f} GB) holds {k} engines") + notes.append("the experts are in RAM once for all engines (shared expert arena)" if shared else + "each engine keeps its own copy of the experts in RAM (this OS cannot share them)") + if not parallel_given and "parallel" not in cfg and not {"--batch", "--slots"} & set(str(a) for a in cfg.get("args") or []): + cfg["parallel"] = 8 if k == 2 else 4 + notes.append(f"{cfg['parallel']} requests at once on each engine (--parallel N changes it)") + return notes diff --git a/tests/terminal/test_replica_profile.py b/tests/terminal/test_replica_profile.py new file mode 100644 index 00000000..fa845126 --- /dev/null +++ b/tests/terminal/test_replica_profile.py @@ -0,0 +1,124 @@ +"""strata_control/replica_profile.py: the opt-in replica serving profile (one engine per card, shared expert arena).""" + +from __future__ import annotations + +import unittest + +from strata_control import replica_profile as P + +R9700 = {"name": "AMD Radeon AI PRO R9700", "vram_gb": 31.9, "vendor": "amd"} + + +def cards(n, **kw): + return [{**R9700, **kw, "index": i} for i in range(n)] + + +def cfg(**kw): + c = {"model_name": "qwen3.8-flash-next-iq3_s", "gpu": [0, 1, 2, 3], "layer_split": "auto", + "args": ["--pack", "/p/iq3_s", "--spec", "4"], "replica_profile": True} + c.update(kw) + return c + + +class Request(unittest.TestCase): + def test_opt_in_only(self): + self.assertFalse(P.requested({}, {})) + self.assertTrue(P.requested({"replica_profile": True}, {})) + self.assertTrue(P.requested({}, {"STRATA_REPLICA_PROFILE": "1"})) + + def test_an_explicit_no_wins(self): + self.assertFalse(P.requested({"replica_profile": False}, {"STRATA_REPLICA_PROFILE": "1"})) + + def test_without_the_request_nothing_changes(self): + c = cfg(replica_profile=None) + c.pop("replica_profile") + before = dict(c) + self.assertEqual(P.apply(cards(4), c, 256.0, platform="linux", environ={}), []) + self.assertEqual(c, before) + + +class Count(unittest.TestCase): + def test_ram_sets_the_number_of_engines(self): + self.assertEqual(P.replicas_for(4, 50.3, 251.0, True), 4) + self.assertEqual(P.replicas_for(4, 50.3, 128.0, True), 4) + self.assertEqual(P.replicas_for(4, 50.3, 96.0, True), 4) + self.assertEqual(P.replicas_for(4, 50.3, 80.0, True), 2) # 80 - 12 - 50.3 = 17.7 -> 2 + self.assertEqual(P.replicas_for(4, 50.3, 64.0, True), 0) # one engine's room only + self.assertEqual(P.replicas_for(2, 50.3, 251.0, False), 2) # private copies: 2 x 58.3 + 12 + self.assertEqual(P.replicas_for(4, 50.3, 200.0, False), 3) + self.assertEqual(P.replicas_for(2, 50.3, 100.0, False), 1) + + +class Apply(unittest.TestCase): + def test_four_cards_with_plenty_of_ram(self): + c = cfg() + notes = P.apply(cards(4), c, 251.0, platform="linux", environ={}) + self.assertTrue(notes) + self.assertEqual(c["replicas"], [{"gpus": [0]}, {"gpus": [1]}, {"gpus": [2]}, {"gpus": [3]}]) + self.assertIs(c["replica_shared_arena"], True) + self.assertEqual(c["parallel"], 4) + self.assertEqual(c["env"]["STRATA_CKPT_POOL"], "1") + self.assertEqual(c["env"]["STRATA_QSA_SLOTS"], "1") + + def test_two_cards_get_eight_slots_each(self): + c = cfg(gpu=[0, 1]) + P.apply(cards(2), c, 128.0, platform="linux", environ={}) + self.assertEqual(c["parallel"], 8) + self.assertEqual(len(c["replicas"]), 2) + + def test_less_ram_drops_cards_and_says_so(self): + c = cfg() + notes = P.apply(cards(4), c, 80.0, platform="linux", environ={}) + self.assertEqual(c["replicas"], [{"gpus": [0]}, {"gpus": [1]}]) + self.assertEqual(c["gpu"], [0, 1]) + self.assertTrue(any("left out" in n for n in notes)) + + def test_too_little_ram_keeps_the_layer_split(self): + c = cfg() + notes = P.apply(cards(4), c, 64.0, platform="linux", environ={}) + self.assertNotIn("replicas", c) + self.assertEqual(c["layer_split"], "auto") + self.assertTrue(notes and "layer split stays" in notes[0]) + + def test_windows_cannot_share_so_the_private_copies_must_fit(self): + c = cfg(gpu=[0, 1]) + self.assertEqual(P.apply(cards(2), c, 100.0, platform="win32", environ={})[0].split(":")[0], "replica profile") + self.assertNotIn("replicas", c) + c = cfg(gpu=[0, 1]) + P.apply(cards(2), c, 160.0, platform="win32", environ={}) + self.assertEqual(len(c["replicas"]), 2) + self.assertNotIn("replica_shared_arena", c) + + def test_what_the_person_chose_wins(self): + for kw in ({"replicas": 2}, {"layer_split": "13,26,39"}, {"args": ["--pack", "/p/iq3_s", "--peer-device", "1"]}, + {"args": ["--pack", "/p/iq3_s", "--layer-split", "20,40"]}, {"arch": "deepseek41"}): + c = cfg(**kw) + before = dict(c) + P.apply(cards(4), c, 251.0, platform="linux", environ={}) + self.assertEqual(c.get("replicas"), before.get("replicas"), kw) + self.assertEqual(c.get("layer_split"), before.get("layer_split"), kw) + + def test_parallel_and_environment_variables_of_the_person_win(self): + c = cfg(parallel=6) + P.apply(cards(4), c, 251.0, platform="linux", environ={"STRATA_CKPT_POOL": "0"}) + self.assertEqual(c["parallel"], 6) + self.assertNotIn("STRATA_CKPT_POOL", c["env"]) + c = cfg() + P.apply(cards(4), c, 251.0, parallel_given=True, platform="linux", environ={}) + self.assertNotIn("parallel", c) + + def test_only_the_measured_hardware_and_model(self): + for kw in ({"model_name": "qwen3.6-35b-a3b-q4"}, {"model_name": "qwen3.8-flash-next-iq2_xs", "args": ["--pack", "/p/x"]}): + c = cfg(**kw) + P.apply(cards(4), c, 251.0, platform="linux", environ={}) + self.assertNotIn("replicas", c, kw) + c = cfg(gpu=[0, 1]) + P.apply(cards(2, vram_gb=16.0), c, 251.0, platform="linux", environ={}) + self.assertNotIn("replicas", c) + c = cfg(gpu=[0]) + P.apply(cards(1), c, 251.0, platform="linux", environ={}) + self.assertNotIn("replicas", c) + + +if __name__ == "__main__": + unittest.main() From a8fba1f5b7cd8f535b835cc62a506b650570ed19 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 03:39:07 +0100 Subject: [PATCH 09/17] multi-prompt PLE: tables, constants and weights from the main session, the history of the prompt's own --- src/prefill/prefill.cpp | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index 58531af2..ef6ae21f 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -6221,13 +6221,13 @@ bool Prefill::lm_run(const std::vector& segs, int depth, const std::vector< // (none without an MMQ layer: the FP16 path has no such table) if (m.ids_identity) mmq::iota(m.ids_identity, m.T * K, m.cs); // ---- the PLE rows, read per chunk on a thread as run_impl does (the same state machine, in chunk order) - const bool ple_on = ssp->ple.ready(); + const bool ple_on = m.ss->ple.ready(); static const bool ple_batch_env = [] { const char* v = std::getenv("STRATA_PLE_BATCH"); return v == nullptr || std::atoi(v) != 0; }(); const bool ple_batch = ple_on && ple_batch_env && strata::kernels::ple_native_postops_enabled() && - (ssp->ple.w.key_bf16 != nullptr || ssp->ple.w.key_native_data != nullptr) && + (m.ss->ple.w.key_bf16 != nullptr || m.ss->ple.w.key_native_data != nullptr) && m.region_bytes / ((uint64_t) (3 * strata::kernels::NG_HC_DIM + N + 4) * 4 + (uint64_t) N * 2 + 4096) >= 64; // ---- the plan. A segment is cut into pieces of at most S tokens (the usual super-chunks), each piece into // sub-chunks of m.T tokens from the piece's start; pieces of several segments are packed into one super-chunk while @@ -6282,12 +6282,12 @@ bool Prefill::lm_run(const std::vector& segs, int depth, const std::vector< int32_t pv[2] = {at(c0), at(c0 + 1)}; for (int64_t t = 0; t < T; ++t) { const int32_t tok = (int32_t) tokens[c0 + t]; - strata::kernels::ngram_rows(&tok, pv, 1, ssq->ple.consts, + strata::kernels::ngram_rows(&tok, pv, 1, m.ss->ple.consts, m.ple_rows[buf].data() + t * strata::kernels::PLE_N_HEADS); pv[0] = pv[1]; pv[1] = tok; } - return ssq->ple.table->gather_batch(m.ple_rows[buf].data(), (size_t) T, m.ple_emb_host[buf], e); + return m.ss->ple.table->gather_batch(m.ple_rows[buf].data(), (size_t) T, m.ple_emb_host[buf], e); }; std::string ple_next_err; std::future ple_next; @@ -6805,7 +6805,7 @@ bool Prefill::lm_run(const std::vector& segs, int depth, const std::vector< if (l == 1 && ple_on && ple_batch) { pt.mark(kPfPle, cs); const auto tp = Clock::now(); - const strata::kernels::PleWeights& pw = ssp->ple.w; + const strata::kernels::PleWeights& pw = m.ss->ple.w; constexpr int64_t HD = strata::kernels::NG_HC_DIM; const uint64_t per_token = (uint64_t) (3 * HD + N + 4) * 4 + (uint64_t) N * (bf16x2(m.f16_io) ? 4 : 2) + 4096; const int64_t SB = std::min(T, (int64_t) (m.region_bytes / per_token)); @@ -6833,7 +6833,7 @@ bool Prefill::lm_run(const std::vector& segs, int depth, const std::vector< m.gemm.bf16(e16, pw.value_bf16, val, nb, N, N); if (e16_lo) m.gemm.bf16(e16_lo, pw.value_bf16, val, nb, N, N, 0, 1.0f); try { - strata::kernels::native_ple_postops_batch(key, m.R + s0 * D, val, ssp->ple.hist, pw, qn, gated, + strata::kernels::native_ple_postops_batch(key, m.R + s0 * D, val, ssp->ple_hist, pw, qn, gated, gate, (int) nb, cs); } catch (const std::exception& e) { err = std::string("prefill PLE: ") + e.what(); return false; } } @@ -6846,10 +6846,10 @@ bool Prefill::lm_run(const std::vector& segs, int depth, const std::vector< po.normalized = m.ple_norm; po.result = m.R + t * D; try { - strata::kernels::ple_block(m.ple_emb + t * N, m.R + t * D, ssp->ple.hist, ssp->ple.w, po, - ssp->ple.scratch, cs); + strata::kernels::ple_block(m.ple_emb + t * N, m.R + t * D, ssp->ple_hist, m.ss->ple.w, po, + m.ss->ple.scratch, cs); } catch (const std::exception& e) { err = std::string("prefill PLE: ") + e.what(); return false; } - strata::kernels::ple_history_advance(ssp->ple.hist, m.ple_norm, cs); + strata::kernels::ple_history_advance(ssp->ple_hist, m.ple_norm, cs); } stats_.ms_ple += ms_since(tp); } From 4f88a2fa8bbc455150f47b9f71968027ac06e7d4 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 04:03:57 +0100 Subject: [PATCH 10/17] multi-prompt read: relaxed stage-chunk rule, sub-chunks at the smallest stage chunk; gated identity variants --- src/prefill/prefill.cpp | 10 +++++++--- tools/par/variants/co_lm3.txt | 3 +++ 2 files changed, 10 insertions(+), 3 deletions(-) create mode 100644 tools/par/variants/co_lm3.txt diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index ef6ae21f..5ce4fd9b 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -6008,7 +6008,7 @@ int64_t Prefill::lm_prepare(int64_t n, std::string& why, int64_t cap) { int64_t S = std::max(min_s, want / m.T * m.T); S = std::min(S, (n + m.T - 1) / m.T * m.T); if (cap > 0) { - if (cap % m.T != 0 || cap < min_s) { why = "the layer split's stages read different chunks"; return 0; } + if ((!lm_force_ && cap % m.T != 0) || cap < min_s) { why = "the layer split's stages read different chunks"; return 0; } S = std::min(S, cap); } const int64_t G = m.T * K; @@ -6233,6 +6233,10 @@ bool Prefill::lm_run(const std::vector& segs, int depth, const std::vector< // sub-chunks of m.T tokens from the piece's start; pieces of several segments are packed into one super-chunk while // their rows fit S. One segment: exactly the usual super-chunks. const int64_t S = m.lm.S; + // the sub-chunk size of a multi-prompt read: what every stage of the chain can take (the stages' chunks may differ) + int64_t tsub = m.T; + if (lm_force_) + for (const Prefill* q = this; q != nullptr; q = q->next_) tsub = std::min(tsub, q->impl_->T); std::vector> supers_own; const std::vector>* supers_p = &supers_own; std::vector> one_given; @@ -6250,11 +6254,11 @@ bool Prefill::lm_run(const std::vector& segs, int depth, const std::vector< cur.clear(); cur_rows = 0; } - for (int64_t c0s = 0; c0s < piece; c0s += m.T) { + for (int64_t c0s = 0; c0s < piece; c0s += tsub) { LmSub sb; sb.seg = (int) si; sb.c0 = p0s + c0s; - sb.T = std::min(m.T, piece - c0s); + sb.T = std::min(tsub, piece - c0s); sb.t0 = cur_rows + c0s; cur.push_back(sb); } diff --git a/tools/par/variants/co_lm3.txt b/tools/par/variants/co_lm3.txt new file mode 100644 index 00000000..a768ae88 --- /dev/null +++ b/tools/par/variants/co_lm3.txt @@ -0,0 +1,3 @@ +cog_ref | 0,1 | 8 | | STRATA_ROUTE_TAIL_SKIP=0;STRATA_PREFILL_CPU_SHARE=0;STRATA_PREFILL_LM_FORCE=1 | --expert-cache 1500 --adapt-every 100000 --pcie-frac 0 --suffix-draft 0 | w2:4@--nonce idg --keep-text --max-tokens 48 --rounds 2;w2:8@--nonce idg --keep-text --max-tokens 48 --rounds 2 +cog_ref2 | 0,1 | 8 | | STRATA_ROUTE_TAIL_SKIP=0;STRATA_PREFILL_CPU_SHARE=0;STRATA_PREFILL_LM_FORCE=1 | --expert-cache 1500 --adapt-every 100000 --pcie-frac 0 --suffix-draft 0 | w2:4@--nonce idg --keep-text --max-tokens 48 --rounds 2;w2:8@--nonce idg --keep-text --max-tokens 48 --rounds 2 +cog_on | 0,1 | 8 | | STRATA_ROUTE_TAIL_SKIP=0;STRATA_PREFILL_CPU_SHARE=0;STRATA_PREFILL_LM_FORCE=1;STRATA_COALESCE_PREFILL=8;STRATA_COALESCE_WAIT_MS=40;STRATA_CRASH_BT=/data/external/niko/par29/res/co/crash.txt | --expert-cache 1500 --adapt-every 100000 --pcie-frac 0 --suffix-draft 0 | w2:4@--nonce idg --keep-text --max-tokens 48 --rounds 2;w2:8@--nonce idg --keep-text --max-tokens 48 --rounds 2 From 301a09f10035185ab25701b5be2d4af954947520 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 03:52:28 +0100 Subject: [PATCH 11/17] STRATA_PREFILL_LM_FORCE (research reference for the multi-prompt read) --- src/prefill/prefill.cpp | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index 5ce4fd9b..310d62ae 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -5882,6 +5882,11 @@ const char* lm_env_text() { // STRATA_PREFILL_LAYER_MAJOR as set (nullptr: uns } bool lm_env_explicit() { const char* e = lm_env_text(); return e != nullptr && e[0] != '\0'; } +bool lm_force_env() { // STRATA_PREFILL_LM_FORCE=1 (research): the layer-major walk even for a prompt under two chunks (the identity reference of a multi-prompt read) + static const bool v = [] { const char* e = std::getenv("STRATA_PREFILL_LM_FORCE"); return e != nullptr && e[0] == '1'; }(); + return v; +} + bool lm_even_env() { // even super-chunks (STRATA_PREFILL_LM_EVEN=0: the largest that fits, then the rest) static const bool v = [] { const char* e = std::getenv("STRATA_PREFILL_LM_EVEN"); return !(e && e[0] == '0'); }(); return v; @@ -6002,7 +6007,7 @@ int64_t Prefill::lm_prepare(int64_t n, std::string& why, int64_t cap) { if (m.fused_bufs) { why = "the fused expert layout"; return 0; } if (!m.GU || !m.Dm || (n_mmq > 0 && !m.ids_identity)) { why = "the MoE buffers are missing"; return 0; } // (an explicit STRATA_PREFILL_LM_TOKENS may be one chunk: the A/B against the default path, experts aside) - const int64_t min_s = lm_force_ ? 1 : (lm_tokens_env() > 0 ? m.T : 2 * m.T); // (lm_force_: run_multi - several prompts, however short) + const int64_t min_s = (lm_force_ || lm_force_env()) ? 1 : (lm_tokens_env() > 0 ? m.T : 2 * m.T); // (lm_force_: run_multi - several prompts, however short) if (n < min_s) { why = "a short prompt"; return 0; } const int64_t want = lm_tokens_env() > 0 ? lm_tokens_env() : (lm_by_default_ ? kLmDefaultOnTokens : kLmDefaultTokens); int64_t S = std::max(min_s, want / m.T * m.T); From 1064ff7783112140c44ff7cbecd3b280906412b1 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 03:56:07 +0100 Subject: [PATCH 12/17] STRATA_COALESCE_WAIT_MS --- serve/server.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/serve/server.py b/serve/server.py index 5a77af9a..c809bfb4 100644 --- a/serve/server.py +++ b/serve/server.py @@ -342,6 +342,7 @@ ADMIT_MAX_WAIT_S = _timeout_env("STRATA_ADMIT_MAX_WAIT_S", 30.0) or 0.0 # read once for all of them, and each waiting request then finds its prompt in its slot. 0 / unset: off. COALESCE_PREFILL = max(0, int(os.environ.get("STRATA_COALESCE_PREFILL") or 0)) COALESCE_TOKENS = max(256, int(os.environ.get("STRATA_COALESCE_TOKENS") or 16384)) # prompt tokens read together at most +COALESCE_WAIT_MS = max(0, int(os.environ.get("STRATA_COALESCE_WAIT_MS") or 0)) # a burst's later requests register within this long STALL_CPU_EPS_S = 0.5 # CPU seconds the engine may use in the whole stall window and still count as frozen STALL_IO_EPS_B = 1 << 20 # bytes read or written in that window @@ -1809,6 +1810,8 @@ class StrataEngine: same layer-major pass as the BGEN's prompt. Each follower finds its slot in cancel._strata_preslot.""" picked = [] total = len(lead_prompt) + if COALESCE_WAIT_MS > 0: + time.sleep(COALESCE_WAIT_MS / 1000.0) # the requests of a burst reach the server a few ms apart with self.slot_cv: for e in list(self.wait_lens): if len(picked) >= COALESCE_PREFILL - 1: From 9f1d0ad0418e88f819bd57cc194eb5556c73d399 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 05:07:50 +0100 Subject: [PATCH 13/17] replicas: a replica that fails to start beside a living one is started again on the same shared arena before the sharing is given up --- serve/server.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/serve/server.py b/serve/server.py index c809bfb4..2d3bdc82 100644 --- a/serve/server.py +++ b/serve/server.py @@ -8586,6 +8586,16 @@ def start_replicas(cfg: dict, groups: list[list[int]], exe: str, effort_end: lis t.start() for t in ts: t.join() + if errors and arena and len(errors) < n: + # a replica that died while it was reading the experts into the shared arena (or any other one that failed) while + # another one is up: the living engine took the read over (the state word in the arena), so start the failed ones + # once more on the same arena instead of giving up the sharing + failed = sorted(i for i, _ in errors) + errors.clear() + print(f"[strata] replica(s) {', '.join(map(str, failed))} did not start; starting again on the shared expert arena", + flush=True) + for i in failed: + one(i) if errors: for e in engines: if e is not None: From 9d8337844eee741cd6525d14affc271b2b1beec8 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 05:24:08 +0100 Subject: [PATCH 14/17] replicas: the memfd arena warns when the RAM will not hold all the engines --- serve/replicas.py | 13 ++++++++++++- serve/test_replicas.py | 16 ++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/serve/replicas.py b/serve/replicas.py index 93309875..f4df027d 100644 --- a/serve/replicas.py +++ b/serve/replicas.py @@ -194,8 +194,19 @@ def shared_arena_plan(cfg: dict, n_replicas: int, *, shm: str = "/dev/shm", plat if fd >= 0: path = f"/proc/{os.getpid()}/fd/{fd}" _MEMFDS[path] = fd + warn = "" + try: # the arena once plus about 7 GB of private RAM per engine (measured); a warning, never a refusal + from strata_control import replica_profile as _rp + q = _rp.quant_of(cfg) + if q: + need = _rp.ARENA_GB[q] + n_replicas * _rp.PER_REPLICA_GB + if avail / 2**30 < need + 4: + warn = (f" WARNING: about {need:.0f} GB are needed for {n_replicas} engines and {avail / 2**30:.0f} GB " + "are free: the last engine may not fit (fewer replicas, or the layer split, is safer)") + except Exception: # noqa: BLE001 + pass return path, (f"one shared expert arena for {n_replicas} replicas (memory object of the server, not limited by " - "the size of /dev/shm): the experts are in RAM once") + f"the size of /dev/shm): the experts are in RAM once.{warn}") if not (os.path.isdir(shm) and os.access(shm, os.W_OK)): return None, f"shared expert arena: {shm} is not usable; each replica keeps its own copy of the experts in RAM" free = free_bytes if free_bytes is not None else _free(shm) diff --git a/serve/test_replicas.py b/serve/test_replicas.py index ae1d3e0c..5f236e6c 100644 --- a/serve/test_replicas.py +++ b/serve/test_replicas.py @@ -384,6 +384,22 @@ class SharedArena(unittest.TestCase): os.fstat(fd) remove_arena(path) # (twice: harmless) + @unittest.skipUnless(hasattr(__import__("os"), "memfd_create") and sys.platform.startswith("linux"), "Linux memfd") + def test_a_memfd_arena_warns_when_the_engines_will_not_all_fit(self): + cfg = {**self.CFG, "model_name": "qwen3.8-flash-next-iq3_s"} + path, note = shared_arena_plan(cfg, 2, platform="linux", environ={}, free_bytes=57 << 30, use_memfd=True) + try: + self.assertTrue(path) + self.assertIn("WARNING", note) + self.assertIn("may not fit", note) + finally: + remove_arena(path) + path, note = shared_arena_plan(cfg, 2, platform="linux", environ={}, free_bytes=120 << 30, use_memfd=True) + try: + self.assertNotIn("WARNING", note) + finally: + remove_arena(path) + def test_a_memfd_needs_free_ram(self): path, note = shared_arena_plan(self.CFG, 2, platform="linux", environ={}, free_bytes=2 << 30, use_memfd=True) self.assertIsNone(path) From 62b042e8436806c0cb3b0d5397748b9a9a056a92 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 05:58:22 +0100 Subject: [PATCH 15/17] co-read: adaptive - the leader waits and the read is shared only when the engine has seen at least 3 requests arrive together (STRATA_COALESCE_MIN_REQS, STRATA_COALESCE_BURST_MS) --- serve/server.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/serve/server.py b/serve/server.py index 2d3bdc82..c71da8c2 100644 --- a/serve/server.py +++ b/serve/server.py @@ -343,6 +343,12 @@ ADMIT_MAX_WAIT_S = _timeout_env("STRATA_ADMIT_MAX_WAIT_S", 30.0) or 0.0 COALESCE_PREFILL = max(0, int(os.environ.get("STRATA_COALESCE_PREFILL") or 0)) COALESCE_TOKENS = max(256, int(os.environ.get("STRATA_COALESCE_TOKENS") or 16384)) # prompt tokens read together at most COALESCE_WAIT_MS = max(0, int(os.environ.get("STRATA_COALESCE_WAIT_MS") or 0)) # a burst's later requests register within this long +# Adaptive: a read is shared (and the leader waits for its followers) only when this engine has seen at least +# COALESCE_MIN_REQS requests arrive within COALESCE_BURST_MS (the leader's own included). Fewer than that - one or two +# requests per engine, as 8 requests on 4 replicas - read plainly, with no wait: nothing to share, and the wait and the +# longer pass only cost (measured 4 x 1 replicas at 8 requests: 168 against 196 tok/s with an unconditional wait). +COALESCE_MIN_REQS = max(2, int(os.environ.get("STRATA_COALESCE_MIN_REQS") or 3)) +COALESCE_BURST_MS = max(10, int(os.environ.get("STRATA_COALESCE_BURST_MS") or 250)) STALL_CPU_EPS_S = 0.5 # CPU seconds the engine may use in the whole stall window and still count as frozen STALL_IO_EPS_B = 1 << 20 # bytes read or written in that window @@ -1201,6 +1207,7 @@ class StrataEngine: self.waiting = 0 # requests waiting for the control lines (ctl) self.admitting = 0 # #1603: ... and those that gave the lines back to wait for a slot self.wait_lens: list[list[int]] = [] # ... their prompt lengths (a long read gives way to short ones) + self.arrivals = collections.deque(maxlen=64) # monotonic times of the last requests that came for the control lines (STRATA_COALESCE_*) self.ctl_epoch = 0 # how often the control lines were taken self.ctl = threading.Lock() # one admission or solo request on the control lines at a time self.gen = self.__dict__.get("gen", 0) + 1 # which engine process this is (a request notes its own) @@ -1743,6 +1750,7 @@ class StrataEngine: with self.slot_cv: self.waiting += 1 self.wait_lens.append(entry) + self.arrivals.append(entry[1]) beat = time.monotonic() try: while True: @@ -1810,6 +1818,11 @@ class StrataEngine: same layer-major pass as the BGEN's prompt. Each follower finds its slot in cancel._strata_preslot.""" picked = [] total = len(lead_prompt) + now = time.monotonic() + with self.slot_cv: + recent = max(sum(1 for t in self.arrivals if now - t <= COALESCE_BURST_MS / 1000.0), 1 + len(self.wait_lens)) + if recent < COALESCE_MIN_REQS: + return # one or two requests on this engine: a plain read, no wait if COALESCE_WAIT_MS > 0: time.sleep(COALESCE_WAIT_MS / 1000.0) # the requests of a burst reach the server a few ms apart with self.slot_cv: From 7aafac422c0c73af8d0cffb4421b33e55de597e1 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 06:02:30 +0100 Subject: [PATCH 16/17] shared arena on Windows: a named pagefile-backed section (CreateFileMapping/MapViewOfFile) with the same header and load state word, pid election on both OSes, server option enabled, pinned_shared_win_test --- CMakeLists.txt | 7 ++ serve/replicas.py | 11 ++- serve/test_replicas.py | 8 +- src/core/expert_source.cpp | 27 +++++- src/core/pinned.cu | 118 ++++++++++++++++++++++++- src/core/pinned_shared_win_test.cpp | 107 ++++++++++++++++++++++ strata_control/replica_profile.py | 2 +- tests/terminal/test_replica_profile.py | 12 ++- 8 files changed, 278 insertions(+), 14 deletions(-) create mode 100644 src/core/pinned_shared_win_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index ee7e9f1b..076249c6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1638,6 +1638,13 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) target_link_libraries(pinned_shared_test PRIVATE strata_core ${_strata_gpu_runtime_target}) add_test(NAME pinned_shared_test COMMAND pinned_shared_test) endif() + if(WIN32) # the replicas' shared expert arena as a named section (src/core/pinned.cu): two arenas, a child process, the last handle + add_executable(pinned_shared_win_test src/core/pinned_shared_win_test.cpp) + target_include_directories(pinned_shared_win_test PRIVATE include) + target_link_libraries(pinned_shared_win_test PRIVATE strata_core ${_strata_gpu_runtime_target}) + add_test(NAME pinned_shared_win_test COMMAND pinned_shared_win_test) + set_tests_properties(pinned_shared_win_test PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 180) + endif() # Plan v0.3 P2: the n-gram table read straight from the SSD. The selftest needs only the CPU and a temp file. add_executable(ple_reader_test src/ngram/ple_reader_test.cpp) diff --git a/serve/replicas.py b/serve/replicas.py index f4df027d..bf1723bf 100644 --- a/serve/replicas.py +++ b/serve/replicas.py @@ -23,7 +23,7 @@ replicas need two copies and four need four. "replica_shared_arena": true (or S backs the experts' arena of all replicas with ONE shared memory object (the engine's --shared-expert-arena, MAP_SHARED): the experts are in RAM once, whatever the number of replicas. On Linux the object is a memfd the server creates (os.memfd_create): it is not a file on a mount, so the size of /dev/shm does not limit it; the engines open it as -/proc//fd/. Without memfd (an old kernel, a sandbox that refuses it) a file in /dev/shm is used instead. +/proc//fd/. On Windows it is a named section (CreateFileMapping) the engines open by name. Without memfd (an old kernel, a sandbox that refuses it) a file in /dev/shm is used instead. The pack hash in the header refuses another model; the first engine to start reads the experts into it and the others wait for it and skip the read (a state word in the header: nobody / a pid reading / filled; an engine whose reader died takes the read over); the server closes the memfd (or unlinks the file) when it ends - an engine that still maps it @@ -171,8 +171,13 @@ def shared_arena_plan(cfg: dict, n_replicas: int, *, shm: str = "/dev/shm", plat want = str(environ.get("STRATA_REPLICA_SHARED_ARENA", "")).strip().lower() in ("1", "true", "on", "yes") if not want or n_replicas < 2: return None, "" + if platform.startswith("win"): + # a named pagefile-backed section the engines open by name (src/core/pinned.cu); it lives while any engine holds it + name = f"Local{chr(92)}strata-replicas-{os.getpid()}" + return f"win:{name}", (f"one shared expert arena for {n_replicas} replicas (a Windows section, {name}): " + "the experts are in RAM once") if not platform.startswith("linux"): - return None, "shared expert arena: Linux only; each replica keeps its own copy of the experts in RAM" + return None, "shared expert arena: Linux and Windows only; each replica keeps its own copy of the experts in RAM" args = [str(a) for a in cfg.get("args") or []] if isinstance(cfg.get("replicas"), list): for e in cfg["replicas"]: @@ -222,7 +227,7 @@ def shared_arena_plan(cfg: dict, n_replicas: int, *, shm: str = "/dev/shm", plat def remove_arena(path: str | None) -> None: """Close the memfd / unlink the arena file (the server's end, or a start that fell back): engines that map it keep their pages.""" - if not path: + if not path or path.startswith("win:"): # (a Windows section goes with its last handle) return fd = _MEMFDS.pop(path, None) if fd is not None: diff --git a/serve/test_replicas.py b/serve/test_replicas.py index 5f236e6c..a85e2c8a 100644 --- a/serve/test_replicas.py +++ b/serve/test_replicas.py @@ -335,9 +335,15 @@ class SharedArena(unittest.TestCase): self.assertTrue(path2) self.assertEqual(path, path2) # (one server: one file) + def test_windows_gets_a_named_section(self): + with tempfile.TemporaryDirectory() as d: + path, note = self.plan(tmp=d, platform="win32") + self.assertTrue(path and path.startswith("win:Local") and "strata-replicas-" in path, path) + self.assertIn("Windows section", note) + remove_arena(path) # (nothing to remove: harmless) + def test_only_where_there_is_an_arena_to_share(self): with tempfile.TemporaryDirectory() as d: - self.assertIsNone(self.plan(tmp=d, platform="win32")[0]) self.assertIsNone(self.plan(tmp=d, platform="darwin")[0]) for flag in ("--mmap-experts", "--resident-experts", "--resident-budget-gib", "--shared-expert-arena"): self.assertIsNone(self.plan({**self.CFG, "args": ["--pack", "/p", flag]}, tmp=d)[0], flag) diff --git a/src/core/expert_source.cpp b/src/core/expert_source.cpp index 3eee2116..f4d41e18 100644 --- a/src/core/expert_source.cpp +++ b/src/core/expert_source.cpp @@ -76,6 +76,27 @@ #endif #endif +// The pid of this process and whether a pid still exists (the shared arena's reader election, both OSes) +static uint64_t arena_this_pid() { +#ifdef _WIN32 + return (uint64_t) GetCurrentProcessId(); +#else + return (uint64_t) getpid(); +#endif +} +static bool arena_pid_alive(uint64_t pid) { +#ifdef _WIN32 + HANDLE h = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, (DWORD) pid); + if (h == nullptr) return GetLastError() == ERROR_ACCESS_DENIED; // (denied: it exists; invalid parameter: it does not) + DWORD code = 0; + const bool alive = GetExitCodeProcess(h, &code) && code == STILL_ACTIVE; + CloseHandle(h); + return alive; +#else + return !(kill((pid_t) pid, 0) != 0 && errno == ESRCH); +#endif +} + namespace strata::core { namespace detail { @@ -6011,9 +6032,8 @@ bool ArenaExpertSource::open(const std::string& pack_dir, int64_t n_layers, int6 constexpr uint64_t kArenaFilled = ~uint64_t{0}; std::atomic* arena_state = padded_ ? nullptr : a->shared_state; bool arena_leader = false, arena_skip = false; -#ifndef _WIN32 if (arena_state != nullptr) { - const uint64_t me = (uint64_t) getpid(); + const uint64_t me = arena_this_pid(); const auto t_begin = std::chrono::steady_clock::now(); bool said = false; for (;;) { @@ -6023,7 +6043,7 @@ bool ArenaExpertSource::open(const std::string& pack_dir, int64_t n_layers, int6 if (arena_state->compare_exchange_strong(cur, me, std::memory_order_acq_rel)) { arena_leader = true; break; } continue; } - const bool alive = !(kill((pid_t) cur, 0) != 0 && errno == ESRCH); + const bool alive = arena_pid_alive(cur); if (!alive) { if (arena_state->compare_exchange_strong(cur, me, std::memory_order_acq_rel)) { std::fprintf(stderr, "strata generate: shared expert arena: engine %llu died while reading it; reading it again\n", @@ -6042,7 +6062,6 @@ bool ArenaExpertSource::open(const std::string& pack_dir, int64_t n_layers, int6 std::this_thread::sleep_for(std::chrono::milliseconds(200)); } } -#endif std::atomic ready{0}; std::thread reg; if (deferred) { diff --git a/src/core/pinned.cu b/src/core/pinned.cu index 468971c4..45be4abb 100644 --- a/src/core/pinned.cu +++ b/src/core/pinned.cu @@ -58,6 +58,12 @@ struct SharedArenaHeader { }; static_assert(sizeof(SharedArenaHeader) <= kSharedArenaHeaderBytes); +#ifdef _WIN32 +// Shared-section views: MapViewOfFile address -> the section handle, so release() unmaps and closes instead of VirtualFree. +std::mutex g_views_mu; +std::map g_views; +#endif + #ifndef _WIN32 // MemAvailable and, in a cgroup v2 group with a limit, the room left under it; ~0 when neither can be read uint64_t free_ram_room() { @@ -137,8 +143,102 @@ void* reserve(uint64_t bytes, PageBacking& got, std::string& note, const std::st #ifdef _WIN32 if (!shared_file.empty()) { - note = "shared-file arena backing is not implemented on Windows"; - return nullptr; + // The shared arena on Windows is a named, pagefile-backed section (CreateFileMappingW) that every engine of the + // replicas opens by name: "win:Local\\strata-replicas-" (a name without the prefix is used as it is). + // The first engine creates it at the full size, the others open it; it lives while any process holds a handle or + // a view, so a replica that restarts joins the same pages. The first 4 KiB are the same header as on Linux (the + // pack hash, the size, the load state word); a new section is all zeros, so "no magic" means "nobody has been + // here" - two engines that both find it empty write the same header. + if (shared_pack_hash == 0) { + note = "shared arena requires a nonzero pack hash"; + return nullptr; + } + if (bytes > UINT64_MAX - kSharedArenaHeaderBytes) { + note = "shared arena size overflows its header"; + return nullptr; + } + const uint64_t file_bytes = bytes + kSharedArenaHeaderBytes; + std::string nm = shared_file.rfind("win:", 0) == 0 ? shared_file.substr(4) : shared_file; + std::wstring wname(nm.begin(), nm.end()); + // the section is committed memory (RAM plus pagefile) from the moment it exists: say so before taking it + MEMORYSTATUSEX ms{}; + ms.dwLength = sizeof ms; + if (GlobalMemoryStatusEx(&ms) && file_bytes + (2ull << 30) > ms.ullAvailPageFile) { + HANDLE probe = OpenFileMappingW(FILE_MAP_ALL_ACCESS, FALSE, wname.c_str()); // (already there: nothing new is committed) + if (probe == nullptr) { + note = "not enough free memory (RAM + pagefile) for the shared arena " + nm + ": it needs " + + std::to_string((unsigned long long) (file_bytes >> 20)) + " MiB and " + + std::to_string((unsigned long long) (ms.ullAvailPageFile >> 20)) + " MiB are available (2 GiB are kept back)"; + return nullptr; + } + CloseHandle(probe); + } + HANDLE h = CreateFileMappingW(INVALID_HANDLE_VALUE, nullptr, PAGE_READWRITE, (DWORD) (file_bytes >> 32), + (DWORD) (file_bytes & 0xffffffffu), wname.c_str()); + if (h == nullptr) { + note = "cannot create shared arena section " + nm + ": Windows error " + std::to_string((unsigned long long) GetLastError()); + return nullptr; + } + void* view = MapViewOfFile(h, FILE_MAP_ALL_ACCESS, 0, 0, 0); // (0: the whole section, as big as it was created) + if (view == nullptr) { + const DWORD e = GetLastError(); + CloseHandle(h); + note = "cannot map shared arena section " + nm + ": Windows error " + std::to_string((unsigned long long) e); + return nullptr; + } + MEMORY_BASIC_INFORMATION mbi{}; + if (VirtualQuery(view, &mbi, sizeof mbi) == 0 || (uint64_t) mbi.RegionSize < file_bytes) { + UnmapViewOfFile(view); + CloseHandle(h); + note = "shared arena section " + nm + " is smaller than the " + std::to_string((unsigned long long) file_bytes) + + " B this engine needs (including its header)"; + return nullptr; + } + auto* hdr = static_cast(view); + bool empty = true; + for (char c : hdr->magic) empty = empty && c == 0; + if (empty) { + SharedArenaHeader fresh{}; + std::memcpy(fresh.magic, kSharedArenaMagic, sizeof fresh.magic); + fresh.version = 1; + fresh.header_bytes = (uint32_t) kSharedArenaHeaderBytes; + fresh.arena_bytes = bytes; + fresh.pack_hash = shared_pack_hash; + // fields after the magic first, the magic last: an engine that sees the magic sees a complete header (the + // state word, reserved[0], stays zero) + std::memcpy(reinterpret_cast(hdr) + sizeof fresh.magic, reinterpret_cast(&fresh) + sizeof fresh.magic, + sizeof fresh - sizeof fresh.magic); + MemoryBarrier(); + std::memcpy(hdr->magic, fresh.magic, sizeof fresh.magic); + } else { + // (an engine that arrives in the instant between another one's header fields and its magic never sees a + // half header: the magic is written last) + if (std::memcmp(hdr->magic, kSharedArenaMagic, sizeof hdr->magic) != 0 || hdr->version != 1 || + hdr->header_bytes != kSharedArenaHeaderBytes || hdr->arena_bytes != bytes) { + UnmapViewOfFile(view); + CloseHandle(h); + note = "shared arena section " + nm + " has an incompatible or missing header"; + return nullptr; + } + if (hdr->pack_hash != shared_pack_hash) { + char b[256]; + std::snprintf(b, sizeof b, "shared arena %s was written for pack hash %016llx, expected %016llx", nm.c_str(), + (unsigned long long) hdr->pack_hash, (unsigned long long) shared_pack_hash); + UnmapViewOfFile(view); + CloseHandle(h); + note = b; + return nullptr; + } + } + { + std::lock_guard lk(g_views_mu); + g_views[view] = h; + } + mapping_base = view; + mapping_bytes = file_bytes; + got = PageBacking::NormalPages; + note = "shared section " + nm + " (pack hash checked)"; + return static_cast(view) + kSharedArenaHeaderBytes; } // MEM_LARGE_PAGES needs SeLockMemoryPrivilege. Having it assigned to the account is not enough: the // PROCESS must enable it in its own token (AdjustTokenPrivileges) before VirtualAlloc, or the call fails. @@ -413,6 +513,20 @@ void release(void* p, uint64_t bytes) { if (!p) return; #ifdef _WIN32 (void) bytes; + HANDLE section = nullptr; + { + std::lock_guard lk(g_views_mu); + auto it = g_views.find(p); + if (it != g_views.end()) { + section = it->second; + g_views.erase(it); + } + } + if (section != nullptr) { // a shared-section view: the pages stay while another engine has the section + UnmapViewOfFile(p); + CloseHandle(section); + return; + } VirtualFree(p, 0, MEM_RELEASE); #else munmap(p, bytes); diff --git a/src/core/pinned_shared_win_test.cpp b/src/core/pinned_shared_win_test.cpp new file mode 100644 index 00000000..006a18d4 --- /dev/null +++ b/src/core/pinned_shared_win_test.cpp @@ -0,0 +1,107 @@ +// Windows selftest for the named-section backing of PinnedArena (the replicas' shared expert arena). +// Parent: two arenas on one section name (same bytes, header refused for another pack / a bigger size, the load state word), +// then a CHILD PROCESS that opens the same section, reads what the parent wrote and writes back; the parent closes its first +// arena while the second still holds the section (the pages must stay), and a third arena joins the live section. After the +// last close, a new arena with the same name starts empty (the section is gone with its last handle). +#ifdef _WIN32 +#include "strata/core/pinned.hpp" + +#include + +#include +#include +#include +#include +#include + +namespace { +constexpr uint64_t kBytes = 64u << 10; +constexpr uint64_t kPackHash = 0x0123456789abcdefull; +} // namespace + +int main(int argc, char** argv) { + const std::vector no_bounds; + if (argc >= 3 && std::strcmp(argv[1], "child") == 0) { + const std::string name = argv[2]; + strata::core::PinnedArena a(kBytes, no_bounds, 0, name, kPackHash); + if (!a.valid()) { std::fprintf(stderr, "child: cannot open: %s\n", a.note.c_str()); return 2; } + if (a.data()[123] != 0x5a) { std::fprintf(stderr, "child: did not see the parent's byte\n"); return 3; } + if (a.shared_state == nullptr || a.shared_state->load() != 4242u) { std::fprintf(stderr, "child: state word is not the parent's\n"); return 4; } + a.data()[456] = 0xa5; + a.shared_state->store(~uint64_t{0}); // (filled) + return 0; + } + + int fail = 0; + const std::string name = "win:Local\\strata-pinned-test-" + std::to_string((unsigned long long) GetCurrentProcessId()); + { + auto* first = new strata::core::PinnedArena(kBytes, no_bounds, 0, name, kPackHash); + if (!first->valid()) { + std::fprintf(stderr, "first section failed: %s\n", first->note.c_str()); + return 1; + } + first->data()[123] = 0x5a; + if (first->shared_state == nullptr) { std::fprintf(stderr, "no state word\n"); fail = 1; } + else if (first->shared_state->load() != 0) { std::fprintf(stderr, "a new section's state word is not 0\n"); fail = 1; } + else first->shared_state->store(4242u); + + auto* second = new strata::core::PinnedArena(kBytes, no_bounds, 0, name, kPackHash); + if (!second->valid()) { std::fprintf(stderr, "second section failed: %s\n", second->note.c_str()); fail = 1; } + else { + if (second->data()[123] != 0x5a) { std::fprintf(stderr, "same name did not expose the same bytes\n"); fail = 1; } + if (second->shared_state == nullptr || second->shared_state->load() != 4242u) { std::fprintf(stderr, "state word not shared\n"); fail = 1; } + } + { + strata::core::PinnedArena wrong_pack(kBytes, no_bounds, 0, name, kPackHash + 1); + if (wrong_pack.valid() || wrong_pack.note.find("pack hash") == std::string::npos) { + std::fprintf(stderr, "pack mismatch was not refused: %s\n", wrong_pack.note.c_str()); + fail = 1; + } + strata::core::PinnedArena wrong_size(kBytes * 2, no_bounds, 0, name, kPackHash); + if (wrong_size.valid()) { std::fprintf(stderr, "a bigger size was not refused\n"); fail = 1; } + } + + // a child process opens the live section + char exe[MAX_PATH]; + GetModuleFileNameA(nullptr, exe, MAX_PATH); + std::string cmd = std::string("\"") + exe + "\" child \"" + name + "\""; + STARTUPINFOA si{}; + si.cb = sizeof si; + PROCESS_INFORMATION pi{}; + std::vector cl(cmd.begin(), cmd.end()); + cl.push_back('\0'); + if (!CreateProcessA(nullptr, cl.data(), nullptr, nullptr, FALSE, 0, nullptr, nullptr, &si, &pi)) { + std::fprintf(stderr, "cannot start the child (error %lu)\n", GetLastError()); + fail = 1; + } else { + WaitForSingleObject(pi.hProcess, 120000); + DWORD code = 99; + GetExitCodeProcess(pi.hProcess, &code); + CloseHandle(pi.hProcess); + CloseHandle(pi.hThread); + if (code != 0) { std::fprintf(stderr, "the child failed with %lu\n", code); fail = 1; } + if (first->data()[456] != 0xa5) { std::fprintf(stderr, "the child's byte did not arrive\n"); fail = 1; } + if (first->shared_state->load() != ~uint64_t{0}) { std::fprintf(stderr, "the child's state word did not arrive\n"); fail = 1; } + } + + delete first; // the creator is gone; the second still holds the section + if (second != nullptr && second->valid()) { + if (second->data()[123] != 0x5a || second->data()[456] != 0xa5) { std::fprintf(stderr, "bytes lost when the first arena closed\n"); fail = 1; } + strata::core::PinnedArena third(kBytes, no_bounds, 0, name, kPackHash); + if (!third.valid() || third.data()[123] != 0x5a) { std::fprintf(stderr, "a late arena did not join the live section\n"); fail = 1; } + } + delete second; + } + { + strata::core::PinnedArena again(kBytes, no_bounds, 0, name, kPackHash); // (every handle was closed: a new, empty section) + if (!again.valid() || again.data()[123] != 0 || again.shared_state->load() != 0) { + std::fprintf(stderr, "the section survived its last handle\n"); + fail = 1; + } + } + std::printf("pinned_shared_win_test: %s\n", fail ? "FAILED" : "OK"); + return fail; +} +#else +int main() { return 77; } // (skipped: the Linux build has pinned_shared_test) +#endif diff --git a/strata_control/replica_profile.py b/strata_control/replica_profile.py index f7a3b516..78a4bf84 100644 --- a/strata_control/replica_profile.py +++ b/strata_control/replica_profile.py @@ -60,7 +60,7 @@ def quant_of(cfg: dict) -> str | None: def shared_arena_possible(platform: str | None = None) -> bool: - return (sys.platform if platform is None else platform).startswith("linux") + return (sys.platform if platform is None else platform).startswith(("linux", "win")) def replicas_for(n_cards: int, arena_gb: float, ram_gb: float, shared: bool) -> int: diff --git a/tests/terminal/test_replica_profile.py b/tests/terminal/test_replica_profile.py index fa845126..1bbff576 100644 --- a/tests/terminal/test_replica_profile.py +++ b/tests/terminal/test_replica_profile.py @@ -80,12 +80,18 @@ class Apply(unittest.TestCase): self.assertEqual(c["layer_split"], "auto") self.assertTrue(notes and "layer split stays" in notes[0]) - def test_windows_cannot_share_so_the_private_copies_must_fit(self): + def test_windows_shares_through_a_section(self): c = cfg(gpu=[0, 1]) - self.assertEqual(P.apply(cards(2), c, 100.0, platform="win32", environ={})[0].split(":")[0], "replica profile") + P.apply(cards(2), c, 80.0, platform="win32", environ={}) + self.assertEqual(len(c["replicas"]), 2) + self.assertIs(c["replica_shared_arena"], True) + + def test_an_os_that_cannot_share_needs_the_private_copies_to_fit(self): + c = cfg(gpu=[0, 1]) + self.assertEqual(P.apply(cards(2), c, 100.0, platform="darwin", environ={})[0].split(":")[0], "replica profile") self.assertNotIn("replicas", c) c = cfg(gpu=[0, 1]) - P.apply(cards(2), c, 160.0, platform="win32", environ={}) + P.apply(cards(2), c, 160.0, platform="darwin", environ={}) self.assertEqual(len(c["replicas"]), 2) self.assertNotIn("replica_shared_arena", c) From 4772f0cd9acd81bb2958a9ebd3c9a8a7a326ec6b Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 06:21:21 +0100 Subject: [PATCH 17/17] co-read adaptive: look 10 ms for the burst's other arrivals before deciding --- serve/server.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/serve/server.py b/serve/server.py index c71da8c2..9b85342d 100644 --- a/serve/server.py +++ b/serve/server.py @@ -1818,11 +1818,16 @@ class StrataEngine: same layer-major pass as the BGEN's prompt. Each follower finds its slot in cancel._strata_preslot.""" picked = [] total = len(lead_prompt) - now = time.monotonic() - with self.slot_cv: - recent = max(sum(1 for t in self.arrivals if now - t <= COALESCE_BURST_MS / 1000.0), 1 + len(self.wait_lens)) + def seen() -> int: + now = time.monotonic() + with self.slot_cv: + return max(sum(1 for t in self.arrivals if now - t <= COALESCE_BURST_MS / 1000.0), 1 + len(self.wait_lens)) + recent = seen() + if recent < COALESCE_MIN_REQS and COALESCE_WAIT_MS > 0: + time.sleep(0.010) # (the first request of a burst holds the lines before its neighbours have arrived) + recent = seen() if recent < COALESCE_MIN_REQS: - return # one or two requests on this engine: a plain read, no wait + return # one or two requests on this engine: a plain read, at most 10 ms lost if COALESCE_WAIT_MS > 0: time.sleep(COALESCE_WAIT_MS / 1000.0) # the requests of a burst reach the server a few ms apart with self.slot_cv: