prefill: STRATA_PREFILL_CPU_SHARE=auto hands a small chunk's least-routed experts to the idle CPU pool

A chunk below stream_all_min() (1,024 tokens: an agent's tool result, a test's
output) streams every routed expert outside VRAM over PCIe while the decode CPU
pool sits idle and the arena holds the same experts in RAM.  Opt-in,
STRATA_PREFILL_CPU_SHARE=auto|x:
- the non-resident experts routed by at most MAXT of the chunk's tokens,
  fewest first, run on the pool on a thread while the GPU streams and runs the
  rest; their unweighted rows go to Dm's tail (as a peer GPU's do) and the
  combine weights them as any other row;
- auto: the share measured, each layer's CPU time per expert and the GPU's per
  streamed expert (CUDA events around its expert work, read after the next
  layer's routing sync), and the next layers hand the CPU g / (c + g), where
  both end together - a small CPU takes a small share;
- Q2_0 layers take the pool's Q2_0 kernels' activations (ActQ, as decode's),
  i-quant layers the native vec_dot type.

Serve without batch slots and generate, one GPU.  Unset (the default): no
expert moves, byte-identical to v0.1.39 (IQ2_XS 600 and 2,100 tokens, fixed
cache, logits and 32 tokens).

RTX 5090 + 9950X3D, 600-token prompts, ms (off / auto):
- UD-Q4_K_XL: 1,548 / 1,369, 1,528 / 1,296 (share 0.65; KL between 0.006);
- Q2_0: 472 / 445, 491 / 448;
- IQ2_XS: 493 / 481, 502 / 460, 527 / 529; 250 tokens 424 / 355, 402 / 382,
  379 / 367 (KL to STRATA_PREFILL_MMQ=0 0.0053 off, 0.0080 auto);
- IQ2_XS on 2 AVX2 workers (emulated): 512 / 487 (share 0.49).
This commit is contained in:
sergiywith authored and Niko1221 committed 2026-10-07 11:51:33 +02:00
1 parent 9476ef1ade
commit 978d355852
3 files changed
+196 -4

No files matched your search

+12
View File
@@ -23,6 +23,7 @@
#include <string>
namespace strata::core { class PeerExperts; }
namespace strata::kernels::cpu { class ExpertPool; }
namespace strata::prefill {
@@ -34,6 +35,8 @@ struct PrefillStats {
int64_t experts_streamed = 0; ///< expert blobs copied host -> device
int64_t experts_dma = 0; ///< ...of which straight from the pinned arena (no CPU copy)
int64_t experts_resident = 0; ///< expert-layer groups served from the VRAM tier
int64_t experts_cpu = 0; ///< ...computed on the CPU pool instead of streamed (STRATA_PREFILL_CPU_SHARE)
double cpu_share = 0; ///< ...the share of the streamed ones it took last (measured by default)
double ms_ple = 0;
};
@@ -162,9 +165,18 @@ public:
/// run `init`.
bool set_stage_helper(Prefill* helper, std::string& err);
/// The CPU expert pool (decode's, idle while a prompt is read). With STRATA_PREFILL_CPU_SHARE set, a chunk below
/// stream_all_min() tokens - an agent's tool output - hands it the non-resident experts routed by at most MAXT of
/// its tokens, fewest first, up to a share of the experts it would stream (`auto`: measured, where both sides end
/// together). The CPU reads them from RAM while the rest come over PCIe; their rows go to Dm's tail, as a peer's
/// do. Not bit-identical to the GPU's rows (the CPU's own activation format). Unset (default) or null: every
/// expert on the GPU. Only for a pool no other thread runs meanwhile (no batch slots). Set before `init`.
void set_cpu_pool(kernels::cpu::ExpertPool* pool);
private:
static uint64_t bytes_needed_impl(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk,
bool owned_pages);
kernels::cpu::ExpertPool* cpu_pool_ = nullptr; ///< set_cpu_pool
// Stage-1 pipeline: intermediate stages return after handing their chunk to
// the direct successor. The public run() drains the chain once at prompt end.
bool run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err);
+178 -2
View File
@@ -9,6 +9,8 @@
#include "strata/core/layout.hpp"
#include "strata/kernels/cpu/expert.hpp"
#include "strata/kernels/native_qsa.hpp"
#include "strata/kernels/cpu/native_expert.hpp"
#include "strata/kernels/cpu/pool.hpp"
#include "strata/kernels/native_qsa_indexer.hpp"
#include "strata/kernels/ngram.hpp"
#include "strata/kernels/ple.hpp"
@@ -138,6 +140,24 @@ inline int64_t stream_all_min() {
static const int64_t v = [] { const char* e = std::getenv("STRATA_PREFILL_STREAM_MIN"); return e ? (int64_t) std::atoll(e) : (int64_t) 1024; }();
return v;
}
// STRATA_PREFILL_CPU_SHARE (opt-in; set_cpu_pool): a chunk below stream_all_min() hands the decode CPU pool - idle
// while a prompt is read - the non-resident experts few of its tokens route to, instead of streaming them over PCIe.
// `auto`: the share measured, each layer's CPU time per expert and the GPU's per streamed expert (CUDA events around
// its expert work, read after the next layer's routing sync - the host reaches the combine long before the GPU does)
// as running means, and the next layers hand the CPU g / (c + g) of them, where both sides end together - so a small
// CPU takes a small share. x: a fixed x. Unset or 0: every expert on the GPU (the default). RTX 5090 + 9950X3D,
// 600-token prompts: UD-Q4_K_XL 1,528-1,548 -> 1,296-1,369 ms (share 0.65), Q2_0 472-491 -> 445-448, IQ2_XS 507 -> 490
// (250 tokens 402 -> 368); IQ2_XS on 2 AVX2 workers 512 -> 487 (share 0.49).
inline double cpu_share_env() { // -1: measured
static const double v = [] {
const char* e = std::getenv("STRATA_PREFILL_CPU_SHARE");
if (e == nullptr) return 0.0;
if (std::strcmp(e, "auto") == 0) return -1.0;
return std::clamp(std::atof(e), 0.0, 1.0);
}();
return v;
}
inline bool cpu_share_on() { return cpu_share_env() != 0.0; }
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) -
@@ -618,6 +638,20 @@ struct Prefill::Impl {
std::vector<int32_t> bounds_host;
std::unique_ptr<mmq::Context> mmq_ctx;
std::vector<int32_t> ids_host, slot_host, src_host, cnt, off;
// set_cpu_pool: a small chunk's CPU experts. The layer's MoE input (T x N, from `mixed`) and the rows the pool
// writes (Dm's tail, in Dm's row order), both pinned; and the pool's per-token activations and jobs.
float *cpu_x = nullptr, *cpu_rows = nullptr;
size_t cpu_x_n = 0, cpu_rows_n = 0;
std::vector<uint8_t> cpu_nact;
std::vector<kernels::cpu::ActQ> cpu_actq; // a Q2_0 layer's activations (the pool's Q2_0 kernels read ActQ)
std::vector<kernels::cpu::ExpertJobMulti> cpu_jobs;
// the measured share: running means of the CPU's ms per expert and the GPU's per streamed expert, and the share
// they balance at (cpu_share_env)
double cpu_c_ms = 0, cpu_g_ms = 0, cpu_share_now = 0.5;
cudaEvent_t cpu_ev[2] = {}; // the GPU's expert work of the last CPU-sharing layer, read once the stream is done
bool cpu_pend = false;
double pend_cpu_ms = 0;
int64_t pend_n_cpu = 0, pend_n_gpu = 0;
// The grouping tables in mapped pinned memory, [ids | slot | src] of T_max * K each, then the MMQ bounds: kernels
// read and write them in place. A cudaMemcpyAsync of them queues behind the expert blobs the copy stream already
// holds (up to `ring` of them, ~70 us each), and the GPU idles meanwhile - measured 4.2 s of a 128K prompt's
@@ -734,6 +768,10 @@ void Prefill::release() {
if (impl_->kv_released) cudaEventDestroy(impl_->kv_released);
if (impl_->kv_ready) cudaEventDestroy(impl_->kv_ready);
if (impl_->grp_host) cudaFreeHost(impl_->grp_host);
for (float* p : {impl_->cpu_x, impl_->cpu_rows})
if (p) cudaFreeHost(p);
for (cudaEvent_t e : impl_->cpu_ev)
if (e) cudaEventDestroy(e);
for (void* p : impl_->owned) cudaFree(p);
}
@@ -1601,6 +1639,7 @@ int64_t Prefill::ring_max_slots() {
}
int64_t Prefill::ring_slots_for(int64_t chunk) { return ring_slots((size_t) chunk); }
void Prefill::set_cpu_pool(kernels::cpu::ExpertPool* pool) { cpu_pool_ = pool; }
bool Prefill::ring_bytes_enabled() { return ring_bytes_on(); }
@@ -2578,6 +2617,12 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
}
size_t n_order = 0; // the routed experts (the debug report; unknown when fused)
bool peer_now = false; // multi-GPU: the peer computed rows of this layer (MMQ path only)
// set_cpu_pool: this layer's CPU experts (rows [T * K - rows_cpu, T * K) of Dm), computed on a thread
// that reads on_cpu: declared first, so the future (which waits for the thread) goes first
std::vector<char> on_cpu;
double cpu_ms = 0; // the thread's time (the measured share)
std::future<bool> cpu_fut;
int64_t rows_cpu = 0, n_cpu = 0, n_stream = 0;
if (fused_l) {
if (static bool said = false; !said) {
said = true;
@@ -2647,12 +2692,39 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
int32_t* src_h = grp_mapped ? m.grp_host + 2 * m.grp_tk : m.src_host.data();
if (grp_mapped) copy_i32(m.grp_dev, m.ids, T * K, m.cs);
else cudaMemcpyAsync(m.ids_host.data(), m.ids, (size_t) T * K * 4, cudaMemcpyDeviceToHost, m.cs);
const bool cpu_maybe = cpu_pool_ != nullptr && cpu_share_on() && !stream_all && m.src != nullptr &&
lay.native && !m.pp && !lay.fmt.empty();
if (cpu_maybe) {
const size_t want = (size_t) T * N;
if (m.cpu_x_n < want) {
if (m.cpu_x) cudaFreeHost(m.cpu_x);
m.cpu_x = nullptr;
m.cpu_x_n = 0;
if (cudaHostAlloc((void**) &m.cpu_x, want * sizeof(float), cudaHostAllocDefault) != cudaSuccess) {
err = "prefill: cannot allocate the CPU experts' activations";
return false;
}
m.cpu_x_n = want;
}
cudaMemcpyAsync(m.cpu_x, m.mixed, want * sizeof(float), cudaMemcpyDeviceToHost, m.cs);
}
// #579: a stall here is the GPU (this layer's attention and router, or the previous layer's
// work), not the host: the watchdog's report says so (only its text changes)
core::progress_at("reading the prompt (batched): waiting for the GPU (attention, router) at layer",
l, p0);
cudaStreamSynchronize(m.cs);
core::progress_at("reading the prompt (batched): layer", l, p0);
if (m.cpu_pend) { // the measured share: the last CPU-sharing layer's GPU time is final now
m.cpu_pend = false;
float g_ms = 0;
if (cudaEventElapsedTime(&g_ms, m.cpu_ev[0], m.cpu_ev[1]) == cudaSuccess) {
constexpr double a = 0.25;
const double c = m.pend_cpu_ms / (double) m.pend_n_cpu, g = g_ms / (double) m.pend_n_gpu;
m.cpu_c_ms = m.cpu_c_ms > 0 ? m.cpu_c_ms + a * (c - m.cpu_c_ms) : c;
m.cpu_g_ms = m.cpu_g_ms > 0 ? m.cpu_g_ms + a * (g - m.cpu_g_ms) : g;
m.cpu_share_now = std::clamp(m.cpu_g_ms / (m.cpu_c_ms + m.cpu_g_ms), 0.05, 0.9);
}
}
pt.fold();
if (pe.on) { // the peer's marks so far are done: the primary waited for its last rows
int pd = 0; cudaGetDevice(&pd); cudaSetDevice(pe.dev); cudaStreamSynchronize(m.pp->s); pe.fold(); cudaSetDevice(pd);
@@ -2668,6 +2740,29 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
if (e < 0 || e >= m.g->n_expert) { err = "prefill: routed id out of range"; return false; }
++m.cnt[(size_t) e];
}
if (cpu_maybe) {
std::vector<std::pair<int32_t, int32_t>> cand; // (tokens, expert)
int64_t nstream = 0;
for (int32_t e = 0; e < m.g->n_expert; ++e) {
const int32_t c = m.cnt[(size_t) e];
if (c == 0 || (m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0))
continue;
++nstream;
if (c <= strata::kernels::cpu::MAXT && m.src->pinned(l, e)) cand.emplace_back(c, e);
}
std::sort(cand.begin(), cand.end());
const double share = cpu_share_env() >= 0.0 ? cpu_share_env() : m.cpu_share_now;
const size_t take = std::min(cand.size(), (size_t) std::llround(share * (double) nstream));
n_stream = nstream;
if (take > 0) {
on_cpu.assign((size_t) m.g->n_expert, 0);
for (size_t i = 0; i < take; ++i) {
on_cpu[(size_t) cand[i].second] = 1;
rows_cpu += cand[i].first;
}
n_cpu = (int64_t) take;
}
}
// multi-GPU: the rows of the experts the peer computes go last, as one block [rows_local, T*K)
const bool pre_mmq = mmq_plan().any && mmq_plan().layer[(size_t) l];
std::vector<char> on_peer;
@@ -2693,12 +2788,18 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
{
int32_t r = 0;
for (int32_t e = 0; e < m.g->n_expert; ++e)
if (on_peer.empty() || !on_peer[(size_t) e]) { m.off[(size_t) e] = r; r += m.cnt[(size_t) e]; }
if ((on_peer.empty() || !on_peer[(size_t) e]) && (on_cpu.empty() || !on_cpu[(size_t) e])) {
m.off[(size_t) e] = r;
r += m.cnt[(size_t) e];
}
// the peer's rows in order_peer's order (peer-held, then peer-streamed): its groups need them contiguous
for (int kind = 1; kind <= 2 && !on_peer.empty(); ++kind)
for (int32_t e = 0; e < m.g->n_expert; ++e)
if (on_peer[(size_t) e] == kind) { m.off[(size_t) e] = r; r += m.cnt[(size_t) e]; }
for (int32_t e = 0; e < m.g->n_expert && !on_cpu.empty(); ++e) // the CPU's block, last
if (on_cpu[(size_t) e]) { m.off[(size_t) e] = r; r += m.cnt[(size_t) e]; }
m.off[(size_t) m.g->n_expert] = r;
rows_local -= rows_cpu;
}
std::vector<int32_t> fill(m.off.begin(), m.off.end() - 1);
for (int64_t i = 0; i < T * K; ++i) {
@@ -2710,7 +2811,8 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
// the experts, in id order: resident ones from VRAM, the others through the staging ring
std::vector<int32_t> order, order_peer;
for (int32_t e = 0; e < m.g->n_expert; ++e)
if (m.cnt[(size_t) e] > 0 && (on_peer.empty() || on_peer[(size_t) e] != 2))
if (m.cnt[(size_t) e] > 0 && (on_peer.empty() || on_peer[(size_t) e] != 2) &&
(on_cpu.empty() || !on_cpu[(size_t) e]))
(!on_peer.empty() && on_peer[(size_t) e] ? order_peer : order).push_back(e);
for (int32_t e = 0; e < m.g->n_expert && !on_peer.empty(); ++e) // the peer-streamed ones last:
if (on_peer[(size_t) e] == 2) order_peer.push_back(e); // their copies get the most time
@@ -2730,6 +2832,65 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
cudaMemcpyAsync(m.src_dev, m.src_host.data(), (size_t) T * K * 4, cudaMemcpyHostToDevice, m.cs);
}
n_order = order.size();
// set_cpu_pool: the CPU's experts on a thread while this one streams and runs the GPU's; their
// unweighted rows land in Dm's row order, so the combine weights them as any other row
if (rows_cpu > 0) {
const size_t want = (size_t) rows_cpu * N;
if (m.cpu_rows_n < want) {
if (m.cpu_rows) cudaFreeHost(m.cpu_rows);
m.cpu_rows = nullptr;
m.cpu_rows_n = 0;
const size_t cap_rows = (size_t) rows_cpu * 2; // headroom: few reallocations
if (cudaHostAlloc((void**) &m.cpu_rows, cap_rows * N * sizeof(float), cudaHostAllocDefault) != cudaSuccess) {
err = "prefill: cannot allocate the CPU experts' rows";
return false;
}
m.cpu_rows_n = cap_rows * N;
}
const int64_t r0c = T * K - rows_cpu;
const strata::kernels::cpu::NativeFmt& cf = lay.fmt[(size_t) l];
if (cpu_share_env() < 0.0) {
if (m.cpu_ev[0] == nullptr) { cudaEventCreate(&m.cpu_ev[0]); cudaEventCreate(&m.cpu_ev[1]); }
cudaEventRecord(m.cpu_ev[0], m.cs);
}
cpu_fut = std::async(std::launch::async, [&m, &on_cpu, &cf, &cpu_ms, src_h, l, T, r0c, this]() -> bool {
const auto t0 = std::chrono::steady_clock::now();
constexpr size_t AB = strata::kernels::cpu::kNativeActBytes;
// a Q2_0 layer (gate/up on the pool's Q2_0 kernels): ActQ activations, as decode's
const bool q2 = strata::kernels::cpu::q2_native_kernels(cf.gu_type);
if (q2 && m.cpu_actq.size() < (size_t) T) m.cpu_actq.resize((size_t) T);
if (!q2 && m.cpu_nact.size() < (size_t) T * AB) m.cpu_nact.resize((size_t) T * AB);
std::vector<char> need((size_t) T, 0);
for (int32_t e = 0; e < m.g->n_expert; ++e) // the tokens the CPU's experts read
if (on_cpu[(size_t) e])
for (int32_t r = 0; r < m.cnt[(size_t) e]; ++r)
need[(size_t) src_h[(size_t) m.off[(size_t) e] + (size_t) r]] = 1;
for (int64_t t = 0; t < T; ++t)
if (need[(size_t) t] && q2)
strata::kernels::cpu::act_quant_any(m.cpu_x + (size_t) t * N, (int) N, m.cpu_actq[(size_t) t]);
else if (need[(size_t) t])
strata::kernels::cpu::native_quant_act(cf, m.cpu_x + (size_t) t * N,
m.cpu_nact.data() + (size_t) t * AB);
m.cpu_jobs.clear();
for (int32_t e = 0; e < m.g->n_expert; ++e) {
if (!on_cpu[(size_t) e]) continue;
strata::kernels::cpu::ExpertJobMulti j;
j.blob = m.src->blob(l, e);
if (j.blob == nullptr) return false;
j.nt = m.cnt[(size_t) e];
for (int r = 0; r < j.nt; ++r) {
const int64_t p = (int64_t) m.off[(size_t) e] + r;
if (q2) j.act[r] = &m.cpu_actq[(size_t) src_h[(size_t) p]];
else j.nact[r] = m.cpu_nact.data() + (size_t) src_h[(size_t) p] * AB;
j.out[r] = m.cpu_rows + (size_t) (p - r0c) * N;
}
m.cpu_jobs.push_back(j);
}
cpu_pool_->run_split_multi_native(cf, m.cpu_jobs.data(), (int) m.cpu_jobs.size());
cpu_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t0).count();
return true;
});
}
const size_t mmq_gub = use_mmq ? mmq::matrix_bytes(mmq_gt, 1280, N) : 0;
const size_t mmq_db = use_mmq ? mmq::matrix_bytes(mmq_dt, N, 640) : 0;
pt.mark(kPfGather, cs);
@@ -3183,6 +3344,21 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
}
}
pt.mark(kPfCombine, cs);
if (cpu_fut.valid()) { // set_cpu_pool: the CPU's rows into Dm's tail
if (cpu_share_env() < 0.0) cudaEventRecord(m.cpu_ev[1], m.cs); // after the GPU's expert work
if (!cpu_fut.get()) { err = "prefill: a CPU expert has no blob"; return false; }
if (cpu_share_env() < 0.0 && n_stream > n_cpu) {
m.cpu_pend = true;
m.pend_cpu_ms = cpu_ms;
m.pend_n_cpu = n_cpu;
m.pend_n_gpu = n_stream - n_cpu;
}
stats_.cpu_share = cpu_share_env() >= 0.0 ? cpu_share_env() : m.cpu_share_now;
const int64_t r0c = T * K - rows_cpu;
cudaMemcpyAsync(m.Dm + (size_t) r0c * N, m.cpu_rows, (size_t) rows_cpu * N * sizeof(float),
cudaMemcpyHostToDevice, m.cs);
stats_.experts_cpu += n_cpu;
}
if (peer_now) { // multi-GPU: the peer's rows are in Dm (or, without P2P, in host memory)
cudaStreamWaitEvent(m.cs, m.pp->ev_done, 0);
if (!m.pp->p2p)
+6 -2
View File
@@ -6329,6 +6329,8 @@ int main(int argc, char** argv) {
}
}
if (multi_gpu) sp.set_stage(0, split_at[0], &stages[0]->sp);
// the pool is idle while a prompt is read unless batch slots decode between its parts
if (!multi_gpu && o.batch <= 0 && !o.no_pool) sp.set_cpu_pool(&pool);
if (!sp.init(wt, g, ss, srcp, &xcache, host_res.data(), o.prefill_chunk, main_cs, err, borrow, borrow_bytes))
return err.find("do not fit") != std::string::npos ? 2 : 1;
return 0;
@@ -10911,6 +10913,7 @@ int main(int argc, char** argv) {
}
if (borrow == nullptr)
std::fprintf(stderr, "strata generate: prompt path allocates its own buffers (no cache slots to borrow)\n");
if (!multi_gpu && !o.no_pool) prefill.set_cpu_pool(&pool);
if (!prefill.init(wt, g, ss, srcp, o.expert_cache > 0 ? &xcache : nullptr,
host_res.empty() ? nullptr : host_res.data(), o.prefill_chunk, main_cs, err, borrow,
borrow_bytes)) {
@@ -10974,10 +10977,11 @@ int main(int argc, char** argv) {
ss.ple_prev[1] = pos_start >= 1 ? (int32_t) o.tokens[(size_t) (pos_start - 1)] : -1;
const strata::prefill::PrefillStats& ps = prefill.stats();
std::fprintf(stderr, "strata generate: prefill %lld tokens in %lld chunks, %.1f ms (%.1f tok/s); experts "
"streamed %lld (%lld by DMA, host %.1f ms), resident %lld; PLE %.1f ms\n",
"streamed %lld (%lld by DMA, host %.1f ms), resident %lld, on the CPU %lld (share %.2f); PLE %.1f ms\n",
(long long) ps.tokens, (long long) ps.chunks, ps.ms_total,
ps.ms_total > 0 ? 1000.0 * (double) ps.tokens / ps.ms_total : 0.0, (long long) ps.experts_streamed,
(long long) ps.experts_dma, ps.ms_experts_host, (long long) ps.experts_resident, ps.ms_ple);
(long long) ps.experts_dma, ps.ms_experts_host, (long long) ps.experts_resident,
(long long) ps.experts_cpu, ps.cpu_share, ps.ms_ple);
}
if (!kvg_started) { // the elastic K/V of a prompt that was not read in batches