mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-10 17:33:53 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0c2e63e930 | ||
|
|
888751a0e1 | ||
|
|
af1c91f245 | ||
|
|
14b22c51c1 | ||
|
|
be028948f4 | ||
|
|
470b3d7c21 | ||
|
|
be0e0c94e3 | ||
|
|
42810ee8bb | ||
|
|
0f8b5d4938 |
No files matched your search
@@ -98,8 +98,8 @@ Not sure? Take **IQ2_XS**. You can add another one later with `START-HERE.bat --
|
||||
- **Pictures:** in the chat page click **Picture**; in `chat.py` type `/image <path>`; in apps just attach them.
|
||||
- **From your phone or another PC:** see the [details](docs/DETAILS.md#using-it) (set an API key first).
|
||||
|
||||
**Good to know:** it answers one request at a time, and it re-reads the whole conversation for every answer. So in
|
||||
very long chats you wait longer before it starts writing: about 1 minute per 30,000 tokens of conversation.
|
||||
**Good to know:** it answers one request at a time. The first message of a chat is read in full (about 1 minute per
|
||||
30,000 tokens); after that it keeps the conversation and reads only what is new, so follow-ups start in seconds.
|
||||
|
||||
## Something went wrong?
|
||||
|
||||
|
||||
+8
-3
@@ -200,9 +200,14 @@ print(r.choices[0].message.content)
|
||||
add `"api_key": "some-long-secret"` to `strata-<model>.json` (or set the `STRATA_API_KEY` environment variable);
|
||||
clients then send it as their API key.
|
||||
|
||||
**Current limits (v1):** one request at a time; greedy decoding (temperature is ignored); every request processes its
|
||||
whole prompt again (no conversation cache yet, so long chats have a long time-to-first-token); images only when set up
|
||||
with them (below); no video.
|
||||
**Conversation cache.** A request that continues a chat reads only the part after what the engine already holds: the
|
||||
live session, or one of the checkpoints it keeps in RAM (up to 6, ~118 MB each, taken at the start of each new
|
||||
assistant turn and every 16K prompt tokens). A checkpoint is used only when the prompt starts with exactly its tokens
|
||||
and pictures. Engine options: `--prompt-cache N` (0 = off), `--prompt-cache-every N`, `--turn-token ID`.
|
||||
|
||||
**Current limits (v1):** one request at a time, and one conversation cached at a time (switching between two chats
|
||||
re-reads the other one); greedy decoding (temperature is ignored); images only when set up with them (below); no
|
||||
video.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -46,6 +46,8 @@ public:
|
||||
/// The prompt's length: prefill() skips the cells the attention window can never reach again.
|
||||
void set_prompt_len(int64_t n) { prompt_len_ = n; }
|
||||
uint64_t vram_bytes() const { return vram_; }
|
||||
/// The draft layer's K/V state (read-only: --serve's STRATA_STATE_HASH check hashes it)
|
||||
const QsaState& kv_state() const { return st_; }
|
||||
/// The main model's embedding and head, and the verify window's final residuals (T rows, hc*n_embd each).
|
||||
bool bind(const WeightTable& wt, const NativeHead* head, const float* window_R, std::string& err);
|
||||
|
||||
|
||||
@@ -23,7 +23,10 @@
|
||||
#include "strata/kernels/cpu/native_expert.hpp"
|
||||
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <condition_variable>
|
||||
#include <cstdint>
|
||||
#include <mutex>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
@@ -136,10 +139,19 @@ public:
|
||||
repark = ms_repark_;
|
||||
}
|
||||
|
||||
/// **A PARKED WORKER SPINS FOR THIS LONG, THEN SLEEPS.** The park is a `_mm_pause` spin because a layer's
|
||||
/// batches are microseconds apart and a wake-up from the OS costs more than that. But a spin that never ends
|
||||
/// keeps every worker's core at 100% while the engine waits for a request - issue #4, "CPU 50% even when
|
||||
/// doing nothing" (7 of the 5700X's 16 threads). Inside a request the gaps between batches are far below
|
||||
/// this, so the token path still never sleeps; between requests the workers block on `sleep_cv_`.
|
||||
static constexpr std::chrono::milliseconds kSpinBeforeSleep{20};
|
||||
|
||||
private:
|
||||
void worker(int id);
|
||||
void drain(int id, ExpertScratch& scratch);
|
||||
void run_phase(int mode, int n_tasks);
|
||||
/// Bump `epoch_`, and wake the workers that went to sleep. Every publish goes through here.
|
||||
void publish();
|
||||
|
||||
int n_ = 0;
|
||||
bool host_works_ = true;
|
||||
@@ -169,6 +181,11 @@ private:
|
||||
alignas(64) std::atomic<uint32_t> parked_{0};
|
||||
alignas(64) std::atomic<uint32_t> epoch_{0};
|
||||
alignas(64) std::atomic<bool> stop_{false};
|
||||
// The sleep after `kSpinBeforeSleep`. `sleepers_` is how `publish` knows whether anyone needs waking, so the
|
||||
// token path pays one uncontended load per publish and never takes the mutex while the workers spin.
|
||||
alignas(64) std::atomic<uint32_t> sleepers_{0};
|
||||
std::mutex sleep_mu_;
|
||||
std::condition_variable sleep_cv_;
|
||||
std::vector<std::thread> threads_;
|
||||
std::vector<ExpertScratch> scratch_; // one per worker: no allocation, no false sharing of the hot data
|
||||
// run_split state: mode 0 = whole experts, 1 = gate/up row parts, 2 = down row parts
|
||||
|
||||
+60
-18
@@ -127,6 +127,10 @@ class StrataEngine:
|
||||
if cancel.is_set():
|
||||
return
|
||||
yield int(line[2:])
|
||||
elif line.startswith("PP "): # prompt progress, one per chunk: also a heartbeat (the
|
||||
if cancel.is_set(): # lines reset the 10 s wait, so without this a long prompt
|
||||
return # would send no keep-alives at all)
|
||||
yield None
|
||||
elif line.startswith("DONE"):
|
||||
self._parse_done(line)
|
||||
done = True
|
||||
@@ -370,7 +374,7 @@ class Service:
|
||||
s["phase"] = "tool call complete"
|
||||
s["tail"] = (s["tail"] + (ev.text or ""))[-600:]
|
||||
|
||||
def _progress(self, last_print, every=15.0):
|
||||
def _progress(self, last_print, every=1.0):
|
||||
"""A progress line in the server window every `every` seconds while a request runs."""
|
||||
now = time.time()
|
||||
if now - last_print < every:
|
||||
@@ -390,6 +394,7 @@ class Service:
|
||||
"""Yields ("event", Event) as text arrives, then ("done", {"finish": .., "completion_tokens": ..})."""
|
||||
parser = OutputParser(thinking=thinking, tools=tools, stream_tools=True)
|
||||
detok, n, finish = Detokenizer(self.tok), 0, "length"
|
||||
raw_ids = [] # every generated id (STRATA_DEBUG: dump raw model text)
|
||||
emb = getattr(self.embeddings, "path", None)
|
||||
with self.status_lock:
|
||||
self.status["queued"] += 1
|
||||
@@ -402,35 +407,66 @@ class Service:
|
||||
last_print = time.time()
|
||||
gen = self.engine.generate(ids, max_new, sampling, cancel, embeddings=emb) if emb else \
|
||||
self.engine.generate(ids, max_new, sampling, cancel)
|
||||
for t in gen:
|
||||
if t is None: # heartbeat while the engine is quiet
|
||||
try:
|
||||
for t in gen:
|
||||
if t is None: # heartbeat while the engine is quiet
|
||||
last_print = self._progress(last_print)
|
||||
yield "ping", None
|
||||
continue
|
||||
n += 1
|
||||
if t in self.stop_ids:
|
||||
finish = "stop"
|
||||
raw_ids.append(t)
|
||||
break
|
||||
raw_ids.append(t)
|
||||
evs = parser.feed(detok.push(t))
|
||||
self._note(n, evs)
|
||||
last_print = self._progress(last_print)
|
||||
yield "ping", None
|
||||
continue
|
||||
n += 1
|
||||
if t in self.stop_ids:
|
||||
finish = "stop"
|
||||
break
|
||||
evs = parser.feed(detok.push(t))
|
||||
self._note(n, evs)
|
||||
last_print = self._progress(last_print)
|
||||
for ev in evs:
|
||||
yield "event", ev
|
||||
if cancel.is_set():
|
||||
finish = "cancel"
|
||||
for ev in evs:
|
||||
yield "event", ev
|
||||
if cancel.is_set():
|
||||
finish = "cancel"
|
||||
finally:
|
||||
gen.close() # STOP+drain to THIS request's DONE while still holding the
|
||||
# fifo, so a stop-token break can't leave the shared engine
|
||||
# queue mid-drain for the next request to read as its own DONE
|
||||
except GeneratorExit: # the client disconnected mid-stream
|
||||
finish = "disconnect"
|
||||
raise
|
||||
finally:
|
||||
if emb:
|
||||
Path(emb).unlink(missing_ok=True)
|
||||
with self.status_lock:
|
||||
if self.status.get("busy"):
|
||||
el = time.time() - self.status.get("started", time.time())
|
||||
print(f"[strata] done: {n} tokens in {el:.0f} s ({finish})", flush=True)
|
||||
now = time.time()
|
||||
el = now - self.status.get("started", now)
|
||||
ft = self.status.get("first_token")
|
||||
rate = n / max(1e-6, now - ft) if ft else 0.0
|
||||
print(f"[strata] done: {n} tokens in {el:.0f} s ({rate:.1f} tok/s) "
|
||||
f"({finish}, cancel={cancel.is_set()})", flush=True)
|
||||
if os.environ.get("STRATA_DEBUG") and raw_ids:
|
||||
print(f"[strata] raw: {self.tok.decode(raw_ids)!r}", flush=True)
|
||||
self.status["busy"] = False
|
||||
for ev in parser.finish():
|
||||
yield "event", ev
|
||||
yield "done", {"finish": finish, "completion_tokens": n}
|
||||
|
||||
|
||||
def _debug_req(api, req, messages, tools, max_new, thinking, prompt_tokens):
|
||||
"""One compact line per request while diagnosing blank/empty turns. Set STRATA_DEBUG=1 to enable."""
|
||||
if not os.environ.get("STRATA_DEBUG"):
|
||||
return
|
||||
last = messages[-1] if messages else {}
|
||||
body = last.get("content")
|
||||
if isinstance(body, list):
|
||||
body = " ".join(p.get("text", "") for p in body if isinstance(p, dict))
|
||||
preview = (str(body or "")[:80]).replace("\n", " ")
|
||||
print(f"[strata] req {api}: msgs={len(messages)} tools={len(tools or [])} "
|
||||
f"max_tokens_raw={req.get('max_tokens')!r}/{req.get('max_completion_tokens')!r} "
|
||||
f"max_new={max_new} thinking={thinking} stream={bool(req.get('stream'))} "
|
||||
f"prompt_tokens={prompt_tokens} last={last.get('role')!r}:{preview!r}", flush=True)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ OpenAI
|
||||
def openai_chunks(svc: Service, req: dict, ids, thinking, tools, max_new, cancel):
|
||||
cid, created = "chatcmpl-" + uuid.uuid4().hex[:24], int(time.time())
|
||||
@@ -668,7 +704,10 @@ def make_handler(svc: Service):
|
||||
def _openai(self, req):
|
||||
messages, tools, kw = openai_to_messages(req)
|
||||
max_new = int(req.get("max_completion_tokens") or req.get("max_tokens") or 1024)
|
||||
if max_new <= 0: # a non-positive budget (some clients send -1) means "unset"
|
||||
max_new = 1024
|
||||
ids, thinking = svc.prepare(messages, tools, kw, max_new)
|
||||
_debug_req("openai", req, messages, tools, max_new, thinking, len(ids))
|
||||
cancel = threading.Event()
|
||||
chunks = openai_chunks(svc, req, ids, thinking, tools, max_new, cancel)
|
||||
if not req.get("stream"):
|
||||
@@ -689,7 +728,10 @@ def make_handler(svc: Service):
|
||||
def _anthropic(self, req):
|
||||
messages, tools, kw = anthropic_to_messages(req)
|
||||
max_new = int(req.get("max_tokens") or 1024)
|
||||
if max_new <= 0: # a non-positive budget (some clients send -1) means "unset"
|
||||
max_new = 1024
|
||||
ids, thinking = svc.prepare(messages, tools, kw, max_new)
|
||||
_debug_req("anthropic", req, messages, tools, max_new, thinking, len(ids))
|
||||
cancel = threading.Event()
|
||||
events = anthropic_events(svc, req, ids, thinking, tools, max_new, cancel)
|
||||
if not req.get("stream"):
|
||||
|
||||
@@ -49,11 +49,12 @@ std::vector<int> physical_cores(bool skip_first) {
|
||||
#else
|
||||
cpu_set_t set;
|
||||
CPU_ZERO(&set);
|
||||
if (sched_getaffinity(0, sizeof set, &set) == 0)
|
||||
if (sched_getaffinity(0, sizeof set, &set) == 0) {
|
||||
for (int i = 0; i < CPU_SETSIZE; ++i)
|
||||
if (CPU_ISSET(i, &set)) cores.push_back(i);
|
||||
else
|
||||
} else {
|
||||
for (unsigned i = 0; i < std::thread::hardware_concurrency(); ++i) cores.push_back((int) i);
|
||||
}
|
||||
#endif
|
||||
if (skip_first && !cores.empty()) cores.erase(cores.begin());
|
||||
return cores;
|
||||
@@ -127,10 +128,23 @@ ExpertPool::ExpertPool(int n_workers, bool pin, bool host_works) : host_works_(h
|
||||
ExpertPool::~ExpertPool() {
|
||||
stop_.store(true, std::memory_order_release);
|
||||
// Bump the epoch so a PARKED worker notices the stop flag rather than sleeping through it.
|
||||
epoch_.fetch_add(1, std::memory_order_release);
|
||||
publish();
|
||||
for (auto& t : threads_) t.join();
|
||||
}
|
||||
|
||||
void ExpertPool::publish() {
|
||||
// Both sides are seq_cst, and that is the whole lost-wakeup argument: a worker going to sleep does
|
||||
// `sleepers_++` and then reads `epoch_`, the host does `epoch_++` and then reads `sleepers_`. In one total
|
||||
// order at least one of them sees the other's write - the worker sees the new epoch and does not sleep, or
|
||||
// the host sees the sleeper and notifies under the mutex the worker holds until it is inside `wait`.
|
||||
// On x86 the fetch_add is a locked xadd either way, so this costs the token path nothing.
|
||||
epoch_.fetch_add(1, std::memory_order_seq_cst);
|
||||
if (sleepers_.load(std::memory_order_seq_cst) != 0) {
|
||||
std::lock_guard<std::mutex> lk(sleep_mu_);
|
||||
sleep_cv_.notify_all();
|
||||
}
|
||||
}
|
||||
|
||||
void ExpertPool::worker(int id) {
|
||||
uint32_t seen = 0;
|
||||
// ARRIVE at the park before the first wait, so `parked_ == n_` is true from construction. Counting only
|
||||
@@ -147,9 +161,22 @@ void ExpertPool::worker(int id) {
|
||||
// - a locked read-modify-write, five workers against one cache line - so the workers spent their wait
|
||||
// invalidating each other's caches and the very line the host writes to publish work. The counter was
|
||||
// diagnostic and nothing branched on it. See the note on the atomics in pool.hpp.
|
||||
//
|
||||
// After `kSpinBeforeSleep` with no work the worker sleeps instead (issue #4). The clock is read once
|
||||
// every 1024 pauses, so the spin itself is unchanged.
|
||||
const auto parked_at = std::chrono::steady_clock::now();
|
||||
uint32_t spins = 0;
|
||||
while (epoch_.load(std::memory_order_acquire) == seen) {
|
||||
if (stop_.load(std::memory_order_relaxed)) return;
|
||||
_mm_pause();
|
||||
if ((++spins & 1023u) != 0) continue;
|
||||
if (std::chrono::steady_clock::now() - parked_at < kSpinBeforeSleep) continue;
|
||||
std::unique_lock<std::mutex> lk(sleep_mu_);
|
||||
sleepers_.fetch_add(1, std::memory_order_seq_cst);
|
||||
sleep_cv_.wait(lk, [&] {
|
||||
return epoch_.load(std::memory_order_seq_cst) != seen || stop_.load(std::memory_order_relaxed);
|
||||
});
|
||||
sleepers_.fetch_sub(1, std::memory_order_relaxed);
|
||||
}
|
||||
if (stop_.load(std::memory_order_acquire)) return;
|
||||
seen = epoch_.load(std::memory_order_relaxed);
|
||||
@@ -245,7 +272,7 @@ void ExpertPool::run_phase(int mode, int n_tasks) {
|
||||
njobs_ = n_tasks;
|
||||
head_.store(0, std::memory_order_relaxed);
|
||||
done_.store(0, std::memory_order_relaxed);
|
||||
epoch_.fetch_add(1, std::memory_order_release);
|
||||
publish();
|
||||
if (host_works_) drain(-1, host_scratch_);
|
||||
while (done_.load(std::memory_order_acquire) != (uint32_t) n_tasks) _mm_pause();
|
||||
while (parked_.load(std::memory_order_acquire) != (uint32_t) n_) _mm_pause();
|
||||
@@ -356,7 +383,7 @@ void ExpertPool::run(ExpertJob* jobs, int n) {
|
||||
mode_ = 0;
|
||||
head_.store(0, std::memory_order_relaxed);
|
||||
done_.store(0, std::memory_order_relaxed);
|
||||
epoch_.fetch_add(1, std::memory_order_release); // release: jobs_/njobs_ are visible before the bump
|
||||
publish(); // a release (seq_cst): jobs_/njobs_ are visible before the bump
|
||||
|
||||
// ---- **THE HOST DRAINS TOO (R2.2), INSTEAD OF SPINNING ON `done_`.**
|
||||
//
|
||||
|
||||
+396
-30
@@ -235,6 +235,19 @@ struct Options {
|
||||
/// The vision path: keep a per-cell (t, h, w) rotary position table so --serve can take GENI requests.
|
||||
bool vision = false;
|
||||
int adapt_swaps = 96;
|
||||
/// --serve: how many conversation checkpoints to keep between requests (0 = every request reads its whole
|
||||
/// prompt again, the v0.1.2 behaviour). One is the GDN recurrence of the 36 layers, the QSA indexer tails and
|
||||
/// the PLE history (~118 MB of host RAM); the KV cache itself is positional and stays where it is.
|
||||
int prompt_cache = 6;
|
||||
/// --serve: also keep a checkpoint every N freshly read prompt tokens (0 = only at the last turn boundary)
|
||||
int64_t prompt_cache_every = 16384;
|
||||
/// --serve: the token that opens a chat turn (<|im_start|>). The last one in a prompt is where the chat's
|
||||
/// history ends and the new assistant turn begins, which is the checkpoint the next request can reuse.
|
||||
int64_t turn_token = 248045;
|
||||
/// --serve: a text part of the prompt of at most N tokens (a chat message, the assistant header, a short tool
|
||||
/// result) goes through the verify windows, S tokens at a time, instead of the batched prompt path (0 = always
|
||||
/// the batched path)
|
||||
int64_t short_read = 64;
|
||||
};
|
||||
|
||||
void usage() {
|
||||
@@ -292,6 +305,12 @@ void usage() {
|
||||
" --shared-late A/B: shared expert after the CPU pool (default: overlapped with it)\n"
|
||||
" --keep-canonical A/B: also load canonical copies of natively served tensors (more VRAM)\n"
|
||||
" --vision --serve takes images too (GENI requests; embeddings from strata-vision)\n"
|
||||
" --prompt-cache N --serve: keep N conversation checkpoints between requests (default 6, ~118 MB\n"
|
||||
" of RAM each; 0 = read every prompt from the start)\n"
|
||||
" --prompt-cache-every N --serve: also checkpoint every N fresh prompt tokens (default 16384, 0 = off)\n"
|
||||
" --turn-token ID --serve: the token that opens a chat turn (default 248045, <|im_start|>)\n"
|
||||
" --short-read N --serve: read at most N fresh text tokens through the decode windows instead\n"
|
||||
" of the batched prompt path (default 64, 0 = off)\n"
|
||||
" --no-token-graph A/B: two graphs per layer (the host launches each) instead of one per token\n"
|
||||
" --no-fused-gr A/B: the six-kernel hyper-connection read and a separate write (native)\n"
|
||||
" --prefill CHUNK batched prompt processing in chunks of CHUNK tokens (needs --native)\n"
|
||||
@@ -442,6 +461,84 @@ int argmax(const std::vector<float>& v) {
|
||||
return best;
|
||||
}
|
||||
|
||||
// ---- --serve's conversation cache. A chat or an agent sends the whole conversation again with every request, and
|
||||
// reading it again is what made a long session wait minutes for every turn. What a sequence leaves behind splits in
|
||||
// two, and only one half needs copying:
|
||||
// * POSITIONAL state - the KV cache of the 12 QSA layers and their pooled indexer keys, the draft layer's KV. A
|
||||
// cell is written once for its position and read only by later positions (the block scores take `dead` for the
|
||||
// block being filled, never its pooled row), so rewinding to a position just means writing from there again.
|
||||
// * RUNNING state - the 36 GDN recurrences and conv histories, each QSA layer's indexer tail (the unfinished
|
||||
// block's raw keys) and the PLE's normalized history. Each describes "everything so far" and cannot be
|
||||
// rewound, so a checkpoint is a copy of exactly these: ~118 MB, the same set the verifier snapshots to roll
|
||||
// back rejected drafts.
|
||||
// A checkpoint is only valid while the positional cells below it still hold ITS tokens, so the serve loop keeps
|
||||
// just the checkpoints that are prefixes of the tokens the session holds now.
|
||||
struct ImgKey {
|
||||
int64_t start = 0; ///< the image's first <|image_pad|> position
|
||||
uint64_t hash = 0; ///< its embeddings and grid: the pad tokens alone are the same for every picture
|
||||
bool operator==(const ImgKey& o) const { return start == o.start && hash == o.hash; }
|
||||
};
|
||||
|
||||
struct ConvCheckpoint {
|
||||
std::vector<int32_t> ids; ///< the tokens this state has consumed
|
||||
std::vector<ImgKey> imgs; ///< the images among them
|
||||
std::vector<uint8_t> gdn, ple, tails;
|
||||
};
|
||||
|
||||
uint64_t fnv1a(const void* data, size_t n, uint64_t h = 1469598103934665603ull) {
|
||||
const uint8_t* p = (const uint8_t*) data;
|
||||
for (size_t i = 0; i < n; ++i) { h ^= p[i]; h *= 1099511628211ull; }
|
||||
return h;
|
||||
}
|
||||
|
||||
struct ConvStateSizes {
|
||||
size_t gdn = 0, ple = 0, tail = 0;
|
||||
};
|
||||
|
||||
ConvStateSizes conv_state_sizes(const strata::core::ModelGeometry& g) {
|
||||
ConvStateSizes z;
|
||||
z.gdn = (size_t) g.n_gdn_layers() *
|
||||
((size_t) g.ssm_state_size * (size_t) g.ssm_v_heads * (size_t) g.ssm_state_size +
|
||||
(size_t) g.ssm_conv_channels * (size_t) (g.ssm_d_conv - 1)) * sizeof(float);
|
||||
z.ple = (size_t) strata::kernels::NG_HIST * (size_t) strata::kernels::NG_HC_DIM * sizeof(float);
|
||||
z.tail = (size_t) (strata::kernels::qsa_real_shapes().idx_block - 1) * (size_t) g.idx_key_dim * sizeof(float);
|
||||
return z;
|
||||
}
|
||||
|
||||
/// Copies the running state out. The caller has synchronized the device.
|
||||
bool checkpoint_save(ConvCheckpoint& c, const strata::core::SessionState& ss, const strata::core::ModelGeometry& g) {
|
||||
const ConvStateSizes z = conv_state_sizes(g);
|
||||
c.gdn.resize(z.gdn);
|
||||
c.ple.resize(ss.ple_hist != nullptr ? z.ple : 0);
|
||||
c.tails.resize(z.tail * (size_t) g.n_qsa_layers());
|
||||
if (cudaMemcpy(c.gdn.data(), ss.gdn_state, z.gdn, cudaMemcpyDeviceToHost) != cudaSuccess) return false;
|
||||
if (!c.ple.empty() && cudaMemcpy(c.ple.data(), ss.ple_hist, z.ple, cudaMemcpyDeviceToHost) != cudaSuccess)
|
||||
return false;
|
||||
for (int64_t i = 0; i < g.n_qsa_layers(); ++i)
|
||||
if (cudaMemcpy(c.tails.data() + (size_t) i * z.tail, ss.qsa_states[i].idx_tail, z.tail, cudaMemcpyDeviceToHost) !=
|
||||
cudaSuccess)
|
||||
return false;
|
||||
return true;
|
||||
}
|
||||
|
||||
/// Puts a checkpoint's running state back; the positional cells below it are the caller's to guarantee.
|
||||
bool checkpoint_restore(const ConvCheckpoint& c, strata::core::SessionState& ss, const strata::core::ModelGeometry& g) {
|
||||
const ConvStateSizes z = conv_state_sizes(g);
|
||||
if (c.gdn.size() != z.gdn || c.tails.size() != z.tail * (size_t) g.n_qsa_layers()) return false;
|
||||
if (cudaMemcpy(ss.gdn_state, c.gdn.data(), z.gdn, cudaMemcpyHostToDevice) != cudaSuccess) return false;
|
||||
if (!c.ple.empty() && cudaMemcpy(ss.ple_hist, c.ple.data(), z.ple, cudaMemcpyHostToDevice) != cudaSuccess)
|
||||
return false;
|
||||
for (int64_t i = 0; i < g.n_qsa_layers(); ++i)
|
||||
if (cudaMemcpy(ss.qsa_states[i].idx_tail, c.tails.data() + (size_t) i * z.tail, z.tail, cudaMemcpyHostToDevice) !=
|
||||
cudaSuccess)
|
||||
return false;
|
||||
// the PLE's token window is the last two tokens, OLDEST FIRST, -1 where there is none (as session_zero leaves it)
|
||||
const size_t L = c.ids.size();
|
||||
ss.ple_prev[0] = L >= 2 ? c.ids[L - 2] : -1;
|
||||
ss.ple_prev[1] = L >= 1 ? c.ids[L - 1] : -1;
|
||||
return cudaDeviceSynchronize() == cudaSuccess;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
@@ -557,6 +654,10 @@ int main(int argc, char** argv) {
|
||||
else if (a == "--pcie-mode") o.pcie_mode = next("--pcie-mode");
|
||||
else if (a == "--serve") o.serve = true;
|
||||
else if (a == "--vision") o.vision = true;
|
||||
else if (a == "--prompt-cache") o.prompt_cache = std::max(0, std::atoi(next("--prompt-cache")));
|
||||
else if (a == "--prompt-cache-every") o.prompt_cache_every = std::max(0LL, std::atoll(next("--prompt-cache-every")));
|
||||
else if (a == "--turn-token") o.turn_token = std::atoll(next("--turn-token"));
|
||||
else if (a == "--short-read") o.short_read = std::max(0LL, std::atoll(next("--short-read")));
|
||||
else if (a == "--no-spec-split") o.spec_split = false;
|
||||
else if (a == "--eos-ids") {
|
||||
std::string e;
|
||||
@@ -1812,11 +1913,15 @@ int main(int argc, char** argv) {
|
||||
//
|
||||
// and each generated token is written to stdout as `T <id>` as soon as its verify window is done, followed by
|
||||
//
|
||||
// DONE <generated> <prompt_tokens> <prompt_ms> <decode_ms> <stop|length>
|
||||
// DONE <generated> <prompt_tokens> <prompt_ms> <decode_ms> <stop|length|cancel> <drafts accepted>
|
||||
// <drafts offered> <prompt tokens reused>
|
||||
//
|
||||
// (`ERR <message>` instead when a request cannot run; `QUIT` ends the process). Every request starts from an
|
||||
// empty sequence (`session_zero`): the prompt goes through the batched prompt path and its last token through
|
||||
// the first verify window - the path all three model files share. Decoding is greedy.
|
||||
// Before that, `RESUME <n>` (n prompt tokens are not read again), `PP <position> <prompt_tokens> <ms> <tok/s>`
|
||||
// after every prompt chunk, and `REUSED <n>` once the prompt is read. (`ERR <message>` instead when a request
|
||||
// cannot run; `STOP` ends the running request at its next step; `QUIT` ends the process.) A request continues
|
||||
// from the live session or the longest conversation checkpoint its prompt starts with (see ConvCheckpoint),
|
||||
// otherwise from an empty sequence (`session_zero`); the rest of the prompt goes through the batched prompt path
|
||||
// and its last token through the first verify window - the path all three model files share. Decoding is greedy.
|
||||
if (o.serve) {
|
||||
if (o.spec < 2 || o.mtp.empty() || o.prefill_chunk <= 0 || thits.d_res == nullptr || host_res.empty()) {
|
||||
std::fprintf(stderr, "strata serve: needs --spec T, --mtp DIR, --prefill CHUNK, --expert-profile P and "
|
||||
@@ -1876,10 +1981,47 @@ int main(int argc, char** argv) {
|
||||
ver.set_pcie_mode(o.pcie_mode == "dma" ? 0 : o.pcie_mode == "direct" ? 1 : o.pcie_mode == "kernel" ? 2
|
||||
: native_pack ? 0 : 2); // auto: DMA for the native packs, the copy kernel for Q2_0
|
||||
std::vector<int64_t> cur;
|
||||
// ---- the conversation cache (see ConvCheckpoint). `live` is what the session holds right now: the tokens
|
||||
// it has consumed, so a request that starts with exactly them continues without any copy. `checks` are the
|
||||
// saved points; every one of them is a prefix of `live` (the loop drops the rest).
|
||||
std::vector<int32_t> live;
|
||||
std::vector<ImgKey> live_imgs, req_imgs;
|
||||
bool live_ok = false;
|
||||
std::vector<ConvCheckpoint> checks;
|
||||
int64_t pp_total = 0, pp_from = 0, pp_next_check = 0;
|
||||
Clock::time_point pp_t0 = Clock::now();
|
||||
auto imgs_below = [&](const std::vector<ImgKey>& all, int64_t L) {
|
||||
std::vector<ImgKey> v;
|
||||
for (const ImgKey& k : all) if (k.start < L) v.push_back(k);
|
||||
return v;
|
||||
};
|
||||
// a checkpoint of the state after `cur[0, L)`; false only when the copy itself failed
|
||||
auto checkpoint_at = [&](int64_t L) -> bool {
|
||||
if (o.prompt_cache <= 0 || L < 1) return true;
|
||||
for (const ConvCheckpoint& c : checks) if ((int64_t) c.ids.size() == L) return true;
|
||||
ConvCheckpoint c;
|
||||
c.ids.assign(cur.begin(), cur.begin() + L);
|
||||
c.imgs = imgs_below(req_imgs, L);
|
||||
if (cudaDeviceSynchronize() != cudaSuccess || !checkpoint_save(c, ss, g)) return false;
|
||||
checks.push_back(std::move(c));
|
||||
while ((int) checks.size() > o.prompt_cache) checks.erase(checks.begin()); // the oldest goes first
|
||||
return true;
|
||||
};
|
||||
sp.on_chunk = [&](const float* R_rows, int64_t T, int64_t p0, std::string& e) -> bool {
|
||||
std::vector<int32_t> nxt((size_t) T);
|
||||
for (int64_t t = 0; t < T; ++t) nxt[(size_t) t] = (int32_t) cur[(size_t) (p0 + t + 1)];
|
||||
return mtp.prefill(R_rows, nxt.data(), T, p0, e);
|
||||
if (!mtp.prefill(R_rows, nxt.data(), T, p0, e)) return false;
|
||||
// progress for the server window: PP <position reached> <prompt tokens> <ms> <fresh tokens/s>
|
||||
const int64_t done = p0 + T;
|
||||
const double ms = std::chrono::duration<double, std::milli>(Clock::now() - pp_t0).count();
|
||||
std::printf("PP %lld %lld %.0f %.1f\n", (long long) done, (long long) pp_total, ms,
|
||||
ms > 0.0 ? 1000.0 * (double) (done - pp_from) / ms : 0.0);
|
||||
std::fflush(stdout);
|
||||
if (o.prompt_cache_every > 0 && done >= pp_next_check) {
|
||||
if (!checkpoint_at(done)) { e = "saving a conversation checkpoint failed"; return false; }
|
||||
pp_next_check = done + o.prompt_cache_every;
|
||||
}
|
||||
return true;
|
||||
};
|
||||
drive.d.plan = ver.plan_sink();
|
||||
drive.d.pcie_num = std::max(0, std::min(256, (int) (o.pcie_frac * 256.0 + 0.5)));
|
||||
@@ -2025,6 +2167,7 @@ int main(int argc, char** argv) {
|
||||
continue;
|
||||
}
|
||||
const int64_t n = (int64_t) ids.size();
|
||||
req_imgs.clear();
|
||||
if (geni && !o.vision) { std::printf("ERR this engine was started without --vision\n"); continue; }
|
||||
if (geni || !mrope_identity) {
|
||||
// positions for every cell this request can reach; the identity again for a text request
|
||||
@@ -2066,6 +2209,12 @@ int main(int argc, char** argv) {
|
||||
if (ids[(size_t) i] != kImagePad) { put(i, p, p, p); ++p; ++i; continue; }
|
||||
if (k >= imgs.size()) { ve = "the prompt has more images than the embeddings file"; break; }
|
||||
const Img& im = imgs[k++];
|
||||
{ // what the conversation cache compares: a picture is its grid and its embeddings
|
||||
const int64_t grid[3] = {im.n, im.nx, im.ny};
|
||||
uint64_t h = fnv1a(grid, sizeof grid);
|
||||
h = fnv1a(img_rows.data() + im.off, (size_t) im.n * (size_t) g.n_embd * sizeof(float), h);
|
||||
req_imgs.push_back({i, h});
|
||||
}
|
||||
for (int64_t j = 0; j < im.n && ve.empty(); ++j)
|
||||
if (i + j >= n || ids[(size_t) (i + j)] != kImagePad)
|
||||
ve = "image " + std::to_string(k) + " has " + std::to_string((long long) im.n) +
|
||||
@@ -2110,43 +2259,183 @@ int main(int argc, char** argv) {
|
||||
if (bad) { std::printf("ERR a token id is outside the vocabulary\n"); continue; }
|
||||
cur = ids;
|
||||
const Clock::time_point r0 = Clock::now();
|
||||
strata::core::session_zero(ss, g, nullptr, main_cs);
|
||||
cudaStreamSynchronize(main_stream);
|
||||
// ---- where this request starts reading: the live session, or a checkpoint, whose tokens AND pictures are
|
||||
// exactly the start of this prompt - at most n - 1 of them, the last token is always the first window
|
||||
auto starts_with = [&](const std::vector<int32_t>& pre, const std::vector<ImgKey>& pre_imgs) -> bool {
|
||||
const int64_t L = (int64_t) pre.size();
|
||||
if (L < 1 || L > n - 1) return false;
|
||||
for (int64_t i = 0; i < L; ++i)
|
||||
if ((int32_t) ids[(size_t) i] != pre[(size_t) i]) return false;
|
||||
return imgs_below(req_imgs, L) == pre_imgs;
|
||||
};
|
||||
int64_t resume = 0;
|
||||
bool from_live = false;
|
||||
if (o.prompt_cache > 0) {
|
||||
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)) {
|
||||
resume = (int64_t) c.ids.size();
|
||||
from_live = false;
|
||||
}
|
||||
}
|
||||
// this request rewrites every cell from `resume` on, so a checkpoint past it (or not on this prompt's
|
||||
// path) no longer has its cells; the ones kept are prefixes of both the old tokens and the new
|
||||
checks.erase(std::remove_if(checks.begin(), checks.end(), [&](const ConvCheckpoint& c) {
|
||||
return (int64_t) c.ids.size() > resume || !starts_with(c.ids, c.imgs);
|
||||
}), checks.end());
|
||||
live_ok = false; // until this request has finished, the session is in between
|
||||
int64_t reread_to = -1; // STRATA_CKPT_REREAD only: read [0, reread_to) again instead of restoring
|
||||
if (resume == 0) {
|
||||
strata::core::session_zero(ss, g, nullptr, main_cs);
|
||||
cudaStreamSynchronize(main_stream);
|
||||
checks.clear();
|
||||
} else if (!from_live) {
|
||||
const ConvCheckpoint* c = nullptr;
|
||||
for (const ConvCheckpoint& k : checks) if ((int64_t) k.ids.size() == resume) c = &k;
|
||||
static const bool reread = std::getenv("STRATA_CKPT_REREAD") != nullptr;
|
||||
if (reread && c != nullptr) {
|
||||
// THE CHECK OF THE CHECKPOINT: instead of restoring it, read its tokens again from position 0 in
|
||||
// one run (below, with the prompt path's slots lent like any read) - the same chunks the request
|
||||
// that saved it read them in, when that request started at 0. With the VRAM expert set fixed
|
||||
// (--adapt-swaps 0) the answer must match the restored one token for token; anything the
|
||||
// checkpoint missed shows up as a difference.
|
||||
strata::core::session_zero(ss, g, nullptr, main_cs);
|
||||
cudaStreamSynchronize(main_stream);
|
||||
reread_to = resume;
|
||||
std::fprintf(stderr, "strata serve: STRATA_CKPT_REREAD: reading %lld tokens again instead of "
|
||||
"restoring\n", (long long) resume);
|
||||
} else if (c == nullptr || !checkpoint_restore(*c, ss, g)) {
|
||||
std::printf("ERR restoring a conversation checkpoint failed\n");
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
tr("request", n, geni ? 1 : 0);
|
||||
mtp.set_prompt_len(n);
|
||||
const int64_t read_from = reread_to > 0 ? 0 : resume;
|
||||
pp_total = n;
|
||||
pp_from = read_from;
|
||||
pp_t0 = r0;
|
||||
pp_next_check = reread_to > 0 ? INT64_MAX : resume + o.prompt_cache_every;
|
||||
std::printf("RESUME %lld\n", (long long) resume); // before reading: this many prompt tokens are reused
|
||||
std::fflush(stdout);
|
||||
// A SHORT PART OF THE PROMPT - the new message of a chat that continues from a checkpoint, the assistant
|
||||
// header - goes through the verify windows, S tokens at a time, as decode reads them. The batched path
|
||||
// costs ~300 ms per run however few tokens it has (it streams every expert the chunk routes to that is
|
||||
// not in VRAM over PCIe), and it borrows slots it must refill after (~180 ms); a window costs ~16 ms a
|
||||
// token, with the misses on the CPU. Each part below is decided on its own, so a long first message is
|
||||
// read batched and its header still goes through the windows. Picture rows need the batched path.
|
||||
// STRATA_CKPT_REREAD compares a restored checkpoint with a batched re-read, so it keeps every read batched.
|
||||
static const bool no_short = std::getenv("STRATA_CKPT_REREAD") != nullptr;
|
||||
auto windows_ok = [&](int64_t a, int64_t b) -> bool {
|
||||
if (no_short || b - a > o.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;
|
||||
return true;
|
||||
};
|
||||
// tokens [a, b) through the windows: commit all of them, then give the draft layer their residuals
|
||||
auto read_windows = [&](int64_t a, int64_t b, std::string& e) -> bool {
|
||||
std::vector<int32_t> win((size_t) S), outw((size_t) S), nxt((size_t) S);
|
||||
for (int64_t q = a; q < b;) {
|
||||
if (stop_req.load()) { e = "cancelled"; return false; }
|
||||
const int T = (int) std::min<int64_t>(S, b - q);
|
||||
for (int t = 0; t < T; ++t) {
|
||||
win[(size_t) t] = (int32_t) cur[(size_t) (q + t)];
|
||||
nxt[(size_t) t] = (int32_t) cur[(size_t) (q + t + 1)];
|
||||
}
|
||||
drive.d.layers = 0;
|
||||
drive.d.experts = 0;
|
||||
drive.d.failed = false;
|
||||
if (!ver.run(T, win.data(), q, &drive_pool_multi, &drive, outw.data(), e) || drive.d.failed) {
|
||||
if (drive.d.failed && drive.d.fail) e = drive.d.fail;
|
||||
return false;
|
||||
}
|
||||
if (!ver.commit(T, e) || !mtp.prefill(ver.final_R_all(), nxt.data(), T, q, e)) return false;
|
||||
q += T;
|
||||
}
|
||||
const double ms = std::chrono::duration<double, std::milli>(Clock::now() - pp_t0).count();
|
||||
std::printf("PP %lld %lld %.0f %.1f\n", (long long) b, (long long) pp_total, ms,
|
||||
ms > 0.0 ? 1000.0 * (double) (b - pp_from) / ms : 0.0);
|
||||
std::fflush(stdout);
|
||||
return true;
|
||||
};
|
||||
// the batched path's slots are lent just before its first run and given back (refilled) before a window
|
||||
// reads - so the windows always see the whole expert cache - or once the prompt is read
|
||||
std::vector<std::pair<int32_t, int32_t>> lent_now;
|
||||
apply_pending(true);
|
||||
if (lend_first >= 0) {
|
||||
auto lend = [&] {
|
||||
if (lend_first < 0 || !lent_now.empty()) return;
|
||||
for (size_t i = 0; i < host_res.size(); ++i)
|
||||
if (host_res[i] >= lend_first) {
|
||||
lent_now.emplace_back((int32_t) i, host_res[i]);
|
||||
host_res[i] = strata::core::kNotResident;
|
||||
}
|
||||
cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice);
|
||||
}
|
||||
};
|
||||
auto refill = [&](std::string& e) -> bool {
|
||||
if (lent_now.empty()) return true;
|
||||
tr("refill start", (long long) lent_now.size());
|
||||
for (const auto& [i, slot] : lent_now) {
|
||||
const uint8_t* b = srcp->blob(i / g.n_expert, i % g.n_expert);
|
||||
if (b == nullptr || !xcache.fill_slot_blocking(slot, b, e,
|
||||
(int64_t) strata::kernels::cpu::expert_layout().blob_bytes(i / g.n_expert)))
|
||||
return false;
|
||||
host_res[(size_t) i] = slot;
|
||||
}
|
||||
cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice);
|
||||
lent_now.clear();
|
||||
return true;
|
||||
};
|
||||
apply_pending(true);
|
||||
bool cancelled = false;
|
||||
tr("prompt start", n - 1);
|
||||
if (n > 1 && !sp.run(ids.data(), n - 1, 0, err)) {
|
||||
if (!stop_req.load()) {
|
||||
std::fprintf(stderr, "strata serve: %s\n", err.c_str());
|
||||
std::printf("ERR %s\n", err.c_str());
|
||||
return 1;
|
||||
}
|
||||
cancelled = true; // stopped while reading the prompt: refill the lent slots below, then DONE cancel
|
||||
}
|
||||
for (const auto& [i, slot] : lent_now) {
|
||||
const uint8_t* b = srcp->blob(i / g.n_expert, i % g.n_expert);
|
||||
if (b == nullptr || !xcache.fill_slot_blocking(slot, b, err,
|
||||
(int64_t) strata::kernels::cpu::expert_layout().blob_bytes(i / g.n_expert))) {
|
||||
// The prompt is read in two parts when it has a turn boundary past `resume`: up to the last <|im_start|>
|
||||
// (the conversation so far), a checkpoint there, then the new turn's header. The next request of the same
|
||||
// chat renders the same history - but not always the same header or the thinking of this reply - so that
|
||||
// checkpoint is the one it reuses.
|
||||
int64_t turn_at = -1;
|
||||
if (o.prompt_cache > 0 && o.turn_token >= 0)
|
||||
for (int64_t i = n - 1; i > resume; --i)
|
||||
if (ids[(size_t) i] == o.turn_token) { turn_at = i; break; }
|
||||
int64_t at = read_from;
|
||||
for (const int64_t to : {reread_to, turn_at, n - 1}) {
|
||||
if (to <= at) continue;
|
||||
const bool win = windows_ok(at, to);
|
||||
if (win && !refill(err)) {
|
||||
std::printf("ERR refilling a lent slot failed: %s\n", err.c_str());
|
||||
return 1;
|
||||
}
|
||||
host_res[(size_t) i] = slot;
|
||||
if (!win) lend();
|
||||
const auto tsp = Clock::now();
|
||||
const bool sp_ok = win ? read_windows(at, to, err) : sp.run(ids.data() + at, to - at, at, err);
|
||||
if (trace) {
|
||||
std::fprintf(stderr, "strata trace: read %lld tokens (%s) in %.1f ms\n", (long long) (to - at),
|
||||
win ? "windows" : "batched",
|
||||
std::chrono::duration<double, std::milli>(Clock::now() - tsp).count());
|
||||
std::fflush(stderr);
|
||||
}
|
||||
if (!sp_ok) {
|
||||
if (!stop_req.load()) {
|
||||
std::fprintf(stderr, "strata serve: %s\n", err.c_str());
|
||||
std::printf("ERR %s\n", err.c_str());
|
||||
return 1;
|
||||
}
|
||||
cancelled = true; // stopped while reading the prompt: refill the lent slots below, then DONE cancel
|
||||
break;
|
||||
}
|
||||
at = to;
|
||||
if (to == turn_at && !checkpoint_at(to)) {
|
||||
std::printf("ERR saving a conversation checkpoint failed\n");
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
if (!refill(err)) {
|
||||
std::printf("ERR refilling a lent slot failed: %s\n", err.c_str());
|
||||
return 1;
|
||||
}
|
||||
if (!lent_now.empty())
|
||||
cudaMemcpy(d_res, host_res.data(), host_res.size() * sizeof(int32_t), cudaMemcpyHostToDevice);
|
||||
tr("prompt done (slots refilled)");
|
||||
const double prompt_ms = std::chrono::duration<double, std::milli>(Clock::now() - r0).count();
|
||||
std::printf("REUSED %lld\n", (long long) resume); // the prompt is read; the first window comes next
|
||||
std::fflush(stdout);
|
||||
// the verify windows: the first holds the last prompt token alone
|
||||
int64_t p = n - 1;
|
||||
int32_t x = (int32_t) ids[(size_t) (n - 1)];
|
||||
@@ -2154,6 +2443,11 @@ int main(int argc, char** argv) {
|
||||
std::vector<float> dprob((size_t) S, 0.0f);
|
||||
bool first_window = true;
|
||||
int64_t produced_n = 0;
|
||||
int64_t draft_offered = 0, draft_accepted = 0;
|
||||
// what the session holds once this request is done: the prompt read so far, then every committed token
|
||||
std::vector<int32_t> consumed;
|
||||
consumed.reserve((size_t) (n + max_new + S));
|
||||
for (int64_t i = 0; i < n - 1; ++i) consumed.push_back((int32_t) ids[(size_t) i]);
|
||||
const char* finish = "length";
|
||||
const Clock::time_point d0 = Clock::now();
|
||||
if (cancelled) finish = "cancel";
|
||||
@@ -2187,6 +2481,10 @@ int main(int argc, char** argv) {
|
||||
std::printf("ERR %s\n", err.c_str());
|
||||
return 1;
|
||||
}
|
||||
// the window's first a + 1 tokens are in the session now (the last output is not: it is next x)
|
||||
for (int i = 0; i <= a; ++i) consumed.push_back(window[(size_t) i]);
|
||||
draft_offered += T - 1;
|
||||
draft_accepted += a;
|
||||
first_window = false;
|
||||
bool eos = false;
|
||||
for (int i = 0; i <= a && produced_n < max_new && !eos; ++i) {
|
||||
@@ -2213,12 +2511,80 @@ int main(int argc, char** argv) {
|
||||
p += a + 1;
|
||||
}
|
||||
const double decode_ms = std::chrono::duration<double, std::milli>(Clock::now() - d0).count();
|
||||
std::printf("DONE %lld %lld %.1f %.1f %s\n", (long long) produced_n, (long long) n, prompt_ms, decode_ms,
|
||||
finish);
|
||||
if (!cancelled) {
|
||||
// a prompt stopped halfway leaves the session somewhere between two chunks: nothing to continue from
|
||||
// (the checkpoints taken while reading it are still good)
|
||||
live.swap(consumed);
|
||||
live_imgs = imgs_below(req_imgs, (int64_t) live.size());
|
||||
live_ok = o.prompt_cache > 0;
|
||||
}
|
||||
static const bool state_hash = std::getenv("STRATA_STATE_HASH") != nullptr;
|
||||
if (state_hash && live_ok) {
|
||||
// DEBUG: a fingerprint of every part of the session over the positions it holds ([0, L)), and
|
||||
// separately of what lies past them in the last KV page (stale cells, fine unless something reads them)
|
||||
cudaDeviceSynchronize();
|
||||
const int64_t L = (int64_t) live.size();
|
||||
const strata::kernels::QsaShapes qs = [&] {
|
||||
strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes();
|
||||
s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_dim = g.idx_key_dim;
|
||||
return s;
|
||||
}();
|
||||
auto hash_dev = [&](const void* p, size_t bytes, uint64_t h) {
|
||||
std::vector<uint8_t> b(bytes);
|
||||
if (bytes) cudaMemcpy(b.data(), p, bytes, cudaMemcpyDeviceToHost);
|
||||
return fnv1a(b.data(), b.size(), h);
|
||||
};
|
||||
// the cells [c0, c1) of one int8 K or V pool ([page][kv_head][page_size][head_dim]), bytes per value `w`
|
||||
auto hash_cells = [&](const void* pool, int64_t per_cell, int64_t c0, int64_t c1, uint64_t h) {
|
||||
const int64_t ps = qs.page_size;
|
||||
for (int64_t pg = c0 / ps; pg * ps < c1; ++pg)
|
||||
for (int64_t hd = 0; hd < qs.n_head_kv; ++hd) {
|
||||
const int64_t a = std::max(c0, pg * ps) - pg * ps, e = std::min(c1, (pg + 1) * ps) - pg * ps;
|
||||
const size_t off = (size_t) (((pg * qs.n_head_kv + hd) * ps + a) * per_cell);
|
||||
h = hash_dev((const uint8_t*) pool + off, (size_t) ((e - a) * per_cell), h);
|
||||
}
|
||||
return h;
|
||||
};
|
||||
const ConvStateSizes z = conv_state_sizes(g);
|
||||
uint64_t h_gdn = hash_dev(ss.gdn_state, z.gdn, 1469598103934665603ull);
|
||||
uint64_t h_ple = hash_dev(ss.ple_hist, z.ple, 1469598103934665603ull);
|
||||
uint64_t h_tail = 1469598103934665603ull, h_pool = h_tail, h_kv = h_tail, h_stale = h_tail;
|
||||
const int64_t kvb = qs.head_dim, scb = (qs.head_dim / 64) * 2;
|
||||
const int64_t end_cell = std::min<int64_t>(((L + qs.page_size - 1) / qs.page_size) * qs.page_size,
|
||||
ss.qsa_states[0].max_cells);
|
||||
for (int64_t i = 0; i < g.n_qsa_layers(); ++i) {
|
||||
const strata::core::QsaState& st = ss.qsa_states[i];
|
||||
h_tail = hash_dev(st.idx_tail, z.tail, h_tail);
|
||||
h_pool = hash_dev(st.idx_pooled, (size_t) (L / qs.idx_block) * qs.idx_dim * 4, h_pool);
|
||||
for (const auto& [pool, w] : {std::pair<const void*, int64_t>{st.k_q, kvb}, {st.v_q, kvb},
|
||||
{st.k_scale, scb}, {st.v_scale, scb}}) {
|
||||
h_kv = hash_cells(pool, w, 0, L, h_kv);
|
||||
h_stale = hash_cells(pool, w, L, end_cell, h_stale);
|
||||
}
|
||||
}
|
||||
const strata::core::QsaState& ms = mtp.kv_state();
|
||||
uint64_t h_mtp = 1469598103934665603ull;
|
||||
const int64_t mL = std::min<int64_t>(L, ms.max_cells);
|
||||
for (const auto& [pool, w] : {std::pair<const void*, int64_t>{ms.k_q, kvb}, {ms.v_q, kvb},
|
||||
{ms.k_scale, scb}, {ms.v_scale, scb}})
|
||||
if (pool != nullptr) h_mtp = hash_cells(pool, w, 0, mL, h_mtp);
|
||||
std::fprintf(stderr, "strata serve: STATE_HASH L=%lld gdn=%016llx ple=%016llx tail=%016llx pooled=%016llx "
|
||||
"kv=%016llx mtp=%016llx stale=%016llx ple_prev=%d,%d\n", (long long) L,
|
||||
(unsigned long long) h_gdn, (unsigned long long) h_ple, (unsigned long long) h_tail,
|
||||
(unsigned long long) h_pool, (unsigned long long) h_kv, (unsigned long long) h_mtp,
|
||||
(unsigned long long) h_stale, ss.ple_prev[0], ss.ple_prev[1]);
|
||||
}
|
||||
// DONE <generated> <prompt> <prompt ms> <decode ms> <finish> <drafts accepted> <drafts offered> <reused>
|
||||
std::printf("DONE %lld %lld %.1f %.1f %s %lld %lld %lld\n", (long long) produced_n, (long long) n, prompt_ms,
|
||||
decode_ms, finish, (long long) draft_accepted, (long long) draft_offered, (long long) resume);
|
||||
std::fflush(stdout);
|
||||
std::fprintf(stderr, "strata serve: %lld prompt tokens in %.0f ms (%.1f tok/s), %lld generated in %.0f ms "
|
||||
"(%.1f tok/s)\n", (long long) n, prompt_ms, prompt_ms > 0 ? 1000.0 * n / prompt_ms : 0.0,
|
||||
(long long) produced_n, decode_ms, decode_ms > 0 ? 1000.0 * produced_n / decode_ms : 0.0);
|
||||
const int64_t fresh = n - resume;
|
||||
std::fprintf(stderr, "strata serve: prompt %lld tokens = %lld reused + %lld read in %.0f ms (%.1f tok/s), "
|
||||
"%lld generated in %.0f ms (%.1f tok/s), drafts accepted %lld of %lld, %zu checkpoints%s\n",
|
||||
(long long) n, (long long) resume, (long long) fresh, prompt_ms,
|
||||
prompt_ms > 0 ? 1000.0 * fresh / prompt_ms : 0.0, (long long) produced_n, decode_ms,
|
||||
decode_ms > 0 ? 1000.0 * produced_n / decode_ms : 0.0, (long long) draft_accepted,
|
||||
(long long) draft_offered, checks.size(), cancelled ? " (cancelled)" : "");
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user