From 2b65710f934726626ed7644a69f28d9a3c0126c5 Mon Sep 17 00:00:00 2001 From: Vop3p <69450118+Vop3p@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:51:07 -0400 Subject: [PATCH] prefill: STRATA_PREFILL_CPU_SHARE_BATCH=1 (opt-in) gives the prompt path the CPU pool beside batch slots, for a read with no slot decoding The share was off whenever --batch was set (share_pool = batch <= 0). With the variable set the pool goes to the prompt path as without batch slots, and before each read the serve loop sets Prefill::set_share_blocked from the slots: a read with a slot decoding takes no share (the slots' windows use the pool), one with none does. Co-Authored-By: Claude Opus 5.5 (1M context) --- include/strata/prefill/prefill.hpp | 3 +++ src/prefill/prefill.cpp | 6 +++++- src/program/generate.cpp | 23 ++++++++++++++++++++++- 3 files changed, 30 insertions(+), 2 deletions(-) diff --git a/include/strata/prefill/prefill.hpp b/include/strata/prefill/prefill.hpp index 7aadc07269..5081888717 100644 --- a/include/strata/prefill/prefill.hpp +++ b/include/strata/prefill/prefill.hpp @@ -245,6 +245,9 @@ class Prefill { /// `by_default`: this path is one the share was measured on (CUDA, one GPU, no batch slots), so with the variable unset /// the share is on for chunks below 1024 tokens (STRATA_PREFILL_CPU_SHARE=0 turns it off). static void arm_cpu_share(bool applies, bool by_default = false); + /// STRATA_PREFILL_CPU_SHARE_BATCH=1 (batch slots beside the pool): set before each read - true while a slot decodes + /// (its windows use the pool), so no chunk of that read takes the share; false lets the read's chunks take it. + static void set_share_blocked(bool blocked); private: static uint64_t bytes_needed_impl(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk, diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index f00f1cc134..1fcdd34274 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -226,6 +226,8 @@ inline int64_t cpu_share_max() { // On a layer split every stage may have the pool (set_cpu_pool), and the stages read different chunks at the same time: // one stage at a time takes it for a chunk, the other reads its chunk with the GPU alone. std::atomic g_cpu_pool_busy{false}; +// STRATA_PREFILL_CPU_SHARE_BATCH: a slot is decoding beside this read (Prefill::set_share_blocked): the pool is its windows' +std::atomic g_share_blocked{false}; double g_pinned_share = 1.0; // The streamed ring: 384 slots when (nearly) every streamed expert is DMA'd from pinned RAM - measured on Q2_0, // 8192-token chunks: 96 slots 1153 tok/s, 384 1294 (the next layer's experts arrive during its attention half) - @@ -2385,6 +2387,7 @@ void Prefill::arm_cpu_share(bool applies, bool by_default) { if (applies && cpu_share_on()) g_stream_min_share = cpu_share_max(); } +void Prefill::set_share_blocked(bool blocked) { g_share_blocked.store(blocked, std::memory_order_release); } bool Prefill::ring_bytes_enabled() { return ring_bytes_on(); } size_t Prefill::slot_guard_bytes() { return mmq::direct_slot_guard(); } @@ -2849,7 +2852,8 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri bool held = false; ~PoolHold() { if (held) g_cpu_pool_busy.store(false, std::memory_order_release); } } pool_hold; - if (cpu_pool_ != nullptr && cpu_share_on() && !stream_all && !pipe_chunked) { + if (cpu_pool_ != nullptr && cpu_share_on() && !stream_all && !pipe_chunked && + !g_share_blocked.load(std::memory_order_acquire)) { bool expect = false; pool_hold.held = g_cpu_pool_busy.compare_exchange_strong(expect, true, std::memory_order_acq_rel); } diff --git a/src/program/generate.cpp b/src/program/generate.cpp index 01b99f06cf..3e770f284e 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -8108,7 +8108,13 @@ int main(int argc, char** argv) { } } #endif - const bool share_pool = o.batch <= 0 && !o.no_pool; + // STRATA_PREFILL_CPU_SHARE_BATCH=1 (opt-in): the prompt path gets the pool beside batch slots too; a read takes the + // share only while no slot decodes (set_share_blocked before each read), since the slots' windows use the pool + static const bool share_batch = [] { const char* v = std::getenv("STRATA_PREFILL_CPU_SHARE_BATCH"); return v != nullptr && v[0] == '1'; }(); + const bool share_pool = (o.batch <= 0 || share_batch) && !o.no_pool; + if (share_batch && o.batch > 0 && !o.no_pool) + std::fprintf(stderr, "strata serve: STRATA_PREFILL_CPU_SHARE_BATCH=1: the prompt path may use the CPU pool beside the " + "%lld batch slots, for a read with no slot decoding\n", (long long) o.batch); bool share_has_work = strata::prefill::share_rules::default_share_has_work( (int64_t) xcache.slots(), (int64_t) (g.n_layers * g.n_expert)); // #1595 if (!share_has_work) @@ -13068,6 +13074,19 @@ int main(int argc, char** argv) { // the request ends with `YIELDED ` + DONE cancel, and the server sends it again later: it // continues from the slot with the same chunks. #656's cooperative preemption, with a slot as the park. int64_t req_chunk = 0, req_chunk_len = 0; // the DONE line's chunk=: the longest batched segment's first chunk + + // STRATA_PREFILL_CPU_SHARE_BATCH: the read takes the CPU share only with no slot decoding (their windows use the pool) + auto share_gate = [&]() { + int busy = 0; + for (const auto& x : bs) busy += x.active ? 1 : 0; + strata::prefill::Prefill::set_share_blocked(busy > 0); + static int said = 0; + static const bool on = [] { const char* v = std::getenv("STRATA_PREFILL_CPU_SHARE_BATCH"); return v != nullptr && v[0] == '1'; }(); + if (on && said < 4) { + ++said; + std::fprintf(stderr, "strata batch: CPU share for this read: %s (%d slot(s) decoding)\n", busy > 0 ? "off" : "on", busy); + } + }; auto read_part = [&](int64_t a0, int64_t b0, std::string& e) -> bool { // --prefill-pipe on a layer split: the segment's chunk is planned ONCE, on its whole length (or taken // from the read it goes on from: resume_plan), and every run of it reads in that chunk (force_chunk) - @@ -13115,6 +13134,7 @@ int main(int argc, char** argv) { } seg_plan_scope{seg_plan}; auto run = [&](int64_t q, int64_t r) -> bool { if (planned) sp.force_chunk(plan.chunk, plan_single, plan.chunked); + share_gate(); return sp.run(ids.data() + q, r - q, q, e); }; // (a layer split reads its stages as a pipeline over one run's chunks: beside slots a chunk at a time, with @@ -13339,6 +13359,7 @@ int main(int argc, char** argv) { std::string read_err; in_prompt_read = true; sp.stop_at_chunk = [&] { return want_yield.load(std::memory_order_acquire) >= 0; }; + share_gate(); // on this thread, before the reader starts std::thread reader([&] { conc_thread_cpus(); try { // (an exception leaving a std::thread ends the process: a failed read instead)