diff --git a/docs/DETAILS.md b/docs/DETAILS.md index e2565af6f98..7b543eec5b4 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -549,6 +549,23 @@ applies at the next start. Network, key, MCP and program keys are not editable t `--parallel`, only a `max_tokens` 1 request gets them: it is the one a batch slot decodes through the single-token path, which writes an LP line. A longer reply there is a 400, since the slot's `BT` lines carry none -- entries for a prefix only, which a client cannot tell from the model's own answer. +- **Canonical reads: the same logprobs for the same prompt (#4059).** A typed decision reads its answer from one + token's logprobs and compares them with its other reads (a label prior, a calibrated threshold), so it needs the same + numbers for the same prompt. By default they move a little from request to request: an expert is computed on the GPU's + hit path, over PCIe or on the CPU depending on what the expert cache holds - which the adaptive tier and the prompt + path's slot loan keep changing - and on the PCIe share and the prompt's CPU share, and a low-ranked missed expert can be + skipped (`STRATA_ROUTE_TAIL_SKIP`); the paths round differently. A request for logprobs whose reply is one token + (`max_tokens` or `max_completion_tokens` 1) is read canonically: its prompt in one batched part with every expert on + the GPU (no chunk shared with the CPU), its one-token window with every routed expert treated as a miss (over PCIe + where its blob is pinned, else on the CPU, the same path every time; none skipped), and no cached prefix reused. Its + logprobs then depend on the prompt alone: the same across requests, with other traffic in between and after a restart + (same build, GPU and settings). Every other request keeps the adaptive paths, so chat decodes as fast as before. + `"strata_deterministic": true` asks for a canonical read on any request, `false` turns it off; the engine gets + `canon=1` on the request line (an older engine skips the key). Cost: a canonical read costs what a first read of its + prompt does, since it reuses nothing. Measured (Qwen3.8-Flash-Next IQ3_S, launch defaults): RTX 5080, a 65-read + decision set identical over 4 runs with long prompts and chat between them and across a restart, a read 363 ms, chat + decode unchanged; GTX 1080 Ti (12 GB/s PCIe), the same, a 101-token read 1.9 s against 1.5 s for its first plain read + and 0.24 s for a plain read that reuses it. A layer split, a peer or remote expert GPU, and batch slots are not covered. - `STRATA_SPEC_PROB=1` with the default `--spec-min-p` gate: the gate tests the token the draft head drew from its distribution q. A draft that fails the gate still counts as a proposal. The verifier keeps q's mass on the tokens that pass the gate (not renormalised; call it mu), accepts an offered token D with probability min(1, p(D)/mu(D)), and draws the replacement diff --git a/include/strata/core/expert_source.hpp b/include/strata/core/expert_source.hpp index 2bcfca39b46..65bba01da52 100644 --- a/include/strata/core/expert_source.hpp +++ b/include/strata/core/expert_source.hpp @@ -669,6 +669,8 @@ void expert_pool_dispatch(void* user, const float* x_f, const int32_t* ids, cons /// STRATA_ROUTE_TAIL_SKIP rank (0: off). Unset in the environment: 0 until the serve start-up sets the default. int tail_skip_rank(); void set_tail_skip_rank(int r); +/// A canonical read (#4059, GEN canon=1): no missed expert is skipped while it runs, whatever the rank above. +void pause_tail_skip(bool paused); void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k, float* out, strata::kernels::cpu::NativeArithmetic arithmetic = strata::kernels::cpu::NativeArithmetic::Default); diff --git a/include/strata/prefill/prefill.hpp b/include/strata/prefill/prefill.hpp index 549aff380a6..7c4ea920577 100644 --- a/include/strata/prefill/prefill.hpp +++ b/include/strata/prefill/prefill.hpp @@ -261,6 +261,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); + /// A canonical read (#4059, GEN canon=1): while paused, no chunk hands experts to the CPU pool - every expert on + /// the GPU, as with STRATA_PREFILL_CPU_SHARE=0 - and the measured share keeps its state for the requests after it. + static void pause_cpu_share(bool paused); /// 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. diff --git a/serve/server.py b/serve/server.py index 8c5213b1f80..bb8f25695e1 100644 --- a/serve/server.py +++ b/serve/server.py @@ -1465,6 +1465,12 @@ def sampling_keys(sampling: dict) -> str: keys += f" pin={pn}" # logprobs: the engine writes an `LP` line after each `T` line (an engine without them ignores # the key, so an older engine answers exactly as it did). Absent = no LP lines at all. + # a canonical read (a typed decision: one token's logprobs, or "strata_deterministic": true): the engine computes + # every expert of this request on the GPU, so the same prompt gives the same logprobs whatever the expert cache, + # the adaptive tier, the PCIe share and the prefill CPU share have done since the start (#4059). An engine + # without it skips the key. + if canonical_read(sampling): + keys += " canon=1" k = top_logprobs(sampling) if k is not None: keys += f" logprobs={k}" @@ -3823,6 +3829,20 @@ def top_logprobs(req: dict): return None +def canonical_read(req: dict) -> bool: + """Is this request a canonical read (#4059)? `"strata_deterministic": true|false` decides when it is sent; else a + request for logprobs whose reply is one token (max_tokens or max_completion_tokens 1) is one: a typed decision read + from one next-token distribution, which a client compares with its other reads and so needs the same numbers for + the same prompt. Longer replies and other requests keep the engine's adaptive paths (and their speed).""" + v = req.get("strata_deterministic") + if isinstance(v, bool): + return v + if top_logprobs(req) is None: + return False + mt = req.get("max_completion_tokens") or req.get("max_tokens") + return isinstance(mt, int) and not isinstance(mt, bool) and mt == 1 + + def allowed_ids(req: dict) -> list: """`allowed_token_ids`: the ids a caller named, or [] (vLLM's semantics, see LogprobEntries).""" ids = req.get("allowed_token_ids") diff --git a/serve/test_logprobs.py b/serve/test_logprobs.py index a21237db782..3d84798a541 100644 --- a/serve/test_logprobs.py +++ b/serve/test_logprobs.py @@ -21,7 +21,7 @@ from serve.frontend import ChatTemplate from serve.server import (ByteTokenizer, LogprobEntries, LPLines, MockEngine, Service, StrataEngine, batch_slots_on, - logprob_span, serve, top_logprobs, allowed_ids) + logprob_span, serve, top_logprobs, allowed_ids, canonical_read) ROOT = Path(__file__).resolve().parent.parent TEMPLATE = ChatTemplate(ROOT / "serve" / "chat_template.jinja") @@ -522,6 +522,35 @@ def test_slots_behind_replicas(self): self.assertTrue(batch_slots_on(SimpleNamespace(batch=4))) +class Canonical(unittest.TestCase): + """#4059: a typed decision (one token's logprobs) is a canonical read - the engine gets canon=1 and computes every + expert of it on the GPU, so the same prompt gives the same logprobs. "strata_deterministic" decides when sent.""" + + def keys(self, req): + return StrataEngine.sampling_keys(req).split() + + def test_a_one_token_logprobs_request_is_canonical(self): + for req in ({"logprobs": True, "max_tokens": 1}, {"logprobs": 0, "max_tokens": 1}, + {"logprobs": True, "top_logprobs": 5, "max_completion_tokens": 1}, + {"logprobs": 1, "max_tokens": 1, "allowed_token_ids": [65, 66]}): + self.assertIn("canon=1", self.keys(req), req) + + def test_other_requests_are_not(self): + for req in ({}, {"max_tokens": 1}, {"logprobs": True, "max_tokens": 2}, {"logprobs": True}, + {"logprobs": True, "max_tokens": True}, {"logprobs": False, "max_tokens": 1}): + self.assertNotIn("canon=1", self.keys(req), req) + + def test_the_field_decides_when_sent(self): + self.assertIn("canon=1", self.keys({"strata_deterministic": True, "max_tokens": 200})) + self.assertIn("canon=1", self.keys({"strata_deterministic": True})) + self.assertNotIn("canon=1", self.keys({"strata_deterministic": False, "logprobs": True, "max_tokens": 1})) + self.assertIn("canon=1", self.keys({"strata_deterministic": "yes", "logprobs": True, "max_tokens": 1})) # not a bool: auto + + def test_canonical_read(self): + self.assertTrue(canonical_read({"logprobs": 3, "max_tokens": 1})) + self.assertFalse(canonical_read({"logprobs": 3, "max_tokens": 0})) + + class Span(unittest.TestCase): def entries(self, *toks): return [{"token": t} for t in toks] diff --git a/src/core/expert_source.cpp b/src/core/expert_source.cpp index 1b5b115524d..3661c28577f 100644 --- a/src/core/expert_source.cpp +++ b/src/core/expert_source.cpp @@ -4808,6 +4808,10 @@ int tail_skip_rank() { return r; } void set_tail_skip_rank(int r) { g_tail_skip_rank.store(r > 0 && r < 10 ? r : 0, std::memory_order_relaxed); } +namespace { +std::atomic g_tail_skip_paused{false}; // pause_tail_skip: a canonical read is running +} // namespace +void pause_tail_skip(bool paused) { g_tail_skip_paused.store(paused, std::memory_order_relaxed); } void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k, float* out, strata::kernels::cpu::NativeArithmetic arithmetic) { @@ -4887,7 +4891,8 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32 // STRATA_ROUTE_TAIL_SKIP=R (default 7 on CUDA serve, =0 off; changes the output): a missed expert that EVERY token of the window routes // only at rank >= R (the smallest weights of the top-10) is neither copied nor computed; its contribution is // zero, without renormalisation (the others keep the weight the router gave them). Less PCIe and CPU work. - const int tail_rank = tail_skip_rank(); + // a canonical read computes every routed expert: which ones are missed depends on the cache, so a skip would too + const int tail_rank = g_tail_skip_paused.load(std::memory_order_relaxed) ? 0 : tail_skip_rank(); bool skip_d[kMaxWindowEntries]; for (int q = 0; q < nd; ++q) skip_d[q] = false; if (tail_rank > 0 && d.peer == nullptr && d.remote_count == 0) { diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index 16a609a5ded..47dfad19006 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -268,6 +268,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}; +// Prefill::pause_cpu_share: a canonical read is running (the share's measured state is left as it is) +std::atomic g_cpu_share_paused{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) - @@ -2622,6 +2624,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; } +void Prefill::pause_cpu_share(bool paused) { g_cpu_share_paused.store(paused, std::memory_order_relaxed); } void Prefill::arm_cpu_share(bool applies, bool by_default) { #if !defined(STRATA_USE_HIP) && !defined(STRATA_HIP_GFX906) if (applies && by_default && cpu_share_explicit() == -2.0 && !g_share_default) { @@ -3335,7 +3338,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_cpu_share_paused.load(std::memory_order_relaxed)) { 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 824bd19dbe6..23528710a04 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -13714,6 +13714,12 @@ int main(int argc, char** argv) { // embedding floats, so a picture encoded again with other bits still matches. Absent: the embeddings' hash. std::vector req_img_keys; bool req_pcie_set = false; // the request names its own share (the measured one holds meanwhile) + // canon=1 (#4059): a canonical read. Every routed expert of this request is computed on the GPU - the missed + // ones over PCIe, none skipped, no prompt chunk shared with the CPU - and no cached prefix is reused, so its + // logits depend on the prompt alone, not on what the expert cache, the adaptive tier, the PCIe share and the + // prefill share have done since the start. For typed decisions (one token's logprobs); other requests keep + // the adaptive paths. + bool req_canon = false; if (endp != nullptr) { // GENI takes the same keys (#75: image requests were always greedy); its // embedding file path is the first token without an = for (;;) { @@ -13753,6 +13759,7 @@ int main(int argc, char** argv) { else if (key == "penalty_present") req_penalty_present = fv; else if (key == "seed") req_seed = std::strtoull(tok.c_str() + eq + 1, nullptr, 10); else if (key == "pcie_frac") { req_pcie_frac = std::clamp((double) fv, 0.0, 1.0); req_pcie_set = true; } + else if (key == "canon") req_canon = std::atoi(tok.c_str() + eq + 1) != 0; else if (key == "spec_min_p") req_spec_min_p = std::clamp((double) fv, 0.0, 1.0); else if (key == "pipe") req_pipe = std::max(0.0, (double) fv); else if (key == "aux_cpus") req_aux_cpus = std::atoi(tok.c_str() + eq + 1) != 0 ? 1 : 0; @@ -13797,6 +13804,21 @@ int main(int argc, char** argv) { std::printf("ERR prefill=1 needs a text prompt of at least 2 tokens and no --batch\n"); continue; } + // canon=1: the two process-wide switches are paused for this request only, and given back however it ends + struct CanonPause { + bool on; + explicit CanonPause(bool c) : on(c) { + if (on) { strata::core::pause_tail_skip(true); strata::prefill::Prefill::pause_cpu_share(true); } + } + ~CanonPause() { + if (on) { strata::core::pause_tail_skip(false); strata::prefill::Prefill::pause_cpu_share(false); } + } + } canon_pause(req_canon); + if (req_canon) { + req_pcie_frac = 1.0; // every missed expert of its one-token window over PCIe to the GPU + req_pcie_set = true; + req_ckpt = false; // nothing reuses it: the prompt is read in one part, not split at the turn + } req_imgs.clear(); if (geni && !o.vision) { std::printf("ERR this engine was started without --vision\n"); continue; } if (geni || !mrope_identity) { @@ -13998,7 +14020,7 @@ int main(int argc, char** argv) { } int64_t resume = 0; bool from_live = false; - if (o.prompt_cache > 0 && want_cvec == cvec_cached) { + if (o.prompt_cache > 0 && !req_canon && want_cvec == cvec_cached) { if (live_ok && starts_with(live, live_imgs)) { resume = (int64_t) live.size(); from_live = true; } for (const ConvCheckpoint& c : checks) if ((int64_t) c.ids.size() > resume && starts_with(c.ids, c.imgs)) { @@ -14012,7 +14034,7 @@ int main(int argc, char** argv) { // of a slot's checkpoint this request moved into `checks` (from the slot it is admitted to), -1: none int lent_slot = -1; int64_t moved_in = -1; - if (slot_ckpt_move_env() && o.prompt_cache > 0 && want_cvec == cvec_cached && live_ok) + if (slot_ckpt_move_env() && o.prompt_cache > 0 && !req_canon && want_cvec == cvec_cached && live_ok) for (int b = 0; b < (int) bs.size(); ++b) { if (bs[(size_t) b].cvec != cvec_cached || bs[(size_t) b].tainted) continue; // (#879) for (const ConvCheckpoint& c : bs[(size_t) b].checks) { @@ -14037,7 +14059,7 @@ int main(int argc, char** argv) { ConvCheckpoint::PromptPlan resume_plan; int64_t resume_plan_at = -1, resume_plan_end = -1; // where it goes on, and the segment end it needs const ConvCheckpoint* slot_ck = nullptr; // the slot's checkpoint the prompt continues from (else its end) - if (o.prompt_cache > 0 && req_imgs.empty()) + if (o.prompt_cache > 0 && !req_canon && req_imgs.empty()) for (int b = 0; b < (int) bs.size(); ++b) { const BSlot& sl = bs[(size_t) b]; const bool idle_source = !sl.active && sl.cached; @@ -14092,7 +14114,8 @@ int main(int argc, char** argv) { k + 1, (long long) (k < now.size() ? now[k].start : -1)); } } - const auto parked = conversations.best(ids, req_imgs, want_cvec); + auto parked = conversations.best(ids, req_imgs, want_cvec); + if (req_canon) parked.tokens = 0; // a canonical read reads its whole prompt (canon=1) std::optional incoming; // A checkpoint inside ANOTHER parked conversation (a sibling subagent's system prompt and tools, say) is // borrowed: its K/V up to there and the checkpoint come over, and that conversation stays parked for its @@ -14579,7 +14602,9 @@ int main(int argc, char** argv) { "whole prompt through the batched path, none of it through the verify windows\n"); } auto windows_ok = [&](int64_t a, int64_t b) -> bool { - if (no_short || b - a > req_short_read) return false; + // a canonical read never reads its prompt through the windows: a window of several tokens can miss more + // experts than its PCIe staging holds, and the rest fall to the CPU by what the cache holds (canon=1) + if (no_short || req_canon || b - a > req_short_read) return false; if (sp.embd_rows != nullptr) for (int64_t i = a; i < b; ++i) if (sp.embd_rows[i] != nullptr) return false; @@ -16652,7 +16677,32 @@ int main(int argc, char** argv) { if (gated_proposal) ver.set_spec_proposal(proposal.n_q ? proposal.mu.data() : nullptr, proposal.n_q); else if (use_mtp && mtp.prob() && !from_sfx && !first_window && T_mtp > 1) ver.set_spec_q(mtp.spec_q(), T_mtp - 1); const auto decode_pool_fn = use_mtp && mtp.prob() ? invariant_win_pool_fn : win_pool_fn; - if (!ver.run(T, window.data(), p, decode_pool_fn, win_pool_user, outv.data(), err) || drive.d.failed) { + // canon=1: the window computes every routed expert as a miss. A resident expert goes through the hit + // path and a missed one through the PCIe (or CPU) path, which round differently, and which experts are + // resident is what the adaptive tier and the prompt path's slot loan change from request to request: + // with every expert a miss, each takes the same path every time (over PCIe where its blob is pinned, + // else the CPU pool), whatever the cache holds. The tables are given back right after the window, + // before the next adaptive round looks at them. + std::vector canon_res; + const bool canon_win = req_canon && !host_res.empty() && d_res != nullptr; + if (canon_win) { + if (aflipped || !pending.empty()) apply_pending(true); // the last swaps land: nothing is in flight + canon_res = host_res; + std::fill(host_res.begin(), host_res.end(), strata::core::kNotResident); + if (res_upload() != cudaSuccess) { + std::printf("ERR canon=1: the residency table could not be cleared for the window\n"); + return 1; + } + } + const bool win_ok = ver.run(T, window.data(), p, decode_pool_fn, win_pool_user, outv.data(), err) && !drive.d.failed; + if (canon_win) { + std::copy(canon_res.begin(), canon_res.end(), host_res.begin()); + if (res_upload() != cudaSuccess) { + std::printf("ERR canon=1: the residency table could not be restored after the window\n"); + return 1; + } + } + if (!win_ok) { std::printf("ERR %s\n", drive.d.failed && drive.d.fail ? drive.d.fail : err.c_str()); return 1; }