mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 14:21:06 -10:00
Conversation cache: the checkpoint pool stays full, the park's small copies gathered, helpers off the host core, park traces
Measured on a CMP 170HX (PCIe Gen2 x4): a conversation switch's park still paid 38-52 ms for the live state - the STRATA_CKPT_POOL pool was empty at every park (a switch takes two checkpoints and the refill added one), so the 118 MB GDN vector was allocated fresh and page-faulted. Now: - STRATA_CKPT_POOL=N: the pool keeps N (default 2) ready checkpoints, pre-sized to the session's state sizes, topped up on a background thread after any take. - STRATA_CKPT_GATHER=1: the 37 small synchronous device-to-host copies of a checkpoint (ple, tails, dead, block_pos) gathered device-to-device into one buffer and fetched with one copy through a pinned buffer. - The STRATA_PARK_ASYNC stage is allocated with 64 MiB of headroom (no cudaFree/cudaMalloc as the conversation grows) and the pointer-attribute query is skipped for a VRAM-resident K/V. - The three background helpers (pool top-up, the checkpoint DMA, the park DMA) leave the host's pinned core. - A parked image whose live GDN state is still being copied joins before its buffers are freed. - STRATA_PARK_TRACE prints the joins, the live-state vectors' readiness and the K/V sub-timers; the d2d time is measured right. Park at a 3K context 64-75 -> 21-40 ms; outputs identical (greedy 11 of 11, the A/B/A conversation check). HIP: ckpt_gather_on() is false under STRATA_HIP_GFX906 as well, and the runtime calls whose result the cache does not check are cast to (void), so this commit carries its own HIP guard and builds alone for gfx1100 and gfx906. Co-Authored-By: Claude Fable 5.1 <[email protected]> Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
1 parent
3e0c141c63
commit
56088e7154
5 files changed
+228
-39
No files matched your search
@@ -171,8 +171,10 @@ struct SavedConversation {
|
||||
// with a layer split, the later stages' own images, one per stage, in stage order
|
||||
std::vector<SavedConversation> stage_images;
|
||||
// STRATA_PARK_ASYNC: set while a background copy fills this image's K/V; its deleter waits for that copy, and as
|
||||
// the last member it runs before any buffer is freed (whichever copy of the image goes last)
|
||||
// one of the last members it runs before any buffer is freed (whichever copy of the image goes last)
|
||||
std::shared_ptr<void> pending;
|
||||
// STRATA_CKPT_ASYNC: the same for the background copy that fills `live.gdn` (the park's running state)
|
||||
std::shared_ptr<void> live_pending;
|
||||
|
||||
/// Holds a pinned shared prefix (see ConversationCheckpoint::pinned): the parked-conversation budget keeps it.
|
||||
bool pinned() const {
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
#include "strata/core/conversation_cache.hpp"
|
||||
#include "strata/core/progress.hpp"
|
||||
#include "strata/platform/aux_cpus.hpp"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
@@ -12,6 +13,10 @@
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
#if defined(__linux__)
|
||||
#include <sched.h>
|
||||
#endif
|
||||
|
||||
namespace strata::core { void conversation_checkpoint_join(const void* only); }
|
||||
namespace strata::core::conversation_detail {
|
||||
|
||||
@@ -36,6 +41,20 @@ inline cudaError_t private_stream(cudaStream_t& s) {
|
||||
return cudaSuccess;
|
||||
}
|
||||
|
||||
/// First thing a background copy thread does (STRATA_CKPT_ASYNC, STRATA_PARK_ASYNC): a thread the host creates
|
||||
/// inherits the host's one CPU (strata/platform/aux_cpus.hpp), where its memcpys would take turns with the decode.
|
||||
/// Wide first: the mask the aux placement keeps as this thread's own is the one restore() puts back (a request with
|
||||
/// aux_cpus=0), and the host's one CPU must not be it. No output depends on where a thread runs.
|
||||
inline void background_thread_placement() {
|
||||
#if defined(__linux__)
|
||||
cpu_set_t all;
|
||||
CPU_ZERO(&all);
|
||||
for (int c = 0; c < CPU_SETSIZE; ++c) CPU_SET(c, &all);
|
||||
(void) sched_setaffinity(0, sizeof all, &all); // the kernel keeps the CPUs the process may use
|
||||
strata::aux_cpus::pin_current_thread(); // (a no-op while the feature is off)
|
||||
#endif
|
||||
}
|
||||
|
||||
/// `queued`: the copy is only queued on the private stream, and the caller synchronizes the device (or orders a
|
||||
/// later launch after it with order_after_queued) before anything reads the target or writes the source.
|
||||
inline cudaError_t copy_nonblocking(void* dst, const void* src, size_t bytes, bool queued = false) {
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <array>
|
||||
#include <chrono>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
@@ -146,6 +147,14 @@ bool on_device(const void* p) {
|
||||
return at.type == cudaMemoryTypeDevice;
|
||||
#endif
|
||||
}
|
||||
// STRATA_PARK_TRACE: where conversation_kv_save's time goes since the park's begin (ms): the pointer-attribute
|
||||
// queries, the host buffer resizes and the copies (into the stage, or the synchronous ones). Read by the park trace.
|
||||
struct KvSaveTrace { double attrs = 0, resize = 0, copy = 0; };
|
||||
KvSaveTrace& kv_save_trace() { static KvSaveTrace t; return t; }
|
||||
bool park_trace_on() {
|
||||
static const bool on = std::getenv("STRATA_PARK_TRACE") != nullptr;
|
||||
return on;
|
||||
}
|
||||
|
||||
bool transfer(void* dst, const void* src, size_t n, std::string& error) {
|
||||
if (!n) return true;
|
||||
@@ -210,11 +219,17 @@ bool conversation_kv_save(ConversationKv& image, const QsaState& st, const Model
|
||||
if (keep > dst[i]->size()) { error = "conversation snapshot: missing reusable prefix"; return false; }
|
||||
// STRATA_PARK_ASYNC: device bytes go to the VRAM stage now and into the buffer on the background thread
|
||||
ParkAsync& a = park_async();
|
||||
const bool trace = park_trace_on();
|
||||
auto tick = [trace] { return trace ? std::chrono::steady_clock::now() : std::chrono::steady_clock::time_point{}; };
|
||||
const auto t0 = tick();
|
||||
// (kv_mode 0: the pools are device memory by construction, so no pointer-attribute query)
|
||||
const bool staged = a.active && src[i] != nullptr && sizes[i] > keep && a.used + (sizes[i] - keep) <= a.cap &&
|
||||
on_device(src[i]);
|
||||
(st.kv_mode == 0 || on_device(src[i]));
|
||||
const auto t1 = tick();
|
||||
if (staged) dst[i]->resize_for_overwrite(sizes[i]);
|
||||
else dst[i]->resize(sizes[i]);
|
||||
if (!dst[i]->visit(keep, sizes[i] - keep, [&](uint8_t* p, size_t n, size_t at) {
|
||||
const auto t2 = tick();
|
||||
const bool copied = dst[i]->visit(keep, sizes[i] - keep, [&](uint8_t* p, size_t n, size_t at) {
|
||||
const uint8_t* s = src[i] ? static_cast<const uint8_t*>(src[i]) + at : nullptr;
|
||||
if (!staged) return transfer(p, s, n, error);
|
||||
const cudaError_t e = conversation_detail::copy_nonblocking(a.stage + a.used, s, n, true);
|
||||
@@ -222,7 +237,14 @@ bool conversation_kv_save(ConversationKv& image, const QsaState& st, const Model
|
||||
a.items.push_back({p, n, a.used});
|
||||
a.used += n;
|
||||
return true;
|
||||
})) return false;
|
||||
});
|
||||
if (trace) {
|
||||
auto ms = [](std::chrono::steady_clock::time_point x, std::chrono::steady_clock::time_point y) {
|
||||
return std::chrono::duration<double, std::milli>(y - x).count(); };
|
||||
KvSaveTrace& t = kv_save_trace();
|
||||
t.attrs += ms(t0, t1); t.resize += ms(t1, t2); t.copy += ms(t2, tick());
|
||||
}
|
||||
if (!copied) return false;
|
||||
if (reused_bytes) *reused_bytes += keep;
|
||||
}
|
||||
return true;
|
||||
@@ -415,7 +437,13 @@ void conversation_park_join() {
|
||||
}
|
||||
}
|
||||
|
||||
void conversation_kv_save_trace(double& attrs_ms, double& resize_ms, double& copy_ms) {
|
||||
const KvSaveTrace& t = kv_save_trace();
|
||||
attrs_ms = t.attrs; resize_ms = t.resize; copy_ms = t.copy;
|
||||
}
|
||||
|
||||
bool conversation_park_async_begin(size_t need) {
|
||||
if (park_trace_on()) kv_save_trace() = {}; // (the sums cover this park's saves)
|
||||
conversation_park_join(); // (the stage is reused)
|
||||
ParkAsync& a = park_async();
|
||||
a.active = false;
|
||||
@@ -447,14 +475,17 @@ bool conversation_park_async_begin(size_t need) {
|
||||
return false;
|
||||
}
|
||||
if (a.stage == nullptr) {
|
||||
// rounded up to 64 MiB: the next parks of a growing conversation reuse it instead of re-allocating (a cudaFree
|
||||
// synchronizes the device, and the checkpoint DMA is in flight at this point)
|
||||
const size_t want = (need + ((size_t) 64 << 20) - 1) & ~(((size_t) 64 << 20) - 1);
|
||||
size_t fr = 0, tot = 0;
|
||||
if (cudaMemGetInfo(&fr, &tot) != cudaSuccess || fr < need + ((size_t) 3 << 29) ||
|
||||
cudaMalloc((void**) &a.stage, need) != cudaSuccess) {
|
||||
if (cudaMemGetInfo(&fr, &tot) != cudaSuccess || fr < want + ((size_t) 3 << 29) ||
|
||||
cudaMalloc((void**) &a.stage, want) != cudaSuccess) {
|
||||
a.stage = nullptr;
|
||||
(void) cudaGetLastError();
|
||||
return false;
|
||||
}
|
||||
a.cap = need;
|
||||
a.cap = want;
|
||||
}
|
||||
a.active = true;
|
||||
return true;
|
||||
@@ -492,6 +523,7 @@ bool conversation_park_async_launch(std::string& error) {
|
||||
std::lock_guard<std::mutex> lk(a.mu);
|
||||
// two pinned pieces in turn: the next piece's DMA runs while this one is copied into its segment
|
||||
a.pending = std::async(std::launch::async, [items = std::move(items), stage, host, st, dev] {
|
||||
conversation_detail::background_thread_placement();
|
||||
(void) cudaSetDevice(dev);
|
||||
cudaEvent_t done[2] = {nullptr, nullptr};
|
||||
cudaError_t e = cudaEventCreateWithFlags(&done[0], cudaEventDisableTiming | cudaEventBlockingSync);
|
||||
|
||||
+125
-20
@@ -213,6 +213,60 @@ bool ckpt_async_on() {
|
||||
return on;
|
||||
#endif
|
||||
}
|
||||
// STRATA_CKPT_GATHER=1 (CUDA, opt-in): the small running-state vectors (ple, then each layer's indexer tail, dead row
|
||||
// and block position: 1 + 3 x layers pieces) are gathered device to device into one buffer, allocated once, and
|
||||
// fetched with ONE copy into a pinned buffer, then copied into the vectors. A small D2H copy issued while the
|
||||
// checkpoint DMA (STRATA_CKPT_ASYNC) is in flight waits behind the 4 MiB piece on the copy engine: one such wait
|
||||
// instead of 37. The same bytes in the same places. One set of buffers: the mutex serializes the saves a layer
|
||||
// split runs on one thread per stage (STRATA_CKPT_THREADS), which otherwise gather through the same buffer.
|
||||
struct CkptGather { std::mutex mu; void* stage = nullptr; void* host = nullptr; size_t cap = 0; int dev = -1; };
|
||||
CkptGather& ckpt_gather() { static CkptGather a; return a; }
|
||||
bool ckpt_gather_on() {
|
||||
#if defined(STRATA_USE_HIP) || defined(STRATA_HIP_GFX906) || defined(STRATA_GPU_METAL)
|
||||
return false;
|
||||
#else
|
||||
static const bool on = [] { const char* v = std::getenv("STRATA_CKPT_GATHER"); return v != nullptr && v[0] == '1'; }();
|
||||
return on;
|
||||
#endif
|
||||
}
|
||||
// 1: the vectors are filled; 0: no buffers (the caller copies them one by one); -1: a copy failed (`error` set)
|
||||
int ckpt_gather_small(ConversationCheckpoint& c, const SessionState& ss, const ConversationStateSizes& z,
|
||||
size_t layers, std::string& error) {
|
||||
const size_t ple = c.ple.size(), tails = c.tails.size(), dead = c.dead.size(), block_pos = c.block_pos.size();
|
||||
const size_t total = ple + tails + dead + block_pos;
|
||||
if (total == 0) return 1;
|
||||
CkptGather& a = ckpt_gather();
|
||||
std::lock_guard<std::mutex> lk(a.mu);
|
||||
int dev = 0;
|
||||
(void) cudaGetDevice(&dev);
|
||||
if (a.stage != nullptr && (a.cap < total || a.dev != dev)) {
|
||||
const int cur = dev;
|
||||
(void) cudaSetDevice(a.dev); (void) cudaFree(a.stage); (void) cudaFreeHost(a.host); (void) cudaSetDevice(cur);
|
||||
a.stage = a.host = nullptr; a.cap = 0;
|
||||
}
|
||||
if (a.stage == nullptr) {
|
||||
if (cudaMalloc(&a.stage, total) == cudaSuccess && cudaMallocHost(&a.host, total) == cudaSuccess) { a.cap = total; a.dev = dev; }
|
||||
else { if (a.stage) (void) cudaFree(a.stage); if (a.host) (void) cudaFreeHost(a.host); a.stage = a.host = nullptr; }
|
||||
(void) cudaGetLastError();
|
||||
if (a.stage == nullptr) return 0;
|
||||
}
|
||||
// the pieces in the vectors' order, queued on the private stream (in order), then the one copy to the host
|
||||
uint8_t* s = (uint8_t*) a.stage;
|
||||
if (!copy(s, ss.ple_hist, ple, error, true)) return -1;
|
||||
for (size_t j = 0; j < layers; ++j) {
|
||||
const auto& st = owned(ss, j);
|
||||
if (!copy(s + ple + j * z.tail, st.idx_tail, z.tail, error, true) ||
|
||||
!copy(s + ple + tails + j * z.dead, st.idx_dead, z.dead, error, true) ||
|
||||
!copy(s + ple + tails + dead + j * z.block_pos, st.idx_block_pos, z.block_pos, error, true)) return -1;
|
||||
}
|
||||
if (!copy(a.host, a.stage, total, error)) return -1;
|
||||
const uint8_t* h = (const uint8_t*) a.host;
|
||||
if (ple) std::memcpy(c.ple.data(), h, ple);
|
||||
if (tails) std::memcpy(c.tails.data(), h + ple, tails);
|
||||
if (dead) std::memcpy(c.dead.data(), h + ple + tails, dead);
|
||||
if (block_pos) std::memcpy(c.block_pos.data(), h + ple + tails + dead, block_pos);
|
||||
return 1;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
void conversation_checkpoint_join(const void* only) {
|
||||
@@ -224,8 +278,12 @@ void conversation_checkpoint_join(const void* only) {
|
||||
}
|
||||
}
|
||||
|
||||
bool conversation_checkpoint_save(ConversationCheckpoint& c, const SessionState& ss,
|
||||
const ModelGeometry& g, std::string& error) {
|
||||
namespace {
|
||||
// conversation_checkpoint_save; `staged` (STRATA_CKPT_ASYNC): the GDN bytes land in `c.gdn` on the background thread
|
||||
// after this returns (conversation_checkpoint_join waits for them)
|
||||
bool checkpoint_save(ConversationCheckpoint& c, const SessionState& ss, const ModelGeometry& g, std::string& error,
|
||||
bool& staged) {
|
||||
staged = false;
|
||||
const auto j0 = std::chrono::steady_clock::now();
|
||||
conversation_checkpoint_join(nullptr); // (the staging buffer is reused)
|
||||
if (std::getenv("STRATA_PARK_TRACE")) std::fprintf(stderr, "strata parktrace: save join %.2f ms\n",
|
||||
@@ -237,9 +295,13 @@ bool conversation_checkpoint_save(ConversationCheckpoint& c, const SessionState&
|
||||
ConversationStateSizes z;
|
||||
if (!checkpoint_targets(ss, g, c.ids.size(), z, error)) return false;
|
||||
const size_t layers = owned_qsa(ss);
|
||||
static const bool atr = std::getenv("STRATA_PARK_TRACE") != nullptr;
|
||||
const bool ready = c.gdn.size() == z.gdn; // (a pooled checkpoint's pages are in; a fresh vector faults them all)
|
||||
const auto r0 = std::chrono::steady_clock::now();
|
||||
c.gdn.resize(z.gdn); c.ple.resize(ss.ple_hist ? z.ple : 0);
|
||||
c.tails.resize(layers * z.tail); c.dead.resize(layers * z.dead); c.block_pos.resize(layers * z.block_pos);
|
||||
bool staged = false;
|
||||
if (atr) std::fprintf(stderr, "strata parktrace: save vectors %s, resize %.2f ms\n", ready ? "ready" : "fresh",
|
||||
std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - r0).count());
|
||||
if (ckpt_async_on() && c.gdn.size() >= ((size_t) 1 << 20)) {
|
||||
CkptAsync& a = ckpt_async();
|
||||
int dev = 0;
|
||||
@@ -256,7 +318,7 @@ bool conversation_checkpoint_save(ConversationCheckpoint& c, const SessionState&
|
||||
else { if (a.stage) (void) cudaFree(a.stage); if (a.host) (void) cudaFreeHost(a.host); a.stage = a.host = nullptr; }
|
||||
}
|
||||
(void) cudaGetLastError();
|
||||
static const bool atr = std::getenv("STRATA_PARK_TRACE") != nullptr;
|
||||
const auto q0 = std::chrono::steady_clock::now();
|
||||
VramSlot* slot = vram_slots_n() > 0 ? vram_slot_for(c.gdn.size(), dev) : nullptr;
|
||||
if (slot != nullptr) slot->tag = 0; // (being rewritten)
|
||||
void* dsrc = slot != nullptr ? slot->p : a.stage;
|
||||
@@ -274,12 +336,13 @@ bool conversation_checkpoint_save(ConversationCheckpoint& c, const SessionState&
|
||||
const uint8_t* stage = (const uint8_t*) dsrc;
|
||||
cudaStream_t st = a.stream;
|
||||
if (atr) std::fprintf(stderr, "strata parktrace: async ckpt d2d %.2f ms\n",
|
||||
std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - q1).count());
|
||||
std::chrono::duration<double, std::milli>(q1 - q0).count());
|
||||
std::lock_guard<std::mutex> lk(a.mu);
|
||||
// the DMA in 4 MiB pieces, one in flight: a D2H copy queues behind the one before it on the copy engine, and
|
||||
// the decode's own small reads would otherwise wait for the whole 100+ MB at the link's ~1.6 GB/s
|
||||
a.dst = dst;
|
||||
a.pending = std::async(std::launch::async, [dst, host, stage, n, st, dev] {
|
||||
conversation_detail::background_thread_placement();
|
||||
(void) cudaSetDevice(dev);
|
||||
cudaEvent_t done = nullptr;
|
||||
cudaError_t e = cudaEventCreateWithFlags(&done, cudaEventDisableTiming | cudaEventBlockingSync);
|
||||
@@ -297,15 +360,33 @@ bool conversation_checkpoint_save(ConversationCheckpoint& c, const SessionState&
|
||||
staged = true;
|
||||
}
|
||||
}
|
||||
if ((!staged && !copy(c.gdn.data(), ss.gdn_state, c.gdn.size(), error)) ||
|
||||
!copy(c.ple.data(), ss.ple_hist, c.ple.size(), error)) return false;
|
||||
for (size_t j = 0; j < layers; ++j) {
|
||||
const auto& st = owned(ss, j);
|
||||
if (!copy(c.tails.data() + j * z.tail, st.idx_tail, z.tail, error) ||
|
||||
!copy(c.dead.data() + j * z.dead, st.idx_dead, z.dead, error) ||
|
||||
!copy(c.block_pos.data() + j * z.block_pos, st.idx_block_pos, z.block_pos, error)) return false;
|
||||
}
|
||||
return true;
|
||||
if (!staged && !copy(c.gdn.data(), ss.gdn_state, c.gdn.size(), error)) return false;
|
||||
const auto s0 = std::chrono::steady_clock::now();
|
||||
const bool ok = [&] {
|
||||
// STRATA_CKPT_GATHER=1: the small vectors through one device gather and one copy (else one copy each)
|
||||
if (ckpt_gather_on()) {
|
||||
const int r = ckpt_gather_small(c, ss, z, layers, error);
|
||||
if (r != 0) return r > 0;
|
||||
}
|
||||
if (!copy(c.ple.data(), ss.ple_hist, c.ple.size(), error)) return false;
|
||||
for (size_t j = 0; j < layers; ++j) {
|
||||
const auto& st = owned(ss, j);
|
||||
if (!copy(c.tails.data() + j * z.tail, st.idx_tail, z.tail, error) ||
|
||||
!copy(c.dead.data() + j * z.dead, st.idx_dead, z.dead, error) ||
|
||||
!copy(c.block_pos.data() + j * z.block_pos, st.idx_block_pos, z.block_pos, error)) return false;
|
||||
}
|
||||
return true;
|
||||
}();
|
||||
if (atr) std::fprintf(stderr, "strata parktrace: save small copies %.2f ms\n",
|
||||
std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - s0).count());
|
||||
return ok;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
bool conversation_checkpoint_save(ConversationCheckpoint& c, const SessionState& ss,
|
||||
const ModelGeometry& g, std::string& error) {
|
||||
bool staged = false;
|
||||
return checkpoint_save(c, ss, g, error, staged);
|
||||
}
|
||||
|
||||
bool conversation_checkpoint_restore(const ConversationCheckpoint& c, SessionState& ss,
|
||||
@@ -430,18 +511,22 @@ bool conversation_snapshot_capture_bytes(const ConversationKvReuse& reuse, const
|
||||
return true;
|
||||
}
|
||||
|
||||
// (conversation_snapshot.cpp) STRATA_PARK_TRACE: conversation_kv_save's time since the park's begin, by part (ms)
|
||||
void conversation_kv_save_trace(double& attrs_ms, double& resize_ms, double& copy_ms);
|
||||
|
||||
bool conversation_snapshot_save(SavedConversation& image, const ConversationView& view,
|
||||
const SessionState& ss, const ModelGeometry& g,
|
||||
const QsaState* draft, std::string& error,
|
||||
ConversationKvReuse reuse, size_t* reused_bytes,
|
||||
std::vector<ConversationCheckpoint>* storage,
|
||||
std::vector<ConversationCheckpoint>* moved) {
|
||||
conversation_checkpoint_join(nullptr);
|
||||
conversation_park_join();
|
||||
static const bool ptr = std::getenv("STRATA_PARK_TRACE") != nullptr;
|
||||
auto now = [] { return std::chrono::steady_clock::now(); };
|
||||
auto ms = [](std::chrono::steady_clock::time_point a, std::chrono::steady_clock::time_point b) {
|
||||
return std::chrono::duration<double, std::milli>(b - a).count(); };
|
||||
const auto pj = now();
|
||||
conversation_checkpoint_join(nullptr); // the last turn checkpoint's DMA (STRATA_CKPT_ASYNC)
|
||||
conversation_park_join(); // the previous park's K/V DMA
|
||||
const auto p0 = now();
|
||||
size_t estimate = 0;
|
||||
if (!conversation_snapshot_capture_bytes(reuse, view, ss, g, draft, estimate, error) || !sync(error)) return false;
|
||||
@@ -467,7 +552,17 @@ bool conversation_snapshot_save(SavedConversation& image, const ConversationView
|
||||
captured.kv = std::move(reuse.kv);
|
||||
const size_t layers = owned_qsa(ss);
|
||||
captured.kv.resize(layers + (draft ? 1 : 0));
|
||||
if (!conversation_checkpoint_save(captured.live, ss, g, error)) return false;
|
||||
bool live_staged = false;
|
||||
const bool live_saved = checkpoint_save(captured.live, ss, g, error, live_staged);
|
||||
// STRATA_CKPT_ASYNC: the running state lands in live.gdn after the save returns; the image waits for that copy
|
||||
// before the vector is freed (an eviction within the copy's ~70 ms), as `pending` does for the K/V. The join is
|
||||
// selective: a no-op once the copy has landed or another checkpoint's is in flight.
|
||||
if (live_staged)
|
||||
captured.live_pending = std::shared_ptr<void>(new char, [dst = (const void*) captured.live.gdn.data()](void* p) {
|
||||
conversation_checkpoint_join(dst);
|
||||
delete (char*) p;
|
||||
});
|
||||
if (!live_saved) return false;
|
||||
const auto p3 = now();
|
||||
const int64_t upto = (int64_t) view.ids.size();
|
||||
// STRATA_PARK_ASYNC: the K/V bytes these saves copy (the stage they need; an over-estimate by the kept prefix)
|
||||
@@ -475,15 +570,25 @@ bool conversation_snapshot_save(SavedConversation& image, const ConversationView
|
||||
for (size_t j = 0; j < layers; ++j) kv_need += conversation_kv_bytes(owned(ss, j), g, upto, true);
|
||||
if (draft) kv_need += conversation_kv_bytes(*draft, g, upto, false);
|
||||
const bool staged = conversation_park_async_begin(kv_need);
|
||||
const auto p4 = now();
|
||||
struct Abandon { bool on; ~Abandon() { if (on) conversation_park_async_abandon(); } } abandon{staged};
|
||||
for (size_t j = 0; j < layers; ++j)
|
||||
double kv_max = 0, kv_sum = 0; // (STRATA_PARK_TRACE: the per-layer save time)
|
||||
for (size_t j = 0; j < layers; ++j) {
|
||||
const auto k0 = now();
|
||||
if (!conversation_kv_save(captured.kv[j], owned(ss, j), g, upto, true, error,
|
||||
unchanged, reused_bytes)) return false;
|
||||
const double kt = ms(k0, now());
|
||||
kv_sum += kt; kv_max = std::max(kv_max, kt);
|
||||
}
|
||||
if (ptr) {
|
||||
size_t kvb = 0;
|
||||
for (const auto& k : captured.kv) kvb += k.bytes();
|
||||
std::fprintf(stderr, "strata parktrace: sync %.1f ms, %zu checkpoints copied %.1f ms, live state %.1f ms, K/V %zu MB %.1f ms\n",
|
||||
ms(p0, p1), captured.checkpoints.size(), ms(p1, p2), ms(p2, p3), kvb >> 20, ms(p3, now()));
|
||||
double attrs_ms = 0, resize_ms = 0, copy_ms = 0;
|
||||
conversation_kv_save_trace(attrs_ms, resize_ms, copy_ms); // (the main layers' saves above)
|
||||
std::fprintf(stderr, "strata parktrace: joins %.1f ms, sync %.1f ms, %zu checkpoints copied %.1f ms, live state %.1f ms, "
|
||||
"begin %.1f ms, K/V %zu MB %.1f ms (layer max %.2f, sum %.2f; attrs %.2f, resize %.2f, copy %.2f)\n",
|
||||
ms(pj, p0), ms(p0, p1), captured.checkpoints.size(), ms(p1, p2), ms(p2, p3), ms(p3, p4), kvb >> 20,
|
||||
ms(p4, now()), kv_max, kv_sum, attrs_ms, resize_ms, copy_ms);
|
||||
}
|
||||
// The draft's final cell may not have been computed when the output cap was
|
||||
// reached. Refresh that page even when the main prefix continued unchanged.
|
||||
|
||||
+43
-12
@@ -210,6 +210,17 @@ bool ckpt_pool_env() {
|
||||
}();
|
||||
return on;
|
||||
}
|
||||
// STRATA_CKPT_POOL=N (N >= 2, read once): how many ready checkpoints the pool keeps (1 or any other value: 2, as
|
||||
// before). Each is a running-state image with its pages in (~118 MB on Qwen3.8-Flash-Next), allocated on a
|
||||
// background thread.
|
||||
size_t ckpt_pool_cap_env() {
|
||||
static const size_t n = [] {
|
||||
const char* v = std::getenv("STRATA_CKPT_POOL");
|
||||
const long k = v ? std::atol(v) : 0;
|
||||
return k >= 2 ? (size_t) std::min(16L, k) : (size_t) 2;
|
||||
}();
|
||||
return n;
|
||||
}
|
||||
// STRATA_CKPT_THREADS=1 (read once; off by default): a layer split's checkpoint saves every stage's part at once, the
|
||||
// first on the calling thread and each later stage's on a thread of its own device (kept for the process: the copies'
|
||||
// private streams are per thread), instead of one stage after the other. Each part is the same copy of the same state
|
||||
@@ -1733,11 +1744,14 @@ bool checkpoint_restore(const ConvCheckpoint& c, strata::core::SessionState& ss,
|
||||
|
||||
/// STRATA_CKPT_POOL: the storage of dropped checkpoints, for the next save or slot copy to reuse. `give` keeps a
|
||||
/// checkpoint's vectors (its tokens, pictures and marks cleared), `take` hands one back (an empty checkpoint when none
|
||||
/// is kept). At most `cap` are kept (a batch admission drops and takes about two), the rest are freed as before.
|
||||
/// is kept). At most `cap` are kept (a batch admission drops and takes about two; STRATA_CKPT_POOL=N), the rest are
|
||||
/// freed as before.
|
||||
class CheckpointPool {
|
||||
public:
|
||||
explicit CheckpointPool(bool on, size_t cap = 2) : on_(on), cap_(cap) {}
|
||||
bool on() const { return on_; }
|
||||
size_t cap() const { return cap_; }
|
||||
size_t held() { std::lock_guard<std::mutex> lk(mu_); return free_.size(); }
|
||||
void give(ConvCheckpoint&& c) {
|
||||
if (!on_ || c.gdn.empty()) return;
|
||||
strata::core::conversation_checkpoint_join(c.gdn.data());
|
||||
@@ -1810,7 +1824,6 @@ public:
|
||||
done_.wait(lk, [&] { return pending_ == 0; });
|
||||
job_ = nullptr;
|
||||
}
|
||||
private:
|
||||
static void spread_current_thread() {
|
||||
#if defined(__linux__)
|
||||
// wide first: the mask the aux placement keeps as this thread's own is the one restore() puts back (a request
|
||||
@@ -1822,6 +1835,7 @@ private:
|
||||
strata::aux_cpus::pin_current_thread(); // (a no-op while the feature is off)
|
||||
#endif
|
||||
}
|
||||
private:
|
||||
void loop(size_t k) {
|
||||
strata::platform::name_this_thread("strata-ckpt", (int) k);
|
||||
spread_current_thread();
|
||||
@@ -9241,7 +9255,7 @@ int main(int argc, char** argv) {
|
||||
std::vector<ConvCheckpoint> checks;
|
||||
// STRATA_CKPT_POOL: the dropped checkpoints' storage, reused. STRATA_SLOT_CKPT_MOVE turns it on too: a moved
|
||||
// checkpoint frees nothing at the next prune, so every save would fault in fresh pages (17 -> 40-60 ms)
|
||||
CheckpointPool ckpt_pool(ckpt_pool_env() || slot_ckpt_move_env());
|
||||
CheckpointPool ckpt_pool(ckpt_pool_env() || slot_ckpt_move_env(), ckpt_pool_cap_env());
|
||||
uint64_t check_clock = 0; // the checkpoints' LRU clock; creation and every use advance it
|
||||
int64_t tail_ckpt_len = -1; // --prompt-cache-tail: the length of the one tail checkpoint alive (-1 = none)
|
||||
// --prefill-pipe: the plan of the prompt segment being read (read_part sets it around its runs), for the
|
||||
@@ -9333,6 +9347,30 @@ int main(int argc, char** argv) {
|
||||
// into the parked image instead of being copied); a background allocation then refills STRATA_CKPT_POOL
|
||||
std::function<bool(const ConvCheckpoint&)> park_drops;
|
||||
std::future<void> pool_refill;
|
||||
// STRATA_CKPT_POOL: the pool holds cap() checkpoints whose vectors already have the session's sizes (the ones
|
||||
// conversation_checkpoint_save resizes to, so a pooled one is never resized) with their pages in: a fresh
|
||||
// 118 MB vector costs ~40 ms of page faults at a save. Topped up on a background thread, off the host's core,
|
||||
// whenever a take leaves it short; one top-up at a time. Allocation only: a save overwrites every byte. Not
|
||||
// with a layer split (its checkpoints carry stage parts; saved as before).
|
||||
const ConvStateSizes pool_z = conv_state_sizes(g, ss);
|
||||
auto pool_topup = [&] {
|
||||
if (!ckpt_pool.on() || !stages.empty() || pool_z.gdn == 0) return;
|
||||
if (pool_refill.valid() && pool_refill.wait_for(std::chrono::seconds(0)) != std::future_status::ready) return;
|
||||
const size_t have = ckpt_pool.held();
|
||||
if (have >= ckpt_pool.cap()) return;
|
||||
pool_refill = std::async(std::launch::async, [&ckpt_pool, want = ckpt_pool.cap() - have, z = pool_z,
|
||||
layers = (size_t) std::max<int64_t>(ss.qsa_alloc, 0),
|
||||
ple = ss.ple_hist != nullptr] {
|
||||
StageThreads::spread_current_thread(); // (a thread the host starts inherits its one CPU: aux_cpus.hpp)
|
||||
for (size_t k = 0; k < want; ++k) {
|
||||
ConvCheckpoint c;
|
||||
c.gdn.resize(z.gdn); c.ple.resize(ple ? z.ple : 0); c.tails.resize(layers * z.tail);
|
||||
c.dead.resize(layers * z.dead); c.block_pos.resize(layers * z.block_pos);
|
||||
ckpt_pool.give(std::move(c));
|
||||
}
|
||||
});
|
||||
};
|
||||
pool_topup(); // the first request's checkpoints find the pool full
|
||||
auto park_current_body = [&](size_t held) -> bool {
|
||||
if (!conversations.enabled() || !live_ok || live.empty()) return true;
|
||||
// #2049: parking copies the outgoing conversation's running checkpoints and K/V into host memory, which
|
||||
@@ -9484,20 +9522,12 @@ int main(int argc, char** argv) {
|
||||
}
|
||||
checks.erase(checks.begin() + (std::ptrdiff_t) w, checks.end());
|
||||
}
|
||||
const size_t moved_gdn = moved.empty() ? 0 : moved.front().gdn.size();
|
||||
const bool saved = strata::core::conversation_snapshot_save(image, view, ss, g, draft0, err,
|
||||
std::move(reuse), &reused_bytes, storage.empty() ? nullptr : &storage,
|
||||
moved.empty() ? nullptr : &moved);
|
||||
for (ConvCheckpoint& c : storage) ckpt_pool.give(std::move(c));
|
||||
if (!saved) return false;
|
||||
// the next checkpoint save would fault in fresh pages for what moved out: one is allocated now
|
||||
if (moved_gdn > 0 && ckpt_pool.on() &&
|
||||
(!pool_refill.valid() || pool_refill.wait_for(std::chrono::seconds(0)) == std::future_status::ready))
|
||||
pool_refill = std::async(std::launch::async, [&ckpt_pool, moved_gdn] {
|
||||
ConvCheckpoint c;
|
||||
c.gdn.resize(moved_gdn);
|
||||
ckpt_pool.give(std::move(c));
|
||||
});
|
||||
pool_topup(); // what the park took (and what moved out) is allocated again, in the background
|
||||
for (size_t k = 0; k < n_st; ++k) { // the later stages (the last one with the draft layer's K/V)
|
||||
auto& st = stages[k];
|
||||
const strata::core::OnDevice on(st->dev);
|
||||
@@ -9605,6 +9635,7 @@ int main(int argc, char** argv) {
|
||||
return true;
|
||||
}
|
||||
ConvCheckpoint c = parts == nullptr ? ckpt_pool.take() : ConvCheckpoint{}; // (STRATA_CKPT_POOL)
|
||||
if (parts == nullptr) pool_topup(); // the next save finds one ready
|
||||
c.ids.assign(cur.begin(), cur.begin() + L);
|
||||
c.imgs = imgs_below(req_imgs, L);
|
||||
if (parts != nullptr) {
|
||||
|
||||
Reference in new issue
Block a user